[fx2trt] support masked_fill, repeat (#29)

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

BERT_pytorch model in torchbench:
1. masked_fill, repeat
2. fix a nit in embedding op

Reviewed By: yinghai, wushirong

Differential Revision: D35034664

fbshipit-source-id: 2c0e66df5a17f6960d71c3bc6e9bdca9a5daf3e5

Author: Wei Wei

fx/converters/acc_ops_converters.py

@@ -1117,15 +1117,12 @@ def acc_ops_embedding(
     embedding_tensor = kwargs["weight"]
 
     # unsupported parameters
-    padding_idx = kwargs["padding_idx"]
+    # ignore padding_idx since it is meaningful for training only
     max_norm = kwargs["max_norm"]
     norm_type = kwargs["norm_type"]
     scale_grad_by_freq = kwargs["scale_grad_by_freq"]
     sparse = kwargs["sparse"]
 
-    if padding_idx is not None:
-        raise RuntimeError(f"Currently we don't support specifying padding_idx, got {padding_idx}.")
-
     if max_norm is not None:
         raise RuntimeError(f"Currently we don't support specifying max_norm, got {max_norm}.")
 
@@ -1684,6 +1681,54 @@ def acc_ops_expand_tensor(
     return layer.get_output(0)
 
 
+@tensorrt_converter(acc_ops.masked_fill, no_implicit_batch_dim=True)
+def acc_ops_masked_fill_tensor(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    input_t = kwargs["input"]
+    mask_t = kwargs["mask"]
+    value_t = kwargs["value"]
+    if network.has_implicit_batch_dimension:
+        raise RuntimeError("We don't support masked_fill with implicit batch dimension due to select layer!")
+
+    shape = list(input_t.shape)
+    mask_shape = list(mask_t.shape)
+
+    assert type(value_t) in (float, int, torch.Tensor), f"value {value_t} is not one of (float, int, torch.Tensor)!"
+
+    if type(mask_t) != TRTTensor:
+        assert mask_t.dtype == torch.bool, "mask dtype is not bool!"
+        if mask_shape != shape:
+            mask_t = mask_t.expand(shape)
+        mask_t = mask_t.to(torch.int32)
+        mask_const = get_trt_tensor(network, mask_t, f"{name}_mask")
+        mask_layer = network.add_identity(mask_const)
+        mask_layer.set_output_type(0, trt.bool)
+        set_layer_name(mask_layer, target, f"{name}_mask")
+        mask_val = mask_layer.get_output(0)
+    else:
+        assert mask_t.dtype == trt.bool, "mask dtype is not bool!"
+        if mask_shape != shape:
+            mask_val = acc_ops_expand_tensor(network, target, None, {"input": mask_t, "sizes": shape}, name=f"{name}_expand")
+        else:
+            mask_val = mask_t
+
+    if type(value_t) is torch.Tensor:
+        value_t = value_t.cpu().numpy()
+    value_t = float(value_t)
+    value_t = torch.ones(shape)*value_t
+
+    input_val = get_trt_tensor(network, input_t, f"{name}_input")
+    value_val = get_trt_tensor(network, value_t, f"{name}_input")
+    layer = network.add_select(mask_val, value_val, input_val)
+    set_layer_name(layer, target, f"{name}_select")
+    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_masked_fill.py

@@ -0,0 +1,66 @@
+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 TestMaskedFill(AccTestCase):
+    @parameterized.expand(
+        [
+            ("same_dims", (2, 3), 5),
+            ("same_dims_tensor", (2, 3), torch.tensor(5)),
+            ("not_same_dims", (2, 1), 5),
+            ("not_same_dims_tensor", (2, 1), torch.tensor(5)),
+        ]
+    )
+    def test_masked_fill(self, _, input_shape, value):
+        class MaskedFill(nn.Module):
+            def __init__(self, input_shape):
+                super().__init__()
+                self.mask = torch.zeros(input_shape)
+                self.mask[0,0] = 1
+                self.mask = self.mask.to(torch.bool)
+                self.value = value
+            def forward(self, x):
+                return x.masked_fill(self.mask, self.value)
+
+        inputs = [torch.ones(*input_shape)]
+        self.run_test(
+            MaskedFill(input_shape),
+            inputs,
+            expected_ops={acc_ops.masked_fill},
+            test_implicit_batch_dim = False
+        )
+
+    @parameterized.expand(
+        [
+            ("same_dims", (2, 3), (2,3), 5),
+             ("expand_first_dims", (2, 3), (1,3), 5),
+             ("expand_second_dims", (2, 3), (2,1), 5),
+             ("expand_third_dims", (2, 3, 4), (2, 3, 1), 5),
+        ]
+    )
+    def test_masked_fill_expand(self, _, input_shape, mask_shape, value):
+        class MaskedFill(nn.Module):
+            def __init__(self, input_shape):
+                super().__init__()
+                self.value = value
+            def forward(self, x, mask_input):
+                return x.masked_fill(mask_input, self.value)
+
+        mask_input = torch.zeros(*mask_shape)
+        index = (0)*len(mask_shape)
+        mask_input[index] = 1
+        mask_input = mask_input.to(torch.bool)
+        inputs = [torch.ones(*input_shape), mask_input]
+        self.run_test(
+            MaskedFill(input_shape),
+            inputs,
+            expected_ops={acc_ops.masked_fill},
+            test_implicit_batch_dim = False
+        )
+
+if __name__ == '__main__':
+    run_tests()

test/tracer/test_acc_tracer.py

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

tracer/acc_tracer/acc_ops.py

@@ -256,6 +256,29 @@ def unsqueeze(*, input, dim):
 def tile(*, input, dims):
     return torch.tile(input=input, dims=dims)
 
+@register_custom_acc_mapper_fn(
+    op_and_target=("call_method", "repeat"),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("*", "sizes"),
+    ],
+)
+def repeat_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
+    """
+    Map repeat to tile.
+    """
+    with node.graph.inserting_before(node):
+        inputs = node.kwargs["input"]
+        dims = node.kwargs["sizes"][0]
+        new_node = node.graph.create_node(
+            "call_function",
+            tile,
+            kwargs={"input": inputs, "dims": dims},
+            name=f"{node.name}_repeat_map",
+        )
+        new_node.meta = node.meta.copy()
+        return new_node
+
 
 @register_custom_acc_mapper_fn(
     op_and_target=("call_function", torch.stack),
@@ -1648,6 +1671,18 @@ def expand(*, input, sizes):
     return input.expand(*sizes)
 
 
+@register_acc_op_mapping(
+    op_and_target=("call_method", "masked_fill"),
+     arg_replacement_tuples=[
+        ("input", "input"),
+        ("mask", "mask"),
+        ("value", "value"),
+    ],
+)
+@register_acc_op
+def masked_fill(*, input, mask, value):
+    return input.masked_fill(mask, value)
+
 @register_acc_op_properties(AccOpProperty.unary)
 @register_acc_op
 def slice_tensor(*, input, dim, start, stop, step):