slice_scatter adding to decomposition group

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py

@@ -174,6 +174,7 @@
     aten.full,
     aten.repeat,
     aten.var_mean,
+    aten.slice_scatter,
 }
 torch_disabled_decompositions: Set[Union[OpOverload, OpOverloadPacket]] = {
     aten._softmax.default,

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -171,23 +171,27 @@ def slice_scatter_decomposition(
     input_tensor: torch.Tensor,
     src_tensor: torch.Tensor,
     dim: int,
-    start: Optional[int],
-    end: Optional[int],
-    step: int,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    step: Optional[int] = None,
 ):
     dim_size = input_tensor.shape[dim]
-    start = get_positive_dim(start, input_tensor.shape)
-    end = get_positive_dim(end, input_tensor.shape)
+    start = get_positive_dim(start, input_tensor.shape[dim])
+    if end is None:
+        end = dim_size
+    end = get_positive_dim(end, input_tensor.shape[dim])
+    if step is None:
+        step = 1
 
     src_dim = src_tensor.shape
     step_dim = (end - start) // step
     end_dim = end
     if step_dim > src_dim[dim]:
         end_dim = src_dim[dim]
     else:
-        indices = torch.arange(0, step_dim)
+        indices = torch.Tensor(torch.arange(0, step_dim))
         indices = indices.to(torch.int32)
-        src = torch.index_select(src, dim, indices)
+        src = torch.index_select(src_tensor, dim, indices)
 
     if start == 0 and end == dim_size and step == 0:
         return input_tensor
@@ -201,8 +205,9 @@ def slice_scatter_decomposition(
     for index in range(start, end_dim, step):
         unbind_input_tensors_list[index] = unbind_source_tensors_list[i]
         i = i + 1
+    output_tensor = torch.stack(tuple(unbind_input_tensors_list), dim)
 
-    return torch.stack(unbind_input_tensors_list, dim)
+    return output_tensor
 
 
 def get_decompositions(

tests/py/dynamo/lowering/test_decompositions.py

@@ -425,22 +425,22 @@ 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)
+            def forward(self, x, src, dim, start, end, step):
+                y = torch.ops.aten.slice_scatter.default(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}
+        unexpected_ops = {torch.ops.aten.slice_scatter}
 
-        inputs = [torch.zeros(8, 8).cuda(), torch.ones(2, 8).cuda(), 0, 6]
+        inputs = [torch.zeros(8, 8).cuda(), torch.ones(2, 8).cuda(), 0, 6, None, 1]
 
         fx_graph = torch.fx.symbolic_trace(sliceScatter())
+
         unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(
             fx_graph,
             inputs,
@@ -490,7 +490,7 @@ 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)
+                y = torch.ops.aten.slice_scatter(x, src, dim, start, end, step)
                 return y
 
         # Operations expected to be removed in the traced graph after decompositions