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. if the returned MappingFn is null, then the node is not supported).
| 942 | // when called. You can use this function to see if a node is supported |
| 943 | // (i.e. if the returned MappingFn is null, then the node is not supported). |
| 944 | NNAPIDelegateKernel::MappingFn NNAPIDelegateKernel::Map( |
| 945 | const TfLiteContext* context, int builtin_code, int version, |
| 946 | int android_sdk_version, const TfLiteNode* node, |
| 947 | bool is_accelerator_specified) { |
| 948 | switch (builtin_code) { |
| 949 | case kTfLiteBuiltinAdd: |
| 950 | if (version <= 2) { |
| 951 | if (!IsFloatOrQuant8Operator(context, node)) { |
| 952 | return nullptr; |
| 953 | } |
| 954 | return [](const NNAPIOpMappingArgs& mapping_args) |
| 955 | -> ANeuralNetworksOperationType { |
| 956 | auto builtin = reinterpret_cast<TfLiteAddParams*>( |
| 957 | mapping_args.node->builtin_data); |
| 958 | mapping_args.builder->AddScalarInt32Operand(builtin->activation); |
| 959 | return ANEURALNETWORKS_ADD; |
| 960 | }; |
| 961 | } |
| 962 | break; |
| 963 | case kTfLiteBuiltinArgMax: |
| 964 | case kTfLiteBuiltinArgMin: |
| 965 | if (version <= 2) { |
| 966 | // Those operators were introduced in NNAPI 1.2. |
| 967 | if (android_sdk_version < kMinSdkVersionForNNAPI12) { |
| 968 | return nullptr; |
| 969 | } |
| 970 | // Only certain input types are supported. |
| 971 | auto input_type = context->tensors[node->inputs->data[0]].type; |
| 972 | if (input_type != kTfLiteFloat16 && input_type != kTfLiteFloat32 && |
| 973 | input_type != kTfLiteInt32 && input_type != kTfLiteUInt8 && |
| 974 | input_type != kTfLiteInt8) { |
| 975 | return nullptr; |
| 976 | } |
| 977 | // NNAPI only supports axis as int32. If the axis type is int64 and |
| 978 | // constant we can convert it to int32 if the value isn't too large. |
| 979 | const auto& axis_tensor = context->tensors[node->inputs->data[1]]; |
| 980 | if (axis_tensor.type == kTfLiteInt64) { |
| 981 | if (axis_tensor.allocation_type != kTfLiteMmapRo || |
| 982 | *axis_tensor.data.i64 > std::numeric_limits<int32_t>::max() || |
| 983 | *axis_tensor.data.i64 < std::numeric_limits<int32_t>::min()) { |
| 984 | return nullptr; |
| 985 | } |
| 986 | } else if (axis_tensor.type != kTfLiteInt32) { |
| 987 | return nullptr; |
| 988 | } |
| 989 | if (builtin_code == kTfLiteBuiltinArgMax) { |
| 990 | // NNAPI only supports int32 output. |
| 991 | auto builtin = |
| 992 | reinterpret_cast<TfLiteArgMaxParams*>(node->builtin_data); |
| 993 | if (builtin->output_type != kTfLiteInt32) { |
| 994 | return nullptr; |
| 995 | } |
| 996 | return BasicMappingFn<ANEURALNETWORKS_ARGMAX>; |
| 997 | } else { |
| 998 | // NNAPI only supports int32 output. |
| 999 | auto builtin = |
| 1000 | reinterpret_cast<TfLiteArgMinParams*>(node->builtin_data); |
| 1001 | if (builtin->output_type != kTfLiteInt32) { |