scatter_add_decomposition

Fixing scatter_add test cases. To do: fix the index collision cases

Index collision cases

Index collision cases- removing the torch.unique checl

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -243,6 +243,37 @@ def empty_strided_decomposition(*args, **kwargs) -> torch.Tensor:
     )
 
 
+@register_torch_trt_decomposition(
+    torch.ops.aten.scatter_add.default, registry=TORCH_TRT_DECOMPOSITIONS
+)
+def scatter_add_decomposition(
+    input_tensor: torch.Tensor,
+    src_tensor: torch.Tensor,
+    dim: int,
+    index: torch.Tensor,
+) -> torch.Tensor:
+    scatter_add_tensor = input_tensor
+    src_copy = src_tensor
+    src_shape = list(src_tensor.shape)
+    del src_shape[dim]
+    select_src_dim = src_copy.shape[dim]
+    to_stack_dummy_src = tuple(torch.empty(src_shape) for _ in range(select_src_dim))
+    for index_src_dim in range(0, select_src_dim, 1):
+        select_tensor_dim = torch.select(src_copy, dim, index_src_dim)
+        to_stack_src = to_stack_dummy_src
+        if(index_src_dim == 0):
+            to_stack_src = (select_tensor_dim.cpu(),) + to_stack_dummy_src[index_src_dim+1:] 
+        elif(index_src_dim == select_src_dim - 1 ):
+            to_stack_src = to_stack_dummy_src[:index_src_dim] + (select_tensor_dim.cpu(),)
+        else:
+            to_stack_src = to_stack_dummy_src[:index_src_dim] + (select_tensor_dim.cpu(),) + to_stack_dummy_src[index_src_dim+1:]
+
+        stacked_src = torch.stack(to_stack_src, dim)
+        input_tensor_to_add = torch.scatter(torch.empty_like(input_tensor, dtype= torch.float32), dim, index, stacked_src.cuda())
+        scatter_add_tensor = torch.add(scatter_add_tensor, input_tensor_to_add)
+    return scatter_add_tensor
+
+    
 def get_decompositions(
     enable_experimental_decompositions: bool = False,
 ) -> Dict[OpOverload, Callable[[Any], Any]]:

tests/py/dynamo/lowering/test_decompositions.py

@@ -2,6 +2,7 @@
 import torch_tensorrt
 from parameterized import parameterized
 from torch.testing._internal.common_utils import TestCase, run_tests
+from parameterized import parameterized
 
 from ..testing_utilities import DECIMALS_OF_AGREEMENT, lower_graph_testing
 
@@ -963,5 +964,60 @@ def forward(self, input):
         )
 
 
+class TestScatterAdd(TestCase):
+    @parameterized.expand(
+        [
+            (
+                "scatter_add_zero_dim_indexOne_constant",
+                0,
+                torch.tensor([[0, 1, 2, 0]]),
+                torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
+            ),
+            (
+                "scatter_add_zero_dim_indexTwo_constant",
+                0,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+                torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.int32),
+            ),
+            (
+                "scatter_add_one_dim_indexOne_constant",
+                1,
+                torch.tensor([[0, 1, 2, 0]]),
+                torch.tensor([[1, 2, 3, 1]], dtype=torch.int32),
+            ),
+            (
+                "scatter_add_one_dim_indexTwo_costant",
+                1,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+                torch.tensor([[1, 2, 3, 1], [5, 6, 5, 5]], dtype=torch.int32),
+            ),
+        ]
+    )
+    def test_scatter_add(self, _, dim, index, src):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                return torch.ops.aten.scatter_add.default(input, dim, index, src)
+
+        # Operations expected to be included in the traced graph after decompositions
+        expected_ops = {torch.ops.aten.scatter.src}
+
+        input = torch.zeros(3, 5, dtype=torch.int32)
+        inputs = [input]
+
+        fx_graph = torch.fx.symbolic_trace(TestModule())
+        _, expected_ops_unseen = lower_graph_testing(
+            fx_graph, inputs, expected_ops=expected_ops, min_block_size=2
+        )
+
+        self.assertEquals(
+            len(expected_ops_unseen),
+            0,
+            f"The following expected ops were not encountered: {expected_ops_unseen}",
+        )
+
+
 if __name__ == "__main__":
     run_tests()