From c3a54d6253d234d5396b9a703c137ec5a7191feb Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 29 Apr 2026 14:22:06 +0300 Subject: [PATCH] ggml-metal: implement async 2D tensor copy functions Add ggml_metal_set_tensor_2d_async and ggml_metal_get_tensor_2d_async functions to the Metal backend, mirroring the CUDA implementation. These functions use MTLBlitCommandEncoder's 2D copy API to efficiently copy tensor data with different strides between host and device memory, enabling row-strided transfers without requiring contiguous layouts. Assisted-by: llama.cpp:local pi --- ggml/src/ggml-metal/ggml-metal-context.h | 2 + ggml/src/ggml-metal/ggml-metal-context.m | 96 ++++++++++++++++++++++++ ggml/src/ggml-metal/ggml-metal.cpp | 18 ++++- 3 files changed, 114 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-context.h b/ggml/src/ggml-metal/ggml-metal-context.h index abf4b06ed2..0f33d3fa10 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.h +++ b/ggml/src/ggml-metal/ggml-metal-context.h @@ -21,6 +21,8 @@ void ggml_metal_synchronize(ggml_metal_t ctx); void ggml_metal_set_tensor_async(ggml_metal_t ctx, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); void ggml_metal_get_tensor_async(ggml_metal_t ctx, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); +void ggml_metal_set_tensor_2d_async(ggml_metal_t ctx, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); +void ggml_metal_get_tensor_2d_async(ggml_metal_t ctx, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); bool ggml_metal_cpy_tensor_async(ggml_metal_t ctx_src, ggml_metal_t ctx_dst, const struct ggml_tensor * src, struct ggml_tensor * dst); enum ggml_status ggml_metal_graph_compute (ggml_metal_t ctx, struct ggml_cgraph * gf); diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 32d97cd5d0..b5101ff718 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -392,6 +392,102 @@ void ggml_metal_get_tensor_async(ggml_metal_t ctx, const struct ggml_tensor * te } } +void ggml_metal_set_tensor_2d_async(ggml_metal_t ctx, struct ggml_tensor * tensor, const void * data, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + @autoreleasepool { + // wrap the source data into a Metal buffer + id device = ggml_metal_device_get_obj(ctx->dev); + id buf_src = [device newBufferWithBytes:data + length:stride_data * n_copies + options:MTLResourceStorageModeShared]; + + GGML_ASSERT(buf_src); + + struct ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(tensor); + if (bid_dst.metal == nil) { + GGML_ABORT("%s: failed to find buffer for tensor '%s'\n", __func__, tensor->name); + } + + bid_dst.offs += offset; + + // queue the copy operation into the queue of the Metal context + id queue = ggml_metal_device_get_queue(ctx->dev); + id cmd_buf = [queue commandBuffer]; + id encoder = [cmd_buf blitCommandEncoder]; + + [encoder copyFromBuffer:buf_src + sourceOffset:0 + sourceBytesPerRow:stride_data + sourceBytesPerImage:0 + sourceHeight:n_copies + sourceDepth:1 + toBuffer:bid_dst.metal + destinationOffset:bid_dst.offs + destinationBytesPerRow:stride_tensor + destinationBytesPerImage:0 + destinationHeight:n_copies + destinationDepth:1 + size:MTLSizeMake(size, 1, 1)]; + + [encoder endEncoding]; + [cmd_buf commit]; + [buf_src release]; + + [ctx->cmd_bufs_ext addObject:cmd_buf]; + ctx->cmd_buf_last = cmd_buf; + + [cmd_buf retain]; + } +} + +void ggml_metal_get_tensor_2d_async(ggml_metal_t ctx, const struct ggml_tensor * tensor, void * data, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + @autoreleasepool { + id device = ggml_metal_device_get_obj(ctx->dev); + id buf_dst = [device newBufferWithBytesNoCopy:data + length:stride_data * n_copies + options:MTLResourceStorageModeShared + deallocator:nil]; + + GGML_ASSERT(buf_dst); + + struct ggml_metal_buffer_id bid_src = ggml_metal_get_buffer_id(tensor); + if (bid_src.metal == nil) { + GGML_ABORT("%s: failed to find buffer for tensor '%s'\n", __func__, tensor->name); + } + + bid_src.offs += offset; + + // queue the copy operation into the queue of the Metal context + id queue = ggml_metal_device_get_queue(ctx->dev); + id cmd_buf = [queue commandBuffer]; + id encoder = [cmd_buf blitCommandEncoder]; + + [encoder copyFromBuffer:bid_src.metal + sourceOffset:bid_src.offs + sourceBytesPerRow:stride_tensor + sourceBytesPerImage:0 + sourceHeight:n_copies + sourceDepth:1 + toBuffer:buf_dst + destinationOffset:0 + destinationBytesPerRow:stride_data + destinationBytesPerImage:0 + destinationHeight:n_copies + destinationDepth:1 + size:MTLSizeMake(size, 1, 1)]; + + [encoder endEncoding]; + [cmd_buf commit]; + [buf_dst release]; + + [ctx->cmd_bufs_ext addObject:cmd_buf]; + ctx->cmd_buf_last = cmd_buf; + + [cmd_buf retain]; + } +} + bool ggml_metal_cpy_tensor_async(ggml_metal_t ctx_src, ggml_metal_t ctx_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) { @autoreleasepool { struct ggml_metal_buffer_id bid_src = ggml_metal_get_buffer_id(src); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 6a836e4590..22901fcfe3 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -507,6 +507,20 @@ static void ggml_backend_metal_get_tensor_async(ggml_backend_t backend, const gg ggml_metal_get_tensor_async(ctx, tensor, data, offset, size); } +static void ggml_backend_metal_set_tensor_2d_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + ggml_metal_t ctx = (ggml_metal_t)backend->context; + + ggml_metal_set_tensor_2d_async(ctx, tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +static void ggml_backend_metal_get_tensor_2d_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, + size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data) { + ggml_metal_t ctx = (ggml_metal_t)backend->context; + + ggml_metal_get_tensor_2d_async(ctx, tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + static bool ggml_backend_metal_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { if (!ggml_backend_is_metal(backend_src) || !ggml_backend_is_metal(backend_dst)) { return false; @@ -567,8 +581,8 @@ static ggml_backend_i ggml_backend_metal_i = { /* .free = */ ggml_backend_metal_free, /* .set_tensor_async = */ ggml_backend_metal_set_tensor_async, /* .get_tensor_async = */ ggml_backend_metal_get_tensor_async, - /* .get_tensor_2d_async = */ NULL, - /* .set_tensor_2d_async = */ NULL, + /* .get_tensor_2d_async = */ ggml_backend_metal_set_tensor_2d_async, + /* .set_tensor_2d_async = */ ggml_backend_metal_get_tensor_2d_async, /* .cpy_tensor_async = */ ggml_backend_metal_cpy_tensor_async, // only needed for multi-GPU setups /* .synchronize = */ ggml_backend_metal_synchronize, /* .graph_plan_create = */ NULL,