diff --git a/test/modules/op/clamp.py b/test/modules/op/clamp.py index fe7311ee..f148141c 100644 --- a/test/modules/op/clamp.py +++ b/test/modules/op/clamp.py @@ -103,6 +103,21 @@ def get_example_inputs(self): return (torch.randn(5, 3) * 10,), {} +class SimpleClampTensorMinMax(TestModuleBase): + """Clamp with scalar tensor buffers as used by Gemma4ClippableLinear.""" + + def __init__(self): + super().__init__() + self.register_buffer("min_value", torch.tensor(-11.0)) + self.register_buffer("max_value", torch.tensor(11.0)) + + def forward(self, x): + return torch.clamp(x, self.min_value, self.max_value) + + def get_example_inputs(self): + return (torch.randn(5, 3) * 20,), {} + + class ClampIntInputFloatMinMax(TestModuleBase): def __init__(self): super().__init__() diff --git a/test/unit_test/passes/test_cast_clamp_mixed_type_args.py b/test/unit_test/passes/test_cast_clamp_mixed_type_args.py index ed183f54..0bd93ea7 100644 --- a/test/unit_test/passes/test_cast_clamp_mixed_type_args.py +++ b/test/unit_test/passes/test_cast_clamp_mixed_type_args.py @@ -106,3 +106,29 @@ def test_pass(self): self.run_value_test(CastClampMixedTypeArgs()) self.assertEqual(num_of_ops(self.exported_program(), ops.aten.clamp), 1) self.assertEqual(num_of_ops(self.exported_program(), ops.aten._to_copy), 1) + + +class CastClampTensorBounds(torch.nn.Module): + """Use lifted scalar tensor buffers as Clamp min/max arguments.""" + + def __init__(self): + super().__init__() + self.register_buffer("min_value", torch.tensor(-10.0)) + self.register_buffer("max_value", torch.tensor(10.0)) + + def forward(self, x): + return torch.clamp(x, self.min_value, self.max_value) + + def get_example_inputs(self): + return (torch.randn(5, 3) * 20,), {} + + +class CastClampTensorBoundsTest(SinglePassValueTest): + def test_pass(self): + self.setup(CastClampTensorBounds()) + self.assertEqual(num_of_ops(self.exported_program(), ops.aten.clamp), 1) + + self.run_value_test(CastClampMixedTypeArgs()) + + self.assertEqual(num_of_ops(self.exported_program(), ops.aten.clamp), 1) + self.assertEqual(num_of_ops(self.exported_program(), ops.aten._to_copy), 0) diff --git a/tico/passes/cast_clamp_mixed_type_args.py b/tico/passes/cast_clamp_mixed_type_args.py index 80288534..f587ce76 100644 --- a/tico/passes/cast_clamp_mixed_type_args.py +++ b/tico/passes/cast_clamp_mixed_type_args.py @@ -100,6 +100,7 @@ def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> boo graph = graph_module.graph # clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor + # clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor args = ClampArgs(*node.args, **node.kwargs) # type: ignore[arg-type] input = args.input @@ -113,20 +114,42 @@ def _convert_arg(arg, arg_name: str): if arg is None: return False - arg_dtype = torch.tensor(arg).dtype arg_idx = node.args.index(arg) - if arg_dtype != output_dtype: - assert output_dtype in [torch.float, torch.int] - if output_dtype == torch.float: - arg = float(arg) - else: - arg = int(arg) - node.update_arg(arg_idx, arg) + + if isinstance(arg, torch.fx.Node): + arg_dtype = extract_torch_dtype(arg) + if arg_dtype == output_dtype: + return False + + with graph.inserting_after(arg): + to_copy = create_node( + graph, + torch.ops.aten._to_copy.default, + (arg,), + {"dtype": output_dtype}, + origin=arg, + ) + set_new_meta_val(to_copy) + node.update_arg(arg_idx, to_copy) + logger.debug( - f"Casting {arg_name} value from {arg_dtype} to {output_dtype} for clamp operation at {node.name}" + f"Inserting cast for {arg_name} from {arg_dtype} to " + f"{output_dtype} for clamp operation at {node.name}" ) return True - return False + + arg_dtype = torch.tensor(arg).dtype + if arg_dtype == output_dtype: + return False + + assert output_dtype in [torch.float, torch.int] + converted_arg = float(arg) if output_dtype == torch.float else int(arg) + node.update_arg(arg_idx, converted_arg) + logger.debug( + f"Casting {arg_name} value from {arg_dtype} to {output_dtype} " + f"for clamp operation at {node.name}" + ) + return True modified |= _convert_arg(min, "min") modified |= _convert_arg(max, "max") diff --git a/tico/utils/validate_args_kwargs.py b/tico/utils/validate_args_kwargs.py index 91b5cc16..d497e281 100644 --- a/tico/utils/validate_args_kwargs.py +++ b/tico/utils/validate_args_kwargs.py @@ -208,11 +208,12 @@ class CircleRMSNormArgs: class ClampArgs: """ clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor + clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor """ input: torch.fx.Node - min: Optional[Union[int, float]] = None - max: Optional[Union[int, float]] = None + min: Optional[Union[torch.fx.Node, int, float]] = None + max: Optional[Union[torch.fx.Node, int, float]] = None @enforce_type