[TORCH] scatter_reduce implementation

python/tvm/relay/frontend/pytorch.py

@@ -2684,6 +2684,57 @@ def scatter_add(self, inputs, input_types):
         src = inputs[3]
         return _op.scatter_add(data, index, src, axis=axis)
 
+    def scatter_reduce(self, inputs, input_types):
+        assert len(inputs) == 5 or len(inputs) == 6, (
+            "scatter_reduce takes 5 or 6 inputs (data, dim, index, src, reduce, include_self), "
+            + "but {} given".format(len(inputs))
+        )
+        data = inputs[0]
+        dim = inputs[1]
+        index = inputs[2]
+        src = inputs[3]
+        reduce = inputs[4]
+        if len(inputs) == 6:
+            include_self = inputs[5]
+            # TODO(vvchernov): support include_self == False
+            assert include_self, "include_self=False has not been suppoted for scatter_reduce yet"
+
+        data_shape = self.infer_shape(inputs[0])
+        data_rank = len(data_shape)
+        index_shape = self.infer_shape(inputs[2])
+        index_rank = len(index_shape)
+        src_shape = self.infer_shape(inputs[3])
+        src_rank = len(src_shape)
+        assert data_rank == index_rank, "Index rank is not the same as data rank"
+        assert data_rank == src_rank, "Src rank is not the same as data rank"
+
+        assert 0 <= dim < data_rank, "Dim is out of bounds"
+
+        for i in range(data_rank):
+            assert index_shape[i] <= src_shape[i], "Index dim size should be less than src one"
+            if i != dim:
+                assert (
+                    index_shape[i] <= data_shape[i]
+                ), "Index dim size should be less than data one"
+
+        red_valids = ["sum", "prod", "mean", "amax", "amin"]
+        assert reduce in red_valids, "Only {} modes are supported, but {} is gotten".format(
+            red_valids, reduce
+        )
+        if reduce == "sum":
+            reduce = "add"
+        elif reduce == "prod":
+            reduce = "mul"
+        elif reduce == "amin":
+            reduce = "min"
+        elif reduce == "amax":
+            reduce = "max"
+        else:  # reduce == "mean"
+            # TODO(vvchernov): support mean reduction
+            raise NotImplementedError("Mean reduction has not been supported yet!")
+
+        return _op.scatter_elements(data, index, src, axis=dim, reduction=reduce)
+
     def cumsum(self, inputs, input_types):
         data = inputs[0]
         dim = inputs[1]
@@ -3785,6 +3836,8 @@ def create_convert_map(self):
             "aten::nonzero": self.nonzero,
             "aten::nonzero_numpy": self.nonzero_numpy,
             "aten::scatter": self.scatter,
+            "aten::scatter_add": self.scatter_add,
+            "aten::scatter_reduce": self.scatter_reduce,
             "aten::index_put": self.index_put,
             "aten::scalar_tensor": self.scalar_tensor,
             "aten::__interpolate": self.interpolate,
@@ -3796,7 +3849,6 @@ def create_convert_map(self):
             "aten::new_empty": self.new_empty,
             "aten::randn": self.randn,
             "aten::bincount": self.bincount,
-            "aten::scatter_add": self.scatter_add,
             "aten::__not__": self.logical_not,
             "aten::hardswish": self.hard_swish,
             "aten::hardsigmoid": self.hard_sigmoid,

python/tvm/topi/scatter_elements.py

@@ -127,9 +127,9 @@ def gen_ir(data_ptr, indices_ptr, updates_ptr, out_ptr):
                 elif reduction == "mul":
                     out[index2] *= updates[index1]
                 elif reduction == "min":
-                    tir.min(out[index2], updates[index1])
+                    out[index2] = tir.min(out[index2], updates[index1])
                 elif reduction == "max":
-                    tir.max(out[index2], updates[index1])
+                    out[index2] = tir.max(out[index2], updates[index1])
                 else:
                     raise NotImplementedError(
                         "scatter_elements reduction not in [update, add, mul, min, max]:",

tests/python/frontend/pytorch/test_forward.py

@@ -4233,6 +4233,32 @@ def test_fn_scatter_add(dim):
     verify_trace_model(test_fn_scatter_add(1), [in_data, in_index, in_src], targets)
 
 
+def test_forward_scatter_reduce():
+    """test_forward_scatter_reduce"""
+    # integer cannot be traced
+    def test_fn_scatter_reduce(dim, reduce):
+        return lambda data, index, src: torch.scatter_reduce(
+            data, dim=dim, index=index, src=src, reduce=reduce
+        )
+
+    in_data = torch.rand(3, 5) - 1
+    in_index = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
+    in_src = torch.rand(2, 5) - 1
+
+    targets = ["llvm", "cuda"]
+    # TODO(vvchernov): support test of mean reduction and include_self=False
+    for reduce in ["sum", "prod", "amin", "amax"]:
+        verify_trace_model(test_fn_scatter_reduce(0, reduce), [in_data, in_index, in_src], targets)
+
+    in_data = torch.rand(2, 4) - 1
+    in_index = torch.tensor([[2], [3]])
+    in_src = torch.rand(2, 1) - 1
+
+    # TODO(vvchernov): support test of mean reduction and include_self=False
+    for reduce in ["sum", "prod", "amin", "amax"]:
+        verify_trace_model(test_fn_scatter_reduce(1, reduce), [in_data, in_index, in_src], targets)
+
+
 def test_forward_index_put():
     """test_forward_index_put"""
     # torch.index_put for 2D tensor and default accumulate (False)