Index converter

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -169,6 +169,24 @@ def aten_ops_gelu(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.index.Tensor)
+def aten_ops_index(
+network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.select.gelu(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.matmul)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.mm.default)  # type: ignore[misc]
 def aten_ops_matmul(

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

@@ -1,4 +1,5 @@
-from typing import Optional, cast
+from typing import Optional, cast, Union, Sequence
+import tensorrt as trt
 
 import numpy as np
 from torch.fx.node import Target
@@ -10,6 +11,7 @@
     to_numpy,
 )
 from torch_tensorrt.fx.types import Shape, TRTNetwork, TRTTensor
+from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 
 
 def select(
@@ -61,3 +63,43 @@ def select(
     if len(out.shape) != 1:
         layer = network.add_shuffle(out)
     return layer.get_output(0)
+
+
+def select(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    index: Union[TRTTensor, Sequence[TRTTensor]]
+) -> TRTTensor:
+    adv_indx_indices = []
+    tensor_indices = []
+
+    for i in len(index):
+        ind = index[i]
+        #FIXME: check if the datatype for the indices needs to be casted to INT32
+        #TRTInterpretor should take care
+        adv_indx_indices.append(i)
+        tensor_indices.append(ind)
+
+    if not tensor_indices:
+        identity_layer = network.add_identity(input)
+        identity_layer.set_output_type(0, trt.int32)
+        set_layer_name(identity_layer, target, name + "_index_identity", source_ir)
+        return identity_layer.get_output(0)
+    elif (len(tensor_indices) == 1):
+        indices_tensor = tensor_indices[0]
+        gather_layer = network.add_gather(input, indices_tensor, adv_indx_indices[0])
+        set_layer_name(gather_layer, target, name + "_index_gather", source_ir)
+        return gather_layer.get_output(0)
+    else:
+        input_shape = input.shape
+        rank = len(input_shape)
+        adv_indx_count = len(adv_indx_indices)
+        input_shape_layer = network.add_shape(input)
+        set_layer_name(input_shape_layer, target, name + "_index_shape", source_ir)
+        input_shape_tensor = input_shape_layer.get_output(0)
+        return input_shape_tensor.get_output(0)
+    
+