From 0a26d9e10a797d62e2299d3daf8e88e30c1ae815 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 14 Sep 2026 09:11:40 -0700 Subject: [PATCH] Preserve Python scalar semantics in Vulkan and XNNPACK quantizers Converting Python scalars to the output tensor dtype can round or overflow them before arithmetic. For example, multiplying an FP16 tensor containing 1e-4 by 100000.0 should produce approximately 10, but materializing the scalar as FP16 produces infinity instead. Limit scalar lifting to FP32 add and multiply operations, preserving the original scalar operands for other dtypes. Cover scalar promotion and rounding across low-precision floats, doubles, integers, bools, and complex values, with and without quantization configuration. Authored with OpenAI Codex. --- .../quantizer/vulkan_quantizer_utils.py | 7 +- .../test/quantizer/test_vulkan_quantizer.py | 66 ++++++++++++------- .../quantizer/xnnpack_quantizer_utils.py | 7 +- .../test/quantizer/test_xnnpack_quantizer.py | 54 +++++++++++---- 4 files changed, 99 insertions(+), 35 deletions(-) diff --git a/backends/vulkan/quantizer/vulkan_quantizer_utils.py b/backends/vulkan/quantizer/vulkan_quantizer_utils.py index a8e914b8ec0..74521523fa0 100644 --- a/backends/vulkan/quantizer/vulkan_quantizer_utils.py +++ b/backends/vulkan/quantizer/vulkan_quantizer_utils.py @@ -201,7 +201,12 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu args = list(n.args) new_args = [] for i in range(len(args)): - if isinstance(args[i], torch.fx.Node): + # Only FP32 binary ops need scalar lifting for quantization. Other + # dtypes must retain Python scalar promotion and rounding semantics. + if ( + isinstance(args[i], torch.fx.Node) + or n.meta["val"].dtype != torch.float32 + ): new_args.append(args[i]) continue prefix = "_tensor_constant_" diff --git a/backends/vulkan/test/quantizer/test_vulkan_quantizer.py b/backends/vulkan/test/quantizer/test_vulkan_quantizer.py index 3f568e55d3b..1d969ee6c48 100644 --- a/backends/vulkan/test/quantizer/test_vulkan_quantizer.py +++ b/backends/vulkan/test/quantizer/test_vulkan_quantizer.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import itertools import unittest import torch @@ -17,9 +18,6 @@ class TestVulkanQuantizer(unittest.TestCase): def test_int64_scalar_add_used_as_index(self): - """Scalars lifted to attrs must keep the op's output dtype; an int64 - add chain used as an index must not be promoted to float32.""" - class M(torch.nn.Module): def forward(self, x): return x[:, torch.arange(4) + 0] @@ -30,20 +28,10 @@ def forward(self, x): example_inputs = (torch.randn(1, 4, 5),) m = export(M(), example_inputs, strict=True).module() m = quantizer.transform_for_annotation(m) - lifted_constants = [ - m.get_buffer(n.target) - for n in m.graph.nodes - if n.op == "get_attr" and n.target.startswith("_tensor_constant_") - ] - self.assertEqual(len(lifted_constants), 1) - self.assertEqual(lifted_constants[0].dtype, torch.int64) m = prepare_pt2e(m, quantizer) - m(*example_inputs) - - def test_int64_scalar_lifted_without_set_global(self): - """Passing through VulkanQuantizer without a config still lifts scalars - and must preserve dtype.""" + torch.testing.assert_close(m(*example_inputs), M()(*example_inputs)) + def test_int64_scalar_add_without_set_global(self): class M(torch.nn.Module): def forward(self, x): return x[:, torch.arange(4) + 0] @@ -51,11 +39,45 @@ def forward(self, x): example_inputs = (torch.randn(1, 4, 5),) m = export(M(), example_inputs, strict=True).module() m = VulkanQuantizer().transform_for_annotation(m) - lifted_constants = [ - m.get_buffer(n.target) - for n in m.graph.nodes - if n.op == "get_attr" and n.target.startswith("_tensor_constant_") + torch.testing.assert_close(m(*example_inputs), M()(*example_inputs)) + + def test_scalar_type_promotion(self): + class M(torch.nn.Module): + def __init__(self, op, scalar): + super().__init__() + self.op = op + self.scalar = scalar + + def forward(self, x): + return self.op(x, self.scalar) + + cases = [ + (torch.float16, 1e-4, 100000.0), + (torch.float16, 1e-4, 100000), + (torch.float16, 10000.0, 1e-8), + (torch.bfloat16, 100.0, 1.0039), + (torch.float64, 1.0, 1.0 + 2**-30), + (torch.int64, 2**54 + 1, 1), + (torch.int32, 1, 1), + (torch.int32, 1, 2**31), + (torch.int8, 1, 256), + (torch.bool, True, False), + (torch.int32, 1, 0.5), + (torch.float32, 1.0, 2), + (torch.complex64, 1j, 1 + 2j), ] - self.assertEqual(len(lifted_constants), 1) - self.assertEqual(lifted_constants[0].dtype, torch.int64) - m(*example_inputs) + for op, (dtype, value, scalar), shape, configured in itertools.product( + (torch.add, torch.mul), cases, ((), (2,)), (False, True) + ): + with self.subTest( + op=op, dtype=dtype, scalar=scalar, shape=shape, configured=configured + ): + model = M(op, scalar) + example_inputs = (torch.full(shape, value, dtype=dtype),) + expected = model(*example_inputs) + quantizer = VulkanQuantizer() + if configured: + quantizer.set_global(get_symmetric_quantization_config()) + m = export(model, example_inputs, strict=True).module() + m = prepare_pt2e(m, quantizer) + torch.testing.assert_close(m(*example_inputs), expected, rtol=0, atol=0) diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index 78a989bd75c..c43fc5da089 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1154,7 +1154,12 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu args = list(n.args) new_args = [] for i in range(len(args)): - if isinstance(args[i], torch.fx.Node): + # Only FP32 binary ops need scalar lifting for quantization. Other + # dtypes must retain Python scalar promotion and rounding semantics. + if ( + isinstance(args[i], torch.fx.Node) + or n.meta["val"].dtype != torch.float32 + ): new_args.append(args[i]) continue prefix = "_tensor_constant_" diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index 27d6a8f65fb..32b94521352 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -1,5 +1,6 @@ # Owner(s): ["oncall: mobile"] import copy +import itertools import operator import torch @@ -1123,9 +1124,6 @@ def forward(self, x): ) def test_int64_scalar_add_used_as_index(self): - """Scalars lifted to attrs must keep the op's output dtype; an int64 - add chain used as an index must not be promoted to float32.""" - class M(torch.nn.Module): def forward(self, x): return x[:, torch.arange(4) + 0] @@ -1136,15 +1134,49 @@ def forward(self, x): example_inputs = (torch.randn(1, 4, 5),) m = export(M(), example_inputs, strict=True).module() m = quantizer.transform_for_annotation(m) - lifted_constants = [ - m.get_buffer(n.target) - for n in m.graph.nodes - if n.op == "get_attr" and n.target.startswith("_tensor_constant_") - ] - self.assertEqual(len(lifted_constants), 1) - self.assertEqual(lifted_constants[0].dtype, torch.int64) m = prepare_pt2e(m, quantizer) - m(*example_inputs) + torch.testing.assert_close(m(*example_inputs), M()(*example_inputs)) + + def test_scalar_type_promotion(self): + class M(torch.nn.Module): + def __init__(self, op, scalar): + super().__init__() + self.op = op + self.scalar = scalar + + def forward(self, x): + return self.op(x, self.scalar) + + cases = [ + (torch.float16, 1e-4, 100000.0), + (torch.float16, 1e-4, 100000), + (torch.float16, 10000.0, 1e-8), + (torch.bfloat16, 100.0, 1.0039), + (torch.float64, 1.0, 1.0 + 2**-30), + (torch.int64, 2**54 + 1, 1), + (torch.int32, 1, 1), + (torch.int32, 1, 2**31), + (torch.int8, 1, 256), + (torch.bool, True, False), + (torch.int32, 1, 0.5), + (torch.float32, 1.0, 2), + (torch.complex64, 1j, 1 + 2j), + ] + for op, (dtype, value, scalar), shape, configured in itertools.product( + (torch.add, torch.mul), cases, ((), (2,)), (False, True) + ): + with self.subTest( + op=op, dtype=dtype, scalar=scalar, shape=shape, configured=configured + ): + model = M(op, scalar) + example_inputs = (torch.full(shape, value, dtype=dtype),) + expected = model(*example_inputs) + quantizer = XNNPACKQuantizer() + if configured: + quantizer.set_global(get_symmetric_quantization_config()) + m = export(model, example_inputs, strict=True).module() + m = prepare_pt2e(m, quantizer) + torch.testing.assert_close(m(*example_inputs), expected, rtol=0, atol=0) def test_cat_same_node(self): """Ensure that concatenating the same node does not cause any unexpected behavior"""