Correcting the slice_scatter case with aten::scatter use

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py

@@ -174,7 +174,6 @@
     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

@@ -1,7 +1,6 @@
 import logging
 from typing import Any, Callable, Dict, List, Optional
 
-import numpy as np
 import torch
 from torch._decomp import register_decomposition
 from torch._ops import OpOverload
@@ -190,14 +189,21 @@ def slice_scatter_decomposition(
         end_dim = src_dim[dim]
     else:
         # In this case src first step_dim need to be selected
-        indices = torch.Tensor(torch.arange(0, step_dim))
-        indices = indices.to(torch.int32)
+        indices = torch.arange(0, step_dim)
         src = torch.index_select(src_tensor, dim, indices)
 
     if start == 0 and end == dim_size and step == 0:
         return input_tensor
 
-    index_tensor = torch.arange(start, end, step_dim)
+    cat_tensors = []
+    index_tensor_shape = []
+    for i, src_each_dim in enumerate(list(src_dim)):
+        if i != dim:
+            index_tensor_shape.append(src_each_dim)
+    for index in range(start, end, step):
+        cat_tensors.append(index * torch.ones(index_tensor_shape))
+    index_tensor = torch.stack(cat_tensors, dim)
+    index_tensor = index_tensor.to(torch.int64).cuda()
     output_tensor = torch.scatter(input_tensor, dim, index_tensor, src)
     return output_tensor
 

tests/py/dynamo/lowering/test_decompositions.py

@@ -431,9 +431,8 @@ def forward(self, x, src, dim, start, end, step):
 
         # 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,
+            torch.ops.aten.scatter.src,
         }
         unexpected_ops = {torch.ops.aten.slice_scatter}
 
@@ -469,6 +468,7 @@ def forward(self, x, src, dim, start, end, step):
             "torch_compile",
             inputs,
             min_block_size=1,
+            truncate_long_and_double=True,
             pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()
@@ -495,14 +495,12 @@ def forward(self, x, src, dim, start=None, end=None, step=1):
 
         # 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,
+            torch.ops.aten.scatter.src,
         }
         unexpected_ops = {torch.ops.aten.select_scatter}
 
-        inputs = [torch.zeros(8, 8).cuda(), torch.ones(2, 8).cuda(), 0, 6]
+        inputs = [torch.zeros(8, 8).cuda(), torch.ones(8, 2).cuda(), 1, 6, None, 1]
 
         fx_graph = torch.fx.symbolic_trace(sliceScatter())
         unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(
@@ -533,6 +531,7 @@ def forward(self, x, src, dim, start=None, end=None, step=1):
             "torch_compile",
             inputs,
             min_block_size=1,
+            truncate_long_and_double=True,
             pass_through_build_failures=True,
         )
         optimized_model_results = optimized_model(*inputs).detach().cpu()