From 357a2c263c9a990817f80500bcd5d146a890e9d2 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Wed, 15 Feb 2023 13:15:21 +0300 Subject: [PATCH 1/7] add scatter_reduce to pytorch front-end --- python/tvm/relay/frontend/pytorch.py | 42 +++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 919ac65f504a..27d5c302e361 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2684,6 +2684,45 @@ 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] + assert include_self, "include_self=False has not been suppoted for scatter_reduce yet" + + data_rank = len(self.infer_shape(inputs[0])) + index_rank = len(self.infer_shape(inputs[2])) + src_rank = len(self.infer_shape(inputs[3])) + 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" + + 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 == "mean": + raise NotImplementedError("Mean reduction has not been supported yet!") + elif reduce == "sum": + reduce = "add" + elif reduce == "prod": + reduce = "mul" + elif reduce == "amin": + reduce = "min" + elif reduce == "amax": + reduce = "max" + + 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 +3824,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 +3837,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, From 224d01b958f029bc024228a990b4598183598c13 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Wed, 15 Feb 2023 13:22:15 +0300 Subject: [PATCH 2/7] test for scatter_reduce was added to pytorch CI --- tests/python/frontend/pytorch/test_forward.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 033ce64b3ac4..fe5ccef07993 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -4233,6 +4233,30 @@ 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.zeros(3, 5) + in_index = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]) + in_src = torch.rand(2, 5) + + targets = ["llvm", "cuda"] + 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.zeros(2, 4) + in_index = torch.tensor([[2], [3]]) + in_src = torch.rand(2, 1) + + 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) From 5bca539792ae86d3c63bd30d6e1fe78f7842d004 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Wed, 15 Feb 2023 18:26:24 +0300 Subject: [PATCH 3/7] update check --- python/tvm/relay/frontend/pytorch.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 27d5c302e361..7f877e1de5e1 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2698,14 +2698,24 @@ def scatter_reduce(self, inputs, input_types): include_self = inputs[5] assert include_self, "include_self=False has not been suppoted for scatter_reduce yet" - data_rank = len(self.infer_shape(inputs[0])) - index_rank = len(self.infer_shape(inputs[2])) - src_rank = len(self.infer_shape(inputs[3])) + 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 From c52490e95d207f9431108b3fc594da7bafb5ce38 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Wed, 15 Feb 2023 18:28:58 +0300 Subject: [PATCH 4/7] add TODOs waiting for other PRs for development --- python/tvm/relay/frontend/pytorch.py | 1 + tests/python/frontend/pytorch/test_forward.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 7f877e1de5e1..6c067422ec62 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2721,6 +2721,7 @@ def scatter_reduce(self, inputs, input_types): red_valids, reduce ) if reduce == "mean": + # TODO(vvchernov): support mean reduction raise NotImplementedError("Mean reduction has not been supported yet!") elif reduce == "sum": reduce = "add" diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index fe5ccef07993..6a6787f3302a 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -4246,6 +4246,7 @@ def test_fn_scatter_reduce(dim, reduce): in_src = torch.rand(2, 5) targets = ["llvm", "cuda"] + # TODO(vvchernov): support test of mean reduction for reduce in ["sum", "prod", "amin", "amax"]: verify_trace_model(test_fn_scatter_reduce(0, reduce), [in_data, in_index, in_src], targets) @@ -4253,6 +4254,7 @@ def test_fn_scatter_reduce(dim, reduce): in_index = torch.tensor([[2], [3]]) in_src = torch.rand(2, 1) + # TODO(vvchernov): support test of mean reduction for reduce in ["sum", "prod", "amin", "amax"]: verify_trace_model(test_fn_scatter_reduce(1, reduce), [in_data, in_index, in_src], targets) From 7cebcc62d0dbdcad7732ff904b7ed54d42e27406 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Fri, 17 Feb 2023 09:59:34 +0300 Subject: [PATCH 5/7] fix lint --- python/tvm/relay/frontend/pytorch.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 6c067422ec62..f87422d6beb0 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2685,10 +2685,10 @@ def scatter_add(self, inputs, input_types): 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)) + 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] @@ -2720,10 +2720,7 @@ def scatter_reduce(self, inputs, input_types): assert reduce in red_valids, "Only {} modes are supported, but {} is gotten".format( red_valids, reduce ) - if reduce == "mean": - # TODO(vvchernov): support mean reduction - raise NotImplementedError("Mean reduction has not been supported yet!") - elif reduce == "sum": + if reduce == "sum": reduce = "add" elif reduce == "prod": reduce = "mul" @@ -2731,6 +2728,9 @@ def scatter_reduce(self, inputs, input_types): 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) From 174f2db506a04c9c0e0dfbc66f9c01016ee67386 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Fri, 17 Feb 2023 14:40:36 +0300 Subject: [PATCH 6/7] fix min-max reduction for cpu --- python/tvm/relay/frontend/pytorch.py | 2 +- python/tvm/topi/scatter_elements.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index f87422d6beb0..f6a623be281d 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2728,7 +2728,7 @@ def scatter_reduce(self, inputs, input_types): reduce = "min" elif reduce == "amax": reduce = "max" - else: # reduce == "mean" + else: # reduce == "mean" # TODO(vvchernov): support mean reduction raise NotImplementedError("Mean reduction has not been supported yet!") 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]:", From 71e5b039543d1fe2b5388c4b7f516785d0bc3b77 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Sat, 18 Feb 2023 21:18:17 +0300 Subject: [PATCH 7/7] final clean code --- python/tvm/relay/frontend/pytorch.py | 1 + tests/python/frontend/pytorch/test_forward.py | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index f6a623be281d..2ea764872c27 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2696,6 +2696,7 @@ def scatter_reduce(self, inputs, input_types): 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]) diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 6a6787f3302a..39d78bd6065c 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -4241,20 +4241,20 @@ def test_fn_scatter_reduce(dim, reduce): data, dim=dim, index=index, src=src, reduce=reduce ) - in_data = torch.zeros(3, 5) + 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) + in_src = torch.rand(2, 5) - 1 targets = ["llvm", "cuda"] - # TODO(vvchernov): support test of mean reduction + # 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.zeros(2, 4) + in_data = torch.rand(2, 4) - 1 in_index = torch.tensor([[2], [3]]) - in_src = torch.rand(2, 1) + in_src = torch.rand(2, 1) - 1 - # TODO(vvchernov): support test of mean reduction + # 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)