[fx2trt] fix setitem (#69)

Summary:
Pull Request resolved: https://github.com/pytorch/fx2trt/pull/69

fix the pass for corner cases：
1. position index is not a slice type but a int in setitem. Change it to a slice type. For ex,  [:, 1] - > [:, 1:2]
2. Similarly, for getitem (input for this setitem) that has integer index like [:, 1], the output will remove the 2nd dimension. It will bring trouble for setitem as we are lack of shape info. I will replace with a slice like [:, 1:2]

Reviewed By: yinghai

Differential Revision: D36136098

fbshipit-source-id: c48f976fe639ff3fa159b44c8c2287f202898901

Author: Wei Wei

fx/passes/lower_basic_pass.py

@@ -1,3 +1,5 @@
+import copy
+import operator
 import operator
 import warnings
 from typing import Any
@@ -348,13 +350,32 @@ def transform_setitem(gm: torch.fx.GraphModule, input: Input):
     """
     map_replace = {}
     for node in gm.graph.nodes:
+        for old_node in map_replace:
+            node.replace_input_with(old_node, map_replace[old_node])
+
         if node.target == operator.setitem:
             input_node = node.args[0]
             sli = node.args[1]
             inp = node.args[2]
 
+            inp_flag = False
+            if inp.target == operator.getitem:
+                new_args = list(copy.deepcopy(inp.args[1]))
+                for ind, val in enumerate(new_args):
+                    if type(val) == int:
+                        inp_flag = True
+                        new_args[ind] = slice(val, val + 1, None)
+                if inp_flag:
+                    with gm.graph.inserting_before(inp):
+                        new_node = gm.graph.call_function(
+                            operator.getitem, args=(inp.args[0], new_args)
+                        )
+                        inp.replace_all_uses_with(new_node)
+                    inp = new_node
+
             if type(sli) is not tuple:
                 sli = [sli]
+            sli = [slice(x, x + 1, None) if type(x) == int else x for x in sli]
             dimension = len(sli)
             with gm.graph.inserting_before(node):
                 if dimension == 1:
@@ -418,9 +439,6 @@ def transform_setitem(gm: torch.fx.GraphModule, input: Input):
                     continue
                 node.replace_all_uses_with(concat_node_0)
                 map_replace[input_node] = concat_node_0
-        else:
-            for old_node in map_replace:
-                node.replace_input_with(old_node, map_replace[old_node])
 
     gm.graph.eliminate_dead_code()
     gm.graph.lint()

test/passes/test_setitem.py

@@ -32,6 +32,55 @@ def transform_fx(gm, example_inputs):
         with optimize_ctx:
             m(*inputs)
 
+    def test_setitem1d_c2(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, x, y):
+                y[:-1] = x
+                y[1:] = x
+                return y
+
+        inputs = [torch.randn(2), torch.randn(3)]
+        m = TestModule()
+
+        inputs = [i.cuda() for i in inputs]
+        m.cuda()
+
+        def transform_fx(gm, example_inputs):
+            gm = transform_setitem(gm, example_inputs)
+            return gm
+
+        optimize_ctx = torchdynamo.optimize(
+            transform_fx,
+            nopython=True,
+        )
+
+        with optimize_ctx:
+            m(*inputs)
+
+    def test_setitem1d_c3(self):
+        class TestModule(torch.nn.Module):
+            def forward(self, x, y):
+                y[1] = x
+                return y
+
+        inputs = [torch.randn(2), torch.randn(3)]
+        m = TestModule()
+
+        inputs = [i.cuda() for i in inputs]
+        m.cuda()
+
+        def transform_fx(gm, example_inputs):
+            gm = transform_setitem(gm, example_inputs)
+            return gm
+
+        optimize_ctx = torchdynamo.optimize(
+            transform_fx,
+            nopython=True,
+        )
+
+        with optimize_ctx:
+            m(*inputs)
+
     @parameterized.expand(
         [
             ("c1", (4, 2), (4, 5), 0, 2),
@@ -96,6 +145,37 @@ def transform_fx(gm, example_inputs):
         with optimize_ctx:
             m(*inputs)
 
+    @parameterized.expand(
+        [
+            ("c1", (4, 2), (4, 2), 0, 1),
+        ]
+    )
+    def test_setitem2d_1v_ex2(self, name, x_shape, y_shape, y_start, y_end):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x, y):
+                y[:, y_start:y_end] = x[:, 0]
+                return y
+
+        inputs = [torch.randn(x_shape), torch.randn(y_shape)]
+        m = TestModule()
+
+        inputs = [i.cuda() for i in inputs]
+        m.cuda()
+
+        def transform_fx(gm, example_inputs):
+            gm = transform_setitem(gm, example_inputs)
+            return gm
+
+        optimize_ctx = torchdynamo.optimize(
+            transform_fx,
+            nopython=True,
+        )
+        with optimize_ctx:
+            m(*inputs)
+
     @parameterized.expand(
         [
             ("c1", (3, 2), (4, 5), 0, 3, 0, 2),
@@ -429,17 +509,17 @@ def transform_fx(gm, example_inputs):
         with optimize_ctx:
             m(*inputs)
 
-    ## test with torchdynamo
+    # test with torchdynamo
     def test_setitem1d_trt(self):
         class TestModule(torch.nn.Module):
             def __init__(self):
                 super().__init__()
 
             def forward(self, x, y):
-                y[0:2] = x
+                y[1] = x
                 return y
 
-        inputs = [torch.randn(2), torch.randn(3)]
+        inputs = [torch.randn(1), torch.randn(3)]
         m = TestModule()
 
         inputs = [i.cuda() for i in inputs]