[fx2trt] matmul, softmax, expand (#28)

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

as titled

Reviewed By: wushirong

Differential Revision: D34997650

fbshipit-source-id: 62ce11b4ca0605f78b9022cb1271582e049f2327

Author: Wei Wei

fx/converters/acc_ops_converters.py

@@ -1654,6 +1654,36 @@ def acc_ops_slice_tensor(
     return layer.get_output(0)
 
 
+@tensorrt_converter(acc_ops.expand)
+def acc_ops_expand_tensor(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    input_t = kwargs["input"]
+    shape = kwargs["sizes"].copy()
+
+    input_val = get_trt_tensor(network, input_t, f"{name}_input")
+
+    if network.has_implicit_batch_dimension:
+        shape = shape[1:]
+
+    ranks = len(input_val.shape)
+    # TRT does not support different dimension size
+    assert len(shape) == ranks
+    shape = [input_val.shape[i] if shape[i] == -1 else shape[i] for i in range(ranks)]
+
+    inshape = tuple(input_val.shape)
+    shape = tuple(shape)
+    start = tuple([0]*ranks)
+    stride = tuple([int(i == o) for i, o in zip(inshape, shape)])  # stride == 1 if dimensions match, 0 otherwise
+    layer = network.add_slice(input_val, start=start, shape=shape, stride=stride)
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)
+
+
 @tensorrt_converter(acc_ops.split, no_explicit_batch_dim=True)
 def acc_ops_split(
     network: TRTNetwork,

test/converters/acc_op/test_expand.py

@@ -0,0 +1,33 @@
+# Owner(s): ["oncall: aiacc"]
+
+import torch
+import fx2trt_oss.tracer.acc_tracer.acc_ops as acc_ops
+import torch.nn as nn
+from torch.testing._internal.common_fx2trt import AccTestCase
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+
+class TestExpandConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            ("2d_dim", (2, 3), (2, 1)),
+            ("3d_dim", (2, 3, 4), (2, 1, 1)),
+            ("4d_dim", (2, 3, 4, 5), (2, 1, 1, 1)),
+            ("keep_dim", (2, 3, -1, -1), (2, 1, 5, 5)),
+        ]
+    )
+    def test_expand(self, _, sizes, init_size):
+        class Expand(nn.Module):
+            def forward(self, x):
+                return x.expand(*sizes)
+
+        inputs = [torch.randn(*init_size)]
+        self.run_test(
+            Expand(),
+            inputs,
+            expected_ops={acc_ops.expand},
+        )
+
+if __name__ == '__main__':
+    run_tests()

test/tracer/test_acc_tracer.py

@@ -2153,5 +2153,6 @@ def test_all_acc_ops_registered(self):
                 acc_ops.rescale_quantize_per_tensor,
                 acc_ops.rescale_quantize_per_channel,
                 acc_ops.nan_to_num,
+                acc_ops.expand,
             },
         )

tracer/acc_tracer/acc_ops.py

@@ -354,9 +354,24 @@ def contiguous(*, input):
 
 
 @register_acc_op_properties(AccOpProperty.unary)
-@register_acc_op_mapping(op_and_target=("call_function", torch.nn.functional.softmax))
+@register_acc_op_mapping(
+    op_and_target=("call_method", "softmax"),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("dim", "dim"),
+        ("dtype", "dtype", this_arg_is_optional),
+    ],
+)
+@register_acc_op_mapping(
+    op_and_target=("call_function", torch.nn.functional.softmax),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("dim", "dim"),
+        ("dtype", "dtype", this_arg_is_optional),
+    ],
+    )
 @register_acc_op
-def softmax(*, input, dim, dtype):
+def softmax(*, input, dim, dtype=None):
     """
     _stacklevel are ignored here.
     """
@@ -471,7 +486,13 @@ def square_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
         new_node.meta = node.meta.copy()
         return new_node
 
-
+@register_acc_op_mapping(
+    op_and_target=("call_function", operator.matmul),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("mat2", "other"),
+    ],
+)
 @register_acc_op_mapping(
     op_and_target=("call_function", torch.bmm),
     arg_replacement_tuples=[
@@ -1614,6 +1635,19 @@ def nan_to_num(*, input, nan=0.0, posinf=None, neginf=None):
     return torch.nan_to_num(input, nan=nan, posinf=posinf, neginf=neginf)
 
 
+@register_acc_op_properties(AccOpProperty.unary)
+@register_acc_op_mapping(
+    op_and_target=("call_method", "expand"),
+     arg_replacement_tuples=[
+        ("input", "input"),
+        ("*", "sizes"),
+    ],
+)
+@register_acc_op
+def expand(*, input, sizes):
+    return input.expand(*sizes)
+
+
 @register_acc_op_properties(AccOpProperty.unary)
 @register_acc_op
 def slice_tensor(*, input, dim, start, stop, step):