Merge branch 'main' into min_cpp_build

pytorch · Apr 16, 2024 · 959f3e5 · 959f3e5
2 parents 22faab8 + b76024d
commit 959f3e5
Show file tree

Hide file tree

Showing 14 changed files with 1,227 additions and 5 deletions.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -47,7 +47,7 @@ RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/
 RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
 RUN apt-get update
 
-RUN apt-get install -y libnvinfer8=${TENSORRT_VERSION}.* libnvinfer-plugin8=${TENSORRT_VERSION}.* libnvinfer-dev=${TENSORRT_VERSION}.* libnvinfer-plugin-dev=${TENSORRT_VERSION}.* libnvonnxparsers8=${TENSORRT_VERSION}.* libnvonnxparsers-dev=${TENSORRT_VERSION}.* libnvparsers8=${TENSORRT_VERSION}.*  libnvparsers-dev=${TENSORRT_VERSION}.*
+RUN apt-get install -y libnvinfer8=${TENSORRT_VERSION}.* libnvinfer-plugin8=${TENSORRT_VERSION}.* libnvinfer-dev=${TENSORRT_VERSION}.* libnvinfer-plugin-dev=${TENSORRT_VERSION}.* libnvonnxparsers8=${TENSORRT_VERSION}.* libnvonnxparsers-dev=${TENSORRT_VERSION}.* libnvparsers8=${TENSORRT_VERSION}.*  libnvparsers-dev=${TENSORRT_VERSION}.* libnvinfer-headers-dev=${TENSORRT_VERSION}.* libnvinfer-headers-plugin-dev=${TENSORRT_VERSION}.*
 
 # Setup Bazel via Bazelisk
 RUN wget -q https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64 -O /usr/bin/bazel &&\

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -1136,6 +1136,23 @@ def aten_ops_exp(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.expm1.default)
+def aten_ops_expm1(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.unary.expm1(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.log.default)
 def aten_ops_log(
     ctx: ConversionContext,
@@ -1391,6 +1408,30 @@ def aten_ops_atanh(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.atan2.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+        1: (TRTTensor,),
+    }
+)
+def aten_ops_atan2(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.elementwise.atan2(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.ceil.default)
 def aten_ops_ceil(
     ctx: ConversionContext,
@@ -1493,6 +1534,23 @@ def aten_ops_isinf(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.isnan.default)
+def aten_ops_isnan(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.unary.isnan(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.add.Tensor)
 @dynamo_tensorrt_converter(torch.ops.aten.add.Scalar)
 def aten_ops_add(
@@ -2185,7 +2243,12 @@ def aten_ops_avg_pool(
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool1d.default)
-def aten_ops_adaptive_avg_pool(
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_adaptive_avg_pool1d(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
@@ -2202,6 +2265,32 @@ def aten_ops_adaptive_avg_pool(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool2d.default)
+@dynamo_tensorrt_converter(torch.ops.aten._adaptive_avg_pool2d.default)
+@dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool3d.default)
+@dynamo_tensorrt_converter(torch.ops.aten._adaptive_avg_pool3d.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_adaptive_avg_poolNd(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pool.adaptive_avg_poolNd(
+        ctx,
+        target,
+        source_ir=SourceIR.ATEN,
+        name=name,
+        input=args[0],
+        output_size=args[1],
+    )
+
+
 def max_pool_param_validator(pool_node: Node) -> bool:
     dilation = args_bounds_check(pool_node.args, 4, 1)
     ceil_mode = args_bounds_check(pool_node.args, 5, False)
@@ -2319,6 +2408,29 @@ def aten_ops_pixel_shuffle(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.pixel_unshuffle.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)
+def aten_ops_pixel_unshuffle(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.shuffle.pixel_unshuffle(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
 @enforce_tensor_types({0: (TRTTensor,)})
 @dynamo_tensorrt_converter(torch.ops.aten.argmax.default)
 def aten_ops_argmax(
@@ -2782,3 +2894,28 @@ def aten_ops_roll(
         args[1],
         args_bounds_check(args, 2, []),
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.index_select.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+        2: (TRTTensor,),
+    }
+)
+def aten_ops_index_select(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.select.index_select(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args[2],
+    )
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py
@@ -9,13 +9,15 @@
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     cast_int_int_div_trt_tensor,
     cast_int_or_float_to_bool,
+    cast_trt_tensor,
     get_trt_tensor,
 )
 from torch_tensorrt.dynamo.conversion.impl.elementwise.base import (
     convert_binary_elementwise,
 )
-from torch_tensorrt.dynamo.conversion.impl.unary import sign
+from torch_tensorrt.dynamo.conversion.impl.unary import atan, sign
 from torch_tensorrt.dynamo.conversion.impl.unary.base import convert_unary
+from torch_tensorrt.fx.converters.converter_utils import broadcast
 from torch_tensorrt.fx.types import TRTTensor
 
 import tensorrt as trt
@@ -214,6 +216,180 @@ def remainder(
     return fmod2_value
 
 
+def atan2(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    other: TRTTensor,
+) -> TRTTensor:
+    """
+    Perform atan2 operation on Tensor, calculating the arctangent of the quotient of input tensors.
+    atan2(x,y) = atan(x/y) if y > 0,
+            = atan(x/y) + π if x ≥ 0 and y < 0,
+            = atan(x/y) - π if x < 0 and y < 0,
+            = π/2 if x > 0 and y = 0,
+            = -π/2 if x < 0 and y = 0,
+            = 0 if x = 0 and y = 0
+
+    Args:
+        ctx: ConversionContext.
+        target: node target
+        source_ir (SourceIR): Source IR calling the function.
+        name: namespace for the op
+        input: Tensor or constant representing the dividend.
+        other: Tensor or constant representing the divisor.
+
+    Returns:
+        A TensorRT tensor representing the result of the atan2 operation.
+    """
+    pi_value = 3.141592653589793
+    pi_tensor = get_trt_tensor(ctx, pi_value, f"{name}_pi")
+
+    if isinstance(input, TRTTensor):
+        input = cast_trt_tensor(ctx, input, trt.float32, f"{name}_input")
+    if isinstance(other, TRTTensor):
+        other = cast_trt_tensor(ctx, other, trt.float32, f"{name}_other")
+
+    input, other = broadcast(ctx.net, input, other, f"{name}_input", f"{name}_other")
+
+    # Calculate x_zero, y_zero (whether inputs are zero)
+    x_zero = eq(ctx, target, source_ir, f"{name}_x_zero", input, 0)
+    y_zero = eq(ctx, target, source_ir, f"{name}_y_zero", other, 0)
+
+    # Get sign of inputs
+    x_positive = gt(ctx, target, source_ir, f"{name}_x_positive", input, 0)
+    x_zero_positive = ge(ctx, target, source_ir, f"{name}_x_zero_positive", input, 0)
+    x_negative = lt(ctx, target, source_ir, f"{name}_x_negative", input, 0)
+    y_positive = gt(ctx, target, source_ir, f"{name}_y_positive", other, 0)
+    y_negative = lt(ctx, target, source_ir, f"{name}_y_negative", other, 0)
+
+    # Calculate atan(x/y)
+    input_div_other = div(
+        ctx, target, source_ir, f"{name}_input_div_other", input, other
+    )
+    atan_val = atan(ctx, target, source_ir, f"{name}_atan", input_div_other)
+
+    # atan(x/y)+π if x≥0 and y<0,
+    atan_add_pi = add(
+        ctx, target, source_ir, f"{name}_atan_add_pi", atan_val, pi_tensor
+    )
+
+    # atan(x/y)-π if x<0 and y<0,
+    atan_sub_pi = sub(
+        ctx, target, source_ir, f"{name}_atan_sub_pi", atan_val, pi_tensor
+    )
+
+    # atan(x/y)+π if x≥0 and y<0,
+    atan_corrected = impl.condition.select(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_atan_corrected",
+        atan_add_pi,
+        atan_val,
+        logical_and(
+            ctx,
+            target,
+            source_ir,
+            f"{name}_x_zero_positive_and_y_negative",
+            x_zero_positive,
+            y_negative,
+        ),
+    )
+
+    # atan(x/y)-π if x<0 and y<0,
+    atan_corrected_2 = impl.condition.select(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_atan_corrected_2",
+        atan_sub_pi,
+        atan_corrected,
+        logical_and(
+            ctx,
+            target,
+            source_ir,
+            f"{name}_x_negative_and_y_negative",
+            x_negative,
+            y_negative,
+        ),
+    )
+
+    # atan(x/y) if y>0
+    atan_output = impl.condition.select(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_atan_output",
+        atan_val,
+        atan_corrected_2,
+        y_positive,
+    )
+
+    # on x or y-axis
+    pi_over_2_tensor = get_trt_tensor(
+        ctx,
+        (pi_value / 2) * np.ones(input.shape, dtype=np.float32),
+        f"{name}_pi_over_2_tensor",
+        dtype=trt.float32,
+    )
+    minus_pi_over_2_tensor = get_trt_tensor(
+        ctx,
+        (-pi_value / 2) * np.ones(input.shape, dtype=np.float32),
+        f"{name}_minus_pi_over_2_tensor",
+        dtype=trt.float32,
+    )
+    zero_tensor = get_trt_tensor(
+        ctx,
+        np.zeros(input.shape, dtype=np.float32),
+        f"{name}_zero_tensor",
+        dtype=trt.float32,
+    )
+
+    # π/2 if x>0 and y=0,
+    pi_over_2_output = impl.condition.select(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_pi_over_2_output",
+        pi_over_2_tensor,
+        atan_output,
+        logical_and(
+            ctx, target, source_ir, f"{name}_x_zero_and_y_positive", x_positive, y_zero
+        ),
+    )
+
+    # -π/2 if x<0 and y=0,
+    minus_pi_over_2_output = impl.condition.select(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_minus_pi_over_2_output",
+        minus_pi_over_2_tensor,
+        pi_over_2_output,
+        logical_and(
+            ctx, target, source_ir, f"{name}_x_zero_and_y_negative", x_negative, y_zero
+        ),
+    )
+
+    # 0 if x=0 and y=0,
+    zero_output = impl.condition.select(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_zero_output",
+        zero_tensor,
+        minus_pi_over_2_output,
+        logical_and(
+            ctx, target, source_ir, f"{name}_x_zero_and_y_zero", y_zero, x_zero
+        ),
+    )
+
+    return zero_output
+
+
 def clamp(
     ctx: ConversionContext,
     target: Target,