diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index fca0f1b14c6..d4b7f956f19 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1157,15 +1157,15 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu prefix = "_tensor_constant_" get_new_attr_name = get_new_attr_name_with_prefix(prefix) tensor_constant_name = get_new_attr_name(model) - float_tensor = torch.tensor(float(args[i])) - model.register_buffer(tensor_constant_name, float_tensor) + scalar_tensor = torch.tensor(args[i], dtype=n.meta["val"].dtype) + model.register_buffer(tensor_constant_name, scalar_tensor) fake_mode = n.meta["val"].fake_mode with model.graph.inserting_before(n): get_attr_node = model.graph.create_node( "get_attr", tensor_constant_name, (), {} ) get_attr_node.meta["val"] = fake_mode.from_tensor( - float_tensor, static_shapes=True + scalar_tensor, static_shapes=True ) new_args.append(get_attr_node) n.args = tuple(new_args) diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index 1e1a473dd59..27d6a8f65fb 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -1122,6 +1122,30 @@ def forward(self, x): node_list, ) + 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] + + quantizer = XNNPACKQuantizer() + quantization_config = get_symmetric_quantization_config(is_per_channel=True) + quantizer.set_global(quantization_config) + 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_cat_same_node(self): """Ensure that concatenating the same node does not cause any unexpected behavior"""