Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 09:15:04 +02:00
Code Issues Packages Projects Releases Wiki Activity

ggml-metal: implement async 2D tensor copy functions

Browse Source 
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
This commit is contained in:
Georgi Gerganov
2026-04-29 14:22:06 +03:00
parent 59237bfbbc
commit c3a54d6253
3 changed files with 114 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-context.h
+2
View File
@@ -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);
ggml/src/ggml-metal/ggml-metal-context.m
+96
View File
@@ -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<MTLDevice> device = ggml_metal_device_get_obj(ctx->dev);
id<MTLBuffer> 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<MTLCommandQueue> queue = ggml_metal_device_get_queue(ctx->dev);
id<MTLCommandBuffer> cmd_buf = [queue commandBuffer];
id<MTLBlitCommandEncoder> 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<MTLDevice> device = ggml_metal_device_get_obj(ctx->dev);
id<MTLBuffer> 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<MTLCommandQueue> queue = ggml_metal_device_get_queue(ctx->dev);
id<MTLCommandBuffer> cmd_buf = [queue commandBuffer];
id<MTLBlitCommandEncoder> 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);
ggml/src/ggml-metal/ggml-metal.cpp
+16 -2
View File
@@ -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,