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

[abdelaziz-mahdy]

🐛 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

[abdelaziz-mahdy]

Update after digging further: the second cause above is stronger than an ordering problem, and it also affects writes.

A view with a storage offset gets its own MTLBuffer from newBufferWithBytesNoCopy over the parent's memory. Standalone Metal program, no ExecuTorch involved: allocate a shared buffer, create a no-copy alias over its tail, then in one serial compute encoder dispatch a kernel that writes through the alias followed by a kernel that copies the parent out.

alias-write then parent-read in one serial encoder: 200 / 200 runs saw stale data
parent+offset write then parent-read in one serial encoder: 0 / 200 runs saw stale data

(macOS 27.0, M2 Pro. The only difference between the two lines is setBuffer:alias offset:0 versus setBuffer:parent offset:<bytes>.)

So the alias is unsound whenever both buffers are used in the same command buffer, not just when something is pending. Inductor relies on both directions:

• reads: aoti_torch_mps_mm_out(buf1, reinterpret_tensor_wrapper(buf0, ..., 96LL), ...) right after buf0 was written;

• writes: the result of a cat is allocated once and filled through views of it,

  auto buf7 = wrap_with_raii_handle_if_needed(reinterpret_tensor_wrapper(buf9, 4, int_array_10, int_array_14, 512LL));  // alias
  auto buf8 = wrap_with_raii_handle_if_needed(reinterpret_tensor_wrapper(buf9, 4, int_array_10, int_array_14, 1024LL));  // alias
  ...kernel writes buf7, buf8...
  ...next kernel reads buf9...

That second pattern is what breaks YOLO: C2f is cv2(cat([a, b, m(b)])). Bisecting yolov8n layer by layer, the first wrong output is the first C2f whose cat result feeds a conv; reducing it, the block is exact if it returns the cat and wrong as soon as anything reads it on the GPU.

#22957 now binds such views to the parent's buffer at an offset instead of aliasing. With it yolo11n, yolov8n and yolov5n match their XNNPACK exports to 0.002 (from 340 to 515 before).

One case is still open and is noted in the PR: a non-packed view (for example a slice of a last-dim cat) is materialized into a copy by reinterpret_tensor, so writes through it never reach the parent.

Standalone repro (clang++ -fobjc-arc -framework Metal -framework Foundation)

#import <Metal/Metal.h>
#import <Foundation/Foundation.h>
#include <cstdio>
// Write through a no-copy alias of a parent buffer's tail, then read the parent,
// all inside one command buffer, using compute kernels in ONE serial encoder
// (the way ExecuTorch's Metal stream coalesces generated kernels).
static const char* kSrc = R"(
#include <metal_stdlib>
using namespace metal;
kernel void fill(device float* out [[buffer(0)]], constant float& v [[buffer(1)]], uint i [[thread_position_in_grid]]) { out[i] = v + i; }
kernel void copy(device float* out [[buffer(0)]], device const float* in [[buffer(1)]], uint i [[thread_position_in_grid]]) { out[i] = in[i]; }
)";
int main() {
  @autoreleasepool {
    id<MTLDevice> dev = MTLCreateSystemDefaultDevice(); id<MTLCommandQueue> q = [dev newCommandQueue];
    NSError* err = nil; id<MTLLibrary> lib = [dev newLibraryWithSource:@(kSrc) options:nil error:&err];
    id<MTLComputePipelineState> fill = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"fill"] error:&err];
    id<MTLComputePipelineState> copy = [dev newComputePipelineStateWithFunction:[lib newFunctionWithName:@"copy"] error:&err];
    const int n = 1536, half = 512;
    int bad_runs = 0;
    for (int run = 0; run < 200; run++) {
      id<MTLBuffer> parent = [dev newBufferWithLength:n * 4 options:MTLResourceStorageModeShared];
      memset([parent contents], 0, n * 4);
      id<MTLBuffer> alias = [dev newBufferWithBytesNoCopy:(float*)[parent contents] + 2 * half length:half * 4
                                                  options:MTLResourceCPUCacheModeWriteCombined | MTLResourceStorageModeShared deallocator:nil];
      id<MTLBuffer> dst = [dev newBufferWithLength:n * 4 options:MTLResourceStorageModeShared];
      id<MTLCommandBuffer> cb = [q commandBuffer]; id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
      float v = 1000.0f;
      [enc setComputePipelineState:fill]; [enc setBuffer:alias offset:0 atIndex:0]; [enc setBytes:&v length:4 atIndex:1];
      [enc dispatchThreads:MTLSizeMake(half, 1, 1) threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
      [enc setComputePipelineState:copy]; [enc setBuffer:dst offset:0 atIndex:0]; [enc setBuffer:parent offset:0 atIndex:1];
      [enc dispatchThreads:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
      [enc endEncoding]; [cb commit]; [cb waitUntilCompleted];
      float* d = (float*)[dst contents]; int wrong = 0;
      for (int i = 0; i < half; i++) if (d[2 * half + i] != 1000.0f + i) wrong++;
      if (wrong) bad_runs++;
    }
    std::printf("alias-write then parent-read in one serial encoder: %d / 200 runs saw stale data\n", bad_runs);
  }
  return 0;
}

[powerofaisinstudy-debug]

Hi @abdelaziz-mahdy , Makes total sense. Without a proper storage offset pointer link or memory fence to the base MTLBuffer, view ops like chunk/split end up reading stale buffer regions before the parent kernel completes writes. Looking forward to seeing PR #22957 land!