From 0eb0c5d1bf738465bb9c858e4a703f24643e333e Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Thu, 9 Feb 2023 16:53:23 +0300
Subject: [PATCH 1/9] update ScatterND attrs class

---
 include/tvm/relay/attrs/transform.h | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h
index 274a421e5719..82403d7c40ee 100644
--- a/include/tvm/relay/attrs/transform.h
+++ b/include/tvm/relay/attrs/transform.h
@@ -168,8 +168,9 @@ struct ScatterNDAttrs : public tvm::AttrsNode<ScatterNDAttrs> {
   String mode;
 
   TVM_DECLARE_ATTRS(ScatterNDAttrs, "relay.attrs.ScatterNDAttrs") {
-    TVM_ATTR_FIELD(mode).describe(
-        "Accumulation mode of the scatter, either \"update\" or \"add\".");
+    TVM_ATTR_FIELD(mode).set_default("update").describe(
+        "Accumulation mode of the ScatterND, "
+        "either \"update\", \"add\", \"mul\", \"min\" or \"max\".");
   }
 };
 

From be1dda1619db83e13d239c05ff57b22d8bb02a55 Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Fri, 10 Feb 2023 08:38:30 +0300
Subject: [PATCH 2/9] check methods were added to ScatterND

---
 python/tvm/relay/frontend/onnx.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 8de5e0e08bd8..8f80fa7565c5 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -2856,8 +2856,37 @@ def _impl_v1(cls, inputs, attr, params):
 class ScatterND(OnnxOpConverter):
     """Operator converter for ScatterND."""
 
+    @classmethod
+    def _inputs_check(cls, inputs):
+        assert (
+            len(inputs) == 3
+        ), "ScatterND takes 3 inputs (data, indices, updates), {} given".format(len(inputs))
+        assert infer_type(inputs[1]).checked_type.dtype == "int64"
+
+        data_rank = len(infer_shape(inputs[0]))
+        assert data_rank > 0, "Data rank higher than 0 is expected"
+        indices_rank = len(infer_shape(inputs[1]))
+        assert indices_rank > 0, "Indices rank higher than 0 is expected"
+        updates_rank = len(infer_shape(inputs[2]))
+        assert (
+            updates_rank == data_rank + indices_rank - infer_shape(inputs[1])[-1] - 1
+        ), "Updates rank should be equal to data_rank + indices_rank - indices_shape[-1] - 1"
+
+    @classmethod
+    def _reduction_check(cls, attr, red_valids=["update"]):
+        reduction = attr.get("reduction", None)
+        if reduction is None:
+            reduction = b"update"
+        reduction = reduction.decode("utf-8")
+        assert reduction in red_valids, "Only {} reductions are supported, but {} is gotten".format(
+            red_valids, reduction
+        )
+
+        return reduction
+
     @classmethod
     def _impl_v11(cls, inputs, attr, params):
+        cls._inputs_check(inputs)
         indices_dim = len(infer_shape(inputs[1]))
         axes = list(range(indices_dim))
         return _op.scatter_nd(

From 051765a263d9953c550f54333218bc86d574353a Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Fri, 10 Feb 2023 08:45:41 +0300
Subject: [PATCH 3/9] support 16 and 18 opset for ScatterND on ONNX front-end

---
 python/tvm/relay/frontend/onnx.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 8f80fa7565c5..682089f313b2 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -2893,6 +2893,28 @@ def _impl_v11(cls, inputs, attr, params):
             inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2], "update"
         )
 
+    @classmethod
+    def _impl_v16(cls, inputs, attr, params):
+        cls._inputs_check(inputs)
+        reduction = cls._reduction_check(attr, ["update", "add", "mul"])
+
+        indices_dim = len(infer_shape(inputs[1]))
+        axes = list(range(indices_dim))
+        return _op.scatter_nd(
+            inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2], reduction
+        )
+
+    @classmethod
+    def _impl_v18(cls, inputs, attr, params):
+        cls._inputs_check(inputs)
+        reduction = cls._reduction_check(attr, ["update", "add", "mul", "min", "max"])
+
+        indices_dim = len(infer_shape(inputs[1]))
+        axes = list(range(indices_dim))
+        return _op.scatter_nd(
+            inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2], reduction
+        )
+
 
 class EyeLike(OnnxOpConverter):
     """Operator converter for EyeLike."""

From 1139780435dcb3fdbfd41628afdd9aaa6a2d61dc Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Fri, 10 Feb 2023 11:55:00 +0300
Subject: [PATCH 4/9] small updates

---
 python/tvm/relay/frontend/onnx.py | 2 +-
 python/tvm/relay/op/transform.py  | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 682089f313b2..a6d4d5a946e1 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -2890,7 +2890,7 @@ def _impl_v11(cls, inputs, attr, params):
         indices_dim = len(infer_shape(inputs[1]))
         axes = list(range(indices_dim))
         return _op.scatter_nd(
-            inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2], "update"
+            inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2]
         )
 
     @classmethod
diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index c7234f340395..782797dadb83 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -420,7 +420,13 @@ def scatter_nd(data, indices, updates, mode="update"):
         The values to update.
 
     mode : string, optional
-        The accumulation mode for scatter. "update" or "add"
+        The accumulation mode for scatter. "update", "add", "mul", "min" or "max"
+        If update, the update values will replace the input data
+        If add, the update values will be added to the input data
+        If mul, the update values will be multiply to the input data
+        If min, there is choice of minimal between the update values and the input data
+        If max, there is choice of maximal between the update values and the input data
+        It is "update" by default
 
     Returns
     -------

From 8dbc01387c890f264f39806f7769a4721aae818d Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Fri, 10 Feb 2023 12:51:11 +0300
Subject: [PATCH 5/9] extend mode types for scatter_nd

---
 python/tvm/topi/cuda/scatter.py | 12 ++++++++++--
 python/tvm/topi/scatter.py      | 12 ++++++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index fa7545cd323a..f104d1e42cb1 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -17,7 +17,7 @@
 # pylint: disable=invalid-name, no-member, too-many-locals, too-many-arguments, too-many-statements, singleton-comparison, unused-argument
 """Scatter operator """
 import tvm
-from tvm import te, autotvm
+from tvm import te, tir, autotvm
 from ..scatter import _verify_scatter_nd_inputs
 from ..generic import schedule_extern
 from .nms import atomic_add
@@ -871,8 +871,16 @@ def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
                             out[index] = updates[i * fused_updates_dimension + j]
                         elif mode == "add":
                             out[index] += updates[i * fused_updates_dimension + j]
+                        elif mode == "mul":
+                            out[index] *= updates[i * fused_updates_dimension + j]
+                        elif mode == "min":
+                            out[index] = tir.min(out[index], updates[i * fused_updates_dimension + j])
+                        elif mode == "max":
+                            out[index] = tir.max(out[index], updates[i * fused_updates_dimension + j])
                         else:
-                            raise NotImplementedError("scatter_nd mode not in [update, add]:", mode)
+                            raise NotImplementedError(
+                                "scatter_nd mode not in [update, add, mul, min, max]:", mode
+                            )
 
         return ib.get()
 
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index e0578aab41b9..2298f09b4a67 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -17,7 +17,7 @@
 # pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks
 """Scatter operator"""
 from ..te import extern, hybrid
-from ..tir import decl_buffer, expr, ir_builder
+from ..tir import decl_buffer, expr, ir_builder, min, max
 
 
 @hybrid.script
@@ -308,8 +308,16 @@ def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
                     out[index] = updates[i * fused_updates_dimension + j]
                 elif mode == "add":
                     out[index] += updates[i * fused_updates_dimension + j]
+                elif mode == "mul":
+                    out[index] *= updates[i * fused_updates_dimension + j]
+                elif mode == "min":
+                    out[index] = min(out[index], updates[i * fused_updates_dimension + j])
+                elif mode == "max":
+                    out[index] = max(out[index], updates[i * fused_updates_dimension + j])
                 else:
-                    raise NotImplementedError("scatter_nd mode not in [update, add]:", mode)
+                    raise NotImplementedError(
+                        "scatter_nd mode not in [update, add, mul, min, max]:", mode
+                    )
 
         return ib.get()
 

From 320d670d69155667b7fc4f3e17f8893ad08eec08 Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Sat, 11 Feb 2023 11:18:06 +0300
Subject: [PATCH 6/9] lint fixes

---
 python/tvm/relay/frontend/onnx.py |  8 ++++----
 python/tvm/topi/cuda/scatter.py   |  8 ++++++--
 python/tvm/topi/scatter.py        | 22 +++++++++++-----------
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index a6d4d5a946e1..aebc6daa5ebe 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -2873,11 +2873,13 @@ def _inputs_check(cls, inputs):
         ), "Updates rank should be equal to data_rank + indices_rank - indices_shape[-1] - 1"
 
     @classmethod
-    def _reduction_check(cls, attr, red_valids=["update"]):
+    def _reduction_check(cls, attr, red_valids=None):
         reduction = attr.get("reduction", None)
         if reduction is None:
             reduction = b"update"
         reduction = reduction.decode("utf-8")
+        if red_valids is None:
+            red_valids = ["update"]
         assert reduction in red_valids, "Only {} reductions are supported, but {} is gotten".format(
             red_valids, reduction
         )
@@ -2889,9 +2891,7 @@ def _impl_v11(cls, inputs, attr, params):
         cls._inputs_check(inputs)
         indices_dim = len(infer_shape(inputs[1]))
         axes = list(range(indices_dim))
-        return _op.scatter_nd(
-            inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2]
-        )
+        return _op.scatter_nd(inputs[0], _op.transpose(inputs[1], axes[-1:] + axes[:-1]), inputs[2])
 
     @classmethod
     def _impl_v16(cls, inputs, attr, params):
diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index f104d1e42cb1..1bdd53156623 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -874,9 +874,13 @@ def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
                         elif mode == "mul":
                             out[index] *= updates[i * fused_updates_dimension + j]
                         elif mode == "min":
-                            out[index] = tir.min(out[index], updates[i * fused_updates_dimension + j])
+                            out[index] = tir.min(
+                                out[index], updates[i * fused_updates_dimension + j]
+                            )
                         elif mode == "max":
-                            out[index] = tir.max(out[index], updates[i * fused_updates_dimension + j])
+                            out[index] = tir.max(
+                                out[index], updates[i * fused_updates_dimension + j]
+                            )
                         else:
                             raise NotImplementedError(
                                 "scatter_nd mode not in [update, add, mul, min, max]:", mode
diff --git a/python/tvm/topi/scatter.py b/python/tvm/topi/scatter.py
index 2298f09b4a67..45629c005f79 100644
--- a/python/tvm/topi/scatter.py
+++ b/python/tvm/topi/scatter.py
@@ -16,11 +16,11 @@
 # under the License.
 # pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks
 """Scatter operator"""
-from ..te import extern, hybrid
-from ..tir import decl_buffer, expr, ir_builder, min, max
+from tvm import te, tir  # hide redefinition of min and max
+from tvm.tir import expr
 
 
-@hybrid.script
+@te.hybrid.script
 def _scatter_1d(data, indices, updates):
     out = output_tensor(data.shape, data.dtype)
     for i in range(data.shape[0]):
@@ -30,7 +30,7 @@ def _scatter_1d(data, indices, updates):
     return out
 
 
-@hybrid.script
+@te.hybrid.script
 def _scatter_2d(data, indices, updates, axis):
     out = output_tensor(data.shape, data.dtype)
     for i in range(data.shape[0]):
@@ -52,7 +52,7 @@ def _scatter_2d(data, indices, updates, axis):
     return out
 
 
-@hybrid.script
+@te.hybrid.script
 def _scatter_3d(data, indices, updates, axis):
     out = output_tensor(data.shape, data.dtype)
     for i in range(data.shape[0]):
@@ -96,7 +96,7 @@ def _scatter_3d(data, indices, updates, axis):
     return out
 
 
-@hybrid.script
+@te.hybrid.script
 def _scatter_4d(data, indices, updates, axis):
     out = output_tensor(data.shape, data.dtype)
     for i in range(data.shape[0]):
@@ -269,7 +269,7 @@ def scatter_nd(data, indices, updates, mode):
 
     def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
         # pylint: disable=invalid-name
-        ib = ir_builder.create()
+        ib = tir.ir_builder.create()
 
         data = ib.buffer_ptr(data_ptr)
         indices = ib.buffer_ptr(indices_ptr)
@@ -311,9 +311,9 @@ def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
                 elif mode == "mul":
                     out[index] *= updates[i * fused_updates_dimension + j]
                 elif mode == "min":
-                    out[index] = min(out[index], updates[i * fused_updates_dimension + j])
+                    out[index] = tir.min(out[index], updates[i * fused_updates_dimension + j])
                 elif mode == "max":
-                    out[index] = max(out[index], updates[i * fused_updates_dimension + j])
+                    out[index] = tir.max(out[index], updates[i * fused_updates_dimension + j])
                 else:
                     raise NotImplementedError(
                         "scatter_nd mode not in [update, add, mul, min, max]:", mode
@@ -321,8 +321,8 @@ def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
 
         return ib.get()
 
-    out_buf = decl_buffer(data.shape, data.dtype, "out_buf")
-    return extern(
+    out_buf = tir.decl_buffer(data.shape, data.dtype, "out_buf")
+    return te.extern(
         [data.shape],
         [data, indices, updates],
         lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0]),

From 12d87c4d63802ebedbe13c20b50f4392c797cf01 Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Sat, 11 Feb 2023 14:08:05 +0300
Subject: [PATCH 7/9] update test_op_level3 with new modes of scatter_nd

---
 tests/python/relay/test_op_level3.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index efd37f2ecd22..225210f4d617 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1983,6 +1983,7 @@ def verify_scatter_nd_with_stack(
         )
         tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=rtol, atol=atol)
 
+    # TODO(vcchernov): check frameworks' int type requirements. ONNX expects int64 only
     for indice_dtype in ["uint8", "uint16", "uint32"]:
         data = np.zeros((2, 2)).astype("int64")
         indices = np.array([[1, 1, 0], [0, 1, 0]]).astype(indice_dtype)
@@ -2009,7 +2010,7 @@ def verify_scatter_nd_with_stack(
         verify_scatter_nd(data, indices, updates, out, mode="add")
         verify_scatter_nd_with_stack(data, indices, updates, out)
 
-        for mode in ["add", "update"]:
+        for mode in ["update", "add", "mul", "min", "max"]:
             indices = np.stack((np.random.randint(2, size=5), np.random.randint(7, size=5))).astype(
                 indice_dtype
             )
@@ -2019,10 +2020,20 @@ def verify_scatter_nd_with_stack(
             out = data.copy()
             for i in range(indices.shape[1]):
                 for j in range(updates.shape[1]):
-                    if mode == "add":
-                        out[indices[0, i], indices[1, i], j] += updates[i, j]
-                    elif mode == "update":
+                    if mode == "update":
                         out[indices[0, i], indices[1, i], j] = updates[i, j]
+                    elif mode == "add":
+                        out[indices[0, i], indices[1, i], j] += updates[i, j]
+                    elif mode == "mul":
+                        out[indices[0, i], indices[1, i], j] *= updates[i, j]
+                    elif mode == "min":
+                        out[indices[0, i], indices[1, i], j] = min(
+                            out[indices[0, i], indices[1, i], j], updates[i, j]
+                        )
+                    elif mode == "max":
+                        out[indices[0, i], indices[1, i], j] = max(
+                            out[indices[0, i], indices[1, i], j], updates[i, j]
+                        )
             verify_scatter_nd(data, indices, updates, out, mode)
             verify_scatter_nd_with_stack(data, indices, updates, out, mode)
 

From c24b906408e1c621dcfc36618ec0ca0ab2f46d8a Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Sat, 11 Feb 2023 14:14:39 +0300
Subject: [PATCH 8/9] the same update for test_topi_scatter

---
 tests/python/topi/python/test_topi_scatter.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/tests/python/topi/python/test_topi_scatter.py b/tests/python/topi/python/test_topi_scatter.py
index 025e44889d63..ccc34837a05a 100644
--- a/tests/python/topi/python/test_topi_scatter.py
+++ b/tests/python/topi/python/test_topi_scatter.py
@@ -61,7 +61,7 @@ def check_scatter_nd(data, indices, updates, out, mode="add"):
     out[0, :] += updates[2, :]
     check_scatter_nd(data, indices, updates, out)
 
-    for mode in ["add", "update"]:
+    for mode in ["update", "add", "mul", "min", "max"]:
         updates = np.ones((5, 3)).astype("float64")
         indices = np.stack((np.random.randint(2, size=5), np.random.randint(7, size=5))).astype(
             "int64"
@@ -71,10 +71,20 @@ def check_scatter_nd(data, indices, updates, out, mode="add"):
         out = data.copy()
         for i in range(indices.shape[1]):
             for j in range(updates.shape[1]):
-                if mode == "add":
-                    out[indices[0, i], indices[1, i], j] += updates[i, j]
-                elif mode == "update":
+                if mode == "update":
                     out[indices[0, i], indices[1, i], j] = updates[i, j]
+                elif mode == "add":
+                    out[indices[0, i], indices[1, i], j] += updates[i, j]
+                elif mode == "mul":
+                    out[indices[0, i], indices[1, i], j] *= updates[i, j]
+                elif mode == "min":
+                    out[indices[0, i], indices[1, i], j] = min(
+                        out[indices[0, i], indices[1, i], j], updates[i, j]
+                    )
+                elif mode == "max":
+                    out[indices[0, i], indices[1, i], j] = max(
+                        out[indices[0, i], indices[1, i], j], updates[i, j]
+                    )
 
         check_scatter_nd(data, indices, updates, out, mode)
 

From a24cd288c4f8fb3a733051600ea7f9df26b22170 Mon Sep 17 00:00:00 2001
From: Valery Chernov <valery.chernov@deelvin.com>
Date: Mon, 13 Feb 2023 08:23:57 +0300
Subject: [PATCH 9/9] enclose scatter_nd tests from ONNX

---
 tests/python/frontend/onnx/test_forward.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index 0a032843267a..470a67e86c93 100644
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -5398,8 +5398,6 @@ def verify_eyelike(indata, dynamic=False):
     "test_reduce_sum_negative_axes_keepdims_random",
     "test_roialign_aligned_true",
     "test_scatter_elements_with_duplicate_indices",
-    "test_scatternd_add",
-    "test_scatternd_multiply",
     "test_sequence_insert_at_back",
     "test_sequence_insert_at_front",
     "test_sequence_map_add_1_sequence_1_tensor",
