Skip to content

Commit 7a4f97d

Browse files
ggerganovggerganov
authored
metal : add diag (ggml-org#19330)
1 parent a498c75 commit 7a4f97d

7 files changed

Lines changed: 110 additions & 1 deletion

File tree

‎ggml/src/ggml-metal/ggml-metal-device.cpp‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_me
176176
return res;
177177
}
178178

179+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag(ggml_metal_library_t lib, const ggml_tensor * op) {
180+
char base[256];
181+
char name[256];
182+
183+
const int n = op->src[0]->ne[0];
184+
185+
snprintf(base, 256, "kernel_diag_%s", ggml_type_name(op->src[0]->type));
186+
snprintf(name, 256, "%s_n=%d", base, n);
187+
188+
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
189+
if (!res.pipeline) {
190+
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
191+
}
192+
193+
res.nsg = 1;
194+
res.smem = 0;
195+
196+
return res;
197+
}
198+
179199
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat(ggml_metal_library_t lib, ggml_type tsrc) {
180200
char base[256];
181201
char name[256];

‎ggml/src/ggml-metal/ggml-metal-device.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_1d
108108
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool);
109109
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows (ggml_metal_library_t lib, enum ggml_type tsrc);
110110
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, enum ggml_type tidx, enum ggml_type tdst);
111+
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag (ggml_metal_library_t lib, const struct ggml_tensor * op);
111112
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat (ggml_metal_library_t lib, enum ggml_type tsrc);
112113
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary (ggml_metal_library_t lib, const struct ggml_tensor * op);
113114
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op);

‎ggml/src/ggml-metal/ggml-metal-device.m‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1152,8 +1152,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
11521152
return has_simdgroup_reduction;
11531153
case GGML_OP_RWKV_WKV6:
11541154
case GGML_OP_RWKV_WKV7:
1155-
case GGML_OP_SOLVE_TRI:
11561155
return true;
1156+
case GGML_OP_SOLVE_TRI:
11571157
case GGML_OP_MUL_MAT:
11581158
case GGML_OP_MUL_MAT_ID:
11591159
return has_simdgroup_reduction;
@@ -1235,6 +1235,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
12351235
return false;
12361236
};
12371237
}
1238+
case GGML_OP_DIAG:
1239+
return true;
12381240
case GGML_OP_OPT_STEP_ADAMW:
12391241
case GGML_OP_OPT_STEP_SGD:
12401242
return has_simdgroup_reduction;

‎ggml/src/ggml-metal/ggml-metal-impl.h‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -792,6 +792,25 @@ typedef struct {
792792
uint64_t nb3;
793793
} ggml_metal_kargs_set_rows;
794794

795+
typedef struct {
796+
int32_t ne00;
797+
int32_t ne01;
798+
int32_t ne02;
799+
int32_t ne03;
800+
uint64_t nb00;
801+
uint64_t nb01;
802+
uint64_t nb02;
803+
uint64_t nb03;
804+
int32_t ne0;
805+
int32_t ne1;
806+
int32_t ne2;
807+
int32_t ne3;
808+
uint64_t nb0;
809+
uint64_t nb1;
810+
uint64_t nb2;
811+
uint64_t nb3;
812+
} ggml_metal_kargs_diag;
813+
795814
typedef struct {
796815
int64_t ne00;
797816
int64_t ne01;

‎ggml/src/ggml-metal/ggml-metal-ops.cpp‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
361361
{
362362
n_fuse = ggml_metal_op_set_rows(ctx, idx);
363363
} break;
364+
case GGML_OP_DIAG:
365+
{
366+
n_fuse = ggml_metal_op_diag(ctx, idx);
367+
} break;
364368
case GGML_OP_L2_NORM:
365369
{
366370
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) {
12591263
return 1;
12601264
}
12611265

1266+
int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
1267+
ggml_tensor * op = ctx->node(idx);
1268+
1269+
ggml_metal_library_t lib = ctx->lib;
1270+
ggml_metal_encoder_t enc = ctx->enc;
1271+
1272+
GGML_TENSOR_LOCALS(int32_t, ne0, op->src[0], ne);
1273+
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
1274+
GGML_TENSOR_LOCALS(int32_t, ne, op, ne);
1275+
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
1276+
1277+
ggml_metal_kargs_diag args = {
1278+
/*.ne00 =*/ne00,
1279+
/*.ne01 =*/ne01,
1280+
/*.ne02 =*/ne02,
1281+
/*.ne03 =*/ne03,
1282+
/*.nb00 =*/nb00,
1283+
/*.nb01 =*/nb01,
1284+
/*.nb02 =*/nb02,
1285+
/*.nb03 =*/nb03,
1286+
/*.ne0 =*/ne0,
1287+
/*.ne1 =*/ne1,
1288+
/*.ne2 =*/ne2,
1289+
/*.ne3 =*/ne3,
1290+
/*.nb0 =*/nb0,
1291+
/*.nb1 =*/nb1,
1292+
/*.nb2 =*/nb2,
1293+
/*.nb3 =*/nb3,
1294+
};
1295+
1296+
auto pipeline = ggml_metal_library_get_pipeline_diag(lib, op);
1297+
1298+
ggml_metal_encoder_set_pipeline(enc, pipeline);
1299+
ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
1300+
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1);
1301+
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2);
1302+
1303+
ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, 32, 1, 1);
1304+
1305+
return 1;
1306+
}
1307+
12621308
int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) {
12631309
ggml_tensor * op = ctx->node(idx);
12641310

‎ggml/src/ggml-metal/ggml-metal-ops.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ int ggml_metal_op_sum_rows (ggml_metal_op_t ctx, int idx);
5656
int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx);
5757
int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx);
5858
int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx);
59+
int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx);
5960
int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx);
6061
int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx);
6162
int ggml_metal_op_ssm_scan (ggml_metal_op_t ctx, int idx);

‎ggml/src/ggml-metal/ggml-metal.metal‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8815,6 +8815,26 @@ kernel void kernel_set_rows_f(
88158815
}
88168816
}
88178817

8818+
kernel void kernel_diag_f32(
8819+
constant ggml_metal_kargs_diag & args,
8820+
device const char * src0,
8821+
device char * dst,
8822+
uint3 tgpig[[threadgroup_position_in_grid]],
8823+
ushort tiitg[[thread_index_in_threadgroup]]) {
8824+
constexpr short NW = N_SIMDWIDTH;
8825+
8826+
const int32_t i3 = tgpig.z;
8827+
const int32_t i2 = tgpig.y;
8828+
const int32_t i1 = tgpig.x;
8829+
8830+
device const float * src0_ptr = (device const float *)(src0 + i2*args.nb02 + i3*args.nb03);
8831+
device float * dst_ptr = (device float *)(dst + i1*args.nb01 + i2*args.nb2 + i3*args.nb3);
8832+
8833+
for (int i0 = tiitg; i0 < args.ne0; i0 += NW) {
8834+
dst_ptr[i0] = i0 == i1 ? src0_ptr[i0] : 0.0f;
8835+
}
8836+
}
8837+
88188838
constant bool FC_mul_mm_bc_inp [[function_constant(FC_MUL_MM + 0)]];
88198839
constant bool FC_mul_mm_bc_out [[function_constant(FC_MUL_MM + 1)]];
88208840

0 commit comments

Comments
 (0)