diff --git a/backends/apple/metal/runtime/ops/common.h b/backends/apple/metal/runtime/ops/common.h index 8f808088454..08050df1960 100644 --- a/backends/apple/metal/runtime/ops/common.h +++ b/backends/apple/metal/runtime/ops/common.h @@ -100,8 +100,14 @@ extern std::unordered_map 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 diff --git a/backends/apple/metal/runtime/ops/common.mm b/backends/apple/metal/runtime/ops/common.mm index e030f750f79..2794e0abf16 100644 --- a/backends/apple/metal/runtime/ops/common.mm +++ b/backends/apple/metal/runtime/ops/common.mm @@ -15,14 +15,38 @@ std::unordered_map graph_cache; CacheStats cache_stats; -id get_mtl_buffer(Tensor* tensor, const char* op_name, const char* tensor_name) { +id 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 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 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 allocate_mtl_buffer(void** data_ptr, size_t size_bytes) { diff --git a/backends/apple/metal/runtime/ops/op_addmm.mm b/backends/apple/metal/runtime/ops/op_addmm.mm index 25410e047f1..8e9ebceaa17 100644 --- a/backends/apple/metal/runtime/ops/op_addmm.mm +++ b/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 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 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 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 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, diff --git a/backends/apple/metal/runtime/ops/op_bmm.mm b/backends/apple/metal/runtime/ops/op_bmm.mm index fd354d2c074..bcb55e4a052 100644 --- a/backends/apple/metal/runtime/ops/op_bmm.mm +++ b/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 self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_bmm_out", "self"); - id mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_bmm_out", "mat2"); - id out_buffer = get_mtl_buffer(out_tensor, "aoti_torch_mps_bmm_out", "out"); + bool settle_aliases = false; + id self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_bmm_out", "self", &settle_aliases); + id mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_bmm_out", "mat2", &settle_aliases); + id 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]); diff --git a/backends/apple/metal/runtime/ops/op_convolution.mm b/backends/apple/metal/runtime/ops/op_convolution.mm index 400787bed4b..4fd836d8c3f 100644 --- a/backends/apple/metal/runtime/ops/op_convolution.mm +++ b/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 input_buffer = get_mtl_buffer(input_tensor, "aoti_torch_mps_convolution", "input"); - id weight_buffer = get_mtl_buffer(weight_tensor, "aoti_torch_mps_convolution", "weight"); + bool settle_aliases = false; + id input_buffer = get_mtl_buffer(input_tensor, "aoti_torch_mps_convolution", "input", &settle_aliases); + id 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 bias_buffer = get_mtl_buffer(bias_tensor, "aoti_torch_mps_convolution", "bias"); + id bias_buffer = get_mtl_buffer(bias_tensor, "aoti_torch_mps_convolution", "bias", &settle_aliases); NSArray* 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]); diff --git a/backends/apple/metal/runtime/ops/op_mm.mm b/backends/apple/metal/runtime/ops/op_mm.mm index 1dab8af7461..87ecd2bd79c 100644 --- a/backends/apple/metal/runtime/ops/op_mm.mm +++ b/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 self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_mm_out", "self"); - id mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_mm_out", "mat2"); - id out_buffer = get_mtl_buffer(out_tensor, "aoti_torch_mps_mm_out", "out"); + bool settle_aliases = false; + id self_buffer = get_mtl_buffer(self_tensor, "aoti_torch_mps_mm_out", "self", &settle_aliases); + id mat2_buffer = get_mtl_buffer(mat2_tensor, "aoti_torch_mps_mm_out", "mat2", &settle_aliases); + id 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]); diff --git a/backends/apple/metal/runtime/ops/op_topk.mm b/backends/apple/metal/runtime/ops/op_topk.mm index 8d1b6722466..089c5b6adf2 100644 --- a/backends/apple/metal/runtime/ops/op_topk.mm +++ b/backends/apple/metal/runtime/ops/op_topk.mm @@ -128,7 +128,8 @@ AOTITorchError aoti_torch_mps_topk( stream->endKernelCoalescing(); - id self_buffer = get_mtl_buffer(self_tensor, "topk", "self"); + bool settle_aliases = false; + id self_buffer = get_mtl_buffer(self_tensor, "topk", "self", &settle_aliases); id values_buffer = ptr_to_mtl_buffer[values_ptr]; id 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]; diff --git a/backends/apple/metal/runtime/shims/et_metal.h b/backends/apple/metal/runtime/shims/et_metal.h index 2aabc9bb9df..66000b72ced 100644 --- a/backends/apple/metal/runtime/shims/et_metal.h +++ b/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 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 diff --git a/backends/apple/metal/runtime/shims/et_metal.mm b/backends/apple/metal/runtime/shims/et_metal.mm index afd75108b1a..9e02edcacfc 100644 --- a/backends/apple/metal/runtime/shims/et_metal.mm +++ b/backends/apple/metal/runtime/shims/et_metal.mm @@ -81,6 +81,35 @@ void dispatch_sync_with_rethrow(dispatch_queue_t queue, void (^block)()) { // Global Metal buffer mapping - accessible for MPS shim std::unordered_map> ptr_to_mtl_buffer; +namespace { +// A view's address mapped to the address of the buffer it lives in, counted +// because several tensors can be views of the same address. A view of CPU +// memory has a no-copy buffer of its own, mapped at its own address, which +// goes away with the view's last handle. +struct MetalView { + void* base; + int32_t count; + bool owns_buffer = false; +}; +std::unordered_map ptr_to_view; +} // namespace + +bool metal_resolve_buffer(void* ptr, id* buffer, size_t* offset) { + void* base = ptr; + auto view = ptr_to_view.find(ptr); + if (view != ptr_to_view.end()) { + base = view->second.base; + } + + auto it = ptr_to_mtl_buffer.find(base); + if (it == ptr_to_mtl_buffer.end()) { + return false; + } + *buffer = it->second; + *offset = static_cast(ptr) - static_cast(base); + return true; +} + // Metal buffer pool with best-fit matching and LRU eviction. // On free, buffers are recycled into a sorted pool. On alloc, the smallest // buffer >= requested size is returned (if within the headroom bound). When the @@ -266,6 +295,7 @@ void metal_cleanup_resources() { [pair.second release]; } ptr_to_mtl_buffer.clear(); + ptr_to_view.clear(); get_metal_buffer_pool().clear(); } @@ -287,8 +317,86 @@ bool metal_buffer_nocopy(void* ptr, size_t nbytes, bool map_ptr_to_buffer) { return true; } +bool metal_register_view(void* view_ptr, void* base_ptr) { + // A view of a view lives in the same buffer as its parent. + auto parent = ptr_to_view.find(base_ptr); + void* base = parent != ptr_to_view.end() ? parent->second.base : base_ptr; + if (ptr_to_mtl_buffer.find(base) == ptr_to_mtl_buffer.end()) { + ET_LOG(Error, "metal_register_view: %p is not inside a Metal buffer", base_ptr); + return false; + } + + auto it = ptr_to_view.find(view_ptr); + if (it == ptr_to_view.end()) { + ptr_to_view[view_ptr] = {base, 1}; + } else { + it->second.base = base; + it->second.count++; + } + return true; +} + +void metal_retain_view(void* view_ptr) { + auto it = ptr_to_view.find(view_ptr); + if (it != ptr_to_view.end()) { + it->second.count++; + } +} + +bool metal_register_cpu_view(void* view_ptr, size_t nbytes) { + auto it = ptr_to_view.find(view_ptr); + if (it != ptr_to_view.end()) { + if (!it->second.owns_buffer) { + ET_LOG(Error, "metal_register_cpu_view: %p is already a view of a Metal buffer", view_ptr); + return false; + } + // Another view at this address, which may reach further than the + // ones before it: the buffer has to cover the longest of them. + id current = ptr_to_mtl_buffer[view_ptr]; + if ([current length] < nbytes) { + if (!metal_buffer_nocopy(view_ptr, nbytes, true)) { + return false; + } + [current release]; + } + it->second.count++; + return true; + } + if (ptr_to_mtl_buffer.find(view_ptr) != ptr_to_mtl_buffer.end()) { + ET_LOG(Error, "metal_register_cpu_view: %p already has a Metal buffer", view_ptr); + return false; + } + if (!metal_buffer_nocopy(view_ptr, nbytes, true)) { + return false; + } + ptr_to_view[view_ptr] = {view_ptr, 1, true}; + return true; +} + +bool metal_is_cpu_view(void* ptr) { + auto it = ptr_to_view.find(ptr); + return it != ptr_to_view.end() && it->second.owns_buffer; +} + +void metal_unregister_view(void* view_ptr) { + auto it = ptr_to_view.find(view_ptr); + if (it == ptr_to_view.end() || --it->second.count > 0) { + return; + } + if (it->second.owns_buffer) { + auto buffer = ptr_to_mtl_buffer.find(view_ptr); + if (buffer != ptr_to_mtl_buffer.end()) { + [buffer->second release]; + ptr_to_mtl_buffer.erase(buffer); + } + } + ptr_to_view.erase(it); +} + bool metal_is_device_pointer(void* ptr) { - return ptr_to_mtl_buffer.find(ptr) != ptr_to_mtl_buffer.end(); + id buffer = nil; + size_t offset = 0; + return metal_resolve_buffer(ptr, &buffer, &offset); } int metal_copy_memory(void* dst, const void* src, size_t nbytes, bool src_is_device, bool dst_is_device) { @@ -300,16 +408,13 @@ int metal_copy_memory(void* dst, const void* src, size_t nbytes, bool src_is_dev @autoreleasepool { // Case 1: Device-to-device copy - use GPU blit encoder (most efficient) if (src_is_device && dst_is_device) { - auto src_it = ptr_to_mtl_buffer.find(const_cast(src)); - auto dst_it = ptr_to_mtl_buffer.find(dst); - - if (src_it != ptr_to_mtl_buffer.end() && dst_it != ptr_to_mtl_buffer.end()) { - id srcBuffer = src_it->second; - id dstBuffer = dst_it->second; + id srcBuffer = nil; + id dstBuffer = nil; + size_t srcOffset = 0; + size_t dstOffset = 0; - // Calculate offsets relative to buffer base - size_t srcOffset = static_cast(src) - static_cast([srcBuffer contents]); - size_t dstOffset = static_cast(dst) - static_cast([dstBuffer contents]); + if (metal_resolve_buffer(const_cast(src), &srcBuffer, &srcOffset) && + metal_resolve_buffer(dst, &dstBuffer, &dstOffset)) { // Use Metal's blit encoder for GPU-accelerated copy ETMetalStream* stream = getCurrentMetalStream(); @@ -324,16 +429,15 @@ int metal_copy_memory(void* dst, const void* src, size_t nbytes, bool src_is_dev } // Case 2: Host-to-device or device-to-host - use memcpy with shared memory - // Since Metal uses shared storage mode, CPU and GPU access the same memory - std::memcpy(dst, src, nbytes); - - // Synchronize only if we need to ensure GPU operations complete before CPU reads - // (device-to-host case where GPU may have written data) - if (src_is_device && !dst_is_device) { - // Ensure any pending GPU writes to source complete before CPU reads + // Since Metal uses shared storage mode, CPU and GPU access the same memory. + // What the GPU still has to do with it must be done before the CPU + // touches it: writes to a device source, and reads or writes of a device + // destination. + if (src_is_device || dst_is_device) { ETMetalStream* stream = getCurrentMetalStream(); stream->synchronize(SyncType::COMMIT_AND_WAIT); } + std::memcpy(dst, src, nbytes); ET_LOG(Debug, "Metal memory copy (memcpy): %zu bytes, src_device=%d, dst_device=%d", nbytes, src_is_device, dst_is_device); @@ -514,11 +618,11 @@ int metal_copy_memory(void* dst, const void* src, size_t nbytes, bool src_is_dev void* data_ptr = tensor.mutable_data_ptr(); size_t totalSize = tensor.numel() * tensor.element_size(); - auto it = ptr_to_mtl_buffer.find(data_ptr); - if (it != ptr_to_mtl_buffer.end()) { - // Use existing Metal buffer - id mtlBuffer = it->second; - [encoder_ setBuffer:mtlBuffer offset:0 atIndex:idx]; + id mtlBuffer = nil; + size_t bufferOffset = 0; + if (metal_resolve_buffer(data_ptr, &mtlBuffer, &bufferOffset)) { + // Use existing Metal buffer; a view binds its parent at an offset + [encoder_ setBuffer:mtlBuffer offset:bufferOffset atIndex:idx]; ET_LOG(Debug, "ETMetalKernelFunction::setArg: Set Metal buffer at index %u (size: %zu)", idx, totalSize); } else { // Handle CPU tensor data @@ -1172,10 +1276,22 @@ static int getDefaultFlushInterval(MTLDevice_t device, const char** outArch) { return !commandBuffer_ && !commandEncoder_; } -void ETMetalStream::executeMPSGraph(MPSGraph* mpsGraph, NSDictionary* feeds, NSDictionary* results, SyncType syncType) { +void ETMetalStream::executeMPSGraph( + MPSGraph* mpsGraph, + NSDictionary* feeds, + NSDictionary* results, + SyncType syncType, + bool settle_aliases) { // Use dispatch_sync_with_rethrow exactly like PyTorch does for MPSGraph execution dispatch_sync_with_rethrow(serialQueue_, ^() { @autoreleasepool { + // An alias (see get_mtl_buffer) is a separate MTLBuffer over memory + // another buffer covers, and Metal orders nothing between the two. + // Settle that memory on both sides of this graph, all within this + // block so that no other work on the stream can come in between. + if (settle_aliases) { + synchronize(SyncType::COMMIT_AND_WAIT); + } endKernelCoalescing(); [mpsGraph encodeToCommandBuffer:commandBuffer() @@ -1183,6 +1299,10 @@ static int getDefaultFlushInterval(MTLDevice_t device, const char** outArch) { targetOperations:nil resultsDictionary:results executionDescriptor:nil]; + + if (settle_aliases) { + synchronize(SyncType::COMMIT_AND_WAIT); + } } }); } diff --git a/backends/apple/metal/runtime/shims/memory.cpp b/backends/apple/metal/runtime/shims/memory.cpp index c6a9b292309..384b5e2c26f 100644 --- a/backends/apple/metal/runtime/shims/memory.cpp +++ b/backends/apple/metal/runtime/shims/memory.cpp @@ -41,8 +41,39 @@ std::unordered_map> tensors; constexpr int32_t NOT_OWN = -1; std::unordered_map memory_to_n_tensor; +// Every handle into memory the runtime owns holds a count on that allocation +// in memory_to_n_tensor. A handle whose own address is not an allocation (a +// view at an offset, or a handle made from one) finds only that address when +// it is deleted, so this maps it to the allocation its count went to. +std::unordered_map view_owner; + namespace { +// The owned allocation that `handle`, with data at `data_ptr`, lives in, or +// null for memory the runtime does not own, such as a model's constants. +void* owning_allocation(Tensor* handle, void* data_ptr) { + auto owner = view_owner.find(handle); + if (owner != view_owner.end()) { + return owner->second; + } + auto memory = memory_to_n_tensor.find(data_ptr); + if (memory != memory_to_n_tensor.end() && memory->second != NOT_OWN) { + return data_ptr; + } + return nullptr; +} + +// Takes a count on `owner`, if any, for a new handle with data at `data_ptr`. +void hold_allocation(Tensor* handle, void* data_ptr, void* owner) { + if (owner == nullptr) { + return; + } + memory_to_n_tensor[owner] += 1; + if (data_ptr != owner) { + view_owner[handle] = owner; + } +} + // Wraps `data` in a tensor whose strides are the ones given. from_blob() does // not keep the strides it is handed: it sorts them into a dim order and derives // the strides again from that. A dimension of size 1 has the same stride as the @@ -249,6 +280,32 @@ AOTITorchError aoti_torch_empty_strided( return Error::Ok; } +// Drops one count on owned memory and frees it when none is left. +static AOTITorchError release_memory(void* data_ptr) { + auto memory_it = memory_to_n_tensor.find(data_ptr); + ET_CHECK_OR_RETURN_ERROR( + memory_it != memory_to_n_tensor.end() && memory_it->second > 0, + Internal, + "Internal error: releasing memory %p that is not owned", + data_ptr); + if (memory_it->second > 1) { + memory_it->second -= 1; + return Error::Ok; + } + if (metal_is_device_pointer(data_ptr)) { + metal_deallocate_buffer(data_ptr); + } else { + // Queued GPU work can still read or write this memory through the no-copy + // buffer of a view of it (metal_register_cpu_view). That buffer does not + // own the memory, so the work has to finish before it is freed. + getCurrentMetalStream()->synchronize(SyncType::COMMIT_AND_WAIT); + free(data_ptr); + ET_LOG(Debug, "aoti_torch_delete_tensor_object: freeing CPU memory"); + } + memory_to_n_tensor.erase(memory_it); + return Error::Ok; +} + AOTITorchError aoti_torch_delete_tensor_object(AOTITensorHandle tensor) { ET_LOG(Debug, "aoti_torch_delete_tensor_object: entered"); @@ -266,30 +323,28 @@ AOTITorchError aoti_torch_delete_tensor_object(AOTITensorHandle tensor) { void* data_ptr = tensor_ptr->mutable_data_ptr(); auto memory_it = memory_to_n_tensor.find(data_ptr); - if (memory_it != memory_to_n_tensor.end()) { - int32_t ref_count = memory_it->second; - - if (ref_count == NOT_OWN) { - tensors.erase(it); - ET_LOG( - Debug, - "aoti_torch_delete_tensor_object: tensor doesn't own memory, skipping free"); - return Error::Ok; - } else if (ref_count == 1) { - if (metal_is_device_pointer(data_ptr)) { - metal_deallocate_buffer(data_ptr); - } else { - free(data_ptr); - ET_LOG(Debug, "aoti_torch_delete_tensor_object: freeing CPU memory"); - } - memory_to_n_tensor.erase(memory_it); - } else if (ref_count > 1) { - memory_to_n_tensor[data_ptr] = ref_count - 1; + ET_CHECK_OR_RETURN_ERROR( + memory_it != memory_to_n_tensor.end(), + Internal, + "Internal error: memory not found during deletion"); + + if (memory_it->second == NOT_OWN) { + // No-op unless this tensor is a view. + metal_unregister_view(data_ptr); + // Give back the count the view held on the allocation it lives in. + auto owner = view_owner.find(tensor); + if (owner != view_owner.end()) { + void* allocation = owner->second; + view_owner.erase(owner); + ET_CHECK_OK_OR_RETURN_ERROR(release_memory(allocation)); } - } else { - ET_CHECK_OR_RETURN_ERROR( - false, Internal, "Internal error: memory not found during deletion"); + tensors.erase(it); + ET_LOG( + Debug, + "aoti_torch_delete_tensor_object: tensor doesn't own memory, skipping free"); + return Error::Ok; } + ET_CHECK_OK_OR_RETURN_ERROR(release_memory(data_ptr)); tensors.erase(it); ET_LOG(Debug, "aoti_torch_delete_tensor_object: successful"); @@ -448,12 +503,14 @@ static void* materialize_packed( if (!dst) return nullptr; - // Ensure pending GPU writes to the source buffer are complete - if (metal_is_device_pointer(src)) { - auto* stream = getCurrentMetalStream(); - if (stream) { - stream->synchronize(SyncType::COMMIT_AND_WAIT); - } + // The copy is made on the CPU, so what the GPU still has to write to the + // source must be there first. That holds for CPU memory too: kernels and + // graphs can write it through the no-copy buffer of a view of it + // (metal_register_cpu_view). The wait also settles any queued work still + // using the buffer `dst` was recycled from. + auto* stream = getCurrentMetalStream(); + if (stream) { + stream->synchronize(SyncType::COMMIT_AND_WAIT); } // Element-by-element strided copy @@ -614,20 +671,40 @@ AOTITorchError aoti_torch__reinterpret_tensor( element_size, adjusted_data); - ET_CHECK_OR_RETURN_ERROR( - metal_buffer_nocopy(adjusted_data, tensor->nbytes(), true), - Internal, - "metal_buffer_nocopy failed for adjusted_data=%p, nbytes=%zu", - adjusted_data, - static_cast(tensor->nbytes())); + if (metal_is_device_pointer(data_ptr) && !metal_is_cpu_view(data_ptr)) { + // The view shares its parent's Metal buffer and is bound at an + // offset. It must not get an MTLBuffer of its own: Metal would treat + // the two as unrelated, and inductor both reads views of a buffer + // another op is still writing and fills a buffer (e.g. the result of + // a cat) by writing through views of it. + ET_CHECK_OR_RETURN_ERROR( + metal_register_view(adjusted_data, data_ptr), + Internal, + "Failed to register adjusted_data=%p as a view of %p", + adjusted_data, + data_ptr); + } else { + // CPU memory has no Metal buffer to be bound into, so the view gets a + // no-copy one of its own, as before. + ET_CHECK_OR_RETURN_ERROR( + metal_register_cpu_view(adjusted_data, tensor->nbytes()), + Internal, + "Failed to wrap adjusted_data=%p, nbytes=%zu in a Metal buffer", + adjusted_data, + static_cast(tensor->nbytes())); + } memory_to_n_tensor[adjusted_data] = NOT_OWN; + } else { + // Another handle at the address of `self`. If `self` is a view, deleting + // either handle must leave the view registered for the other one. + metal_retain_view(data_ptr); } - // Increment the reference count for this memory address only if it is owned - if (memory_to_n_tensor[data_ptr] != NOT_OWN) { - memory_to_n_tensor[data_ptr] += 1; - } + // The new handle keeps the allocation it lives in alive, including when + // `self` is itself a view. + hold_allocation( + tensor.get(), adjusted_data, owning_allocation(self, data_ptr)); } ET_LOG(Debug, "aoti_torch__reinterpret_tensor: successful"); @@ -715,11 +792,13 @@ AOTITorchError aoti_torch_new_tensor_handle( *new_handle = tensor.get(); - // Increment the reference count for this memory address only if it is owned - // by tensor - memory_to_n_tensor[data_ptr] = memory_to_n_tensor[data_ptr] == NOT_OWN - ? NOT_OWN - : memory_to_n_tensor[data_ptr] + 1; + // If the original is a view into a Metal buffer, the new handle is one too, + // and deleting either must leave the view registered for the other one. + metal_retain_view(data_ptr); + + // The new handle keeps the allocation the original lives in alive. + hold_allocation( + tensor.get(), data_ptr, owning_allocation(orig_handle, data_ptr)); ET_LOG(Debug, "aoti_torch_new_tensor_handle: successful"); return Error::Ok; @@ -748,6 +827,7 @@ void cleanup_memory() { // anymore, and a stale entry would make the next model fail to load as soon // as its constants land on an address used before. memory_to_n_tensor.clear(); + view_owner.clear(); // Clean up Metal resources metal_cleanup_resources(); diff --git a/backends/apple/metal/runtime/test/test_memory.cpp b/backends/apple/metal/runtime/test/test_memory.cpp index 04f96d40dd5..c7167437d98 100644 --- a/backends/apple/metal/runtime/test/test_memory.cpp +++ b/backends/apple/metal/runtime/test/test_memory.cpp @@ -8,9 +8,11 @@ #include +#include #include #include +#include #include #include #include @@ -25,9 +27,16 @@ namespace { constexpr int32_t kFloat32 = 6; // DeviceType::MPS. constexpr int32_t kDeviceMps = 13; +// DeviceType::CPU. +constexpr int32_t kDeviceCpu = 0; } // namespace +extern "C" AOTITorchError aoti_torch_mps_mm_out( + AOTITensorHandle out, + AOTITensorHandle self, + AOTITensorHandle mat2); + class MetalMemoryTest : public ::testing::Test { protected: void SetUp() override { @@ -58,6 +67,57 @@ class MetalMemoryTest : public ::testing::Test { /*opaque_metadata_size=*/0); } + // Allocates an 8-element Metal buffer and a view of its last 4 elements. + void createBaseAndView(AOTITensorHandle* base, AOTITensorHandle* view) { + const int64_t base_size = 8; + ASSERT_EQ( + aoti_torch_empty_strided( + 1, &base_size, &kStride, kFloat32, kDeviceMps, 0, base), + Error::Ok); + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + *base, 1, &kViewSize, &kStride, /*storage_offset=*/4, view), + Error::Ok); + ASSERT_NE((*view)->mutable_data_ptr(), (*base)->mutable_data_ptr()); + ASSERT_TRUE(metal_is_device_pointer((*view)->mutable_data_ptr())); + } + + // A second handle at the address of `view`, made the way inductor's wrapper + // makes one: by copying the handle, or by reinterpreting at offset 0. + Error createAlias( + AOTITensorHandle view, + bool by_reinterpret, + AOTITensorHandle* alias) { + if (by_reinterpret) { + return aoti_torch__reinterpret_tensor( + view, 1, &kViewSize, &kStride, /*storage_offset=*/0, alias); + } + return aoti_torch_new_tensor_handle(view, alias); + } + + // Deleting one of two handles to the same view must leave the view bound to + // its parent's Metal buffer for the other handle. + void expectViewOutlivesDeletedHandle(bool by_reinterpret, bool delete_view) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createBaseAndView(&base, &view); + AOTITensorHandle alias = nullptr; + ASSERT_EQ(createAlias(view, by_reinterpret, &alias), Error::Ok); + void* view_ptr = view->mutable_data_ptr(); + ASSERT_EQ(alias->mutable_data_ptr(), view_ptr); + + ASSERT_EQ( + aoti_torch_delete_tensor_object(delete_view ? view : alias), Error::Ok); + EXPECT_TRUE(metal_is_device_pointer(view_ptr)); + + ASSERT_EQ( + aoti_torch_delete_tensor_object(delete_view ? alias : view), Error::Ok); + EXPECT_FALSE(metal_is_device_pointer(view_ptr)); + } + + static constexpr int64_t kViewSize = 4; + static constexpr int64_t kStride = 1; + std::vector blob_ = std::vector(4, 1.0f); }; @@ -93,6 +153,357 @@ TEST_F(MetalMemoryTest, CleanupLeavesNoTrackedMemory) { EXPECT_TRUE(memory_to_n_tensor.empty()); } +TEST_F(MetalMemoryTest, ViewOutlivesDeletedOriginalOfCopiedHandle) { + expectViewOutlivesDeletedHandle( + /*by_reinterpret=*/false, /*delete_view=*/true); +} + +TEST_F(MetalMemoryTest, ViewOutlivesDeletedCopiedHandle) { + expectViewOutlivesDeletedHandle( + /*by_reinterpret=*/false, /*delete_view=*/false); +} + +TEST_F(MetalMemoryTest, ViewOutlivesDeletedOriginalOfSameAddressReinterpret) { + expectViewOutlivesDeletedHandle( + /*by_reinterpret=*/true, /*delete_view=*/true); +} + +TEST_F(MetalMemoryTest, ViewOutlivesDeletedSameAddressReinterpret) { + expectViewOutlivesDeletedHandle( + /*by_reinterpret=*/true, /*delete_view=*/false); +} + +// A view at another address counts towards its parent's memory, and the +// parent must not be freed under it. Once the last handle on that memory is +// gone, view or parent, the buffer has to be released. +TEST_F(MetalMemoryTest, ParentFreedAfterViewDeletedLast) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createBaseAndView(&base, &view); + void* base_ptr = base->mutable_data_ptr(); + + ASSERT_EQ(aoti_torch_delete_tensor_object(base), Error::Ok); + EXPECT_TRUE(metal_is_device_pointer(base_ptr)); + EXPECT_TRUE(metal_is_device_pointer(view->mutable_data_ptr())); + + ASSERT_EQ(aoti_torch_delete_tensor_object(view), Error::Ok); + EXPECT_FALSE(metal_is_device_pointer(base_ptr)); + EXPECT_EQ(memory_to_n_tensor.count(base_ptr), 0u); +} + +TEST_F(MetalMemoryTest, ParentFreedAfterParentDeletedLast) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createBaseAndView(&base, &view); + void* base_ptr = base->mutable_data_ptr(); + + ASSERT_EQ(aoti_torch_delete_tensor_object(view), Error::Ok); + EXPECT_TRUE(metal_is_device_pointer(base_ptr)); + + ASSERT_EQ(aoti_torch_delete_tensor_object(base), Error::Ok); + EXPECT_FALSE(metal_is_device_pointer(base_ptr)); + EXPECT_EQ(memory_to_n_tensor.count(base_ptr), 0u); +} + +// A view of a view lives in the same allocation, and keeps it alive after the +// base and the first view are gone. +TEST_F(MetalMemoryTest, NestedViewKeepsParentAlive) { + for (int64_t nested_offset : {0, 2}) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createBaseAndView(&base, &view); + void* base_ptr = base->mutable_data_ptr(); + + const int64_t nested_size = 2; + AOTITensorHandle nested = nullptr; + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + view, 1, &nested_size, &kStride, nested_offset, &nested), + Error::Ok); + void* nested_ptr = nested->mutable_data_ptr(); + + ASSERT_EQ(aoti_torch_delete_tensor_object(base), Error::Ok); + ASSERT_EQ(aoti_torch_delete_tensor_object(view), Error::Ok); + EXPECT_TRUE(metal_is_device_pointer(base_ptr)) << nested_offset; + EXPECT_TRUE(metal_is_device_pointer(nested_ptr)) << nested_offset; + + ASSERT_EQ(aoti_torch_delete_tensor_object(nested), Error::Ok); + EXPECT_FALSE(metal_is_device_pointer(base_ptr)) << nested_offset; + EXPECT_EQ(memory_to_n_tensor.count(base_ptr), 0u) << nested_offset; + } +} + +// A handle copied from a view keeps the allocation alive after the base and +// the original view are gone. +TEST_F(MetalMemoryTest, CopiedViewHandleKeepsParentAlive) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createBaseAndView(&base, &view); + void* base_ptr = base->mutable_data_ptr(); + AOTITensorHandle alias = nullptr; + ASSERT_EQ(aoti_torch_new_tensor_handle(view, &alias), Error::Ok); + + ASSERT_EQ(aoti_torch_delete_tensor_object(base), Error::Ok); + ASSERT_EQ(aoti_torch_delete_tensor_object(view), Error::Ok); + EXPECT_TRUE(metal_is_device_pointer(base_ptr)); + EXPECT_TRUE(metal_is_device_pointer(alias->mutable_data_ptr())); + + ASSERT_EQ(aoti_torch_delete_tensor_object(alias), Error::Ok); + EXPECT_FALSE(metal_is_device_pointer(base_ptr)); + EXPECT_EQ(memory_to_n_tensor.count(base_ptr), 0u); +} + +class MetalGraphViewTest : public MetalMemoryTest { + protected: + // An 8-element tensor holding 1..8 on `device_type`, and the 2x2 view of + // its last four elements. + void createOffsetMatrix( + int32_t device_type, + AOTITensorHandle* base, + AOTITensorHandle* view) { + const int64_t base_size = 8; + ASSERT_EQ( + aoti_torch_empty_strided( + 1, &base_size, &kStride, kFloat32, device_type, 0, base), + Error::Ok); + auto* data = static_cast((*base)->mutable_data_ptr()); + for (int i = 0; i < 8; i++) { + data[i] = static_cast(i + 1); + } + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + *base, 2, kMatrixSizes, kMatrixStrides, /*storage_offset=*/4, view), + Error::Ok); + } + + // A 2x2 Metal tensor holding the identity matrix. + void createIdentity(AOTITensorHandle* identity) { + ASSERT_EQ( + aoti_torch_empty_strided( + 2, kMatrixSizes, kMatrixStrides, kFloat32, kDeviceMps, 0, identity), + Error::Ok); + auto* data = static_cast((*identity)->mutable_data_ptr()); + std::fill_n(data, 4, 0.0f); + data[0] = data[3] = 1.0f; + } + + void createMatrix(int32_t device_type, AOTITensorHandle* matrix) { + ASSERT_EQ( + aoti_torch_empty_strided( + 2, kMatrixSizes, kMatrixStrides, kFloat32, device_type, 0, matrix), + Error::Ok); + } + + static constexpr int64_t kMatrixSizes[2] = {2, 2}; + static constexpr int64_t kMatrixStrides[2] = {2, 1}; +}; + +// A graph fed an offset view through an alias settles the stream itself: its +// result is there as soon as the op returns. +TEST_F(MetalGraphViewTest, AliasedGraphSettlesItsOwnWork) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createOffsetMatrix(kDeviceMps, &base, &view); + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + AOTITensorHandle out = nullptr; + createMatrix(kDeviceMps, &out); + + ASSERT_EQ(aoti_torch_mps_mm_out(out, view, identity), Error::Ok); + EXPECT_TRUE(getCurrentMetalStream()->isEmpty()); + const auto* got = static_cast(out->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{5, 6, 7, 8})); +} + +// An op that fails after taking an alias for one of its inputs must not leave +// a wait behind for the next, unrelated graph. +TEST_F(MetalGraphViewTest, FailedAliasedGraphLeavesNoWaitBehind) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createOffsetMatrix(kDeviceMps, &base, &view); + AOTITensorHandle unmapped = nullptr; + createMatrix(kDeviceCpu, &unmapped); + AOTITensorHandle out = nullptr; + createMatrix(kDeviceMps, &out); + EXPECT_NE(aoti_torch_mps_mm_out(out, view, unmapped), Error::Ok); + + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + getCurrentMetalStream()->synchronize(SyncType::COMMIT_AND_WAIT); + ASSERT_EQ(aoti_torch_mps_mm_out(out, identity, identity), Error::Ok); + // Left pending on the stream, as a graph fed no alias always is. + EXPECT_FALSE(getCurrentMetalStream()->isEmpty()); + getCurrentMetalStream()->synchronize(SyncType::COMMIT_AND_WAIT); +} + +// A view of CPU memory gets a no-copy Metal buffer of its own, which MPSGraph +// ops can read. +TEST_F(MetalGraphViewTest, CpuBackedOffsetViewIsUsableByMm) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createOffsetMatrix(kDeviceCpu, &base, &view); + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + AOTITensorHandle out = nullptr; + createMatrix(kDeviceMps, &out); + + ASSERT_EQ(aoti_torch_mps_mm_out(out, view, identity), Error::Ok); + getCurrentMetalStream()->synchronize(SyncType::COMMIT_AND_WAIT); + const auto* got = static_cast(out->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{5, 6, 7, 8})); +} + +// The Metal buffer of a view of CPU memory stays while any handle to the view +// does, and goes with the last one, before the CPU memory can be freed. +TEST_F(MetalGraphViewTest, CpuBackedViewBufferGoesWithLastHandle) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + createOffsetMatrix(kDeviceCpu, &base, &view); + void* view_ptr = view->mutable_data_ptr(); + AOTITensorHandle alias = nullptr; + ASSERT_EQ(aoti_torch_new_tensor_handle(view, &alias), Error::Ok); + + ASSERT_EQ(aoti_torch_delete_tensor_object(view), Error::Ok); + EXPECT_TRUE(metal_is_device_pointer(view_ptr)); + ASSERT_EQ(aoti_torch_delete_tensor_object(alias), Error::Ok); + EXPECT_FALSE(metal_is_device_pointer(view_ptr)); +} + +// A second view of CPU memory at the same address can be longer than the +// first; the Metal buffer they share has to cover it. +TEST_F(MetalGraphViewTest, CpuBackedViewBufferCoversLongerViewAtSameAddress) { + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + const int64_t short_size = 2; + AOTITensorHandle short_view = nullptr; + const int64_t base_size = 8; + ASSERT_EQ( + aoti_torch_empty_strided( + 1, &base_size, &kStride, kFloat32, kDeviceCpu, 0, &base), + Error::Ok); + auto* data = static_cast(base->mutable_data_ptr()); + for (int i = 0; i < 8; i++) { + data[i] = static_cast(i + 1); + } + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + base, 1, &short_size, &kStride, /*storage_offset=*/4, &short_view), + Error::Ok); + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + base, 2, kMatrixSizes, kMatrixStrides, /*storage_offset=*/4, &view), + Error::Ok); + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + AOTITensorHandle out = nullptr; + createMatrix(kDeviceMps, &out); + + ASSERT_EQ(aoti_torch_mps_mm_out(out, view, identity), Error::Ok); + getCurrentMetalStream()->synchronize(SyncType::COMMIT_AND_WAIT); + const auto* got = static_cast(out->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{5, 6, 7, 8})); +} + +// A copy to the CPU sees what a graph still pending on the stream writes. +TEST_F(MetalGraphViewTest, CopyToHostWaitsForPendingWrites) { + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + AOTITensorHandle out = nullptr; + createMatrix(kDeviceMps, &out); + std::fill_n(static_cast(out->mutable_data_ptr()), 4, -1.0f); + AOTITensorHandle host = nullptr; + createMatrix(kDeviceCpu, &host); + + ASSERT_EQ(aoti_torch_mps_mm_out(out, identity, identity), Error::Ok); + ASSERT_EQ(aoti_torch_copy_(host, out, 0), Error::Ok); + const auto* got = static_cast(host->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{1, 0, 0, 1})); +} + +// A copy from the CPU does not overwrite what a graph still pending on the +// stream is to read. +TEST_F(MetalGraphViewTest, CopyToDeviceWaitsForPendingReads) { + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + AOTITensorHandle input = nullptr; + createMatrix(kDeviceMps, &input); + auto* input_data = static_cast(input->mutable_data_ptr()); + for (int i = 0; i < 4; i++) { + input_data[i] = static_cast(i + 1); + } + AOTITensorHandle out = nullptr; + createMatrix(kDeviceMps, &out); + AOTITensorHandle host = nullptr; + createMatrix(kDeviceCpu, &host); + std::fill_n(static_cast(host->mutable_data_ptr()), 4, 9.0f); + + ASSERT_EQ(aoti_torch_mps_mm_out(out, input, identity), Error::Ok); + ASSERT_EQ(aoti_torch_copy_(input, host, 0), Error::Ok); + getCurrentMetalStream()->synchronize(SyncType::COMMIT_AND_WAIT); + const auto* got = static_cast(out->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{1, 2, 3, 4})); +} + +// CPU memory under a view is not freed while queued GPU work still reads it +// through the view's no-copy buffer, which does not own the memory. +TEST_F(MetalGraphViewTest, QueuedCpuViewKeepsBackingStorageAlive) { + auto* stream = getCurrentMetalStream(); + AOTITensorHandle base = nullptr; + AOTITensorHandle view = nullptr; + AOTITensorHandle identity = nullptr; + AOTITensorHandle out = nullptr; + createOffsetMatrix(kDeviceCpu, &base, &view); + createIdentity(&identity); + createMatrix(kDeviceMps, &out); + + ASSERT_EQ(aoti_torch_mps_mm_out(out, view, identity), Error::Ok); + EXPECT_FALSE(stream->isEmpty()); + ASSERT_EQ(aoti_torch_delete_tensor_object(base), Error::Ok); + ASSERT_EQ(aoti_torch_delete_tensor_object(view), Error::Ok); + stream->synchronize(SyncType::COMMIT_AND_WAIT); + + const auto* got = static_cast(out->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{5, 6, 7, 8})); +} + +// A graph can write CPU memory through the no-copy buffer of a view of it. A +// view of that memory which is not densely packed is copied on the CPU, and +// the copy has to see the write. +TEST_F(MetalGraphViewTest, MaterializingCpuMemoryWaitsForGpuWrites) { + const int64_t base_size = 12; + AOTITensorHandle base = nullptr; + ASSERT_EQ( + aoti_torch_empty_strided( + 1, &base_size, &kStride, kFloat32, kDeviceCpu, 0, &base), + Error::Ok); + std::fill_n(static_cast(base->mutable_data_ptr()), 12, 0.0f); + AOTITensorHandle target = nullptr; + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + base, 2, kMatrixSizes, kMatrixStrides, /*storage_offset=*/4, &target), + Error::Ok); + AOTITensorHandle input = nullptr; + createMatrix(kDeviceMps, &input); + auto* input_data = static_cast(input->mutable_data_ptr()); + for (int i = 0; i < 4; i++) { + input_data[i] = static_cast(i + 1); + } + AOTITensorHandle identity = nullptr; + createIdentity(&identity); + ASSERT_EQ(aoti_torch_mps_mm_out(target, input, identity), Error::Ok); + EXPECT_FALSE(getCurrentMetalStream()->isEmpty()); + + // Elements 4, 5, 8 and 9: not densely packed, so copied on the CPU. + const int64_t strides[2] = {4, 1}; + AOTITensorHandle read = nullptr; + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + base, 2, kMatrixSizes, strides, /*storage_offset=*/4, &read), + Error::Ok); + const auto* got = static_cast(read->const_data_ptr()); + EXPECT_EQ(std::vector(got, got + 4), (std::vector{1, 2, 0, 0})); +} + class MetalStrideTest : public MetalMemoryTest { protected: AOTITensorHandle make( diff --git a/backends/apple/metal/tests/test_modules.py b/backends/apple/metal/tests/test_modules.py index 07fae650175..b8ba5b0402e 100644 --- a/backends/apple/metal/tests/test_modules.py +++ b/backends/apple/metal/tests/test_modules.py @@ -303,6 +303,124 @@ def forward(self, x: torch.Tensor): } +# ------------------------------------------------------------------------- +# Views with a storage offset. The chunks below are views into the first +# linear's output; the second chunk starts partway into that buffer. The cat +# variants also write through such views. +# ------------------------------------------------------------------------- +class LinearChunkLastDim(nn.Module): + """Chunking the last dim gives a non-packed view (its row stride is still + the parent's), which reinterpret_tensor has to materialize.""" + + 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) + + +MODULE_REGISTRY["linear_chunk_last_dim"] = { + "model_class": LinearChunkLastDim, + "input_shapes": [(12, 7)], + "description": "Linear on the second last-dim chunk of another linear's output", +} + + +# ------------------------------------------------------------------------- +class LinearChunkFirstDim(nn.Module): + """Chunking the first dim gives a packed view that only differs from its + parent by the storage offset.""" + + 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) + + +MODULE_REGISTRY["linear_chunk_first_dim"] = { + "model_class": LinearChunkFirstDim, + "input_shapes": [(12, 7)], + "description": "Linear on the second first-dim chunk of another linear's output", +} + + +# ------------------------------------------------------------------------- +class LinearChunkCatLastDim(nn.Module): + """Inductor builds the cat result by writing through views of it, then the + second linear reads the whole buffer.""" + + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(7, 16, bias=False) + self.linear2 = nn.Linear(24, 5, bias=False) + + def forward(self, x): + first, second = self.linear1(x).chunk(2, dim=-1) + return self.linear2( + torch.cat([first, second, torch.relu(second) * 2.0], dim=-1) + ) + + +MODULE_REGISTRY["linear_chunk_cat_last_dim"] = { + "model_class": LinearChunkCatLastDim, + "input_shapes": [(12, 7)], + "description": "Linear on a last-dim cat assembled from chunks of another linear's output", + # The slices of a last-dim cat are non-packed views, which + # aoti_torch__reinterpret_tensor materializes into a copy. Writes through + # them land in the copy and never reach the cat buffer. + "skip": "Writes through a non-packed view are lost (view is materialized)", +} + + +# ------------------------------------------------------------------------- +class LinearChunkCatFirstDim(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): + first, second = self.linear1(x).chunk(2, dim=0) + return self.linear2(torch.cat([first, second, torch.relu(second) * 2.0], dim=0)) + + +MODULE_REGISTRY["linear_chunk_cat_first_dim"] = { + "model_class": LinearChunkCatFirstDim, + "input_shapes": [(12, 7)], + "description": "Linear on a first-dim cat assembled from chunks of another linear's output", +} + + +# ------------------------------------------------------------------------- +class LinearNestedChunk(nn.Module): + """A view of a view: the last quarter of the first linear's output, taken + as the second chunk of its second chunk.""" + + 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) + _, last = second.chunk(2, dim=0) + return self.linear2(last) + + +MODULE_REGISTRY["linear_nested_chunk"] = { + "model_class": LinearNestedChunk, + "input_shapes": [(12, 7)], + "description": "Linear on a chunk of a chunk of another linear's output", +} + + # ------------------------------------------------------------------------- class LinearNoBiasInt4(nn.Module): def __init__(self):