From eff71c38bfd9299f8a13176267885845dd4985c5 Mon Sep 17 00:00:00 2001
From: zhengzhouheng <zhengzhouheng@outlook.com>
Date: Thu, 15 Apr 2021 16:09:24 +0800
Subject: [PATCH] add argmin and argmax op

---
 include/tim/vx/ops/arg.h | 50 ++++++++++++++++++++++++++++++++++++++++
 src/tim/vx/ops/README.md |  4 ++--
 src/tim/vx/ops/arg.cc    | 45 ++++++++++++++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 2 deletions(-)
 create mode 100755 include/tim/vx/ops/arg.h
 create mode 100755 src/tim/vx/ops/arg.cc

diff --git a/include/tim/vx/ops/arg.h b/include/tim/vx/ops/arg.h
new file mode 100755
index 0000000..7cd6036
--- /dev/null
+++ b/include/tim/vx/ops/arg.h
@@ -0,0 +1,50 @@
+/****************************************************************************
+*
+*    Copyright (c) 2021 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_ARG_H_
+#define TIM_VX_OPS_ARG_H_
+#include "tim/vx/operation.h"
+
+namespace tim {
+namespace vx {
+namespace ops {
+
+#define DECLARE_ARG_OP(NAME)               \
+  class Arg##NAME : public Operation {     \
+   public:                                 \
+    Arg##NAME(Graph* graph, int32_t axis); \
+                                           \
+   protected:                              \
+    int32_t axis_;                         \
+  };
+
+DECLARE_ARG_OP(Min);
+DECLARE_ARG_OP(Max);
+
+#undef DECLARE_ARG_OP
+
+} // namespace ops
+} // namespace vx
+} // namespace tim
+
+#endif /* TIM_VX_OPS_ARG_H_ */
\ No newline at end of file
diff --git a/src/tim/vx/ops/README.md b/src/tim/vx/ops/README.md
index 2804930..687aa2f 100644
--- a/src/tim/vx/ops/README.md
+++ b/src/tim/vx/ops/README.md
@@ -30,7 +30,7 @@ Relu|RELU|Mapped
 L2Normalization|L2_NORMALIZE|Mapped
 FullyConnected|FCL2|Mapped
 ||POOLWITHARGMAX|Unmapped
-||ARGMAX|Unmapped
+ArgMax|ARGMAX|Mapped
 Maximum|MAXIMUM|Mapped
 L2Normalization|L2NORMALIZESCALE|Mapped
 ||CROP|Unmapped
@@ -111,7 +111,7 @@ Gather|GATHER|Mapped
 LogicalNot|LOGICAL_NOT|Mapped
 Sin|SIN|Mapped
 Log|LOG|Mapped
-||ARGMIN|Unmapped
+ArgMin|ARGMIN|Mapped
 ||ROI_ALIGN|Unmapped
 ||HEATMAP_MAX_KEYPOINT|Unmapped
 ||AXIS_ALIGNED_BBOX_TRANSFORM|Unmapped
diff --git a/src/tim/vx/ops/arg.cc b/src/tim/vx/ops/arg.cc
new file mode 100755
index 0000000..b71791f
--- /dev/null
+++ b/src/tim/vx/ops/arg.cc
@@ -0,0 +1,45 @@
+/****************************************************************************
+*
+*    Copyright (c) 2021 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/arg.h"
+
+#include "vsi_nn_pub.h"
+
+#include "operation_private.h"
+
+namespace tim {
+namespace vx {
+namespace ops {
+
+#define DEFINE_ARG_OP(NAME, VSI_OP_TYPE, OP_PARAM)                  \
+  Arg##NAME::Arg##NAME(Graph* graph, int32_t axis)                  \
+      : Operation(graph, VSI_NN_OP_ARG##VSI_OP_TYPE), axis_(axis) { \
+    this->impl()->node()->nn_param.arg##OP_PARAM.axis = axis_;      \
+  }
+
+DEFINE_ARG_OP(Max, MAX, max);
+DEFINE_ARG_OP(Min, MIN, min);
+#undef DEFINE_ARG_OP
+}  // namespace ops
+}  // namespace vx
+}  // namespace tim