1 /*
2  * Copyright (C) 2018 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 "Operations"
18 
19 #include "OperationResolver.h"
20 #include "OperationsUtils.h"
21 #include "Tracing.h"
22 
23 namespace android {
24 namespace nn {
25 namespace gather {
26 
27 constexpr char kOperationName[] = "GATHER";
28 
29 constexpr uint32_t kNumInputs = 3;
30 constexpr uint32_t kInputTensor = 0;
31 constexpr uint32_t kInputAxis = 1;
32 constexpr uint32_t kInputIndices = 2;
33 
34 constexpr uint32_t kNumOutputs = 1;
35 constexpr uint32_t kOutputTensor = 0;
36 
37 namespace {
38 
39 template <typename T>
eval(const T * inputData,const Shape & inputShape,int32_t axis,const int32_t * indicesData,const Shape & indicesShape,T * outputData)40 inline bool eval(const T* inputData, const Shape& inputShape, int32_t axis,
41                  const int32_t* indicesData, const Shape& indicesShape, T* outputData) {
42     const auto outerSize = getNumberOfElements(inputShape, 0, axis);
43     const auto axisSize = getSizeOfDimension(inputShape, axis);
44     const auto innerSize =
45             getNumberOfElements(inputShape, axis + 1, getNumberOfDimensions(inputShape));
46     const auto indicesCount = getNumberOfElements(indicesShape);
47     for (uint32_t outer = 0; outer < outerSize; ++outer) {
48         for (uint32_t outputIndex = 0; outputIndex < indicesCount; ++outputIndex) {
49             const auto inputIndex = static_cast<uint32_t>(indicesData[outputIndex]);
50             NN_RET_CHECK_LE(0u, inputIndex);
51             NN_RET_CHECK_LT(inputIndex, axisSize);
52             std::memcpy(outputData + (outer * indicesCount + outputIndex) * innerSize,
53                         inputData + (outer * axisSize + inputIndex) * innerSize,
54                         sizeof(T) * innerSize);
55         }
56     }
57     return true;
58 }
59 
60 }  // namespace
61 
validate(const IOperationValidationContext * context)62 Result<Version> validate(const IOperationValidationContext* context) {
63     NN_RET_CHECK_EQ(context->getNumInputs(), kNumInputs);
64     NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs);
65     OperandType inputType = context->getInputType(kInputTensor);
66     NN_RET_CHECK(inputType == OperandType::TENSOR_FLOAT16 ||
67                  inputType == OperandType::TENSOR_FLOAT32 ||
68                  inputType == OperandType::TENSOR_INT32 ||
69                  inputType == OperandType::TENSOR_QUANT8_ASYMM ||
70                  inputType == OperandType::TENSOR_QUANT8_ASYMM_SIGNED)
71             << "Unsupported tensor type for operation " << kOperationName;
72     NN_RET_CHECK(validateInputTypes(context,
73                                     {inputType, OperandType::INT32, OperandType::TENSOR_INT32}));
74     NN_RET_CHECK(validateOutputTypes(context, {inputType}));
75     if (inputType == OperandType::TENSOR_QUANT8_ASYMM_SIGNED) {
76         return Version::ANDROID_R;
77     } else {
78         return Version::ANDROID_Q;
79     }
80 }
81 
prepare(IOperationExecutionContext * context)82 bool prepare(IOperationExecutionContext* context) {
83     Shape input = context->getInputShape(kInputTensor);
84     int32_t axis = context->getInputValue<int32_t>(kInputAxis);
85     NN_RET_CHECK(handleNegativeAxis(input, &axis));
86     Shape indices = context->getInputShape(kInputIndices);
87     Shape output = context->getOutputShape(kOutputTensor);
88 
89     output.dimensions.clear();
90     output.dimensions.reserve(getNumberOfDimensions(input) + getNumberOfDimensions(indices) - 1);
91     output.dimensions.insert(output.dimensions.end(), input.dimensions.begin(),
92                              input.dimensions.begin() + axis);
93     output.dimensions.insert(output.dimensions.end(), indices.dimensions.begin(),
94                              indices.dimensions.end());
95     output.dimensions.insert(output.dimensions.end(), input.dimensions.begin() + axis + 1,
96                              input.dimensions.end());
97 
98     return context->setOutputShape(kOutputTensor, output);
99 }
100 
execute(IOperationExecutionContext * context)101 bool execute(IOperationExecutionContext* context) {
102     int32_t axis = context->getInputValue<int32_t>(kInputAxis);
103     NN_RET_CHECK(handleNegativeAxis(context->getInputShape(kInputTensor), &axis));
104     switch (context->getInputType(kInputTensor)) {
105         case OperandType::TENSOR_FLOAT16:
106             return eval(context->getInputBuffer<_Float16>(kInputTensor),
107                         context->getInputShape(kInputTensor), axis,
108                         context->getInputBuffer<int32_t>(kInputIndices),
109                         context->getInputShape(kInputIndices),
110                         context->getOutputBuffer<_Float16>(kOutputTensor));
111         case OperandType::TENSOR_FLOAT32:
112             return eval(context->getInputBuffer<float>(kInputTensor),
113                         context->getInputShape(kInputTensor), axis,
114                         context->getInputBuffer<int32_t>(kInputIndices),
115                         context->getInputShape(kInputIndices),
116                         context->getOutputBuffer<float>(kOutputTensor));
117         case OperandType::TENSOR_INT32:
118             return eval(context->getInputBuffer<int32_t>(kInputTensor),
119                         context->getInputShape(kInputTensor), axis,
120                         context->getInputBuffer<int32_t>(kInputIndices),
121                         context->getInputShape(kInputIndices),
122                         context->getOutputBuffer<int32_t>(kOutputTensor));
123         case OperandType::TENSOR_QUANT8_ASYMM:
124             return eval(context->getInputBuffer<uint8_t>(kInputTensor),
125                         context->getInputShape(kInputTensor), axis,
126                         context->getInputBuffer<int32_t>(kInputIndices),
127                         context->getInputShape(kInputIndices),
128                         context->getOutputBuffer<uint8_t>(kOutputTensor));
129         case OperandType::TENSOR_QUANT8_ASYMM_SIGNED:
130             return eval(context->getInputBuffer<int8_t>(kInputTensor),
131                         context->getInputShape(kInputTensor), axis,
132                         context->getInputBuffer<int32_t>(kInputIndices),
133                         context->getInputShape(kInputIndices),
134                         context->getOutputBuffer<int8_t>(kOutputTensor));
135         default:
136             NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName;
137     }
138 }
139 
140 }  // namespace gather
141 
142 NN_REGISTER_OPERATION(GATHER, gather::kOperationName, gather::validate, gather::prepare,
143                       gather::execute);
144 
145 }  // namespace nn
146 }  // namespace android
147