Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

metal : add diag (#19330)

Browse Source
This commit is contained in:
Georgi Gerganov
2026-02-05 10:08:45 +02:00
committed by GitHub
parent a498c75ad1
commit 7a4f97d196
7 changed files with 110 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-ops.cpp
+46
View File
@@ -361,6 +361,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
{
n_fuse = ggml_metal_op_set_rows(ctx, idx);
} break;
case GGML_OP_DIAG:
{
n_fuse = ggml_metal_op_diag(ctx, idx);
} break;
case GGML_OP_L2_NORM:
{
n_fuse = ggml_metal_op_l2_norm(ctx, idx);
@@ -1259,6 +1263,48 @@ int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) {
return 1;
}
int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
ggml_metal_library_t lib = ctx->lib;
ggml_metal_encoder_t enc = ctx->enc;
GGML_TENSOR_LOCALS(int32_t, ne0, op->src[0], ne);
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
GGML_TENSOR_LOCALS(int32_t, ne, op, ne);
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
ggml_metal_kargs_diag args = {
/*.ne00 =*/ne00,
/*.ne01 =*/ne01,
/*.ne02 =*/ne02,
/*.ne03 =*/ne03,
/*.nb00 =*/nb00,
/*.nb01 =*/nb01,
/*.nb02 =*/nb02,
/*.nb03 =*/nb03,
/*.ne0 =*/ne0,
/*.ne1 =*/ne1,
/*.ne2 =*/ne2,
/*.ne3 =*/ne3,
/*.nb0 =*/nb0,
/*.nb1 =*/nb1,
/*.nb2 =*/nb2,
/*.nb3 =*/nb3,
};
auto pipeline = ggml_metal_library_get_pipeline_diag(lib, op);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1);
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2);
ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, 32, 1, 1);
return 1;
}
int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);