chore: Remove dependency from fx.converters.converter_utils

Author: keehyuna

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -15,13 +15,8 @@
     ConverterRegistry,
     DynamoConverterImplSignature,
 )
-from torch_tensorrt.fx.converters.converter_utils import (  # noqa: F401
-    broadcast,
-    get_axes_for_reduce_op,
-    prepend_ones,
-    set_layer_name,
-)
-from torch_tensorrt.fx.types import TRTDataType, TRTTensor
+
+from ..types import Shape, TRTDataType, TRTLayer, TRTNetwork, TRTTensor
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -218,11 +213,6 @@ def broadcast_to_same_shape(
     return lhs_val, rhs_val
 
 
-get_axes_for_reduce_op = functools.partial(
-    get_axes_for_reduce_op, has_implicit_batch_dimension=False
-)
-
-
 def extend_attr_to_tuple(
     val: Any,
     num_elem: int,
@@ -694,3 +684,183 @@ def calculate_strides(shape: Sequence[int]) -> Sequence[int]:
     for i in range(len(shape) - 2, -1, -1):
         strides[i] = strides[i + 1] * shape[i + 1]
     return strides
+
+
+def broadcast(
+    network: TRTNetwork,
+    a: TRTTensor,
+    b: TRTTensor,
+    a_name: str,
+    b_name: str,
+    preset_diff: int = 0,
+) -> Tuple[TRTTensor, TRTTensor]:
+    """
+    Broadcast two TensorRT tensors to the same number of dimensions by
+    prepending 1s to the tensor with less number of dimensions.
+
+    Args:
+        network (TRTNetwork): TensorRT network object.
+        a (TRTTensor): A TensorRT ITensor.
+        b (TRTTensor): A TensorRT ITensor.
+        a_name (str): Name of tensor a.
+        b_name (str): Name of tensor b.
+        preset_diff (int): The difference of number of dimensions after broadcast.
+            A positive number means after broadcast, tensor `a` would have `preset_diff`
+            more dimensions than `b`. This is used in matmul, since we need to broadcast
+            tensors but not always to the same number of dimension. The reason is that
+            matmul supports Matrix x Vector and in this case broadcasted vector should
+            have 1 less number of dimensions than the matrix tensor.
+
+    Returns:
+        Two TensorRT ITensors that are broadcasted to the same number of dimensions.
+    """
+    a_shape = tuple(a.shape)
+    b_shape = tuple(b.shape)
+
+    diff = len(a_shape) - len(b_shape) - preset_diff
+    if diff > 0:
+        b = prepend_ones(network, b, f"{b_name}_broadcast", diff)
+    elif diff < 0:
+        a = prepend_ones(network, a, f"{a_name}_broadcast", -diff)
+
+    return a, b
+
+
+def get_axes_for_reduce_op(
+    dim: Union[int, Sequence[int]],
+    has_implicit_batch_dimension: bool = False,
+) -> int:
+    """
+    TensorRT reduce layer relies on the binary representation of axes to
+    determine which dims to reduce. For example, if we want to reduce on
+    dim 1 and 2 then axes should be 6(110).
+
+    Args:
+        dim (Union[int, Sequence[int]]): An integer or a sequence of integers
+            that will be used to generate axes for TensorRT.
+        has_implicit_batch_dimension (bool): Whether the TensorRT network is
+            using implicit batch dimension.
+
+    Returns:
+        An integer which binary form can be used as axes for TensorRT reduce
+        layer.
+    """
+    if isinstance(dim, int):
+        dim = (dim,)
+
+    if has_implicit_batch_dimension:
+        assert 0 not in dim, "Can't reduce over batch dimension when it's implicit."
+
+    axes = 0
+    for d in dim:
+        axes |= 1 << (d - (1 if has_implicit_batch_dimension else 0))
+
+    return axes
+
+
+def has_dynamic_shape(shape: Shape) -> bool:
+    """
+    Determine if the given shape has dynamic dim. i.e. if there're -1 in shape.
+
+    Args:
+        shape (Shape): Shape of a tensor. Essentially is a sequence of integers.
+
+    Returns:
+        A boolean value indicates whether there's dynamic dim in the shape.
+    """
+    count = 0
+    for s in shape:
+        count += 1 if s == -1 else 0
+    return count > 0
+
+
+def type_cast(
+    network: TRTNetwork,
+    target: Target,
+    name: str,
+    input: TRTTensor,
+    cast_type: TRTDataType,
+) -> TRTTensor:
+    """
+    This function helps to cast the input type to cast_type
+    """
+    layer_i = network.add_identity(input)
+    layer_i.set_output_type(0, cast_type)
+    set_layer_name(layer_i, target, f"{name}_dtype_change")
+    return layer_i.get_output(0)
+
+
+def prepend_ones(
+    network: TRTNetwork,
+    tensor: TRTTensor,
+    name: str,
+    num_prepend_ones: int,
+) -> TRTTensor:
+    """
+    Prepend 1s to the shape of TensorRT ITensor `tensor`.
+
+    Args:
+        network (TRTNetwork): The TensorRT network that `tensor`
+            belongs to.
+        tensor (TRTTensor): A TensorRT tensor.
+        name (str): Name of the TensorRT Shuffle layer which is used to prepend
+            1s.
+        num_prepend_ones (int): Number of 1s that will be prepend.
+
+    Returns:
+        A Tensorrt ITensor which contains the same value as `tensor` but with
+        more 1s prepended to the beginning of `tensor` shape.
+    """
+    layer = network.add_shuffle(tensor)
+
+    # If there're dynamic dim in tensor's shape, we need to use shape layer to
+    # compute the final shape.
+    if has_dynamic_shape(tensor.shape):
+        tensor_shape_layer = network.add_shape(tensor)
+        tensor_shape = tensor_shape_layer.get_output(0)
+        tensor_shape = type_cast(
+            network, "shape", name + "shape_casted", tensor_shape, trt.int32
+        )
+        tensor_shape_layer.name = f"{name}_broadcast_orig_shape"
+        prepend_shape_layer = network.add_constant(
+            (num_prepend_ones,), np.ones((num_prepend_ones,), dtype=np.int32)
+        )
+        prepend_shape_layer.name = f"{name}_broadcast_prepend_ones"
+        reshape_dim_layer = network.add_concatenation(
+            [prepend_shape_layer.get_output(0), tensor_shape]
+        )
+        reshape_dim_layer.axis = 0
+        reshape_dim_layer.name = f"{name}_broadcast_final_shape"
+        layer.set_input(1, reshape_dim_layer.get_output(0))
+    else:
+        layer.reshape_dims = (1,) * num_prepend_ones + tuple(tensor.shape)
+
+    layer.name = name
+    return layer.get_output(0)
+
+
+def set_layer_name(
+    layer: TRTLayer,
+    target: Union[Target, torch.nn.Module, str],
+    name: str,
+    source_ir: Optional[SourceIR] = None,
+) -> None:
+    """
+    Set the TensorRT layer name to "[TensorRT Layer Type]_[Original Op Name]_[FX Node Name with Suffix]"
+
+    Args:
+        layer (TRTLayer): A TensorRT layer of which we want to set the name.
+        target (Target): A fx node.target or submodule. For call_function node, it's the function that
+            the node represents.
+        name (str): Consists of fx node.name with optional suffix.
+        source_ir: (Optional[SourceIR]): The IR producing the op.
+    """
+
+    source_ir = source_ir if source_ir is not None else SourceIR.UNKNOWN
+
+    target_name = (
+        f"{source_ir}_ops.{target}"
+        if isinstance(target, str)
+        else f"{source_ir}_ops.{target.__name__}"
+    )
+    layer.name = f"[{layer.type.name}]-[{target_name}]-[{name}]"

py/torch_tensorrt/dynamo/types.py

@@ -0,0 +1,24 @@
+from typing import Sequence, Tuple
+
+# @manual=//deeplearning/trt/python:py_tensorrt
+import tensorrt as trt
+
+if hasattr(trt, "__version__"):
+    TRTNetwork = trt.INetworkDefinition
+    TRTTensor = trt.tensorrt.ITensor
+    TRTLayer = trt.ILayer
+    TRTPluginFieldCollection = trt.PluginFieldCollection
+    TRTPlugin = trt.IPluginV2
+    TRTDataType = trt.DataType
+    TRTElementWiseOp = trt.ElementWiseOperation
+else:
+    TRTNetwork = "trt.INetworkDefinition"
+    TRTTensor = "trt.tensorrt.ITensor"
+    TRTLayer = "trt.ILayer"
+    TRTPluginFieldCollection = "trt.PluginFieldCollection"
+    TRTPlugin = "trt.IPluginV2"
+    TRTDataType = "trt.DataType"
+    TRTElementWiseOp = "trt.ElementWiseOperation"
+
+Shape = Sequence[int]
+ShapeRange = Tuple[Shape, Shape, Shape]