aten converter- matmul, tanh, gelu, slice select

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -1980,53 +1980,8 @@ def acc_ops_slice_tensor(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"slice_tensor received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    ranks = len(input_val.shape) + (1 if network.has_implicit_batch_dimension else 0)
-    dim = get_positive_dim(cast(int, kwargs["dim"]), ranks)
-    dynamic_shape = has_dynamic_shape(input_val.shape)
-    if network.has_implicit_batch_dimension:
-        if dim == 0:
-            raise RuntimeError(
-                f"We do not support slice_tensor at batch dim when it's implicit, got {dim}!"
-            )
-        dim = dim - 1
-    else:
-        if dynamic_shape:
-            # Check whether slice target dim is dynamic shape dim
-            assert input_val.shape[dim] != -1, "Can't chunk on dynamic shape dimension!"
-
-    start_int = cast(int, kwargs["start"])
-    stop_int = cast(int, kwargs["stop"])
-    step_int = cast(int, kwargs["step"])
-    start = [0] * len(input_val.shape)
-    start[dim] = start_int
-    stride = [1] * len(start)
-    stride[dim] = step_int
-    output_shape = list(input_val.shape)
-    output_shape[dim] = (stop_int - start_int) // step_int
-
-    if dynamic_shape > 0:
-        output_shape = get_shape_with_dynamic_shape(
-            network, output_shape, input_val, target, name
-        )
-    layer = network.add_slice(
-        input_val,
-        start=start,
-        shape=[] if dynamic_shape else output_shape,
-        stride=stride,
-    )
-    if dynamic_shape:
-        layer.set_input(2, output_shape)
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)
-
+    return operator.add_slice(network, target, kwargs, name)
+    
 
 @tensorrt_converter(acc_ops.expand)
 def acc_ops_expand_tensor(
@@ -2036,29 +1991,8 @@ def acc_ops_expand_tensor(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_t = kwargs["input"]
-    shape = list(kwargs["sizes"])
-
-    input_val = get_trt_tensor(network, input_t, f"{name}_input")
-
-    if network.has_implicit_batch_dimension:
-        shape = shape[1:]
-
-    ranks = len(input_val.shape)
-    # TRT does not support different dimension size
-    assert len(shape) == ranks
-    shape = [input_val.shape[i] if shape[i] == -1 else shape[i] for i in range(ranks)]
-
-    inshape = tuple(input_val.shape)
-    shape = tuple(shape)
-    start = tuple([0] * ranks)
-    stride = tuple(
-        [int(i == o) for i, o in zip(inshape, shape)]
-    )  # stride == 1 if dimensions match, 0 otherwise
-    layer = network.add_slice(input_val, start=start, shape=shape, stride=stride)
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)
-
+    return operator.add_expand(network, target, kwargs, name)
+    
 
 @tensorrt_converter(acc_ops.where)
 def acc_ops_where(
@@ -2754,34 +2688,8 @@ def acc_ops_gelu(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-    approximate = kwargs["approximate"]
-    if approximate != "none":
-        raise RuntimeError("GeLU converter currently doesn't support fast gelu compute")
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"GELU received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-    if network.has_implicit_batch_dimension:
-        raise RuntimeError(
-            "GeLU converter currently doesn't support implicit batch dimension"
-        )
-
-    plugin_name = "CustomGeluPluginDynamic"
-    # type_id 0 for float32, 1 for  float16
-    type_id = trt.PluginField(
-        "type_id", np.array(0, dtype=np.int32), trt.PluginFieldType.INT32
-    )
-    field_collection = TRTPluginFieldCollection([type_id])
-    plugin_version = "1"
-
-    plugin = get_trt_plugin(plugin_name, field_collection, plugin_version)
-
-    layer = network.add_plugin_v2([input_val], plugin)
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)
-
+    return activation.add_gelu(network, target, kwargs, name)
+    
 
 @tensorrt_converter(acc_ops.chunk)
 def acc_ops_chunk(

py/torch_tensorrt/fx/converters/activation.py

@@ -12,6 +12,7 @@
 
 from .converter_utils import mark_as_int8_layer
 from .converter_utils import set_layer_name
+from .converter_utils import get_trt_plugin
 
 from ..types import (
     Shape,
@@ -106,6 +107,35 @@ def add_tanh(network, target, kwargs, name):
     operation_type = trt.ActivationType.TANH
     return add_activation_layer(network, input_val, operation_type, target, name)
 
+def add_gelu(network, target, kwargs, name):
+    input_val = kwargs["input"]
+    approximate = kwargs["approximate"]
+    if approximate != "none":
+        raise RuntimeError("GeLU converter currently doesn't support fast gelu compute")
+    if not isinstance(input_val, TRTTensor):
+        raise RuntimeError(
+            f"GELU received input {input_val} that is not part "
+            "of the TensorRT region!"
+        )
+    if network.has_implicit_batch_dimension:
+        raise RuntimeError(
+            "GeLU converter currently doesn't support implicit batch dimension"
+        )
+
+    plugin_name = "CustomGeluPluginDynamic"
+    # type_id 0 for float32, 1 for  float16
+    type_id = trt.PluginField(
+        "type_id", np.array(0, dtype=np.int32), trt.PluginFieldType.INT32
+    )
+    field_collection = TRTPluginFieldCollection([type_id])
+    plugin_version = "1"
+
+    plugin = get_trt_plugin(plugin_name, field_collection, plugin_version)
+
+    layer = network.add_plugin_v2([input_val], plugin)
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)
+
 def add_hard_sigmoid(network, target, kwargs, name):
     input_val = kwargs["input"]
 

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -254,19 +254,6 @@ def aten_ops_mul(
     }
     return operator.add_mul(network, target, None, kwargs_new, name)
 
-@tensorrt_converter(torch.ops.aten.mul.Tensor)
-def aten_ops_mul(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-        "other": args[1],
-    }
-    return operator.add_mul(network, target, None, kwargs_new, name)
 
 @tensorrt_converter(torch.ops.aten.matmul.Tensor)
 def aten_ops_matmul(
@@ -283,7 +270,6 @@ def aten_ops_matmul(
     return operator.add_matmul(network, target, None, kwargs_new, name)
 
 
-
 @tensorrt_converter(torch.ops.aten.pow.Tensor_Scalar)
 @tensorrt_converter(torch.ops.aten.pow.Tensor_Tensor)
 def aten_ops_pow(
@@ -392,10 +378,8 @@ def aten_ops_expand(
         "input": args[0],
         "sizes": args[1],
     }
-    return acc_ops_converters.acc_ops_expand_tensor(
-        network, target, None, kwargs_new, name
-    )
-
+    return operator.add_expand(network, target, kwargs_new, name)
+     
 
 @tensorrt_converter(operator.floordiv)
 def aten_ops_operator_floordiv(
@@ -497,8 +481,42 @@ def aten_ops_sym_size(
     set_layer_name(slice_layer, target, "_slice_layer")
     return slice_layer.get_output(0)
 
+
+@tensorrt_converter(torch.ops.aten.slice.Tensor)
+def aten_ops_slice(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input" : args[0],
+        "dim" : args[1],
+        "start" : args[2],
+        "stop" : args[3],
+        "step" : args[4],
+    }
+    return operator.add_slice(network, target. kwargs_new, name)
+
+@tensorrt_converter(torch.ops.aten.select.Tensor)
+def aten_ops_select(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    kwargs_new = {
+        "input" : args[0],
+        "dim" : args[1],
+        "index" : args[2],
+    }
+    return operator.add_select(network, target. kwargs_new, name)
+
+
 @tensorrt_converter(torch.ops.aten.leaky_relu.default)
-def aten_ops_relu(
+def aten_ops_leaky_relu(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
@@ -510,8 +528,9 @@ def aten_ops_relu(
     }
     return activation.add_leaky_relu(network, target, kwargs_new, name)
 
+
 @tensorrt_converter(torch.ops.aten.elu.default)
-def aten_ops_relu(
+def aten_ops_elu(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
@@ -521,7 +540,8 @@ def aten_ops_relu(
     kwargs_new = {
         "input": args[0],
     }
-    return activation.elu(network, target, kwargs_new, name)
+    return activation.add_elu(network, target, kwargs_new, name)
+
 
 @tensorrt_converter(torch.ops.aten.selu.default)
 def aten_ops_selu(
@@ -534,10 +554,11 @@ def aten_ops_selu(
     kwargs_new = {
         "input": args[0],
     }
-    return activation.add_selu(network, target, kwargs_new, name)
+    return activation.selu(network, target, kwargs_new, name)
 
-@tensorrt_converter(torch.ops.aten.selu.default)
-def aten_ops_selu(
+
+@tensorrt_converter(torch.ops.aten.gelu.default)
+def aten_ops_gelu(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
@@ -547,7 +568,8 @@ def aten_ops_selu(
     kwargs_new = {
         "input": args[0],
     }
-    return activation.add_selu(network, target, kwargs_new, name)
+    return activation.add_gelu(network, target, kwargs_new, name)
+
 
 @tensorrt_converter(torch.ops.aten.softsign.default)
 def aten_ops_softsign(
@@ -562,6 +584,7 @@ def aten_ops_softsign(
     }
     return activation.add_softsign(network, target, kwargs_new, name)
 
+
 @tensorrt_converter(torch.ops.aten.tanh.default)
 def aten_ops_tanh(
     network: TRTNetwork,
@@ -588,6 +611,7 @@ def aten_ops_softsign(
     }
     return activation.add_softsign(network, target, kwargs_new, name)
 
+
 @tensorrt_converter(torch.ops.aten.softsign.default)
 def aten_ops_hard_sigmoid(
     network: TRTNetwork,
@@ -601,6 +625,7 @@ def aten_ops_hard_sigmoid(
     }
     return activation.add_hard_sigmoid(network, target, kwargs_new, name)
 
+
 @tensorrt_converter(torch.ops.aten.sigmoid.default)
 def aten_ops_hard_tanh(
     network: TRTNetwork,
@@ -614,6 +639,7 @@ def aten_ops_hard_tanh(
     }
     return activation.add_hard_tanh(network, target, kwargs_new, name)
 
+
 @tensorrt_converter(torch.ops.aten.sigmoid.default)
 def aten_ops_sigmoid(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -396,7 +396,7 @@ def get_shape_with_dynamic_shape(
     )
     set_layer_name(zero_layer, target, f"{name}_zeros")
 
-    condition_val = add_binary_elementwise_layer(
+    condition_val = operator.add_binary_elementwise_layer(
         network,
         scale_res,
         zero_layer.get_output(0),