Change ARGMAX definition (#5021)

Chagne ARGMAX definition same with current NNAPI
- Allow one axis reduce
- Output rank = input rank - 1

Signed-off-by: Hyeongseok Oh <hseok82.oh@samsung.com>

13 files changed, 35 insertions, 18 deletions

diff --git a/include/NeuralNetworksEx.h b/include/NeuralNetworksEx.h
index 2ec0f069e..6c5e5aa3d 100644
--- a/include/NeuralNetworksEx.h
+++ b/include/NeuralNetworksEx.h
@@ -412,24 +412,23 @@ typedef enum {
   ANEURALNETWORKS_UNPACK_EX = 50014,
 
   /**
-   * Find index with the largest value across axes of a input tensor.
-   *
-   * Reduces the input tensor along the given dimensions to reduce. The reduced
-   * dimensions are retained with length 1.
+   * Returns the index of the largest element along an axis.
    *
    * Supported tensor {@link OperandCode}:
    * * {@link ANEURALNETWORKS_TENSOR_FLOAT32}
    * * {@link ANEURALNETWORKS_TENSOR_INT32}
    * * {@link ANEURALNETWORKS_TENSOR_QUANT8_ASYMM}
    *
+   * Supported tensor rank: from 1
+   *
    * Inputs:
    * * 0: An n-D tensor, specifying the input.
-   * * 1: A 1-D Tensor of {@link ANEURALNETWORKS_TENSOR_INT32}. The dimensions
-   *      to reduce. Must be in the range [-rank(input_tensor), rank(input_tensor)).
+   * * 1: An {@link ANEURALNETWORKS_INT32} scalar specifying the axis to
+   *      reduce across. Negative index is used to specify axis from the
+   *      end (e.g. -1 for the last axis). Must be in the range [-n, n).
    *
    * Outputs:
-   * * 0: A output tensor of {@link ANEURALNETWORKS_TENSOR_INT32}.
-   *      The rank of output tensor should be same rank of input0.
+   * * 0: An (n - 1)-D {@link ANEURALNETWORKS_TENSOR_INT32} tensor.
    */
   ANEURALNETWORKS_ARGMAX_EX = 50015,
 
diff --git a/libs/ARMComputeEx/src/core/CL/kernels/CLArgOperationKernel.cpp b/libs/ARMComputeEx/src/core/CL/kernels/CLArgOperationKernel.cpp
index fd0789864..7f4b5b0df 100644
--- a/libs/ARMComputeEx/src/core/CL/kernels/CLArgOperationKernel.cpp
+++ b/libs/ARMComputeEx/src/core/CL/kernels/CLArgOperationKernel.cpp
@@ -43,7 +43,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, c
                                         DataType::QASYMM8);
   ARM_COMPUTE_ERROR_ON_DATA_TYPE_NOT_IN(output, DataType::S32);
 
-  ARM_COMPUTE_RETURN_ERROR_ON_MSG(input->tensor_shape().num_dimensions() !=
+  ARM_COMPUTE_RETURN_ERROR_ON_MSG((input->tensor_shape().num_dimensions() - 1) !=
                                       output->tensor_shape().num_dimensions(),
                                   "Input's rank is not same with output");
 
diff --git a/libs/tflite/src/ext/nnapi_delegate.cpp b/libs/tflite/src/ext/nnapi_delegate.cpp
index 2aaa9521b..b530dd725 100644
--- a/libs/tflite/src/ext/nnapi_delegate.cpp
+++ b/libs/tflite/src/ext/nnapi_delegate.cpp
@@ -758,6 +758,7 @@ TfLiteStatus AddOpsAndParams(
             reinterpret_cast<uint32_t*>(node.outputs->data)));
         continue;
       case tflite::BuiltinOperator_ARG_MAX:
+        check_arg_max_input(node.builtin_data);
         CHECK_NN(ANeuralNetworksModel_addOperationEx(
             nn_model, ANEURALNETWORKS_ARGMAX_EX,
             static_cast<uint32_t>(augmented_inputs.size()),
diff --git a/libs/tflite/src/ext/nnapi_delegate_ex_AddOpsAndParams_lambda.inc b/libs/tflite/src/ext/nnapi_delegate_ex_AddOpsAndParams_lambda.inc
index cd06de0e9..7924bb2bc 100644
--- a/libs/tflite/src/ext/nnapi_delegate_ex_AddOpsAndParams_lambda.inc
+++ b/libs/tflite/src/ext/nnapi_delegate_ex_AddOpsAndParams_lambda.inc
@@ -98,3 +98,20 @@
       auto builtin = reinterpret_cast<TfLiteSplitParams*>(data);
       add_scalar_int32(builtin->num_splits);
     };
+
+    auto check_arg_max_input = [&interpreter, &augmented_inputs](void *data) {
+      auto params = reinterpret_cast<TfLiteArgMaxParams*>(data);
+      if (params->output_type != kTfLiteInt32)
+      {
+        FATAL("Cannot handle output type in NNAPI");
+      }
+
+      TfLiteTensor* axis_tensor = interpreter->tensor(augmented_inputs.back());
+      assert(axis_tensor->type == kTfLiteInt32);
+
+      int64_t count = 1;
+      for (int i = 0; i < axis_tensor->dims->size; ++i) {
+        count *= axis_tensor->dims->data[i];
+      }
+      assert(count == 1);
+    };
diff --git a/runtimes/neurun/backend/acl_cl/StageGenerator.cc b/runtimes/neurun/backend/acl_cl/StageGenerator.cc
index de7c1aaa3..ef35f565d 100644
--- a/runtimes/neurun/backend/acl_cl/StageGenerator.cc
+++ b/runtimes/neurun/backend/acl_cl/StageGenerator.cc
@@ -2716,7 +2716,7 @@ void StageGenerator::visit(const model::operation::ArgMaxNode &node)
   assert(_ctx.at(axis_index).hasData());
   // Axis dimension is always 1.
   assert(axis_shape.rank() == 1);
-  assert(ifm_shape.rank() == ofm_shape.rank());
+  assert((ifm_shape.rank() - 1) == ofm_shape.rank());
 
   _tensor_builder->dimCorrection(ofm_index, false);
   _tensor_builder->dimCorrection(ifm_index, false);
diff --git a/runtimes/pure_arm_compute/src/compilation.cc b/runtimes/pure_arm_compute/src/compilation.cc
index fd200203f..4506992be 100644
--- a/runtimes/pure_arm_compute/src/compilation.cc
+++ b/runtimes/pure_arm_compute/src/compilation.cc
@@ -3722,7 +3722,7 @@ void Planner::visit(const ::internal::tflite::op::ArgMax::Node &node)
   assert(_ctx.at(axis_index).hasData());
   // Axis dimension is always 1.
   assert(axis_shape.rank() == 1);
-  assert(ifm_shape.rank() == ofm_shape.rank());
+  assert((ifm_shape.rank() - 1) == ofm_shape.rank());
 
   _builder.addShapeConstr(ofm_index, asTensorInfo(asTensorShape(_ctx.at(ofm_index).shape(), false),
                                                   _ctx.at(ofm_index).type()));
diff --git a/tests/nnapi/specs/Ex/argmax_ex_float_1.mod.py b/tests/nnapi/specs/Ex/argmax_ex_float_1.mod.py
index e2255e004..dc29fb358 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_float_1.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_float_1.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_FLOAT32", "{1, 2, 2, 1}")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [1])
-output = Output("output", "TENSOR_INT32", "{1, 1, 2, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 2, 1}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)
 
diff --git a/tests/nnapi/specs/Ex/argmax_ex_float_2.mod.py b/tests/nnapi/specs/Ex/argmax_ex_float_2.mod.py
index 6f06bfd0b..8b6ed565f 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_float_2.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_float_2.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_FLOAT32", "{1, 2, 2, 2}")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [3])
-output = Output("output", "TENSOR_INT32", "{1, 2, 2, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 2, 2}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)
 
diff --git a/tests/nnapi/specs/Ex/argmax_ex_int32.mod.py b/tests/nnapi/specs/Ex/argmax_ex_int32.mod.py
index f7f98afe2..b5d01dd85 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_int32.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_int32.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_INT32", "{1, 2, 2, 1}")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [1])
-output = Output("output", "TENSOR_INT32", "{1, 1, 2, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 2, 1}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)
 
diff --git a/tests/nnapi/specs/Ex/argmax_ex_neg_axis_float.mod.py b/tests/nnapi/specs/Ex/argmax_ex_neg_axis_float.mod.py
index b29cf1eb6..977cac061 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_neg_axis_float.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_neg_axis_float.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_FLOAT32", "{1, 2, 4, 1}")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [-3])
-output = Output("output", "TENSOR_INT32", "{1, 1, 4, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 4, 1}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)
 
diff --git a/tests/nnapi/specs/Ex/argmax_ex_neg_axis_int32.mod.py b/tests/nnapi/specs/Ex/argmax_ex_neg_axis_int32.mod.py
index ddc27b35a..9f448e08d 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_neg_axis_int32.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_neg_axis_int32.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_INT32", "{1, 2, 4, 1}")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [-3])
-output = Output("output", "TENSOR_INT32", "{1, 1, 4, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 4, 1}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)
 
diff --git a/tests/nnapi/specs/Ex/argmax_ex_quant8.mod.py b/tests/nnapi/specs/Ex/argmax_ex_quant8.mod.py
index fd3c5230b..c3131ef91 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_quant8.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_quant8.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_QUANT8_ASYMM", "{1, 2, 2, 1}, 0.5f, 2")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [1])
-output = Output("output", "TENSOR_INT32", "{1, 1, 2, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 2, 1}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)
 
diff --git a/tests/nnapi/specs/Ex/argmax_ex_quant8_neg_axis.mod.py b/tests/nnapi/specs/Ex/argmax_ex_quant8_neg_axis.mod.py
index c7ddc0e5d..9146d8f14 100644
--- a/tests/nnapi/specs/Ex/argmax_ex_quant8_neg_axis.mod.py
+++ b/tests/nnapi/specs/Ex/argmax_ex_quant8_neg_axis.mod.py
@@ -1,7 +1,7 @@
 model = Model()
 i1 = Input("input", "TENSOR_QUANT8_ASYMM", "{1, 2, 4, 1}, 0.5f, 5")
 axis = Parameter("axis", "TENSOR_INT32", "{1}", [-3])
-output = Output("output", "TENSOR_INT32", "{1, 1, 4, 1}")
+output = Output("output", "TENSOR_INT32", "{1, 4, 1}")
 
 model = model.Operation("ARGMAX_EX", i1, axis).To(output)