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 #include "CpuOperationUtils.h" 18 #include "OperationResolver.h" 19 20 #include "tensorflow/lite/kernels/internal/optimized/legacy_optimized_ops.h" 21 #include "tensorflow/lite/kernels/internal/reference/legacy_reference_ops.h" 22 23 #include "Tracing.h" 24 25 namespace android { 26 namespace nn { 27 namespace concatenation { 28 29 constexpr char kOperationName[] = "CONCATENATION"; 30 31 constexpr uint32_t kNumOutputs = 1; 32 constexpr uint32_t kOutputTensor = 0; 33 34 namespace { 35 36 template <typename T> 37 bool concatenation(const std::vector<const T*>& inputDataPtrs, 38 const std::vector<Shape>& inputShapes, int32_t axis, T* outputData, 39 const Shape& outputShape) { 40 NNTRACE_TRANS("concatenation"); 41 int num_inputs = inputShapes.size(); 42 std::vector<tflite::Dims<4>*> inputDimsPtr(num_inputs); 43 std::vector<tflite::Dims<4> > inputDims(num_inputs); 44 for (int i = 0; i < num_inputs; i++) { 45 inputDims[i] = convertShapeToDims(inputShapes[i]); 46 inputDimsPtr[i] = &inputDims[i]; 47 } 48 NNTRACE_COMP_SWITCH("optimized_ops::Concatenation"); 49 tflite::optimized_ops::Concatenation<tflite::FusedActivationFunctionType::kNone, T>( 50 getNumberOfDimensions(outputShape) - axis - 1, inputDataPtrs.data(), 51 inputDimsPtr.data(), num_inputs, outputData, convertShapeToDims(outputShape)); 52 53 return true; 54 } 55 56 template <> 57 bool concatenation<uint8_t>(const std::vector<const uint8_t*>& inputDataPtrs, 58 const std::vector<Shape>& inputShapes, int32_t axis, 59 uint8_t* outputData, const Shape& outputShape) { 60 NNTRACE_TRANS("concatenationQuant8"); 61 int num_inputs = inputShapes.size(); 62 std::vector<float> inputScales(num_inputs); 63 std::vector<int32> inputOffsets(num_inputs); 64 std::vector<tflite::Dims<4>*> inputDimsPtr(num_inputs); 65 std::vector<tflite::Dims<4> > inputDims(num_inputs); 66 for (int i = 0; i < num_inputs; i++) { 67 inputScales[i] = inputShapes[i].scale; 68 inputOffsets[i] = inputShapes[i].offset; 69 inputDims[i] = convertShapeToDims(inputShapes[i]); 70 inputDimsPtr[i] = &inputDims[i]; 71 } 72 73 NNTRACE_COMP_SWITCH("reference_ops::Concatenation"); 74 tflite::reference_ops::Concatenation( 75 getNumberOfDimensions(outputShape) - axis - 1, inputDataPtrs.data(), 76 inputDimsPtr.data(), inputOffsets.data(), inputScales.data(), num_inputs, outputData, 77 convertShapeToDims(outputShape), outputShape.offset, outputShape.scale); 78 79 return true; 80 } 81 82 template <typename T> 83 inline bool concatenation(IOperationExecutionContext* context) { 84 uint32_t inputCount = context->getNumInputs() - 1; 85 std::vector<const T*> inputDatas; 86 std::vector<Shape> inputShapes; 87 for (uint32_t i = 0; i < inputCount; ++i) { 88 const T* buffer = context->getInputBuffer<T>(i); 89 if (buffer == nullptr) continue; 90 inputDatas.push_back(buffer); 91 inputShapes.push_back(context->getInputShape(i)); 92 } 93 return concatenation(inputDatas, inputShapes, context->getInputValue<int32_t>(inputCount), 94 context->getOutputBuffer<T>(kOutputTensor), 95 context->getOutputShape(kOutputTensor)); 96 } 97 98 } // namespace 99 100 bool validate(const IOperationValidationContext* context) { 101 uint32_t inputCount = context->getNumInputs(); 102 NN_RET_CHECK_GE(inputCount, 2); 103 NN_RET_CHECK_EQ(context->getNumOutputs(), kNumOutputs); 104 const OperandType inputType = context->getInputType(0); 105 if (inputType == OperandType::TENSOR_FLOAT32 || inputType == OperandType::TENSOR_QUANT8_ASYMM) { 106 NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_0)); 107 } else if (inputType == OperandType::TENSOR_FLOAT16) { 108 NN_RET_CHECK(validateHalVersion(context, HalVersion::V1_2)); 109 } else { 110 NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName; 111 } 112 std::vector<OperandType> inExpectedTypes(inputCount - 1, inputType); 113 inExpectedTypes.push_back(OperandType::INT32); 114 if (context->getHalVersion() < HalVersion::V1_2 && 115 inputType == OperandType::TENSOR_QUANT8_ASYMM) { 116 const Shape& output = context->getOutputShape(kOutputTensor); 117 for (uint32_t i = 0; i < inputCount - 1; ++i) { 118 const Shape& input = context->getInputShape(i); 119 NN_RET_CHECK_EQ(input.scale, output.scale); 120 NN_RET_CHECK_EQ(input.offset, output.offset); 121 } 122 } 123 return validateInputTypes(context, inExpectedTypes) && 124 validateOutputTypes(context, {inputType}); 125 } 126 127 bool prepare(IOperationExecutionContext* context) { 128 uint32_t numInputs = context->getNumInputs(); 129 NN_RET_CHECK_GE(numInputs, 2); 130 const Shape& input0 = context->getInputShape(0); 131 uint32_t numDimensions = getNumberOfDimensions(input0); 132 int32_t axis = context->getInputValue<int32_t>(numInputs - 1); 133 NN_RET_CHECK_GE(axis, 0); 134 NN_RET_CHECK_LT(axis, numDimensions); 135 136 uint32_t sumAxis = getSizeOfDimension(input0, axis); 137 for (uint32_t i = 1; i < numInputs - 1; ++i) { 138 const Shape& input = context->getInputShape(i); 139 NN_RET_CHECK_EQ(getNumberOfDimensions(input), numDimensions); 140 NN_RET_CHECK(input.type == input0.type); 141 for (uint32_t d = 0; d < numDimensions; ++d) { 142 if (d == axis) { 143 sumAxis += getSizeOfDimension(input, axis); 144 } else { 145 NN_RET_CHECK_EQ(getSizeOfDimension(input0, d), getSizeOfDimension(input, d)); 146 } 147 } 148 } 149 150 Shape output = context->getOutputShape(kOutputTensor); 151 output.type = input0.type; 152 output.dimensions = input0.dimensions; 153 output.dimensions[axis] = sumAxis; 154 return context->setOutputShape(kOutputTensor, output); 155 } 156 157 bool execute(IOperationExecutionContext* context) { 158 // Bypass execution in the case of zero-sized input. 159 if (getNumberOfElements(context->getOutputShape(kOutputTensor)) == 0) return true; 160 switch (context->getInputType(0)) { 161 case OperandType::TENSOR_FLOAT16: 162 return concatenation<_Float16>(context); 163 case OperandType::TENSOR_FLOAT32: 164 return concatenation<float>(context); 165 case OperandType::TENSOR_QUANT8_ASYMM: 166 return concatenation<uint8_t>(context); 167 default: 168 NN_RET_CHECK_FAIL() << "Unsupported tensor type for operation " << kOperationName; 169 } 170 } 171 172 } // namespace concatenation 173 174 NN_REGISTER_OPERATION(CONCATENATION, concatenation::kOperationName, concatenation::validate, 175 concatenation::prepare, concatenation::execute, .allowZeroSizedInput = true); 176 177 } // namespace nn 178 } // namespace android 179
