Refactoring, and avoid calling function twice in CHECK_NN.
PiperOrigin-RevId: 225461929
This commit is contained in:
A. Unique TensorFlower
TensorFlower Gardener
parent
bc99c3db7f
commit
352a08f34f
@ -37,11 +37,15 @@ namespace {
|
||||
|
||||
// TODO(b/80621585): Consider printing error string, but don't for now to
|
||||
// minimize binary size.
|
||||
#define CHECK_NN(context, code) \
|
||||
if (code != ANEURALNETWORKS_NO_ERROR) { \
|
||||
context->ReportError(context, "NN API returned error (%d).\n", code); \
|
||||
return kTfLiteError; \
|
||||
}
|
||||
#define CHECK_NN(context, code) \
|
||||
do { \
|
||||
const auto _code = (code); \
|
||||
if (_code != ANEURALNETWORKS_NO_ERROR) { \
|
||||
context->ReportError(context, "NN API returned error (%d, line %d).\n", \
|
||||
_code, __LINE__); \
|
||||
return kTfLiteError; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
namespace {
|
||||
int32_t GetAndroidSdkVersion() {
|
||||
@ -349,19 +353,18 @@ class NNAPIOpBuilder {
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
// TfLiteContext for error handling. Must be named context for macros to
|
||||
// work.
|
||||
TfLiteContext* context_;
|
||||
// TfLiteContext for error handling.
|
||||
TfLiteContext* const context_;
|
||||
|
||||
// Tracks relationship between indices
|
||||
// Tracks relationship between indices.
|
||||
OperandMapping* operand_mapping_;
|
||||
|
||||
// The model
|
||||
ANeuralNetworksModel* nn_model_;
|
||||
// The NNAPI model.
|
||||
ANeuralNetworksModel* const nn_model_;
|
||||
|
||||
// Inputs and outputs for the current op. These are augmented in the sense
|
||||
// that NN API uses operands for all arguments, not just tensors, unlike
|
||||
// TensorFlow lite.
|
||||
// TensorFlow Lite.
|
||||
std::vector<uint32_t> augmented_inputs_;
|
||||
std::vector<uint32_t> augmented_outputs_;
|
||||
};
|
||||
@ -374,6 +377,14 @@ struct NNAPIOpMappingArgs {
|
||||
std::vector<int>* model_state_tfl_inputs;
|
||||
};
|
||||
|
||||
// Mapping function simply returning the operation type without adding any
|
||||
// additional parameter.
|
||||
template <ANeuralNetworksOperationType OperationType>
|
||||
ANeuralNetworksOperationType BasicMappingFn(
|
||||
const NNAPIOpMappingArgs& mapping_args) {
|
||||
return OperationType;
|
||||
}
|
||||
|
||||
// The kernel that represents the node sub set of TF Lite being run on NN API.
|
||||
class NNAPIDelegateKernel {
|
||||
public:
|
||||
@ -385,8 +396,8 @@ class NNAPIDelegateKernel {
|
||||
// Return a function that knows how to translate a node into its operands
|
||||
// when called. You can use this function to see if a node is supported
|
||||
// (i.e. that MappingFn is not nullptr).
|
||||
MappingFn Map(TfLiteContext* context, int builtin_code, int version,
|
||||
TfLiteNode* node) {
|
||||
static MappingFn Map(TfLiteContext* context, int builtin_code, int version,
|
||||
TfLiteNode* node) {
|
||||
switch (builtin_code) {
|
||||
case kTfLiteBuiltinAdd:
|
||||
if (version == 1) {
|
||||
@ -397,8 +408,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_ADD;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinMul:
|
||||
@ -410,8 +419,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_MUL;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinAveragePool2d:
|
||||
@ -422,8 +429,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.node->builtin_data);
|
||||
return ANEURALNETWORKS_AVERAGE_POOL_2D;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinMaxPool2d:
|
||||
@ -434,8 +439,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.node->builtin_data);
|
||||
return ANEURALNETWORKS_MAX_POOL_2D;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinL2Pool2d:
|
||||
@ -446,8 +449,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.node->builtin_data);
|
||||
return ANEURALNETWORKS_L2_POOL_2D;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinConv2d:
|
||||
@ -469,8 +470,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_CONV_2D;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinDepthwiseConv2d:
|
||||
@ -487,8 +486,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_DEPTHWISE_CONV_2D;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinFullyConnected:
|
||||
@ -500,8 +497,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_FULLY_CONNECTED;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinSoftmax:
|
||||
@ -513,18 +508,11 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarFloat32Operand(builtin->beta);
|
||||
return ANEURALNETWORKS_SOFTMAX;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinReshape:
|
||||
if (version == 1 && node->inputs->size == 2) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_RESHAPE;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_RESHAPE>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinSqueeze:
|
||||
@ -540,20 +528,15 @@ class NNAPIDelegateKernel {
|
||||
static_cast<uint32_t>(builtin->num_squeeze_dims));
|
||||
return ANEURALNETWORKS_SQUEEZE;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinL2Normalization: {
|
||||
auto builtin =
|
||||
reinterpret_cast<TfLiteL2NormParams*>(node->builtin_data);
|
||||
if (builtin->activation != kTfLiteActNone) {
|
||||
// NNAPI does not support activations
|
||||
return nullptr;
|
||||
if (builtin->activation == kTfLiteActNone) {
|
||||
return BasicMappingFn<ANEURALNETWORKS_L2_NORMALIZATION>;
|
||||
}
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_L2_NORMALIZATION;
|
||||
};
|
||||
break;
|
||||
}
|
||||
case kTfLiteBuiltinLocalResponseNormalization:
|
||||
if (version == 1) {
|
||||
@ -567,10 +550,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarFloat32Operand(builtin->beta);
|
||||
return ANEURALNETWORKS_LOCAL_RESPONSE_NORMALIZATION;
|
||||
};
|
||||
} else {
|
||||
// TODO(miaowang): clean-up code and return early in the unsupported
|
||||
// case.
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinLshProjection:
|
||||
@ -587,8 +566,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->type);
|
||||
return ANEURALNETWORKS_LSH_PROJECTION;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinConcatenation:
|
||||
@ -599,7 +576,7 @@ class NNAPIDelegateKernel {
|
||||
// NNAPI only support concatenating quantized tensor of the same
|
||||
// scale and offset.
|
||||
auto first_param = context->tensors[node->inputs->data[0]].params;
|
||||
for (int i = 0; i < node->inputs->size; i++) {
|
||||
for (int i = 1; i < node->inputs->size; i++) {
|
||||
auto curr_param = context->tensors[node->inputs->data[i]].params;
|
||||
if (curr_param.scale != first_param.scale ||
|
||||
curr_param.zero_point != first_param.zero_point) {
|
||||
@ -614,68 +591,36 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->axis);
|
||||
return ANEURALNETWORKS_CONCATENATION;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinDequantize:
|
||||
if (version == 1) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_DEQUANTIZE;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_DEQUANTIZE>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinFloor:
|
||||
if (version == 1) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_FLOOR;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_FLOOR>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinRelu:
|
||||
if (version == 1) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_RELU;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_RELU>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinReluN1To1:
|
||||
if (version == 1) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_RELU1;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_RELU1>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinRelu6:
|
||||
if (version == 1) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_RELU6;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_RELU6>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinLogistic:
|
||||
if (version == 1) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_LOGISTIC;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_LOGISTIC>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinTanh:
|
||||
@ -683,12 +628,7 @@ class NNAPIDelegateKernel {
|
||||
if (version == 1 &&
|
||||
context->tensors[node->inputs->data[0]].type == kTfLiteFloat32) {
|
||||
// NNAPI only support float tanh.
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_TANH;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_TANH>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinSub:
|
||||
@ -702,8 +642,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_SUB;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinDiv:
|
||||
@ -717,8 +655,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_DIV;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinPad:
|
||||
@ -728,22 +664,12 @@ class NNAPIDelegateKernel {
|
||||
// NNAPI does not support specifying the padding value.
|
||||
// NNAPI pads physical zero for quantized tensors, so only delegate
|
||||
// float pad to NNAPI.
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_PAD;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_PAD>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinSpaceToBatchNd:
|
||||
if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_SPACE_TO_BATCH_ND;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_SPACE_TO_BATCH_ND>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinStridedSlice:
|
||||
@ -758,8 +684,6 @@ class NNAPIDelegateKernel {
|
||||
builtin->shrink_axis_mask);
|
||||
return ANEURALNETWORKS_STRIDED_SLICE;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinTranspose:
|
||||
@ -771,12 +695,7 @@ class NNAPIDelegateKernel {
|
||||
(node->inputs->size > 1) &&
|
||||
(context->tensors[node->inputs->data[1]].allocation_type ==
|
||||
kTfLiteMmapRo)) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_TRANSPOSE;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_TRANSPOSE>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinRnn:
|
||||
@ -799,8 +718,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_RNN;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinSvdf:
|
||||
@ -827,8 +744,6 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(builtin->activation);
|
||||
return ANEURALNETWORKS_SVDF;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinLstm:
|
||||
@ -870,8 +785,6 @@ class NNAPIDelegateKernel {
|
||||
|
||||
return ANEURALNETWORKS_LSTM;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinMean:
|
||||
@ -888,36 +801,27 @@ class NNAPIDelegateKernel {
|
||||
mapping_args.builder->AddScalarInt32Operand(keep_dims);
|
||||
return ANEURALNETWORKS_MEAN;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinEmbeddingLookup:
|
||||
// NNAPI only support float32 values.
|
||||
if (version == 1 &&
|
||||
context->tensors[node->inputs->data[1]].type == kTfLiteFloat32) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_EMBEDDING_LOOKUP;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_EMBEDDING_LOOKUP>;
|
||||
}
|
||||
break;
|
||||
case kTfLiteBuiltinHashtableLookup:
|
||||
// NNAPI only support float32 output.
|
||||
if (version == 1 &&
|
||||
context->tensors[node->outputs->data[0]].type == kTfLiteFloat32) {
|
||||
return [](const NNAPIOpMappingArgs& mapping_args)
|
||||
-> ANeuralNetworksOperationType {
|
||||
return ANEURALNETWORKS_HASHTABLE_LOOKUP;
|
||||
};
|
||||
} else {
|
||||
return nullptr;
|
||||
return BasicMappingFn<ANEURALNETWORKS_HASHTABLE_LOOKUP>;
|
||||
}
|
||||
break;
|
||||
default:
|
||||
// All other operators are not mapped.
|
||||
return nullptr;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// Initialize the kernel (a NN model).
|
||||
@ -1090,7 +994,7 @@ class NNAPIDelegateKernel {
|
||||
outputs.reserve(output_tensors->size);
|
||||
|
||||
size_t total_input_byte_size = 0;
|
||||
// Make the TensorFlow lite inputs and outputs to ann_indices.
|
||||
// Make the TensorFlow Lite inputs and outputs to ann_indices.
|
||||
for (int i : TfLiteIntArrayView(input_tensors)) {
|
||||
// Constant tensors are not NNAPI inputs.
|
||||
if (i != kOptionalTensor &&
|
||||
@ -1149,12 +1053,14 @@ TfLiteDelegate* NnApiDelegate() {
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
// Allocate one element in vector already since TensorFlow Lite uses
|
||||
// the first value as the number of nodes. The actual value will be set
|
||||
// later, after the vector has been filled.
|
||||
std::vector<int> supported_nodes(1);
|
||||
// We don't care about all nodes_, we only care about ones in the
|
||||
// current plan.
|
||||
TfLiteIntArray* plan;
|
||||
TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan));
|
||||
int total_supported_nodes = 0;
|
||||
|
||||
// Check for every node if it is supported
|
||||
// TODO(b/80625235): Fix this to do more careful checking of versioning.
|
||||
@ -1163,14 +1069,12 @@ TfLiteDelegate* NnApiDelegate() {
|
||||
TfLiteRegistration* registration;
|
||||
TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(
|
||||
context, node_index, &node, ®istration));
|
||||
NNAPIDelegateKernel dummy_kernel;
|
||||
if (dummy_kernel.Map(context, registration->builtin_code,
|
||||
registration->version, node)) {
|
||||
if (NNAPIDelegateKernel::Map(context, registration->builtin_code,
|
||||
registration->version, node)) {
|
||||
supported_nodes.push_back(node_index);
|
||||
}
|
||||
total_supported_nodes += 1;
|
||||
}
|
||||
// Put the size at the beginning of the array.
|
||||
// First element in vector must be the number of actual nodes.
|
||||
supported_nodes[0] = supported_nodes.size() - 1;
|
||||
|
||||
// NN API Delegate Registration (the pseudo kernel that will invoke NN
|
||||
@ -1208,11 +1112,10 @@ TfLiteDelegate* NnApiDelegate() {
|
||||
|
||||
// Request TFLite to partition the graph and make kernels
|
||||
// for each independent node sub set a new nnapi_delegate_kernel.
|
||||
context->ReplaceNodeSubsetsWithDelegateKernels(
|
||||
return context->ReplaceNodeSubsetsWithDelegateKernels(
|
||||
context, nnapi_delegate_kernel,
|
||||
reinterpret_cast<TfLiteIntArray*>(supported_nodes.data()),
|
||||
delegate);
|
||||
return kTfLiteOk;
|
||||
}};
|
||||
|
||||
return &delegate;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user