Serialize only a view constant's own bytes in XNNPACK and Vulkan

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:

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(),

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(

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()

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(

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