slice scatter changes

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -164,7 +164,7 @@ def var_decomposition(
 
 
 @register_torch_trt_decomposition(
-    torch.ops.slice_scatter, registry=TORCH_TRT_DECOMPOSITIONS
+    torch.ops.aten.slice_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS
 )
 def slice_scatter_decomposition(
     input_tensor: torch.Tensor,
@@ -175,7 +175,6 @@ def slice_scatter_decomposition(
     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:
@@ -185,22 +184,29 @@ def slice_scatter_decomposition(
     if end is None:
         end = dim_size
 
-    src_dim = list(src_tensor.shape())
+    src_dim = src_tensor.shape
     step_dim = torch.floor_divide(end - start, step)
-    # src = torch.expand(src, src_dim)
     end_dim = end
     if step_dim > src_dim[dim]:
         end_dim = src_dim[dim]
     else:
-        src_tensor = src_tensor[0:step_dim]
+        indices = torch.Tensor(np.arange(0, step_dim))
+        indices = indices.to(torch.int32)
+        src = torch.index_select(src, dim, indices)
 
     if start == 0 and end == dim_size and step == 0:
         return input_tensor
-    index_tensor = np.arange[start, end_dim, step]
+    index_tensor = np.arange(start, end_dim, step)
 
-    unbind_tensors = torch.unbind(input_tensor, dim)
-    unbind_tensors[index_tensor] = src_tensor
-    return torch.cat(unbind_tensors, dim)
+    unbind_input_tensors = torch.unbind(input_tensor, dim)
+    unbind_input_tensors_list = list(unbind_input_tensors)
+    unbind_source_tensors = torch.unbind(src, dim)
+    unbind_source_tensors_list = list(unbind_source_tensors)
+
+    for i, index in enumerate(index_tensor):
+        unbind_input_tensors_list[index] = unbind_source_tensors_list[i]
+
+    return torch.stack(unbind_input_tensors_list, dim)
 
 
 def get_decompositions(

tests/py/dynamo/lowering/test_decompositions.py

@@ -420,30 +420,91 @@ def forward(self, x):
             f"MaxPool3d TRT outputs don't match with the original model.",
         )
 
-
-    def test_lowering_select_scatter_module(self):
-        class selectScatter(torch.nn.Module):
+    def test_lowering_slice_scatter_dimZero_module(self):
+        class sliceScatter(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)
+            def forward(self, x, src, dim, start=None, end=None, step=1):
+                y = self.slice_scatter(x, src, dim, start, end, step)
                 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,
+            torch.ops.aten.slice.Tensor,
+            torch.ops.aten.squeeze.dim,
+            torch.ops.aten.cat.default,
+            torch.ops.aten.index.Tensor,
+        }
+        unexpected_ops = {torch.ops.aten.select_scatter}
+
+        inputs = [torch.zeros(8, 8).cuda(), torch.ones(2, 8).cuda(), 0, 6]
+
+        fx_graph = torch.fx.symbolic_trace(sliceScatter())
+        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"Slice_scatter TRT outputs don't match with the original model.",
+        )
+
+    def test_lowering_slice_scatter_dimOne_module(self):
+        class sliceScatter(torch.nn.Module):
+            def __init__(self, *args, **kwargs) -> None:
+                super().__init__(*args, **kwargs)
+
+            def forward(self, x, src, dim, start=None, end=None, step=1):
+                y = self.slice_scatter(x, src, dim, start, end, step)
+                return y
+
+        # Operations expected to be removed in the traced graph after decompositions
+        expected_ops = {
+            torch.ops.aten.slice.Tensor,
+            torch.ops.aten.squeeze.dim,
+            torch.ops.aten.cat.default,
+            torch.ops.aten.index.Tensor,
         }
         unexpected_ops = {torch.ops.aten.select_scatter}
 
-        inputs = [torch.randn(2, 2), torch.ones(2)]
+        inputs = [torch.zeros(8, 8).cuda(), torch.ones(2, 8).cuda(), 0, 6]
 
-        fx_graph = torch.fx.symbolic_trace(selectScatter())
+        fx_graph = torch.fx.symbolic_trace(sliceScatter())
         unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(
             fx_graph,
             inputs,
@@ -484,8 +545,9 @@ def forward(self, x, src, dim, start):
             max_diff,
             0,
             DECIMALS_OF_AGREEMENT,
-            f"Select_scatter TRT outputs don't match with the original model.",
+            f"Slice_scatter TRT outputs don't match with the original model.",
         )
 
+
 if __name__ == "__main__":
     run_tests()