Metal backend: ops read a view with a storage offset before the GPU has written it (silently wrong results)

Start with backends/apple/metal/runtime/shims/memory.cpp and trace a_reinterpret_tensor handling for non-zero storage offsets and packed views. Run backends/apple/metal/tests/run_metal_test.sh --build, then the unittest command for TestMetalBackendModules -k linear_chunk. Done means both named tests produce correct outputs for float32 and bfloat16 rather than consuming zeros.

🐛 Describe the bug

On the Metal backend, an op that consumes a view starting partway into another op's output reads that memory before the GPU has written it. The result is silently wrong (zeros in the minimal cases below); nothing fails. chunk, split and unbind all produce such views, so this shows up in real models: the YOLO family exported for Metal runs but returns box coordinates that are off by hundreds.

Repro. Two modules for MODULE_REGISTRY in backends/apple/metal/tests/test_modules.py. Both feed the second chunk of a linear's output into another linear:

class LinearChunkLastDim(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(7, 16, bias=False)
        self.linear2 = nn.Linear(8, 5, bias=False)

    def forward(self, x):
        _, second = self.linear1(x).chunk(2, dim=-1)
        return self.linear2(second)


class LinearChunkFirstDim(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(7, 16, bias=False)
        self.linear2 = nn.Linear(16, 5, bias=False)

    def forward(self, x):
        _, second = self.linear1(x).chunk(2, dim=0)
        return self.linear2(second)

with "input_shapes": [(12, 7)].

backends/apple/metal/tests/run_metal_test.sh --build
python -m unittest backends.apple.metal.tests.test_modules.TestMetalBackendModules -k linear_chunk

linear_chunk_last_dim (float32): Output mismatch - max_atol=0.24710503220558167, max_rtol=1.0
linear_chunk_first_dim (float32): Output mismatch - max_atol=0.857249915599823, max_rtol=1.0

max_rtol=1.0 in all four cases (float32 and bfloat16): the consumer saw zeros. Using the first chunk instead (offset 0) gives the right answer.

Cause. Inductor hands the view to the second mm as reinterpret_tensor_wrapper(buf0, ..., storage_offset), and aoti_torch__reinterpret_tensor in backends/apple/metal/runtime/shims/memory.cpp has two paths, both of which miss the pending write to buf0:

1. Last-dim chunk: sizes [12, 8], strides [16, 1], offset 8. Not packed, so materialize_packed copies it on the CPU. It does wait for the GPU first, but only under

   if (metal_is_device_pointer(src)) {
     stream->synchronize(SyncType::COMMIT_AND_WAIT);
   }
   

   and src is data_ptr + storage_offset * element_size. Only a buffer's base address is a key in ptr_to_mtl_buffer, so for any non-zero offset the check is false, the wait is skipped, and the copy reads memory the first mm has only been encoded to write.

2. First-dim chunk: sizes [6, 16], strides [16, 1], offset 96. Packed, so the view gets its own buffer through metal_buffer_nocopy(adjusted_data, ...). That is a second MTLBuffer over the parent's memory, and Metal tracks hazards per buffer object, so nothing orders the work reading the alias after the pending work writing the parent. Waiting for the stream before creating the alias makes the result correct, which is what points at ordering rather than addressing (a no-copy buffer over the unaligned pointer does read the right bytes once the memory is settled).

I have a fix with these two tests and will open a PR.

Versions

ExecuTorch: main @ 11120c8dff (also reproduced with the 1.5.0 release: PyPI wheel for export, runtime built from the v1.5.0 tag)
PyTorch version: 2.14.0
torchao: built from source with TORCHAO_BUILD_EXPERIMENTAL_MPS=1
OS: macOS 27.0 (arm64), Apple M2 Pro
Clang version: 21.0.0 (clang-2100.1.1.101)
CMake version: 4.4.3
Python version: 3.10.11