[acc_ops] Support Tensor.view torch.Shape input correctly (#22)

Summary:
Pull Request resolved: pytorch/fx2trt#22

Alternative to D34911848

Reviewed By: frank-wei

Differential Revision: D34930874

fbshipit-source-id: 0da1081cac09feb188d2624ba04c86d5f3ffef6c

Author: jfix71

test/tracer/test_acc_tracer.py

@@ -1597,6 +1597,7 @@ def test_view(self):
         """
 
         self._make_acc_op_function_test(acc_ops.reshape, lambda x: x.view(1, -1))
+        self._make_acc_op_function_test(acc_ops.reshape, lambda x: x.view([1, -1]))
 
     def test_narrow(self):
         """

tracer/acc_tracer/acc_normalizer.py

@@ -286,7 +286,8 @@ def get_normalized_kwargs(
         orig_kwargs_names, new_kwarg_name, is_optional = replacement_tuple
 
         # Check if this is a varg and if so break/process the rest outside the loop.
-        if len(orig_kwargs_names) == 1 and orig_kwargs_names[0] == "*":
+        if "*" in orig_kwargs_names:
+            assert len(orig_kwargs_names) == 1
             assert i == len(arg_replacement_tuples) - 1
             final_arg_is_varg = True
             break

tracer/acc_tracer/acc_ops.py

@@ -1668,14 +1668,6 @@ def custom_narrow_mapper(node: torch.fx.Node, mod: nn.Module) -> torch.fx.Node:
     ],
     kwargs_to_move_to_acc_out_ty=[("shape", "shape")],
 )
-@register_acc_op_mapping(
-    op_and_target=("call_method", "view"),
-    arg_replacement_tuples=[
-        ("input", "input"),
-        ("*", "shape"),
-    ],
-    kwargs_to_move_to_acc_out_ty=[("shape", "shape")],
-)
 @register_acc_op
 def reshape(*, input, acc_out_ty=None):
     assert acc_out_ty is not None
@@ -1689,11 +1681,18 @@ def reshape(*, input, acc_out_ty=None):
         ("*", "shape"),
     ],
 )
+@register_custom_acc_mapper_fn(
+    op_and_target=("call_method", "view"),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("*", "shape"),
+    ],
+)
 def custom_tensor_reshape_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
     """
-    For Tensor.reshape node, args could be (input, 1, 2, 3) or (input, (1, 2, 3)).
-    Here we do some special handling with the `shape` arg in order to map it to
-    acc_ops.reshape. It also handles the case when `shape` is a list instead of
+    For Tensor.reshape and Tensor.view nodes, args could be (input, 1, 2, 3) or (input,
+    (1, 2, 3)).  Here we do some special handling with the `shape` arg in order to map
+    it to acc_ops.reshape. It also handles the case when `shape` is a list instead of
     tuple.
     """
     input_node = node.kwargs["input"]