diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index 0c49880f927..532d1d5b58f 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -162,6 +162,10 @@ def maybe_add_constant_tensor(self, node: Node) -> int: # Serialize tensor data to bytes tensor = tensor.contiguous() + # The bytes are read from the start of the storage, so a view of a + # larger tensor needs its own copy. + if tensor.nbytes != tensor.untyped_storage().nbytes(): + tensor = tensor.clone() size = tensor.untyped_storage().nbytes() if size > 0: diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index b20073493af..0871cb455a3 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -283,6 +283,8 @@ def serialize_constant_tensors( ): vk_graph.constants.append(VkBytes(current_offset, 0)) elif isinstance(tensor, torch.Tensor): + # The bytes are read from the start of the storage in memory order. + tensor = tensor.clone(memory_format=torch.contiguous_format) array_type = ctypes.c_char * tensor.untyped_storage().nbytes() array = ctypes.cast( tensor.untyped_storage().data_ptr(), diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index 8b95b1df2a1..2eb739a760e 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -6,7 +6,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import ctypes import random import unittest from types import SimpleNamespace @@ -267,6 +266,8 @@ def test_serialize_vulkan_binary(self) -> None: shaders=[], ) const_tensors = self._generate_random_const_tensors(5) + const_tensors.append(torch.randn(3, 4, 4)[1]) + const_tensors.append(torch.randn(3, 4, 4)[1].t()) serialized_binary = serialize_vulkan_graph(vk_graph, const_tensors, []) @@ -305,14 +306,7 @@ def test_serialize_vulkan_binary(self) -> None: constant_data_bytes = constant_data_payload[offset : offset + length] - array_type = ctypes.c_char * tensor.untyped_storage().nbytes() - array = ctypes.cast( - tensor.untyped_storage().data_ptr(), - ctypes.POINTER(array_type), - ).contents - - tensor_bytes = bytes(array) - self.assertEqual(constant_data_bytes, tensor_bytes) + self.assertEqual(constant_data_bytes, tensor.numpy().tobytes()) def test_serialize_deserialize_vkgraph(self) -> None: in_vk_graph = VkGraph( diff --git a/backends/vulkan/test/test_vulkan_graph_builder.py b/backends/vulkan/test/test_vulkan_graph_builder.py index 419759785bc..3ac3a5ff682 100644 --- a/backends/vulkan/test/test_vulkan_graph_builder.py +++ b/backends/vulkan/test/test_vulkan_graph_builder.py @@ -106,5 +106,26 @@ def forward(self, first, second): self.assertEqual(len(builder.input_ids), 2) +class TestVkGraphBuilderConstants(unittest.TestCase): + def test_view_constant_serializes_only_its_own_bytes(self) -> None: + class AddViewBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("b", torch.randn(3, 4)[1]) + + def forward(self, x): + return x + self.b + + module = AddViewBuffer() + edge = to_edge(torch.export.export(module, (torch.randn(4),), strict=True)) + program = apply_passes(edge.exported_program(), [SpecPropPass()]) + builder = VkGraphBuilder( + program, DelegateMappingBuilder(generated_identifiers=True) + ) + builder.build_graph() + (data,) = builder.named_data_store.get_named_data_store_output().buffers + self.assertEqual(bytes(data), module.b.numpy().tobytes()) + + if __name__ == "__main__": unittest.main() diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index 1963e4494fd..238bb992784 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -288,6 +288,10 @@ def get_quant_params( if quant_params.per_channel_group: scale = scale.to(torch.bfloat16) + # The bytes are read from the start of the storage in memory order. + # Scales are small, so always copy them into storage of their own. + scale = scale.clone(memory_format=torch.contiguous_format) + num_bytes = scale.untyped_storage().nbytes() scale_array = ctypes.cast( scale.untyped_storage().data_ptr(), @@ -619,6 +623,11 @@ def get_serialized_buffer_index( if quant_params is not None and quant_params.is_qc4w: const_val = self.convert_to_qc4w(const_val) + # The bytes are read from the start of the storage, so a view of a + # larger tensor needs its own copy. + if const_val.nbytes != const_val.untyped_storage().nbytes(): + const_val = const_val.clone() + size = const_val.untyped_storage().nbytes() array_type = ctypes.c_char * size array = ctypes.cast( diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index 154fa2fa559..e309bc514a7 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -781,6 +781,25 @@ def test_fp32_addmm(self): uses_bias=True, ) + def test_fp32_linear_view_weight(self): + class ViewWeightLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.randn(3, 8, 8)[1]) + + def forward(self, x): + return torch.nn.functional.linear(x, self.weight) + + ( + Tester(ViewWeightLinear(), (torch.randn(2, 8),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_fp32_linear_fused_relu(self): for use_bias in (True, False): for num_batch_dims in range(1, 3): @@ -898,6 +917,40 @@ def test_linear_qd8_per_token_groupwise_unsupported_groupsize(self): lin_mod, inputs, group_size=30, use_bias=False, atol=1e-2 ) + @unittest.skipIf(not torchao_installed, "Requires TorchAO") + def test_qd8_linear_transposed_scale(self): + mod = torch.nn.Linear(64, 4, bias=False) + with torch.no_grad(): + # Rows with very different scales make a misordered scale obvious. + mod.weight.mul_(torch.arange(1.0, 5.0).unsqueeze(1)) + quantize_( + mod, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(32) + ), + ) + unwrap_tensor_subclass(mod) + # XNNPACK stores groupwise scales as bf16, so round them for an exact comparison. + scale = mod.weight.scale.to(torch.bfloat16).float() + mod.weight.scale = scale.t().contiguous().t() + ( + Tester(mod, (torch.randn(2, 64),)) + .export() + .to_edge_transform_and_lower( + ToEdgeTransformAndLower( + [ + XnnpackPartitioner( + config_precisions=ConfigPrecisionType.DYNAMIC_QUANT + ) + ] + ) + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_qd8_per_channel_linear_parallel(self): in_size = 2 input_size = 4