diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 919ac65f504a..2ea764872c27 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/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, diff --git a/python/tvm/topi/scatter_elements.py b/python/tvm/topi/scatter_elements.py index e7948f6cb6c7..bfa765855b0e 100644 --- a/python/tvm/topi/scatter_elements.py +++ b/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]:", diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 033ce64b3ac4..39d78bd6065c 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/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)