slice_scatter decomposition

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -162,6 +162,44 @@ def var_decomposition(
     return variance
 
 
+@register_torch_trt_decomposition(
+    torch.ops.slice_scatter, registry=TORCH_TRT_DECOMPOSITIONS
+)
+def slice_scatter_decomposition(
+    input_tensor: torch.Tensor,
+    src_tensor: torch.Tensor,
+    dim: int,
+    start: Optional[int],
+    end: Optional[int],
+    step: int,
+):
+    dim_size = input_tensor.shape[dim]
+    input_tensor_shape = input_tensor.shape
+    if start is not None and start < 0:
+        start = start + dim_size
+    if end is not None and end < 0:
+        end = end + dim_size
+    if start is None:
+        start = 0
+    if end is None:
+        end = dim_size
+
+    src_dim = list(src_tensor.shape())
+    src_dim[dim] = torch.floor_divide(end - start, step)
+    src = torch.expand(src, src_dim)
+
+    if (start == 0 and end == dim_size and step == 0):
+        return input_tensor
+    mask = []
+    if start != 0:
+        mask.append(torch.ge(input_tensor_shape, start))
+    if end != dim_size:
+        mask.append(torch.ge(input_tensor_shape, end))
+    if step != 1:
+        mask.append(torch.eq(src_dim, 0))
+    src_val = torch.masked(mask, src_dim, 0)
+    return torch.where(mask, src_val,input_tensor)
+
 def get_decompositions(
     enable_experimental_decompositions: bool = False,
 ) -> Dict[OpOverload, Callable[[Any], Any]]:

tests/py/dynamo/lowering/test_decompositions.py

@@ -421,5 +421,71 @@ def forward(self, x):
         )
 
 
+    def test_lowering_select_scatter_module(self):
+        class selectScatter(torch.nn.Module):
+            def __init__(self, *args, **kwargs) -> None:
+                super().__init__(*args, **kwargs)
+
+            def forward(self, x, src, dim, start):
+                y = self.slice_scatter(x, src, dim, start)
+                return y
+
+        # Operations expected to be removed in the traced graph after decompositions
+        expected_ops = {
+            torch.ops.aten.lt.default,
+            torch.ops.aten.lt.default,
+            torch.ops.aten.expand.default,
+            torch.ops.aten.eq.default,
+            torch.ops.aten.where.default,
+
+        }
+        unexpected_ops = {torch.ops.aten.select_scatter}
+
+        inputs = [torch.randn(2, 2), torch.ones(2)]
+
+        fx_graph = torch.fx.symbolic_trace(selectScatter())
+        unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(
+            fx_graph,
+            inputs,
+            expected_ops=expected_ops,
+            unexpected_ops=unexpected_ops,
+            min_block_size=1,
+        )
+
+        self.assertEquals(
+            len(unexpected_ops_seen),
+            0,
+            f"The following unexpected ops were encountered: {unexpected_ops_seen}",
+        )
+
+        self.assertEquals(
+            len(expected_ops_unseen),
+            0,
+            f"The following expected ops were not encountered: {expected_ops_unseen}",
+        )
+
+        torch._dynamo.reset()
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        optimized_model_results = optimized_model(*inputs).detach().cpu()
+        torch_model_results = fx_graph(*inputs).detach().cpu()
+
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            DECIMALS_OF_AGREEMENT,
+            f"Select_scatter TRT outputs don't match with the original model.",
+        )
+
 if __name__ == "__main__":
     run_tests()