Added gather_elements & unit test (#363)

Signed-off-by: Chen Xin <jack.chen@verisilicon.com>

Co-authored-by: Chen Xin <jack.chen@verisilicon.com>

include/tim/vx/ops/gather_elements.h

@@ -0,0 +1,57 @@
+/****************************************************************************
+*
+*    Copyright (c) 2022 Vivante Corporation
+*
+*    Permission is hereby granted, free of charge, to any person obtaining a
+*    copy of this software and associated documentation files (the "Software"),
+*    to deal in the Software without restriction, including without limitation
+*    the rights to use, copy, modify, merge, publish, distribute, sublicense,
+*    and/or sell copies of the Software, and to permit persons to whom the
+*    Software is furnished to do so, subject to the following conditions:
+*
+*    The above copyright notice and this permission notice shall be included in
+*    all copies or substantial portions of the Software.
+*
+*    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+*    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+*    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+*    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+*    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+*    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+*    DEALINGS IN THE SOFTWARE.
+*
+*****************************************************************************/
+#ifndef TIM_VX_OPS_GATHER_H_
+#define TIM_VX_OPS_GATHER_H_
+#include "tim/vx/direct_map_op.h"
+
+namespace tim {
+namespace vx {
+namespace ops {
+
+/**
+ * ## Gather_elements
+ *
+ * Gather_elements slices from input, **axis** according to **indices**.
+ * out[i][j][k] = input[index[i][j][k]][j][k] if axis = 0,
+ * out[i][j][k] = input[i][index[i][j][k]][k] if axis = 1,
+ * out[i][j][k] = input[i][j][index[i][j][k]] if axis = 2,
+ * https://github.com/onnx/onnx/blob/main/docs/Operators.md#GatherElements
+ */
+
+class Gather_elements : public DirectMapOp {
+ public:
+  Gather_elements(Graph* Graph, int axis);
+
+  std::shared_ptr<Operation> Clone(
+      std::shared_ptr<Graph>& graph) const override;
+
+ protected:
+  int axis_;
+};
+
+}  // namespace ops
+}  // namespace vx
+}  // namespace tim
+
+#endif /* TIM_VX_OPS_GATHER_H_ */

src/tim/vx/ops/README.md

@@ -20,6 +20,7 @@ Relu|RELU|Mapped|[tf.nn.relu](https://tensorflow.google.cn/api_docs/python/tf/nn
 Reorg|REORG|Mapped|[darknet.reorg](https://github.com/pjreddie/darknet/blob/master/src/reorg_layer.c)
 L2Normalization|L2_NORMALIZE|Mapped|[tf.math.l2_normalize](https://tensorflow.google.cn/api_docs/python/tf/math/l2_normalize)
 FullyConnected|FCL2|Mapped|[tf.keras.layers.Dense](https://tensorflow.google.cn/api_docs/python/tf/keras/layers/Dense)
+Dense|FCL|Mapped|[tf.keras.layers.Dense](https://tensorflow.google.cn/api_docs/python/tf/keras/layers/Dense)
 MaxpoolWithArgmax|POOLWITHARGMAX|Mapped|[tf.nn.max_pool_with_argmax](https://tensorflow.google.cn/api_docs/python/tf/nn/max_pool_with_argmax)
 ArgMax|ARGMAX|Mapped|[tf.math.argmax](https://tensorflow.google.cn/api_docs/python/tf/math/argmax)
 Maximum|MAXIMUM|Mapped|[tf.math.maximum](https://tensorflow.google.cn/api_docs/python/tf/math/maximum)
@@ -76,6 +77,7 @@ Exp|EXP|Mapped|[tf.math.exp](https://tensorflow.google.cn/api_docs/python/tf/mat
 Clip|CLIP|Mapped|[tf.clip_by_value](https://tensorflow.google.cn/api_docs/python/tf/clip_by_value)
 AddN|ADDN|Mapped|[tf.math.add_n](https://tensorflow.google.cn/api_docs/python/tf/math/add_n)
 Gather|GATHER|Mapped|[tf.gather](https://tensorflow.google.cn/api_docs/python/tf/gather)
+Gather_elements|GATHER_ELEMENTS|Mapped|[onnx/GatherElements](https://github.com/onnx/onnx/blob/main/docs/Operators.md#gatherelements)
 LogicalNot|LOGICAL_NOT|Mapped|[tf.math.logical_not](https://tensorflow.google.cn/api_docs/python/tf/math/logical_not)
 Sin|SIN|Mapped|[tf.math.sin](https://tensorflow.google.cn/api_docs/python/tf/math/sin)
 Log|LOG|Mapped|[tf.math.log](https://tensorflow.google.cn/api_docs/python/tf/math/log)

src/tim/vx/ops/gather_elements.cc

@@ -0,0 +1,45 @@
+/****************************************************************************
+*
+*    Copyright (c) 2022 Vivante Corporation
+*
+*    Permission is hereby granted, free of charge, to any person obtaining a
+*    copy of this software and associated documentation files (the "Software"),
+*    to deal in the Software without restriction, including without limitation
+*    the rights to use, copy, modify, merge, publish, distribute, sublicense,
+*    and/or sell copies of the Software, and to permit persons to whom the
+*    Software is furnished to do so, subject to the following conditions:
+*
+*    The above copyright notice and this permission notice shall be included in
+*    all copies or substantial portions of the Software.
+*
+*    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+*    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+*    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+*    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+*    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+*    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+*    DEALINGS IN THE SOFTWARE.
+*
+*****************************************************************************/
+#include "tim/vx/ops/gather_elements.h"
+
+#include "direct_map_op_impl.h"
+#include "vsi_nn_pub.h"
+
+namespace tim {
+namespace vx {
+namespace ops {
+
+Gather_elements::Gather_elements(Graph* graph, int axis)
+    : DirectMapOp(graph, VSI_NN_OP_GATHER_ELEMENTS), axis_(axis) {
+  this->impl()->node()->nn_param.gather_elements.axis = axis_;
+}
+
+std::shared_ptr<Operation> Gather_elements::Clone(
+    std::shared_ptr<Graph>& graph) const {
+  return graph->CreateOperation<Gather_elements>(this->axis_);
+}
+
+}  // namespace ops
+}  // namespace vx
+}  // namespace tim

src/tim/vx/ops/gather_elements_test.cc

@@ -0,0 +1,170 @@
+/****************************************************************************
+*
+*    Copyright (c) 2022 Vivante Corporation
+*
+*    Permission is hereby granted, free of charge, to any person obtaining a
+*    copy of this software and associated documentation files (the "Software"),
+*    to deal in the Software without restriction, including without limitation
+*    the rights to use, copy, modify, merge, publish, distribute, sublicense,
+*    and/or sell copies of the Software, and to permit persons to whom the
+*    Software is furnished to do so, subject to the following conditions:
+*
+*    The above copyright notice and this permission notice shall be included in
+*    all copies or substantial portions of the Software.
+*
+*    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+*    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+*    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+*    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+*    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+*    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+*    DEALINGS IN THE SOFTWARE.
+*
+*****************************************************************************/
+#include "tim/vx/context.h"
+#include "tim/vx/graph.h"
+#include "tim/vx/ops/gather_elements.h"
+#include <iostream>
+#include "gtest/gtest.h"
+#include "test_utils.h"
+
+TEST(Gather_elements, shape_3_2_1_int32_axis_0) {
+  auto ctx = tim::vx::Context::Create();
+  auto graph = ctx->CreateGraph();
+
+  tim::vx::ShapeType in_shape({3, 2, 1});
+  tim::vx::ShapeType indices_shape({2, 2, 1});
+  tim::vx::ShapeType out_shape({2, 2, 1});
+  tim::vx::TensorSpec input_spec(tim::vx::DataType::INT32, in_shape,
+                                 tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec indices_spec(tim::vx::DataType::INT32, indices_shape,
+                                   tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec output_spec(tim::vx::DataType::INT32, out_shape,
+                                  tim::vx::TensorAttribute::OUTPUT);
+
+  auto input_tensor = graph->CreateTensor(input_spec);
+  auto indices_tensor = graph->CreateTensor(indices_spec);
+  auto output_tensor = graph->CreateTensor(output_spec);
+
+  std::vector<int32_t> in_data = {
+      1, 2, 3, 4, 5, 6,
+  };
+  //The index value greater than rank-1 is regarded as rank-1
+  std::vector<int32_t> indices = {
+      1,
+      2,
+      0,
+      2,
+  };
+  std::vector<int32_t> golden = {
+      2, 3, 4, 6,
+  };
+
+  EXPECT_TRUE(
+      input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * 4));
+  EXPECT_TRUE(
+      indices_tensor->CopyDataToTensor(indices.data(), indices.size() * 4));
+  auto op = graph->CreateOperation<tim::vx::ops::Gather_elements>(0);
+  (*op).BindInputs({input_tensor, indices_tensor}).BindOutputs({output_tensor});
+
+  EXPECT_TRUE(graph->Compile());
+  EXPECT_TRUE(graph->Run());
+
+  std::vector<int32_t> output(golden.size());
+  EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
+  EXPECT_EQ(golden, output);
+}
+
+TEST(Gather_elements, shape_3_2_1_int32_axis_1) {
+  auto ctx = tim::vx::Context::Create();
+  auto graph = ctx->CreateGraph();
+
+  tim::vx::ShapeType in_shape({3, 2, 1});
+  tim::vx::ShapeType indices_shape({2, 2, 1});
+  tim::vx::ShapeType out_shape({2, 2, 1});
+  tim::vx::TensorSpec input_spec(tim::vx::DataType::INT32, in_shape,
+                                 tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec indices_spec(tim::vx::DataType::INT32, indices_shape,
+                                   tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec output_spec(tim::vx::DataType::INT32, out_shape,
+                                  tim::vx::TensorAttribute::OUTPUT);
+
+  auto input_tensor = graph->CreateTensor(input_spec);
+  auto indices_tensor = graph->CreateTensor(indices_spec);
+  auto output_tensor = graph->CreateTensor(output_spec);
+
+  std::vector<int32_t> in_data = {
+      1, 2, 3, 4, 5, 6,
+  };
+  //The index value greater than rank-1 is regarded as rank-1
+  std::vector<int32_t> indices = {
+      1,
+      2,
+      0,
+      2,
+  };
+  std::vector<int32_t> golden = {
+      4, 5, 1, 5,
+  };
+
+  EXPECT_TRUE(
+      input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * 4));
+  EXPECT_TRUE(
+      indices_tensor->CopyDataToTensor(indices.data(), indices.size() * 4));
+  auto op = graph->CreateOperation<tim::vx::ops::Gather_elements>(1);
+  (*op).BindInputs({input_tensor, indices_tensor}).BindOutputs({output_tensor});
+
+  EXPECT_TRUE(graph->Compile());
+  EXPECT_TRUE(graph->Run());
+
+  std::vector<int32_t> output(golden.size());
+  EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
+  EXPECT_EQ(golden, output);
+}
+
+TEST(Gather_elements, shape_3_2_1_float32_axis_2) {
+  auto ctx = tim::vx::Context::Create();
+  auto graph = ctx->CreateGraph();
+
+  tim::vx::ShapeType in_shape({3, 2, 1});
+  tim::vx::ShapeType indices_shape({2, 2, 1});
+  tim::vx::ShapeType out_shape({2, 2, 1});
+  tim::vx::TensorSpec input_spec(tim::vx::DataType::INT32, in_shape,
+                                 tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec indices_spec(tim::vx::DataType::INT32, indices_shape,
+                                   tim::vx::TensorAttribute::INPUT);
+  tim::vx::TensorSpec output_spec(tim::vx::DataType::INT32, out_shape,
+                                  tim::vx::TensorAttribute::OUTPUT);
+
+  auto input_tensor = graph->CreateTensor(input_spec);
+  auto indices_tensor = graph->CreateTensor(indices_spec);
+  auto output_tensor = graph->CreateTensor(output_spec);
+
+  std::vector<float> in_data = {
+      1, 2, 3, 4, 5, 6,
+  };
+  //The index value greater than rank-1 is regarded as rank-1
+  std::vector<int32_t> indices = {
+      1,
+      2,
+      0,
+      2,
+  };
+  std::vector<float> golden = {
+      1, 2, 3, 4,
+  };
+
+  EXPECT_TRUE(
+      input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * 4));
+  EXPECT_TRUE(
+      indices_tensor->CopyDataToTensor(indices.data(), indices.size() * 4));
+  auto op = graph->CreateOperation<tim::vx::ops::Gather_elements>(2);
+  (*op).BindInputs({input_tensor, indices_tensor}).BindOutputs({output_tensor});
+
+  EXPECT_TRUE(graph->Compile());
+  EXPECT_TRUE(graph->Run());
+
+  std::vector<float> output(golden.size());
+  EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
+  EXPECT_EQ(golden, output);
+}

src/tim/vx/ops/max_pool_test.cc

@@ -1,3 +1,26 @@
+/****************************************************************************
+*
+*    Copyright (c) 2022 Vivante Corporation
+*
+*    Permission is hereby granted, free of charge, to any person obtaining a
+*    copy of this software and associated documentation files (the "Software"),
+*    to deal in the Software without restriction, including without limitation
+*    the rights to use, copy, modify, merge, publish, distribute, sublicense,
+*    and/or sell copies of the Software, and to permit persons to whom the
+*    Software is furnished to do so, subject to the following conditions:
+*
+*    The above copyright notice and this permission notice shall be included in
+*    all copies or substantial portions of the Software.
+*
+*    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+*    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+*    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+*    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+*    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+*    FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+*    DEALINGS IN THE SOFTWARE.
+*
+*****************************************************************************/
 #include "tim/vx/context.h"
 #include "tim/vx/graph.h"
 #include "tim/vx/ops/pool2d.h"