scatter_value and scatter_src converter

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -691,6 +691,44 @@ def aten_ops_clamp(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.scatter.value)
+def aten_ops_scatter_value(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.select.scatter_value(
+        ctx, target, SourceIR.ATEN, name, args[0], args[1], args[2]
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.scatter.src)
+def aten_ops_scatter_src(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.select.scatter_src(
+        ctx, target, SourceIR.ATEN, name, args[0], args[1], args[2]
+    )
+
+
+def aten_ops_select(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.select.select(
+        ctx, target, SourceIR.ATEN, name, args[0], args[1], args[2]
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.select.int)
 def aten_ops_select(
     ctx: ConversionContext,

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -400,20 +400,25 @@ def scatter_value(
     input: TRTTensor,
     dim: Shape,
     index: Shape,
-    value: TRTTensor,
+    value: float,
 ) -> TRTTensor:
     if not isinstance(input, TRTTensor):
         raise RuntimeError(
             f"scatter_tensor received input {input} that is not part "
             "of the TensorRT region!"
         )
-
-    ranks = len(input.shape)
+    input_shape = input.shape
+    index_shape = index.shape
+    if (len(input_shape) != len(index_shape)):
+        raise RuntimeError(
+                f"The no of dimensions of input and index should be equal"
+            )
+    ranks = len(input_shape)
     dim = get_positive_dim(cast(int, dim), ranks)
     dynamic_shape = has_dynamic_shape(input.shape)
     if dynamic_shape:
         # Check whether slice target dim is dynamic shape dim
-        assert input.shape[dim] != -1, "Can't select on negative shape dimension!"
+        assert input.shape[dim] != -1, "Can't scatter on negative shape dimension!"
 
     input_dims = len(input.shape)
     for i in range(0, input_dims):
@@ -437,22 +442,32 @@ def scatter_src(
     input: TRTTensor,
     dim: Shape,
     index: Shape,
-    src: float,
+    src: TRTTensor,
 ) -> TRTTensor:
     if not isinstance(input, TRTTensor):
         raise RuntimeError(
             f"scatter_tensor received input {input} that is not part "
             "of the TensorRT region!"
         )
-
-    ranks = len(input.shape)
-    dim = get_positive_dim(cast(int, dim), ranks)
+    input_shape = input.shape
+    index_shape = index.shape
+    src_shape = src.shape
+    if (len(input_shape) != len(index_shape)):
+        raise RuntimeError(
+                f"The no of dimensions of input and index should be equal"
+            )
+    if (len(index_shape) != len(src_shape)):
+        raise RuntimeError(
+                f"The no of dimensions of src and index should be equal"
+            )
+    
+    input_dims = len(input_shape)
+    dim = get_positive_dim(cast(int, dim), input_dims)
     dynamic_shape = has_dynamic_shape(input.shape)
     if dynamic_shape:
         # Check whether slice target dim is dynamic shape dim
-        assert input.shape[dim] != -1, "Can't select on negative shape dimension!"
+        assert input.shape[dim] != -1, "Can't scatter on negative shape dimension!"
 
-    input_dims = len(input.shape)
     for i in range(0, input_dims):
         if index[i] >= input.shape[i]:
             raise RuntimeError(

tests/py/dynamo/conversion/test_scatter_aten.py

@@ -0,0 +1,59 @@
+import torch
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+from .harness import DispatchTestCase
+
+
+class TestScatterValueConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("scatter_zero_dim_indexOne_value", 0, [[0, 1, 2, 0]], 1),
+            ("scatter_zero_dim_indexTwo_value", 0, [[0, 1, 2, 0], [1, 2, 1, 1]], 1),
+            ("scatter_one_dim_indexOne_value", 1, [[0, 1, 2, 0]], 1),
+            ("scatter_one_dim_indexTwo_value", 1, [[0, 1, 2, 0], [1, 2, 1, 1]], 1),
+        ]
+    )
+    def test_scatter(self, _, dim, index, value):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input, src):
+                return torch.ops.aten.scatter.value(input, dim, index, value)
+
+        input = [torch.zeros(3, 5, dtype = torch.int32)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+
+class TestScatterSrcConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("scatter_zero_dim_indexOne", 0, [[0, 1, 2, 0]]),
+            ("scatter_zero_dim_indexTwo", 0, [[0, 1, 2, 0], [1, 2, 1, 1]]),
+            ("scatter_one_dim_indexOne", 1, [[0, 1, 2, 0]]),
+            ("scatter_one_dim_indexTwo", 1, [[0, 1, 2, 0], [1, 2, 1, 1]]),
+        ]
+    )
+    def test_scatter(self, _, dim, index):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input, src):
+                return torch.ops.aten.scatter.src(input, dim, index, src)
+        
+        src = [torch.arange(1, 11).reshape((2,5))]
+        input = torch.zeros(3, 5, dtype = src.dtype)
+        inputs = [input, src]
+        self.run_test(
+            TestModule(),
+            inputs,
+        )
+
+
+