Metal backend: free only the destroyed delegate's memory

backends/apple/metal/CMakeLists.txt

@@ -140,6 +140,20 @@ install(
   DESTINATION lib
 )
 
+# Runs two programs in one process for tests/test_modules.py. Only built on
+# request, by tests/run_metal_test.sh --build.
+if(EXECUTORCH_BUILD_EXECUTOR_RUNNER AND EXECUTORCH_BUILD_EXTENSION_MODULE)
+  add_executable(
+    metal_multi_delegate_runner EXCLUDE_FROM_ALL
+    runtime/test/multi_delegate_runner.cpp
+  )
+  target_link_libraries(
+    metal_multi_delegate_runner
+    PRIVATE extension_module_static extension_tensor portable_ops_lib
+            metal_backend
+  )
+endif()
+
 if(BUILD_TESTING)
   include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)
 

backends/apple/metal/runtime/metal_backend.cpp

@@ -363,6 +363,18 @@ class ET_EXPERIMENTAL MetalBackend final
     handle->so_handle = so_handle;
     handle->so_path = so_path.string();
 
+    // The shims record what this delegate creates while it loads, so that
+    // destroy() frees only this delegate's memory. If loading fails, what it
+    // made so far is freed here, since destroy() will not be called.
+    metal_begin_delegate(handle);
+    MetalDelegateScope delegate_scope(handle, /*loading=*/true);
+    bool initialized = false;
+    executorch::backends::aoti::ScopeGuard release_on_failure([&]() noexcept {
+      if (!initialized) {
+        metal_release_delegate(handle);
+      }
+    });
+
     // Load function pointers specific to this handle's shared library
     ET_CHECK_OK_OR_RETURN_ERROR(
         load_function_pointers_into_handle(so_handle, handle));
@@ -393,6 +405,7 @@ class ET_EXPERIMENTAL MetalBackend final
     }
 
     ET_LOG(Info, "MetalBackend::init - Initialization completed successfully");
+    initialized = true;
 
 #ifdef EXECUTORCH_METAL_COLLECT_STATS
     // Accumulate init timing statistics
@@ -447,6 +460,9 @@ class ET_EXPERIMENTAL MetalBackend final
     });
 
     AOTIDelegateHandle* handle = (AOTIDelegateHandle*)handle_;
+    // Handles created by this run, e.g. folded constants that outlive it,
+    // belong to this delegate.
+    MetalDelegateScope delegate_scope(handle, /*loading=*/false);
 
     ET_LOG(Debug, "MetalBackend Handle generated");
 
@@ -679,9 +695,14 @@ class ET_EXPERIMENTAL MetalBackend final
       }
     }
 
+    // Free this delegate's handles and constants. The shim state, including
+    // the buffer pool its recycled buffers went to, is shared by every
+    // delegate in the process, so it is only reset once none is left.
+    if (metal_release_delegate(handle)) {
+      cleanup_memory();
+      executorch::backends::aoti::cleanup_tensor_metadata();
+    }
     delete handle;
-    cleanup_memory();
-    executorch::backends::aoti::cleanup_tensor_metadata();
     ET_LOG(Debug, "MetalBackend handle %p destroy", handle_);
   }
 };

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,40 @@
 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;
+  // An empty tensor reads and writes nothing, so it needs no buffer of its
+  // own at its offset, and Metal makes none of length 0.
+  if (offset == 0 || tensor->nbytes() == 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

@@ -149,9 +149,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) {
@@ -278,7 +279,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

@@ -503,8 +503,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);
@@ -524,7 +525,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
@@ -558,7 +559,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]);
@@ -650,6 +651,8 @@ AOTITorchError aoti_torch_mps_convolution(
 
       if (actual_numel != expected_numel) {
         ET_LOG(Error, "aoti_torch_mps_convolution: Tensor size mismatch. Expected %zu, got %zu", expected_numel, actual_numel);
+        // The handle goes first: it does not own the memory yet.
+        aoti_torch_delete_tensor_object(output_tensor_handle);
         aoti_torch_mps_free(tensor_data);  // Free the allocated GPU memory on failure
         throw std::runtime_error("Tensor size mismatch");
       }

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_sdpa.mm

@@ -336,6 +336,12 @@ static AOTITorchError sdpa_mps_impl(
     return Error::Internal;
   }
 
+  // Freed on an error until a handle owns them; then the handle is deleted
+  // instead, which frees the memory.
+  void* out_contents_ptr = nullptr;
+  void* attn_contents_ptr = nullptr;
+  AOTITensorHandle out_tensor_handle = nullptr;
+
   try {
     @autoreleasepool {
       // Convert AOTITensorHandle to ExecuTorch tensors
@@ -454,10 +460,8 @@ static AOTITorchError sdpa_mps_impl(
       size_t out_size_bytes = batchSize * num_heads * qSize * headSize * element_size;
       size_t attn_size_bytes = batchSize * num_heads * qSize * kvSeqLength * element_size;
 
-      void* out_contents_ptr = nullptr;
       allocate_mtl_buffer(&out_contents_ptr, out_size_bytes);
 
-      void* attn_contents_ptr = nullptr;
       allocate_mtl_buffer(&attn_contents_ptr, attn_size_bytes);
 
       // Use MLX-style Metal kernels instead of MPSGraph
@@ -499,7 +503,6 @@ static AOTITorchError sdpa_mps_impl(
       }
 
       // Create output tensor handle first so we can use it in the kernel
-      AOTITensorHandle out_tensor_handle = nullptr;
       AOTITorchError create_out_result = aoti_torch_create_tensor_from_blob_v2(
           out_contents_ptr,
           4,  // ndim
@@ -517,14 +520,13 @@ static AOTITorchError sdpa_mps_impl(
 
       if (create_out_result != Error::Ok || !out_tensor_handle) {
         ET_LOG(Error, "aoti_torch_mps__scaled_dot_product_attention_math_for_mps: Failed to create output tensor");
-        aoti_torch_mps_free(out_contents_ptr);
-        aoti_torch_mps_free(attn_contents_ptr);
         throw std::runtime_error("Failed to create output tensor");
       }
 
       // Mark that we own the memory
       extern std::unordered_map<void*, int32_t> memory_to_n_tensor;
       memory_to_n_tensor[out_contents_ptr] = 1;
+      out_contents_ptr = nullptr;
 
       auto* out_tensor = reinterpret_cast<Tensor*>(out_tensor_handle);
 
@@ -680,11 +682,11 @@ static AOTITorchError sdpa_mps_impl(
 
       if (create_attn_result != Error::Ok || !attn_tensor_handle) {
         ET_LOG(Error, "aoti_torch_mps__scaled_dot_product_attention_math_for_mps: Failed to create attention weights tensor");
-        aoti_torch_mps_free(attn_contents_ptr);
         throw std::runtime_error("Failed to create attention weights tensor");
       }
 
       memory_to_n_tensor[attn_contents_ptr] = 1;
+      attn_contents_ptr = nullptr;
 
       // Set output tensor handles
       *ret0 = out_tensor_handle;
@@ -699,9 +701,15 @@ static AOTITorchError sdpa_mps_impl(
 
   } catch (const std::exception& e) {
     ET_LOG(Error, "aoti_torch_mps__scaled_dot_product_attention_math_for_mps exception: %s", e.what());
+    if (out_tensor_handle) aoti_torch_delete_tensor_object(out_tensor_handle);
+    if (out_contents_ptr) aoti_torch_mps_free(out_contents_ptr);
+    if (attn_contents_ptr) aoti_torch_mps_free(attn_contents_ptr);
     return Error::Internal;
   } catch (...) {
     ET_LOG(Error, "aoti_torch_mps__scaled_dot_product_attention_math_for_mps: unknown exception");
+    if (out_tensor_handle) aoti_torch_delete_tensor_object(out_tensor_handle);
+    if (out_contents_ptr) aoti_torch_mps_free(out_contents_ptr);
+    if (attn_contents_ptr) aoti_torch_mps_free(attn_contents_ptr);
     return Error::Internal;
   }
 }