diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc index 299df7ae112..5eda8214a3a 100644 --- a/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc +++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate.cc @@ -557,7 +557,14 @@ class NNAPIOpBuilder { TfLiteStatus AddVectorInt32Operand(const int32_t* values, uint32_t num_values) { return AddVectorOperand(values, num_values, - ANEURALNETWORKS_TENSOR_INT32); + ANEURALNETWORKS_TENSOR_INT32, + /*scale=*/0.f, /*zero_point=*/0); + } + + TfLiteStatus AddVectorInt32Operand(const int32_t* values, uint32_t num_values, + float scale, int32_t zero_point) { + return AddVectorOperand( + values, num_values, ANEURALNETWORKS_TENSOR_INT32, scale, zero_point); } TfLiteStatus AddVectorFloat32Operand(const float* values, @@ -781,9 +788,13 @@ class NNAPIOpBuilder { template TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values, - int32_t nn_type) { - ANeuralNetworksOperandType operand_type{ - .type = nn_type, .dimensionCount = 1, .dimensions = &num_values}; + int32_t nn_type, float scale, + int32_t zero_point) { + ANeuralNetworksOperandType operand_type{.type = nn_type, + .dimensionCount = 1, + .dimensions = &num_values, + .scale = scale, + .zeroPoint = zero_point}; RETURN_TFLITE_ERROR_IF_NN_ERROR( context_, @@ -797,6 +808,13 @@ class NNAPIOpBuilder { return kTfLiteOk; } + template + TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values, + int32_t nn_type) { + return AddVectorOperand(values, num_values, nn_type, /*scale=*/0.f, + /*zero_point=*/0); + } + TfLiteStatus AddFloat32OutputTensor(uint32_t dimension_count, const uint32_t* dimension_data, int* ann_index_out) { @@ -1812,6 +1830,91 @@ class NNAPIDelegateKernel { return BasicMappingFn; } break; + case kTfLiteBuiltinTransposeConv: + if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12) { + return [](const NNAPIOpMappingArgs& mapping_args) + -> ANeuralNetworksOperationType { + const bool hybrid_op = IsHybridOperator(mapping_args.context, + kTfLiteBuiltinTransposeConv, + mapping_args.node); + mapping_args.builder->AddTensorInput(/*kDataInputTensor*/ 2, + hybrid_op); + mapping_args.builder->AddTensorInput(/*kWeightsTensor*/ 1, + hybrid_op); + + // NNAPI requires a bias tensor, so we allocate a new tensor to fill + // it with zeroes. It is deleted with other tensors in the context + // during subgraph destructor call. + int bias_index = -1; + mapping_args.context->AddTensors(mapping_args.context, 1, + &bias_index); + TfLiteTensor* bias_tensor = + &mapping_args.context->tensors[bias_index]; + const auto input_type = + mapping_args.context + ->tensors[mapping_args.node->inputs + ->data[/*kDataInputTensor*/ 2]] + .type; + if (input_type == kTfLiteFloat32) { + bias_tensor->type = kTfLiteFloat32; + } else { + bias_tensor->type = kTfLiteInt32; + } + + // Create an array with a required bias shape and resize the bias + // tensor. + TfLiteIntArray* bias_shape = TfLiteIntArrayCreate(1); + const TfLiteTensor& output_shape = + mapping_args.context->tensors + [mapping_args.node->inputs->data[/*kOutputShapeTensor*/ 0]]; + const int output_depth = output_shape.data.i32[3]; + bias_shape->data[0] = output_depth; + bias_tensor->allocation_type = kTfLiteDynamic; + mapping_args.context->ResizeTensor(mapping_args.context, + bias_tensor, bias_shape); + + // Set tensor's values to zeroes and add it using AddVector*, so + // that the values are copied to NNAPI. We don't use the AddTensor + // function because it doesn't copy values and the tensor we just + // created is not in the node->inputs. + if (input_type == kTfLiteFloat32) { + memset(bias_tensor->data.f, 0, output_depth * sizeof(float)); + mapping_args.builder->AddVectorFloat32Operand(bias_tensor->data.f, + output_depth); + } else { + memset(bias_tensor->data.i32, 0, output_depth * sizeof(int)); + const TfLiteTensor& input_tensor = + mapping_args.context->tensors + [mapping_args.node->inputs->data[/*kDataInputTensor*/ 2]]; + const TfLiteTensor& filter_tensor = + mapping_args.context->tensors + [mapping_args.node->inputs->data[/*kWeightsTensor*/ 1]]; + // NNAPI requires bias scale to be a product of an input scale and + // a filter scale. + bias_tensor->params.scale = + input_tensor.params.scale * filter_tensor.params.scale; + mapping_args.builder->AddVectorInt32Operand( + bias_tensor->data.i32, output_depth, + input_tensor.params.scale * filter_tensor.params.scale, + /*zero_point=*/0); + } + + mapping_args.builder->AddTensorInput(/*kOutputShapeTensor*/ 0, + hybrid_op); + + auto builtin = reinterpret_cast( + mapping_args.node->builtin_data); + mapping_args.builder->AddScalarInt32Operand(builtin->padding); + mapping_args.builder->AddScalarInt32Operand(builtin->stride_width); + mapping_args.builder->AddScalarInt32Operand(builtin->stride_height); + mapping_args.builder->AddScalarInt32Operand( + /*ANEURALNETWORKS_FUSED_NONE*/ 0); + // Use NHWC layout for input and output + mapping_args.builder->AddScalarBoolOperand(false); + return ANEURALNETWORKS_TRANSPOSE_CONV; + }; + } + break; case kTfLiteBuiltinSqrt: if (version == 1 && android_sdk_version >= kMinSdkVersionForNNAPI12 && IsFloat(context->tensors[node->inputs->data[0]].type)) { @@ -3087,13 +3190,17 @@ class NNAPIDelegateKernel { continue; } } - if ((reg->builtin_code == kTfLiteBuiltinSplit) && (input_index == node->inputs->data[0])) { // Skip the axis input tensor; it will be added as a scalar operand // by the Map() mapping. continue; } + if (reg->builtin_code == kTfLiteBuiltinTransposeConv) { + // Everything is added during Map since input tensors + // have different order. + continue; + } // Pad and Padv2 have an optional parameter for a pad value which has // to be converted to a scalar type in NN API. diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD index 6fba580b718..b80062e56ed 100644 --- a/tensorflow/lite/kernels/BUILD +++ b/tensorflow/lite/kernels/BUILD @@ -1460,6 +1460,7 @@ cc_test( name = "transpose_conv_test", size = "small", srcs = ["transpose_conv_test.cc"], + tags = ["tflite_nnapi"], deps = [ ":builtin_ops", ":test_main", diff --git a/tensorflow/lite/nnapi/NeuralNetworksTypes.h b/tensorflow/lite/nnapi/NeuralNetworksTypes.h index 0a92fc2cb99..b4ec12ee14d 100644 --- a/tensorflow/lite/nnapi/NeuralNetworksTypes.h +++ b/tensorflow/lite/nnapi/NeuralNetworksTypes.h @@ -130,6 +130,7 @@ enum { ANEURALNETWORKS_SQRT = 88, ANEURALNETWORKS_TILE = 89, ANEURALNETWORKS_TOPK_V2 = 90, + ANEURALNETWORKS_TRANSPOSE_CONV = 91, ANEURALNETWORKS_UNIDIRECTIONAL_SEQUENCE_LSTM = 92, ANEURALNETWORKS_UNIDIRECTIONAL_SEQUENCE_RNN = 93, ANEURALNETWORKS_RESIZE_NEAREST_NEIGHBOR = 94,