gather converter

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -751,6 +751,25 @@ def aten_ops_clamp(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.gather.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+        2: (TRTTensor,),
+    }
+)
+def aten_ops_gather(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.select.gather(
+        ctx, target, SourceIR.ATEN, name, args[0], args[1], args[2]
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.scatter.src)
 @dynamo_tensorrt_converter(torch.ops.aten.scatter.value)
 @enforce_tensor_types(

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -387,7 +387,7 @@ def index_select(
     dim = get_positive_dim(dim, len(input.shape))
     gather_layer = ctx.net.add_gather(input, index, axis=dim)
 
-    set_layer_name(gather_layer, target, f"{name}_gather", source_ir)
+    set_layer_name(gather_layer, target, f"{name}_gather_layer_default", source_ir)
 
     return gather_layer.get_output(0)
 
@@ -428,3 +428,21 @@ def scatter(
     set_layer_name(scatter_layer, target, name + "_scatter_layer", source_ir)
     out = scatter_layer.get_output(0)
     return out
+
+
+def gather(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dim: int,
+    index: Union[TRTTensor, np.ndarray, torch.Tensor],
+) -> TRTTensor:
+    input_shape = input.shape
+    dim = get_positive_dim(dim, len(input_shape))
+    gather_layer = ctx.net.add_gather(input, index, axis=dim)
+    gather_layer.mode = trt.GatherMode.ELEMENT
+    set_layer_name(gather_layer, target, name + "_gather_layer_element", source_ir)
+    out = gather_layer.get_output(0)
+    return out

tests/py/dynamo/conversion/test_gather_aten.py

@@ -0,0 +1,72 @@
+import torch
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+from .harness import DispatchTestCase
+
+
+class TestGatherValueConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            (
+                "gather_zero_dim_indexOne_constant_value",
+                0,
+                torch.tensor([[0, 1, 2, 0]]),
+            ),
+            (
+                "gather_zero_dim_indexTwo_constant_value",
+                0,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+            ),
+            (
+                "gather_one_dim_indexOne_constant_value",
+                1,
+                torch.tensor([[0, 1, 2, 0]]),
+            ),
+            (
+                "gather_one_dim_indexTwo_costant_value",
+                1,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+            ),
+        ]
+    )
+    def test_gather_index_constant(self, _, dim, index):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input):
+                return torch.ops.aten.gather.default(input, dim, index)
+
+        input = torch.zeros(3, 5, dtype=torch.int32)
+        inputs = [input]
+        self.run_test(TestModule(), inputs)
+
+    @parameterized.expand(
+        [
+            ("gather_zero_dim_indexOne_value", 0, torch.tensor([[0, 1, 2, 0]])),
+            (
+                "gather_zero_dim_indexTwo_value",
+                0,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+            ),
+            ("gather_one_dim_indexOne_value", 1, torch.tensor([[0, 1, 2, 0]])),
+            (
+                "gather_one_dim_indexTwo_value",
+                1,
+                torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]),
+            ),
+        ]
+    )
+    def test_gather_index_input(self, _, dim, index):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, input, index):
+                return torch.ops.aten.gather.default(input, dim, index)
+
+        input = torch.zeros(3, 5, dtype=torch.int32)
+        inputs = [input, index]
+        self.run_test(TestModule(), inputs)