Embedding operator in dynamo

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -4,7 +4,7 @@
 import tensorrt as trt
 from torch_tensorrt.fx.converters import acc_ops_converters
 from ..converter_registry import dynamo_tensorrt_converter
-from torch.fx.node import Argument, Target
+from torch.fx.node import Argument, Target, Node
 
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 from torch_tensorrt.dynamo.conversion import SourceIR, impl
@@ -71,6 +71,62 @@ def aten_ops_div(
         )
 
 
+def embedding_param_validator(embedding_node: Node):
+
+    max_norm = args_bounds_check(embedding_node.args, 2)
+    norm_type = args_bounds_check(embedding_node.args, 3)
+    scale_grad_by_freq = args_bounds_check(embedding_node.args, 4)
+    sparse = args_bounds_check(embedding_node.args, 5)
+
+    if max_norm is not None:
+        _LOGGER.debug(
+            f"Currently we don't support specifying max_norm, got {max_norm}."
+        )
+        return False
+
+    if norm_type is not None and norm_type != 2.0:
+        _LOGGER.debug(
+            f"Currently we don't support specifying norm_type, got {norm_type}."
+        )
+        return False
+
+    if scale_grad_by_freq is not None:
+        _LOGGER.debug(
+            f"Currently we don't support specifying scale gradient by word frequency, got {scale_grad_by_freq}."
+        )
+        return False
+
+    if sparse is not None:
+        _LOGGER.debug(f"Currently we don't support sparse gradient, got {sparse}.")
+        return False
+
+    return True
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.embedding.default, capability_validator=embedding_param_validator
+)
+def aten_ops_embedding(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.embedding.embedding(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[1],
+        weight=args[0],
+        max_norm=or_none(args, 2),
+        norm_type=or_none(args, 3),
+        scale_grad_by_freq=or_none(args, 4),
+        sparse=or_none(args, 5),
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.fmod.Scalar)
 @dynamo_tensorrt_converter(torch.ops.aten.fmod.Tensor)
 def aten_ops_fmod(

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -1,6 +1,7 @@
 from torch_tensorrt.fx.converters.impl import convolution
 from . import condition
 from . import elementwise
+from . import embedding
 from . import normalization
 from . import slice
 from . import unary

py/torch_tensorrt/dynamo/conversion/impl/embedding.py

@@ -0,0 +1,73 @@
+import operator
+import warnings
+from typing import Optional, cast, Any
+
+import numpy as np
+
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    set_layer_name,
+)
+
+from torch_tensorrt.fx.converters.converter_utils import get_trt_tensor
+
+
+def embedding(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    weight: TRTTensor,
+    max_norm: None,
+    norm_type: None,
+    scale_grad_by_freq: bool,
+    sparse: bool,
+) -> TRTTensor:
+
+    if network.has_implicit_batch_dimension:
+        raise RuntimeError(
+            "The `embedding` function should be called with explicit batch dimension."
+        )
+
+    indices_tensor = input
+    embedding_tensor = weight
+    if isinstance(indices_tensor, torch.Tensor) and indices_tensor.dtype == torch.int64:
+        raise RuntimeError(
+            "The `embedding` op has indices_tensor dtype=int64. This is incorrect since it has to be int32 to run on TRT."
+        )
+    indices_tensor = get_trt_tensor(network, indices_tensor, f"{name}_indices_tensor")
+    embedding_tensor = get_trt_tensor(
+        network, embedding_tensor, f"{name}_embedding_tensor"
+    )
+    # unsupported parameters
+    # ignore padding_idx since it is meaningful for training only
+
+    if max_norm is not None:
+        raise RuntimeError(
+            f"Currently we don't support specifying max_norm, got {max_norm}."
+        )
+
+    if norm_type is not None and 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)

tests/py/dynamo/converters/test_embedding_aten.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.dynamo.test_utils import DispatchTestCase
+from parameterized import param, parameterized
+from torch_tensorrt import Input
+
+
+class TestEmbeddingConverter(DispatchTestCase):
+    @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={torch.ops.aten.embedding.default},
+        )
+
+    def test_embedding_with_dynamic_shape_four_dimensions(
+        self,
+        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, input, weights):
+                return torch.nn.functional.embedding(
+                    input=input,
+                    weight=weights,
+                    padding_idx=padding_idx,
+                    max_norm=max_norm,
+                    norm_type=norm_type,
+                    scale_grad_by_freq=scale_grad_by_freq,
+                    sparse=sparse,
+                )
+
+        input_specs = [
+            Input(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 1), (2, 3, 4, 5), (2, 3, 10, 10))],
+            ),
+            Input(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 1), (2, 3, 4, 5), (2, 3, 10, 10))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestEmbedding(),
+            input_specs,
+            expected_ops={torch.ops.aten.embedding.default},
+        )
+
+
+if __name__ == "__main__":
+    run_tests()