support select_last_index for argmin/max

include/tvm/relay/attrs/reduce.h

@@ -61,6 +61,42 @@ struct ReduceAttrs : public tvm::AttrsNode<ReduceAttrs> {
   }
 };
 
+/*! \brief Attributes for Reduce operators which reduce by finding a single element. E.g. argmin */
+struct OneElementReduceAttrs : public tvm::AttrsNode<ReduceAttrs> {
+  Array<Integer> axis;
+  bool keepdims;
+  bool select_last_index;
+  bool exclude;
+
+  TVM_DECLARE_ATTRS(ReduceAttrs, "relay.attrs.ReduceAttrs") {
+    TVM_ATTR_FIELD(axis)
+        .set_default(NullValue<Array<Integer>>())
+        .describe(R"code(The axis or axes along which to perform the reduction.
+
+      The default, `axis=()`, will compute over all elements into a
+      scalar array with shape `(1,)`.
+
+      If `axis` is int, a reduction is performed on a particular axis.
+
+      If `axis` is a tuple of ints, a reduction is performed on all the axes
+      specified in the tuple.
+
+      If `exclude` is true, reduction will be performed on the axes that are
+      NOT in axis instead.)code");
+
+    TVM_ATTR_FIELD(keepdims).set_default(false).describe(
+        "If this is set to `True`, the reduced axes are left "
+        "in the result as dimension with size one.");
+    TVM_ATTR_FIELD(select_last_index)
+        .set_default(false)
+        .describe(
+            "Whether to select the last index if the target element appears multiple times, else "
+            "select the first index which the target element appears");
+    TVM_ATTR_FIELD(exclude).set_default(false).describe(
+        "Whether to perform reduction on axis that are NOT in axis instead.");
+  }
+};
+
 struct VarianceAttrs : public tvm::AttrsNode<VarianceAttrs> {
   Array<Integer> axis;
   bool keepdims;

include/tvm/relay/expr_functor.h

@@ -241,9 +241,9 @@ class MixedModeVisitor : public ::tvm::relay::ExprVisitor {
    * \brief VisitExpr is finalized to preserve call expansion of dataflow regions
    */
   void VisitExpr(const Expr& expr) final;
-  void VisitExpr_(const CallNode* op) override;
-  void VisitExpr_(const TupleNode* op) override;
-  void VisitExpr_(const TupleGetItemNode* op) override;
+  virtual void VisitExpr_(const CallNode* op) override;
+  virtual void VisitExpr_(const TupleNode* op) override;
+  virtual void VisitExpr_(const TupleGetItemNode* op) override;
 
  protected:
   /*!

include/tvm/topi/reduction.h

@@ -431,6 +431,40 @@ inline Tensor max(const Tensor& data, const Array<Integer>& axis, bool keepdims
   return CommReduce(data, axis, MaxOp, keepdims, atleast1d);
 }
 
+inline FCommReduce MakeSinglePassReducer(
+    std::function<PrimExpr(Var, Var)> comparison_op,
+    std::function<PrimExpr(const DataType&)> initial_value_generator, String name) {
+  // Create a Commutative Reducer with a comparison operation, and method to get the initial value.
+  auto fcombine = [&](Array<Var> lhs, Array<Var> rhs) {
+    Array<PrimExpr> result;
+    result.push_back(tvm::tir::Select(comparison_op(lhs[1], rhs[1]), lhs[0], rhs[0]));  // idx
+    result.push_back(tvm::tir::Select(comparison_op(lhs[1], rhs[1]), lhs[1], rhs[1]));  // val
+    return result;
+  };
+  auto fidentity = [&](std::vector<DataType> types) {
+    Array<PrimExpr> result;
+    result.push_back(tvm::tir::make_const(types[0], -1));  // idx
+    result.push_back(initial_value_generator(types[1]));   // val
+    return result;
+  };
+  return MakeCommReducer(fcombine, fidentity, name);
+}
+
+inline FCommReduce MakeArgminReducer(bool select_last_index = false) {
+  std::function<PrimExpr(Var, Var)> comparison_op;
+  if (select_last_index) {
+    comparison_op = [](Var lhs, Var rhs) { return lhs <= rhs; };
+  } else {
+    comparison_op = [](Var lhs, Var rhs) { return lhs < rhs; };
+  }
+
+  std::function<PrimExpr(const DataType&)> initial_value_generator = [](const DataType& data_type) {
+    return tvm::max_value(data_type);
+  };
+
+  return MakeSinglePassReducer(comparison_op, initial_value_generator, "argmin");
+}
+
 /*!
  * \brief Creates an operation that finds the indices of the minimum
  * values over a given axis.
@@ -442,41 +476,30 @@ inline Tensor max(const Tensor& data, const Array<Integer>& axis, bool keepdims
  * left in the result as dimensions with size one. This enables the result
  * to broadcast correctly against the input array.
  * \param atleast1d Whether the output need to be atleast1d.
+ * \param select_last_index Whether to select the last index if the minimum element
+ * appears multiple times, else select the first index.
  *
  * \return A Tensor whose op member is the argmin operation
  */
 inline Tensor argmin(const Tensor& data, const Array<Integer>& axis, bool keepdims = false,
-                     bool atleast1d = false) {
-  auto fcombine = [](Array<Var> lhs, Array<Var> rhs) {
-    Array<PrimExpr> result;
-    result.push_back(tvm::tir::Select(lhs[1] <= rhs[1], lhs[0], rhs[0]));  // idx
-    result.push_back(tvm::tir::Select(lhs[1] <= rhs[1], lhs[1], rhs[1]));  // val
-    return result;
-  };
-  auto fidentity = [](std::vector<DataType> types) {
-    Array<PrimExpr> result;
-    result.push_back(tvm::tir::make_const(types[0], -1));  // idx
-    result.push_back(tvm::max_value(types[1]));            // val
-    return result;
-  };
-  auto func = MakeCommReducer(fcombine, fidentity, "argmin");
-  return CommReduceIdx(data, axis, func, keepdims, atleast1d);
+                     bool atleast1d = false, bool select_last_index = false) {
+  auto reducer = MakeArgminReducer(select_last_index);
+  return CommReduceIdx(data, axis, reducer, keepdims, atleast1d);
 }
 
-inline FCommReduce MakeArgmaxReducer() {
-  auto fcombine = [](Array<Var> lhs, Array<Var> rhs) {
-    Array<PrimExpr> result;
-    result.push_back(tvm::tir::Select(lhs[1] >= rhs[1], lhs[0], rhs[0]));  // idx
-    result.push_back(tvm::tir::Select(lhs[1] >= rhs[1], lhs[1], rhs[1]));  // val
-    return result;
-  };
-  auto fidentity = [](std::vector<DataType> types) {
-    Array<PrimExpr> result;
-    result.push_back(tvm::tir::make_const(types[0], -1));  // idx
-    result.push_back(tvm::min_value(types[1]));            // val
-    return result;
+inline FCommReduce MakeArgmaxReducer(bool select_last_index = false) {
+  std::function<PrimExpr(Var, Var)> comparison_op;
+  if (select_last_index) {
+    comparison_op = [](Var lhs, Var rhs) { return lhs >= rhs; };
+  } else {
+    comparison_op = [](Var lhs, Var rhs) { return lhs > rhs; };
+  }
+
+  std::function<PrimExpr(const DataType&)> initial_value_generator = [](const DataType& data_type) {
+    return tvm::min_value(data_type);
   };
-  return MakeCommReducer(fcombine, fidentity, "argmax");
+
+  return MakeSinglePassReducer(comparison_op, initial_value_generator, "argmax");
 }
 
 /*!
@@ -490,12 +513,13 @@ inline FCommReduce MakeArgmaxReducer() {
  * left in the result as dimensions with size one. This enables the result
  * to broadcast correctly against the input array.
  * \param atleast1d Whether the output need to be atleast1d.
- *
+ * \param select_last_index Whether to select the last index if the maximum element
+ * appears multiple times, else select the first index.
  * \return A Tensor whose op member is the argmax operation
  */
 inline Tensor argmax(const Tensor& data, const Array<Integer>& axis, bool keepdims = false,
-                     bool atleast1d = false) {
-  auto reducer = MakeArgmaxReducer();
+                     bool atleast1d = false, bool select_last_index = false) {
+  auto reducer = MakeArgmaxReducer(select_last_index);
   return CommReduceIdx(data, axis, reducer, keepdims, atleast1d);
 }
 

python/tvm/relay/frontend/onnx.py

@@ -1846,6 +1846,17 @@ def _impl_v1(cls, inputs, attr, params):
         return _op.cast(AttrCvt("argmin")(inputs, attr), "int64")
 
 
+    @classmethod
+    def _impl_v13(cls, inputs, attr, params):
+        if "select_last_index" in attr:
+            raise NotImplementedError("select_last_index not supported in ArgMin")
+        axis = attr.get("axis", 0)
+        keepdims = attr.get("keepdims", True)
+        select_last_index = attr.get("select_last_index", False)
+        attr = {"axis": axis, "keepdims": keepdims}
+        return _op.cast(AttrCvt("argmin")(inputs, attr), "int64")
+        # return _op.argmin()
+
 class Softmax(OnnxOpConverter):
     """Operator converter for Softmax."""
 

src/relay/op/tensor/reduce.cc

@@ -203,13 +203,34 @@ Array<te::Tensor> ReduceCompute(const Attrs& attrs, const Array<te::Tensor>& inp
   auto axes = param->axis;
   if (param->exclude) {
     axes = GetExcludeAxes(inputs[0]->shape.size(), param->axis);
-    if (axes.size() == 0) {
-      return {topi::identity(inputs[0])};
-    }
   }
+
+  if (axes.size() == 0) {
+    return {topi::identity(inputs[0])};
+  }
+
   return {f(inputs[0], axes, param->keepdims, false)};
 }
 
+template <typename F>
+Array<te::Tensor> OneElementReduceCompute(const Attrs& attrs, const Array<te::Tensor>& inputs,
+                                          const Type& out_type, F f) {
+  const OneElementReduceAttrs* param = attrs.as<OneElementReduceAttrs>();
+  ICHECK(param != nullptr);
+  if (inputs[0]->shape.size() == 0) {
+    return {topi::identity(inputs[0])};
+  }
+  auto axes = param->axis;
+  if (param->exclude) {
+    axes = GetExcludeAxes(inputs[0]->shape.size(), param->axis);
+  }
+
+  if (axes.size() == 0) {
+    return {topi::identity(inputs[0])};
+  }
+  return {f(inputs[0], axes, param->keepdims, false, param->select_last_index)};
+}
+
 /*!
  * \brief ReduceShapeImpl get the outshape for the reduction operator
  * \param in_shape Shape of input data.
@@ -333,31 +354,31 @@ Expr MakeReduce(Expr data, Array<Integer> axis, bool keepdims, bool exclude, Str
 
 Array<te::Tensor> ArgMaxCompute(const Attrs& attrs, const Array<te::Tensor>& inputs,
                                 const Type& out_type) {
-  return ReduceCompute(attrs, inputs, out_type, topi::argmax);
+  return OneElementReduceCompute(attrs, inputs, out_type, topi::argmax);
 }
 
 RELAY_REGISTER_REDUCE_OP("argmax")
     .describe(R"code(Creates an operation that finds the indices of the maximum
 values over a given axis.
 
 )code" TVM_ADD_FILELINE)
-    .set_attrs_type<ReduceAttrs>()
+    .set_attrs_type<OneElementReduceAttrs>()
     .set_support_level(4)
     .add_type_rel("ArgReduce", ArgReduceRel)
     .set_attr<FTVMCompute>("FTVMCompute", ArgMaxCompute)
     .set_attr<TOpPattern>("TOpPattern", kCommReduce);
 
 Array<te::Tensor> ArgMinCompute(const Attrs& attrs, const Array<te::Tensor>& inputs,
                                 const Type& out_type) {
-  return ReduceCompute(attrs, inputs, out_type, topi::argmin);
+  return OneElementReduceCompute(attrs, inputs, out_type, topi::argmin);
 }
 
 RELAY_REGISTER_REDUCE_OP("argmin")
     .describe(R"code(Creates an operation that finds the indices of the minimum
 values over a given axis.
 
 )code" TVM_ADD_FILELINE)
-    .set_attrs_type<ReduceAttrs>()
+    .set_attrs_type<OneElementReduceAttrs>()
     .set_support_level(4)
     .add_type_rel("ArgReduce", ArgReduceRel)
     .set_attr<FTVMCompute>("FTVMCompute", ArgMinCompute)