Metal backend: bind offset views to their parent's buffer instead of aliasing it

backends/apple/metal/runtime/ops/common.h

@@ -100,8 +100,14 @@ extern std::unordered_map<GraphCacheKey, CachedGraph, GraphCacheKeyHash>
     graph_cache;
 extern CacheStats cache_stats;
 
-MTLBuffer_t
-get_mtl_buffer(Tensor* tensor, const char* op_name, const char* tensor_name);
+// The buffer to feed a graph for `tensor`. A tensor that starts partway into
+// its buffer gets an aliasing buffer, and `*settle_aliases` is then set: pass
+// it to the executeMPSGraph call that runs the graph this buffer feeds.
+MTLBuffer_t get_mtl_buffer(
+    Tensor* tensor,
+    const char* op_name,
+    const char* tensor_name,
+    bool* settle_aliases);
 MTLBuffer_t allocate_mtl_buffer(void** data_ptr, size_t size_bytes);
 
 } // namespace metal

backends/apple/metal/runtime/ops/common.mm

@@ -15,14 +15,38 @@
 std::unordered_map<GraphCacheKey, CachedGraph, GraphCacheKeyHash> graph_cache;
 CacheStats cache_stats;
 
-id<MTLBuffer> get_mtl_buffer(Tensor* tensor, const char* op_name, const char* tensor_name) {
+id<MTLBuffer> get_mtl_buffer(
+    Tensor* tensor,
+    const char* op_name,
+    const char* tensor_name,
+    bool* settle_aliases) {
   void* data_ptr = tensor->mutable_data_ptr();
-  auto it = ptr_to_mtl_buffer.find(data_ptr);
-  if (it == ptr_to_mtl_buffer.end()) {
+  id<MTLBuffer> buffer = nil;
+  size_t offset = 0;
+  if (!metal_resolve_buffer(data_ptr, &buffer, &offset)) {
     ET_LOG(Error, "%s: %s tensor not found in Metal buffer mapping", op_name, tensor_name);
     throw std::runtime_error(std::string(tensor_name) + " tensor not found in Metal buffer mapping");
   }
-  return it->second;
+  if (offset == 0) {
+    return buffer;
+  }
+
+  // The tensor is a view that starts partway into `buffer`. MPSGraphTensorData
+  // cannot address into a buffer, so the graph needs an MTLBuffer that begins at
+  // the view, over the same memory. Metal does not relate that alias to
+  // `buffer`, and work using one does not see pending work on the other, so the
+  // graph has to run with the memory settled on both sides of it. That is asked
+  // of the one graph this buffer is for, through executeMPSGraph.
+  id<MTLBuffer> alias = [get_metal_device() newBufferWithBytesNoCopy:data_ptr
+                                                              length:tensor->nbytes()
+                                                             options:MTLResourceStorageModeShared
+                                                         deallocator:nil];
+  if (!alias) {
+    ET_LOG(Error, "%s: failed to wrap the %s view in a Metal buffer", op_name, tensor_name);
+    throw std::runtime_error(std::string(tensor_name) + " view could not be wrapped in a Metal buffer");
+  }
+  *settle_aliases = true;
+  return [alias autorelease];
 }
 
 id<MTLBuffer> allocate_mtl_buffer(void** data_ptr, size_t size_bytes) {

backends/apple/metal/runtime/ops/op_addmm.mm

@@ -112,14 +112,15 @@ AOTITorchError aoti_torch_mps_addmm_out(
         throw std::runtime_error("Failed to get Metal device");
       }
 
+      bool settle_aliases = false;
       id<MTLBuffer> bias_buffer =
-          get_mtl_buffer(bias_tensor, "aoti_torch_mps_addmm_out", "self");
+          get_mtl_buffer(bias_tensor, "aoti_torch_mps_addmm_out", "self", &settle_aliases);
       id<MTLBuffer> mat1_buffer =
-          get_mtl_buffer(mat1_tensor, "aoti_torch_mps_addmm_out", "mat1");
+          get_mtl_buffer(mat1_tensor, "aoti_torch_mps_addmm_out", "mat1", &settle_aliases);
       id<MTLBuffer> mat2_buffer =
-          get_mtl_buffer(mat2_tensor, "aoti_torch_mps_addmm_out", "mat2");
+          get_mtl_buffer(mat2_tensor, "aoti_torch_mps_addmm_out", "mat2", &settle_aliases);
       id<MTLBuffer> out_buffer =
-          get_mtl_buffer(out_tensor, "aoti_torch_mps_addmm_out", "out");
+          get_mtl_buffer(out_tensor, "aoti_torch_mps_addmm_out", "out", &settle_aliases);
 
       stream->endKernelCoalescing();
 
@@ -272,7 +273,7 @@ AOTITorchError aoti_torch_mps_addmm_out(
       NSDictionary* results = @{addmmOutput : outputData};
 
       @try {
-        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT);
+        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT, settle_aliases);
       } @catch (NSException* exception) {
         ET_LOG(
             Error,

backends/apple/metal/runtime/ops/op_bmm.mm

@@ -148,9 +148,10 @@ AOTITorchError aoti_torch_mps_bmm_out(
       }
 
       // Get Metal buffers for input and output tensors
-      id<MTLBuffer> self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_bmm_out", "self");
-      id<MTLBuffer> mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_bmm_out", "mat2");
-      id<MTLBuffer> out_buffer = get_mtl_buffer(out_tensor, "aoti_torch_mps_bmm_out", "out");
+      bool settle_aliases = false;
+      id<MTLBuffer> self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_bmm_out", "self", &settle_aliases);
+      id<MTLBuffer> mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_bmm_out", "mat2", &settle_aliases);
+      id<MTLBuffer> out_buffer = get_mtl_buffer(out_tensor, "aoti_torch_mps_bmm_out", "out", &settle_aliases);
 
       // Validate buffers are non-null
       if (!self_buffer || !mat2_buffer || !out_buffer) {
@@ -277,7 +278,7 @@ AOTITorchError aoti_torch_mps_bmm_out(
 
       // Execute the batched matrix multiplication
       @try {
-        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT);
+        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT, settle_aliases);
       } @catch (NSException *exception) {
         ET_LOG(Error, "aoti_torch_mps_bmm_out: NSException caught during executeMPSGraph: %s - %s",
               [[exception name] UTF8String], [[exception reason] UTF8String]);

backends/apple/metal/runtime/ops/op_convolution.mm

@@ -368,8 +368,9 @@ AOTITorchError aoti_torch_mps_convolution(
       NSMutableDictionary* feeds = [NSMutableDictionary dictionary];
 
       // Get Metal buffers from tensors
-      id<MTLBuffer> input_buffer = get_mtl_buffer(input_tensor, "aoti_torch_mps_convolution", "input");
-      id<MTLBuffer> weight_buffer = get_mtl_buffer(weight_tensor, "aoti_torch_mps_convolution", "weight");
+      bool settle_aliases = false;
+      id<MTLBuffer> input_buffer = get_mtl_buffer(input_tensor, "aoti_torch_mps_convolution", "input", &settle_aliases);
+      id<MTLBuffer> weight_buffer = get_mtl_buffer(weight_tensor, "aoti_torch_mps_convolution", "weight", &settle_aliases);
 
       ET_LOG(Debug, "aoti_torch_mps_convolution: Using existing Metal buffers - input=%p, weight=%p",
               input_buffer, weight_buffer);
@@ -389,7 +390,7 @@ AOTITorchError aoti_torch_mps_convolution(
 
       // Add bias data to feeds if provided
       if (bias_tensor && biasPlaceholder) {
-        id<MTLBuffer> bias_buffer = get_mtl_buffer(bias_tensor, "aoti_torch_mps_convolution", "bias");
+        id<MTLBuffer> bias_buffer = get_mtl_buffer(bias_tensor, "aoti_torch_mps_convolution", "bias", &settle_aliases);
 
         NSArray<NSNumber*>* biasShape = @[@(C_out)];
         biasData = [[MPSGraphTensorData alloc] initWithMTLBuffer:bias_buffer
@@ -421,7 +422,7 @@ AOTITorchError aoti_torch_mps_convolution(
 
       @try {
         // Use stream helper to encode and synchronize correctly
-        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT);
+        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT, settle_aliases);
       } @catch (NSException *exception) {
         ET_LOG(Error, "aoti_torch_mps_convolution: NSException caught during executeMPSGraph: %s - %s",
               [[exception name] UTF8String], [[exception reason] UTF8String]);

backends/apple/metal/runtime/ops/op_mm.mm

@@ -97,9 +97,10 @@ AOTITorchError aoti_torch_mps_mm_out(
       }
 
       // Get Metal buffers for input and output tensors
-      id<MTLBuffer> self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_mm_out", "self");
-      id<MTLBuffer> mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_mm_out", "mat2");
-      id<MTLBuffer> out_buffer = get_mtl_buffer(out_tensor, "aoti_torch_mps_mm_out", "out");
+      bool settle_aliases = false;
+      id<MTLBuffer> self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_mm_out", "self", &settle_aliases);
+      id<MTLBuffer> mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_mm_out", "mat2", &settle_aliases);
+      id<MTLBuffer> out_buffer = get_mtl_buffer(out_tensor, "aoti_torch_mps_mm_out", "out", &settle_aliases);
 
       ET_LOG(Debug, "aoti_torch_mps_mm_out: Using existing Metal buffers - self=%p, mat2=%p, out=%p",
              self_buffer, mat2_buffer, out_buffer);
@@ -262,7 +263,7 @@ AOTITorchError aoti_torch_mps_mm_out(
 
       @try {
         // Use stream helper to encode and synchronize correctly
-        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT);
+        stream->executeMPSGraph(mpsGraph, feeds, results, SyncType::COMMIT, settle_aliases);
       } @catch (NSException *exception) {
         ET_LOG(Error, "aoti_torch_mps_mm_out: NSException caught during executeMPSGraph: %s - %s",
               [[exception name] UTF8String], [[exception reason] UTF8String]);

backends/apple/metal/runtime/ops/op_topk.mm

@@ -128,7 +128,8 @@ AOTITorchError aoti_torch_mps_topk(
 
       stream->endKernelCoalescing();
 
-      id<MTLBuffer> self_buffer = get_mtl_buffer(self_tensor, "topk", "self");
+      bool settle_aliases = false;
+      id<MTLBuffer> self_buffer = get_mtl_buffer(self_tensor, "topk", "self", &settle_aliases);
       id<MTLBuffer> values_buffer = ptr_to_mtl_buffer[values_ptr];
       id<MTLBuffer> indices_buffer = ptr_to_mtl_buffer[indices_ptr];
 
@@ -151,7 +152,7 @@ AOTITorchError aoti_torch_mps_topk(
         };
 
         @try {
-          stream->executeMPSGraph(cached.graph, feeds, results, SyncType::COMMIT);
+          stream->executeMPSGraph(cached.graph, feeds, results, SyncType::COMMIT, settle_aliases);
         } @catch (NSException* e) {
           ET_LOG(Error, "aoti_torch_mps_topk: ObjC exception: %s - %s",
                   e.name.UTF8String, e.reason.UTF8String);
@@ -218,7 +219,7 @@ AOTITorchError aoti_torch_mps_topk(
             indices_out: indicesData,
           };
 
-          stream->executeMPSGraph(graph, feeds, results, SyncType::COMMIT);
+          stream->executeMPSGraph(graph, feeds, results, SyncType::COMMIT, settle_aliases);
 
           [selfData release];
           [valuesData release];

backends/apple/metal/runtime/shims/et_metal.h

@@ -290,12 +290,15 @@ class ETMetalStream {
 
   void endKernelCoalescing();
 
-  // MPSGraph execution
+  // MPSGraph execution. `settle_aliases` is for a graph fed an aliasing buffer
+  // (see get_mtl_buffer): the stream then waits for the GPU both before and
+  // after encoding the graph, as one step.
   void executeMPSGraph(
       MPSGraph_t mpsGraph,
       NSDictionary_t feeds,
       NSDictionary_t results,
-      SyncType syncType = SyncType::COMMIT_ADAPTIVE);
+      SyncType syncType = SyncType::COMMIT_ADAPTIVE,
+      bool settle_aliases = false);
 
   // Command buffer lifecycle management
   void commitCommandBuffer(MTLCommandBuffer_t commandBuffer);
@@ -389,6 +392,25 @@ int metal_copy_memory(
 void metal_cleanup_resources();
 bool metal_buffer_nocopy(void* ptr, size_t nbytes, bool map_ptr_to_buffer);
 
+// Records that `view_ptr` points inside the Metal buffer that owns `base_ptr`,
+// so the view is bound as that buffer plus an offset. Giving a view its own
+// MTLBuffer over the same memory does not work: Metal treats the two buffers as
+// unrelated, and a write through one is not seen by a read of the other in the
+// same command buffer. Registrations are counted: every tensor handle at
+// `view_ptr` holds one, taken with metal_register_view when the view is created
+// or with metal_retain_view when another handle is made for the same address,
+// and gives it back with metal_unregister_view. metal_retain_view does nothing
+// for an address that is not a registered view.
+bool metal_register_view(void* view_ptr, void* base_ptr);
+void metal_retain_view(void* view_ptr);
+void metal_unregister_view(void* view_ptr);
+
+// A view of CPU memory that Metal kernels are to use gets a no-copy buffer of
+// its own, mapped at `view_ptr`. It is counted and released like a view of a
+// Metal buffer, and the buffer goes with its last handle.
+bool metal_register_cpu_view(void* view_ptr, size_t nbytes);
+bool metal_is_cpu_view(void* ptr);
+
 // Helper functions to access Metal objects
 MTLDevice_t get_metal_device();
 MTLCommandQueue_t get_metal_command_queue();
@@ -399,6 +421,11 @@ MTLCommandQueue_t get_metal_command_queue();
 // C++ only - expose the Metal buffer mapping
 #ifdef __OBJC__
 extern std::unordered_map<void*, MTLBuffer_t> ptr_to_mtl_buffer;
+
+// Finds the Metal buffer holding `ptr` and how far into it `ptr` is. Handles
+// both a buffer's own address and a registered view. Returns false for memory
+// Metal does not own.
+bool metal_resolve_buffer(void* ptr, MTLBuffer_t* buffer, size_t* offset);
 #endif
 
 #endif