[fx2trt][bootcamp] Add support for torch.nn.functional.embedding (#27)

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

Follows the instruction in T93104604 to add support for torch.nn.functional.embedding in fx2trt.

Reviewed By: frank-wei

Differential Revision: D34945232

fbshipit-source-id: c0ed25e2b7585bfb11736be7dad60ddbbb065050

Author: Zhengping Zhou

fx/converters/acc_ops_converters.py

@@ -1101,6 +1101,49 @@ def acc_ops_fmod(
     )
     return sub_value
 
+
+@tensorrt_converter(acc_ops.embedding, no_implicit_batch_dim=True)
+def acc_ops_embedding(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    if network.has_implicit_batch_dimension:
+        raise RuntimeError("The `embedding` function should be called with explicit batch dimension.")
+
+    indices_tensor = kwargs["input"]
+    embedding_tensor = kwargs["weight"]
+
+    # unsupported parameters
+    padding_idx = kwargs["padding_idx"]
+    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}.")
+
+    if norm_type != 2.0:
+        raise RuntimeError(f"Currently we don't support specifying max_norm, got {norm_type} for norm_type.")
+
+    if scale_grad_by_freq:
+        raise RuntimeError("Currently we don't support scale gradient by word frequency.")
+
+    if sparse:
+        raise RuntimeError("Currently we don't support sparse gradient.")
+
+    # Implement embedding lookup with gather layer
+    gather_layer = network.add_gather(embedding_tensor, indices_tensor, axis=0)
+    set_layer_name(gather_layer, target, name + "_gather")
+    return gather_layer.get_output(0)
+
+
 @tensorrt_converter(acc_ops.max_pool1d, no_explicit_batch_dim=True)
 def acc_ops_max_pool1d(
     network: TRTNetwork,

test/converters/acc_op/test_embedding.py

@@ -0,0 +1,63 @@
+# Owner(s): ["oncall: aiacc"]
+
+import fx2trt_oss.tracer.acc_tracer.acc_ops as acc_ops
+import torch
+from parameterized import parameterized, param
+from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
+
+
+class TestEmbeddingConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            param(
+                test_name="1d_indices",
+                indices_tensor=torch.tensor([3, 1, 2]),
+                weights_tensor=torch.randn(5, 10),
+            ),
+            param(
+                test_name="2d_indices",
+                indices_tensor=torch.tensor([[3, 1, 2], [4, 1, 3]]),
+                weights_tensor=torch.randn(5, 10),
+            ),
+            param(
+                test_name="3d_indices",
+                indices_tensor=torch.tensor([[[0, 1], [2, 3]], [[3, 4], [4, 0]]]),
+                weights_tensor=torch.randn(5, 10),
+            ),
+        ]
+    )
+    def test_embedding(
+        self,
+        test_name,
+        indices_tensor,
+        weights_tensor,
+        padding_idx=None,
+        max_norm=None,
+        norm_type=2.0,
+        scale_grad_by_freq=False,
+        sparse=False,
+    ):
+        class TestEmbedding(torch.nn.Module):
+            def forward(self, indices, weights):
+                return torch.nn.functional.embedding(
+                    input=indices,
+                    weight=weights,
+                    padding_idx=padding_idx,
+                    max_norm=max_norm,
+                    norm_type=norm_type,
+                    scale_grad_by_freq=scale_grad_by_freq,
+                    sparse=sparse,
+                )
+
+        self.run_test(
+            TestEmbedding(),
+            inputs=[indices_tensor.int(), weights_tensor.float()],
+            expected_ops={acc_ops.embedding},
+            test_implicit_batch_dim=False,
+            test_explicit_batch_dim=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

test/tracer/test_acc_tracer.py

@@ -2063,6 +2063,7 @@ def test_all_acc_ops_registered(self):
             acc_normalizer._acc_ops,
             {
                 acc_ops.linear,
+                acc_ops.embedding,
                 acc_ops.max_pool1d,
                 acc_ops.max_pool2d,
                 acc_ops.flatten,

tracer/acc_tracer/acc_ops.py

@@ -80,6 +80,14 @@ def squeeze(*, input, dim=None):
     return input.squeeze(dim=dim)
 
 
+@register_acc_op_mapping(op_and_target=("call_function", nn.functional.embedding))
+@register_acc_op
+def embedding(
+    *, input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse
+):
+    return torch.nn.functional.embedding(**locals())
+
+
 @register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool1d))
 @register_acc_op
 def max_pool1d(
@@ -95,6 +103,7 @@ def max_pool1d(
         return_indices=return_indices,
     )
 
+
 @register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool2d))
 @register_acc_op
 def max_pool2d(
@@ -118,24 +127,17 @@ def max_pool2d(
 def adaptive_avg_pool2d(*, input, output_size):
     return nn.functional.adaptive_avg_pool2d(input=input, output_size=output_size)
 
+
 @register_acc_op_mapping(op_and_target=("call_function", nn.functional.avg_pool1d))
 @register_acc_op
-def avg_pool1d(
-    *,
-    input,
-    kernel_size,
-    stride,
-    padding,
-    ceil_mode,
-    count_include_pad
-):
+def avg_pool1d(*, input, kernel_size, stride, padding, ceil_mode, count_include_pad):
     return nn.functional.avg_pool1d(
         input=input,
         kernel_size=kernel_size,
         stride=stride,
         padding=padding,
         ceil_mode=ceil_mode,
-        count_include_pad=count_include_pad
+        count_include_pad=count_include_pad,
     )
 
 
@@ -499,6 +501,7 @@ def dropout_mapper(node: torch.fx.Node, mod: nn.Module):
 
 try:
     from torchvision.ops import stochastic_depth
+
     assert callable(stochastic_depth)
 except Exception as e:
     warnings.warn(f"Unable to import torchvision related libraries.: {e}")
@@ -903,6 +906,7 @@ def prod(*, input, dim=None, keepdim=False, dtype=None):
     else:
         return input.prod(dtype=dtype)
 
+
 @register_custom_acc_mapper_fn(
     op_and_target=("call_method", "prod"),
     arg_replacement_tuples=[