feat: support aten.clamp.Tensor and update aten.clamp.default dynamo converters

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -683,6 +683,7 @@ def aten_ops_where(
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.clamp.default)
+@dynamo_tensorrt_converter(torch.ops.aten.clamp.Tensor)
 def aten_ops_clamp(
     ctx: ConversionContext,
     target: Target,

py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py

@@ -1,6 +1,5 @@
 from typing import Optional, Union
 
-import numpy as np
 import tensorrt as trt
 import torch
 import torch_tensorrt.dynamo.conversion.impl as impl
@@ -17,7 +16,6 @@
 )
 from torch_tensorrt.dynamo.conversion.impl.unary import sign
 from torch_tensorrt.dynamo.conversion.impl.unary.base import convert_unary
-from torch_tensorrt.fx.converters.converter_utils import set_layer_name, squeeze_left
 from torch_tensorrt.fx.types import TRTTensor
 from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
@@ -186,63 +184,22 @@ def clamp(
     source_ir: Optional[SourceIR],
     name: str,
     input_val: TRTTensor,
-    min_val: Optional[float] = None,
-    max_val: Optional[float] = None,
+    min_val: Optional[Union[int, float, TRTTensor]] = None,
+    max_val: Optional[Union[int, float, TRTTensor]] = None,
 ) -> TRTTensor:
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"Clamp received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    def _add_layer(
-        ctx: ConversionContext,
-        input: TRTTensor,
-        val: float,
-        op: trt.ElementWiseOperation,
-        name: str,
-    ) -> (
-        trt.ILayer
-    ):  # TODO: Simplify and merge implementations, should just be max and min stacked
-        if not len(input.shape):
-            # clamping scalar
-            acc_ops_clamp_trt = get_trt_tensor(
-                ctx,
-                squeeze_left(
-                    np.array(
-                        [val],
-                        dtype=unified_dtype_converter(input.dtype, Frameworks.NUMPY),
-                    )
-                ),
-                f"{name}_clamp_{val}",
-            )
-        else:
-            acc_ops_clamp_shape = (1,) * len(input.shape)  # broadcast all dimensions
-            acc_ops_clamp_tensor = np.full(
-                acc_ops_clamp_shape,
-                val,
-                dtype=unified_dtype_converter(input.dtype, Frameworks.NUMPY),
-            )
-            acc_ops_clamp_trt = ctx.net.add_constant(
-                acc_ops_clamp_shape, acc_ops_clamp_tensor
-            ).get_output(0)
-        layer = ctx.net.add_elementwise(input, acc_ops_clamp_trt, op)
-        return layer
-
-    if min_val is not None:
-        clamp_min_layer = _add_layer(
-            ctx, input_val, min_val, trt.ElementWiseOperation.MAX, name
-        )
-        set_layer_name(clamp_min_layer, target, f"{name}_clamp_min")
-        input_val = clamp_min_layer.get_output(0)
-    if max_val is not None:
-        clamp_max_layer = _add_layer(
-            ctx, input_val, max_val, trt.ElementWiseOperation.MIN, name
-        )
-        set_layer_name(clamp_max_layer, target, f"{name}_clamp_max")
-        input_val = clamp_max_layer.get_output(0)
+    if min_val is None:
+        min_val = float("-inf")
+    if max_val is None:
+        max_val = float("inf")
 
-    return input_val
+    return impl.elementwise.min(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_min",
+        impl.elementwise.max(ctx, target, source_ir, f"{name}_max", input_val, min_val),
+        max_val,
+    )
 
 
 def add(

tests/py/dynamo/conversion/test_clamp_aten.py

@@ -49,7 +49,7 @@ def forward(self, x):
 
         class TestScalarModule(torch.nn.Module):
             def forward(self, x):
-                y = torch.ops.aten.mean.default(x)
+                y = torch.ops.aten.mean.dim(x, None, True)
                 return torch.ops.aten.clamp.default(y, min, max)
 
         input_specs = [
@@ -63,6 +63,30 @@ def forward(self, x):
         self.run_test_with_dynamic_shape(TestModule(), input_specs)
         self.run_test_with_dynamic_shape(TestScalarModule(), input_specs)
 
+    @parameterized.expand(
+        [
+            param("default", min=-1 * torch.randn(3, 4), max=0 * torch.randn(3, 4)),
+            param("min", min=0.5 * torch.randn(3, 4)),
+            param("max", max=0.5 * torch.randn(3, 4)),
+            param(
+                "minBiggerThanMax", min=1 * torch.randn(3, 4), max=0 * torch.randn(3, 4)
+            ),
+            param("float32Boundary", min=-3.4028234663852886e38 * torch.randn(3, 4)),
+        ]
+    )
+    def test_clamp_tensor(
+        self,
+        test_name,
+        min=None,
+        max=None,
+    ):
+        class TestModule(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.clamp.Tensor(x, min, max)
+
+        inputs = [torch.randn(3, 4)]
+        self.run_test(TestModule(), inputs)
+
 
 if __name__ == "__main__":
     run_tests()