1/*
2 * Copyright (C) 2017 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#define LOG_TAG "ModelBuilder"
18
19#include "ModelBuilder.h"
20
21#include "CompilationBuilder.h"
22#include "Utils.h"
23
24#include <map>
25#include <utility>
26
27namespace android {
28namespace nn {
29
30// The maximum number of operands and operations that a model may have.
31const uint32_t MAX_NUMBER_OF_OPERANDS = 0xFFFFFFFE;
32const uint32_t MAX_NUMBER_OF_OPERATIONS = 0xFFFFFFFE;
33
34int ModelBuilder::addOperand(const ANeuralNetworksOperandType& type) {
35    if (mCompletedModel) {
36        LOG(ERROR) << "ANeuralNetworksModel_addOperand can't modify after model finished";
37        return ANEURALNETWORKS_BAD_DATA;
38    }
39    int n = validateOperandType(type, "ANeuralNetworksModel_addOperand", true);
40    if (n != ANEURALNETWORKS_NO_ERROR) {
41        return n;
42    }
43    size_t idx = mOperands.size();
44    if (idx >= MAX_NUMBER_OF_OPERANDS) {
45        LOG(ERROR) << "ANeuralNetworksModel_addOperand exceed max operands";
46        return ANEURALNETWORKS_BAD_DATA;
47    }
48    mOperands.resize(idx + 1);
49    auto& operand = mOperands[idx];
50    operand.type = static_cast<OperandType>(type.type);
51    setFromIntList(&operand.dimensions, type.dimensionCount, type.dimensions);
52    operand.numberOfConsumers = 0;
53    operand.scale = type.scale;
54    operand.zeroPoint = type.zeroPoint;
55    operand.lifetime = OperandLifeTime::TEMPORARY_VARIABLE;
56    operand.location = {.poolIndex = 0, .offset = 0, .length = 0};
57    return ANEURALNETWORKS_NO_ERROR;
58}
59
60int ModelBuilder::setOperandValue(uint32_t index, const void* buffer, size_t length) {
61    VLOG(MODEL) << __func__ << " for operand " << index << " size " << length;
62    if (index >= operandCount()) {
63        LOG(ERROR) << "ANeuralNetworksModel_setOperandValue setting operand " << index << " of "
64                   << operandCount();
65        return ANEURALNETWORKS_BAD_DATA;
66    }
67    Operand& operand = mOperands[index];
68    if (buffer == nullptr) {
69        if (length) {
70            LOG(ERROR) << "ANeuralNetworksModel_setOperandValue buffer is nullptr but length is "
71                          "not 0";
72            return ANEURALNETWORKS_BAD_DATA;
73        }
74        operand.lifetime = OperandLifeTime::NO_VALUE;
75        // The location is unused and is set to zeros.
76        operand.location = {.poolIndex = 0,
77                            .offset = 0,
78                            .length = 0};
79    } else {
80        if (length > 0xFFFFFFFF) {
81            LOG(ERROR) << "ANeuralNetworksModel_setOperandValue value length of " << length
82                       << " exceeds max size";
83            return ANEURALNETWORKS_BAD_DATA;
84        }
85        uint32_t valueLength = static_cast<uint32_t>(length);
86        uint32_t neededLength = sizeOfData(operand.type, operand.dimensions);
87        if (neededLength != valueLength) {
88            LOG(ERROR) << "ANeuralNetworksModel_setOperandValue setting " << valueLength
89                       << " bytes when needing " << neededLength;
90            return ANEURALNETWORKS_BAD_DATA;
91        }
92        if (valueLength <= ANEURALNETWORKS_MAX_SIZE_OF_IMMEDIATELY_COPIED_VALUES) {
93            uint32_t existingSize = static_cast<uint32_t>(mSmallOperandValues.size());
94            uint32_t extraBytes = alignBytesNeeded(existingSize, valueLength);
95            mSmallOperandValues.resize(existingSize + extraBytes + valueLength);
96            operand.lifetime = OperandLifeTime::CONSTANT_COPY;
97            operand.location = {
98                .poolIndex = 0, .offset = existingSize + extraBytes, .length = neededLength};
99            memcpy(&mSmallOperandValues[operand.location.offset], buffer, valueLength);
100            VLOG(MODEL) << "Copied small value to offset " << operand.location.offset;
101        } else {
102            VLOG(MODEL) << "Saving large value";
103            operand.lifetime = OperandLifeTime::CONSTANT_REFERENCE;
104            // The values for poolIndex and offset will be set when the model is finished.
105            operand.location = {.poolIndex = 0, .offset = 0, .length = valueLength};
106            // We keep track of the buffers. We'll allocate the shared memory only
107            // once we know the total size, to avoid needless copies.
108            mLargeOperandValues.push_back(LargeValue{.operandIndex = index, .buffer = buffer});
109        }
110    }
111    return ANEURALNETWORKS_NO_ERROR;
112}
113
114int ModelBuilder::copyLargeValuesToSharedMemory() {
115    VLOG(MODEL) << __func__ << " has " << mLargeOperandValues.size() << " values.";
116    if (!mLargeOperandValues.empty()) {
117        // Calculate the size of the shared memory needed for all the large values.
118        // Also sets the offset for each value within the memory.
119        size_t poolSize = 0;
120        for (LargeValue& l: mLargeOperandValues) {
121            Operand& operand = mOperands[l.operandIndex];
122            nnAssert(operand.lifetime == OperandLifeTime::CONSTANT_REFERENCE);
123            poolSize += alignBytesNeeded(poolSize, operand.location.length);
124            operand.location.offset = poolSize;
125            poolSize += operand.location.length;
126        }
127
128        // Allocated the shared memory.
129        int n = mLargeValueMemory.create(poolSize);
130        if (n != ANEURALNETWORKS_NO_ERROR) {
131            return n;
132        }
133        uint8_t* memoryPointer = nullptr;
134        n = mLargeValueMemory.getPointer(&memoryPointer);
135        if (n != ANEURALNETWORKS_NO_ERROR) {
136            return n;
137        }
138        uint32_t poolIndex = mMemories.add(&mLargeValueMemory);
139        VLOG(MODEL) << "Allocated large value pool of size " << poolSize << " at index "
140                    << poolIndex;
141
142        // Copy the values to this memory.
143        for (LargeValue& l: mLargeOperandValues) {
144            Operand& operand = mOperands[l.operandIndex];
145            operand.location.poolIndex = poolIndex;
146            memcpy(memoryPointer + operand.location.offset, l.buffer, operand.location.length);
147        }
148    }
149    return ANEURALNETWORKS_NO_ERROR;
150}
151
152int ModelBuilder::setOperandValueFromMemory(uint32_t index, const Memory* memory, uint32_t offset,
153                                            size_t length) {
154    VLOG(MODEL) << __func__ << " for operand " << index << " offset " << offset << " size " << length;
155    if (index >= operandCount()) {
156        LOG(ERROR) << "ANeuralNetworksModel_setOperandValueFromMemory setting operand " << index
157                   << " of " << operandCount();
158        return ANEURALNETWORKS_BAD_DATA;
159    }
160    Operand& operand = mOperands[index];
161    uint32_t neededLength = sizeOfData(operand.type, operand.dimensions);
162    if (neededLength != length) {
163        LOG(ERROR) << "ANeuralNetworksModel_setOperandValueFromMemory setting " << length
164                   << " bytes when needing " << neededLength;
165        return ANEURALNETWORKS_BAD_DATA;
166    }
167    // TODO validate does not exceed length of memory
168    operand.lifetime = OperandLifeTime::CONSTANT_REFERENCE;
169    operand.location = {
170                .poolIndex = mMemories.add(memory), .offset = offset, .length = neededLength};
171    return ANEURALNETWORKS_NO_ERROR;
172}
173
174int ModelBuilder::addOperation(ANeuralNetworksOperationType type, uint32_t inputCount,
175                               const uint32_t* inputs, uint32_t outputCount,
176                               const uint32_t* outputs) {
177    if (mCompletedModel) {
178        LOG(ERROR) << "ANeuralNetworksModel_addOperation can't modify after model finished";
179        return ANEURALNETWORKS_BAD_DATA;
180    }
181    if (!validCode(kNumberOfOperationTypes, kNumberOfOperationTypesOEM, type)) {
182        LOG(ERROR) << "ANeuralNetworksModel_addOperation invalid operations type " << type;
183        return ANEURALNETWORKS_BAD_DATA;
184    }
185    int n = validateOperandList(inputCount, inputs, operandCount(),
186                                "ANeuralNetworksModel_addOperation inputs");
187    if (n != ANEURALNETWORKS_NO_ERROR) {
188        return n;
189    }
190    n = validateOperandList(outputCount, outputs, operandCount(),
191                            "ANeuralNetworksModel_addOperation outputs");
192    if (n != ANEURALNETWORKS_NO_ERROR) {
193        return n;
194    }
195
196    uint32_t operationIndex = operationCount();
197    if (operationIndex >= MAX_NUMBER_OF_OPERATIONS) {
198        LOG(ERROR) << "ANeuralNetworksModel_addOperation exceed max operations";
199        return ANEURALNETWORKS_BAD_DATA;
200    }
201    mOperations.resize(operationIndex + 1);
202    auto& entry = mOperations[operationIndex];
203    entry.type = static_cast<OperationType>(type);
204
205    setFromIntList(&entry.inputs, inputCount, inputs);
206    setFromIntList(&entry.outputs, outputCount, outputs);
207    for (uint32_t i : entry.inputs) {
208        mOperands[i].numberOfConsumers++;
209        // TODO mOperands[i].consumers.push_back(operationIndex);
210    }
211    return ANEURALNETWORKS_NO_ERROR;
212}
213
214int ModelBuilder::identifyInputsAndOutputs(uint32_t inputCount, const uint32_t* inputs,
215                                      uint32_t outputCount, const uint32_t* outputs) {
216    if (mCompletedModel) {
217        LOG(ERROR) << "ANeuralNetworksModel_identifyInputsAndOutputs can't modify after model finished";
218        return ANEURALNETWORKS_BAD_DATA;
219    }
220    int n = validateOperandList(inputCount, inputs, operandCount(),
221                                "ANeuralNetworksModel_identifyInputsAndOutputs inputs");
222    if (n != ANEURALNETWORKS_NO_ERROR) {
223        return n;
224    }
225    n = validateOperandList(outputCount, outputs, operandCount(),
226                            "ANeuralNetworksModel_identifyInputsAndOutputs outputs");
227    if (n != ANEURALNETWORKS_NO_ERROR) {
228        return n;
229    }
230
231    // Makes a copy of the index list, validates the arguments, and changes
232    // the lifetime info of the corresponding operand.
233    auto setArguments = [&](std::vector<uint32_t>* indexVector, uint32_t indexCount,
234                            const uint32_t* indexList, OperandLifeTime lifetime) -> bool {
235        indexVector->resize(indexCount);
236        for (uint32_t i = 0; i < indexCount; i++) {
237            const uint32_t operandIndex = indexList[i];
238            if (operandIndex >= mOperands.size()) {
239                LOG(ERROR) << "ANeuralNetworksModel_identifyInputsAndOutputs Can't set input or output "
240                              "to be "
241                           << operandIndex << " as this exceeds the numbe of operands "
242                           << mOperands.size();
243                return false;
244            }
245            (*indexVector)[i] = operandIndex;
246            Operand& operand = mOperands[operandIndex];
247            if (operand.lifetime != OperandLifeTime::TEMPORARY_VARIABLE) {
248                LOG(ERROR) << "ANeuralNetworksModel_identifyInputsAndOutputs Can't set operand "
249                           << operandIndex
250                           << " to be an input or output.  Check that it's not a constant or "
251                              "already an input or output";
252                return false;
253            }
254            operand.lifetime = lifetime;
255        }
256        return true;
257    };
258
259    if (!setArguments(&mInputIndexes, inputCount, inputs, OperandLifeTime::MODEL_INPUT) ||
260        !setArguments(&mOutputIndexes, outputCount, outputs, OperandLifeTime::MODEL_OUTPUT)) {
261        return ANEURALNETWORKS_BAD_DATA;
262    }
263
264    return ANEURALNETWORKS_NO_ERROR;
265}
266
267int ModelBuilder::createCompilation(CompilationBuilder** compilation) {
268    if (!mCompletedModel) {
269        LOG(ERROR) << "ANeuralNetworksCompilation_create passed an unfinished model";
270        *compilation = nullptr;
271        return ANEURALNETWORKS_BAD_STATE;
272    }
273    *compilation = new CompilationBuilder(this);
274    return (*compilation ? ANEURALNETWORKS_NO_ERROR : ANEURALNETWORKS_OUT_OF_MEMORY);
275}
276
277int ModelBuilder::finish() {
278    if (mCompletedModel) {
279        LOG(ERROR) << "ANeuralNetworksModel_finish called more than once";
280        return ANEURALNETWORKS_BAD_STATE;
281    }
282
283    int n = copyLargeValuesToSharedMemory();
284    if (n != ANEURALNETWORKS_NO_ERROR) {
285        return n;
286    }
287
288    // We sort the operations so that they will be in the appropriate
289    // order for a single-threaded, op at a time execution.
290    // TODO: we don't need this if we always run the partitioner.
291    sortIntoRunOrder();
292    mCompletedModel = true;
293    return ANEURALNETWORKS_NO_ERROR;
294}
295
296void ModelBuilder::sortIntoRunOrder() {
297    // Tracks the operations that can be executed.
298    std::vector<uint32_t> opsReadyToRun;
299    std::vector<Operation> runOrder;
300
301    // Tracks how many inputs are needed for each operation to be ready to run.
302    std::multimap<uint32_t, uint32_t> operandToOperations;
303    std::vector<uint32_t> unknownInputCount(operationCount());
304    for (uint32_t operationIndex = 0; operationIndex < operationCount(); operationIndex++) {
305        uint32_t& count = unknownInputCount[operationIndex];
306        count = 0;
307        for (uint32_t operandIndex : mOperations[operationIndex].inputs) {
308            auto lifetime = mOperands[operandIndex].lifetime;
309            if (lifetime == OperandLifeTime::TEMPORARY_VARIABLE ||
310                lifetime == OperandLifeTime::MODEL_OUTPUT) {
311                count++;
312                operandToOperations.insert(
313                            std::pair<uint32_t, uint32_t>(operandIndex, operationIndex));
314            }
315        }
316        if (count == 0) {
317            opsReadyToRun.push_back(operationIndex);
318        }
319    }
320
321    while (opsReadyToRun.size() > 0) {
322        // Execute the next op
323        int opIndex = opsReadyToRun.back();
324        opsReadyToRun.pop_back();
325        const Operation& operation = mOperations[opIndex];
326
327        runOrder.push_back(mOperations[opIndex]);
328
329        // Mark all its outputs as known.
330        for (uint32_t operandIndex : operation.outputs) {
331            auto range = operandToOperations.equal_range(operandIndex);
332            for (auto i = range.first; i != range.second; i++) {
333                uint32_t& count = unknownInputCount[i->second];
334                if (--count == 0) {
335                    opsReadyToRun.push_back(i->second);
336                }
337            }
338        }
339    }
340    mOperations = runOrder;
341}
342
343void ModelBuilder::setHidlModel(Model* model) const {
344    model->operands = mOperands;
345    model->operations = mOperations;
346    model->inputIndexes = mInputIndexes;
347    model->outputIndexes = mOutputIndexes;
348    model->operandValues = mSmallOperandValues;
349
350    uint32_t count = mMemories.size();
351    model->pools.resize(count);
352    for (uint32_t i = 0; i < count; i++) {
353        model->pools[i] = mMemories[i]->getHidlMemory();
354    }
355}
356
357}  // namespace nn
358}  // namespace android
359