Skip to content

Commit f955e39

Browse files
LamotheLamothe
authored
ggml: add f16 out_prod support for CPU and out_prod op for Vulkan (ggml-org#23997)
1 parent 33a75f4 commit f955e39

6 files changed

Lines changed: 167 additions & 7 deletions

File tree

‎ggml/src/ggml-cpu/ggml-cpu.c‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2859,7 +2859,8 @@ struct ggml_cplan ggml_graph_plan(
28592859
} break;
28602860
case GGML_OP_OUT_PROD:
28612861
{
2862-
if (ggml_is_quantized(node->src[0]->type)) {
2862+
if (ggml_is_quantized(node->src[0]->type) ||
2863+
node->src[0]->type == GGML_TYPE_F16) {
28632864
cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
28642865
}
28652866
} break;

‎ggml/src/ggml-cpu/ggml-cpu.cpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -462,11 +462,12 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
462462
return max_bias == 0.0f;
463463
}
464464
case GGML_OP_IM2COL_BACK:
465-
return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32;
465+
return src0->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
466466
case GGML_OP_GET_ROWS_BACK:
467467
return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16;
468468
case GGML_OP_OUT_PROD:
469-
return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
469+
return (src0->type == GGML_TYPE_F32 ||
470+
((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
470471
src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
471472
default:
472473
return true;

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

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4449,6 +4449,70 @@ static void ggml_compute_forward_out_prod_q_f32(
44494449
}
44504450
}
44514451

4452+
static void ggml_compute_forward_out_prod_f16_f32(
4453+
const ggml_compute_params * params,
4454+
ggml_tensor * dst) {
4455+
4456+
const ggml_tensor * src0 = dst->src[0];
4457+
const ggml_tensor * src1 = dst->src[1];
4458+
4459+
GGML_TENSOR_BINARY_OP_LOCALS;
4460+
4461+
const int ith = params->ith;
4462+
const int nth = params->nth;
4463+
4464+
GGML_ASSERT(src0->type == GGML_TYPE_F16);
4465+
GGML_ASSERT(src1->type == GGML_TYPE_F32);
4466+
GGML_ASSERT(dst->type == GGML_TYPE_F32);
4467+
4468+
GGML_ASSERT(ne02 == ne12);
4469+
GGML_ASSERT(ne03 == ne13);
4470+
GGML_ASSERT(ne2 == ne12);
4471+
GGML_ASSERT(ne3 == ne13);
4472+
4473+
GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
4474+
GGML_ASSERT(nb0 == sizeof(float));
4475+
4476+
GGML_ASSERT(ne0 == ne00);
4477+
GGML_ASSERT(ne1 == ne10);
4478+
GGML_ASSERT(ne2 == ne02);
4479+
GGML_ASSERT(ne3 == ne03);
4480+
4481+
if (ith == 0) {
4482+
ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0);
4483+
}
4484+
ggml_barrier(params->threadpool);
4485+
4486+
const int64_t nr = ne1*ne2*ne3;
4487+
const int64_t dr = (nr + nth - 1)/nth;
4488+
const int64_t ir0 = dr*ith;
4489+
const int64_t ir1 = MIN(ir0 + dr, nr);
4490+
4491+
float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith;
4492+
4493+
for (int64_t ir = ir0; ir < ir1; ++ir) {
4494+
const int64_t i3 = ir/(ne2*ne1);
4495+
const int64_t i2 = (ir - i3*ne2*ne1)/ne1;
4496+
const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1);
4497+
4498+
const int64_t i02 = i2;
4499+
const int64_t i03 = i3;
4500+
4501+
const int64_t i12 = i2;
4502+
const int64_t i13 = i3;
4503+
4504+
float * d = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb3));
4505+
4506+
for (int64_t i01 = 0; i01 < ne01; ++i01) {
4507+
const int64_t i11 = i01;
4508+
ggml_fp16_t * s0 = (ggml_fp16_t *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03));
4509+
float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13));
4510+
ggml_fp16_to_fp32_row(s0, wdata, ne0);
4511+
ggml_vec_mad_f32(ne0, d, wdata, *s1);
4512+
}
4513+
}
4514+
}
4515+
44524516
void ggml_compute_forward_out_prod(
44534517
const ggml_compute_params * params,
44544518
ggml_tensor * dst) {
@@ -4486,9 +4550,8 @@ void ggml_compute_forward_out_prod(
44864550
} break;
44874551
case GGML_TYPE_F16:
44884552
{
4489-
GGML_ABORT("fatal error"); // todo
4490-
// ggml_compute_forward_out_prod_f16_f32(params, dst);
4491-
}
4553+
ggml_compute_forward_out_prod_f16_f32(params, dst);
4554+
} break;
44924555
case GGML_TYPE_F32:
44934556
{
44944557
ggml_compute_forward_out_prod_f32(params, dst);
@@ -6469,7 +6532,7 @@ void ggml_compute_forward_im2col_back_f32(
64696532
const ggml_tensor * src1 = dst->src[1]; // convolution kernel
64706533

64716534
GGML_ASSERT(src0->type == GGML_TYPE_F32);
6472-
GGML_ASSERT(src1->type == GGML_TYPE_F32);
6535+
GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
64736536
GGML_ASSERT( dst->type == GGML_TYPE_F32);
64746537

64756538
GGML_TENSOR_BINARY_OP_LOCALS;

‎ggml/src/ggml-vulkan/ggml-vulkan.cpp‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -961,6 +961,7 @@ struct vk_device_struct {
961961
vk_pipeline pipeline_col2im_1d_f32;
962962
vk_pipeline pipeline_col2im_1d_f16;
963963
vk_pipeline pipeline_col2im_1d_bf16;
964+
vk_pipeline pipeline_out_prod_f32;
964965
vk_pipeline pipeline_snake_f32;
965966
vk_pipeline pipeline_snake_f16;
966967
vk_pipeline pipeline_snake_bf16;
@@ -5479,6 +5480,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
54795480
ggml_vk_create_pipeline(device, device->pipeline_col2im_1d_f16, "col2im_1d_f16", col2im_1d_f16_len, col2im_1d_f16_data, "main", 2, sizeof(vk_op_col2im_1d_push_constants), {256, 1, 1}, {}, 1, true);
54805481
ggml_vk_create_pipeline(device, device->pipeline_col2im_1d_bf16, "col2im_1d_bf16", col2im_1d_bf16_len, col2im_1d_bf16_data, "main", 2, sizeof(vk_op_col2im_1d_push_constants), {256, 1, 1}, {}, 1, true);
54815482

5483+
ggml_vk_create_pipeline(device, device->pipeline_out_prod_f32, "out_prod_f32", out_prod_f32_len, out_prod_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {256, 1, 1}, {}, 1);
5484+
54825485
ggml_vk_create_pipeline(device, device->pipeline_snake_f32, "snake_f32", snake_f32_len, snake_f32_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
54835486
ggml_vk_create_pipeline(device, device->pipeline_snake_f16, "snake_f16", snake_f16_len, snake_f16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
54845487
ggml_vk_create_pipeline(device, device->pipeline_snake_bf16, "snake_bf16", snake_bf16_len, snake_bf16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
@@ -10745,6 +10748,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
1074510748
return ctx->device->pipeline_add_id_f32;
1074610749
}
1074710750
return nullptr;
10751+
case GGML_OP_OUT_PROD:
10752+
if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
10753+
return ctx->device->pipeline_out_prod_f32;
10754+
}
10755+
return nullptr;
1074810756
case GGML_OP_CONCAT: {
1074910757
if (src0->type != src1->type || src0->type != dst->type) {
1075010758
return nullptr;
@@ -11701,6 +11709,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
1170111709
case GGML_OP_DIV:
1170211710
case GGML_OP_MUL:
1170311711
case GGML_OP_ADD1:
11712+
case GGML_OP_OUT_PROD:
1170411713
case GGML_OP_ARANGE:
1170511714
case GGML_OP_FILL:
1170611715
case GGML_OP_SCALE:
@@ -12014,6 +12023,24 @@ static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const
1201412023
});
1201512024
}
1201612025

12026+
static void ggml_vk_out_prod(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
12027+
const uint32_t src0_type_size = ggml_type_size(src0->type);
12028+
const uint32_t src1_type_size = ggml_type_size(src1->type);
12029+
const uint32_t dst_type_size = ggml_type_size(dst->type);
12030+
12031+
ggml_vk_op_f32<vk_op_binary_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_OUT_PROD, {
12032+
(uint32_t)ggml_nelements(dst),
12033+
(uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3],
12034+
(uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size,
12035+
(uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3],
12036+
(uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size,
12037+
(uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3],
12038+
(uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size,
12039+
0,
12040+
0.0f, 0.0f, 0,
12041+
});
12042+
}
12043+
1201712044
static void ggml_vk_sub(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
1201812045
const uint32_t src0_type_size = ggml_type_size(src0->type);
1201912046
const uint32_t src1_type_size = ggml_type_size(src1->type);
@@ -14801,6 +14828,9 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
1480114828
ggml_vk_add(ctx, compute_ctx, src0, src1, node);
1480214829
}
1480314830
break;
14831+
case GGML_OP_OUT_PROD:
14832+
ggml_vk_out_prod(ctx, compute_ctx, src0, src1, node);
14833+
break;
1480414834
case GGML_OP_SUB:
1480514835
ggml_vk_sub(ctx, compute_ctx, src0, src1, node);
1480614836

@@ -17655,6 +17685,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
1765517685
case GGML_OP_OPT_STEP_ADAMW:
1765617686
case GGML_OP_OPT_STEP_SGD:
1765717687
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
17688+
case GGML_OP_OUT_PROD:
17689+
return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
17690+
&& ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
17691+
&& op->type == GGML_TYPE_F32;
1765817692
case GGML_OP_LOG:
1765917693
case GGML_OP_TRI:
1766017694
case GGML_OP_DIAG:

‎ggml/src/ggml-vulkan/vulkan-shaders/out_prod.comp‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
#version 450
2+
3+
#extension GL_EXT_shader_16bit_storage : require
4+
5+
layout (push_constant) uniform parameter
6+
{
7+
uint ne;
8+
uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
9+
uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
10+
uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23;
11+
uint misalign_offsets;
12+
float param1; float param2; int param3;
13+
} p;
14+
15+
layout (binding = 0) readonly buffer A {float data_a[];};
16+
layout (binding = 1) readonly buffer B {float data_b[];};
17+
layout (binding = 2) writeonly buffer D {float data_d[];};
18+
19+
uint get_idx() {
20+
return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
21+
}
22+
23+
uint get_aoffset() { return p.misalign_offsets >> 16; }
24+
uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; }
25+
uint get_doffset() { return p.misalign_offsets & 0xFF; }
26+
27+
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
28+
29+
void main() {
30+
uint idx = get_idx();
31+
if (idx >= p.ne) {
32+
return;
33+
}
34+
35+
uint tmp = idx;
36+
uint i0 = tmp % p.ne20; tmp /= p.ne20;
37+
uint i1 = tmp % p.ne21; tmp /= p.ne21;
38+
uint i2 = tmp % p.ne22; tmp /= p.ne22;
39+
uint i3 = tmp;
40+
41+
uint a_i0 = i0 % p.ne00;
42+
uint a_i2 = i2 / (p.ne22 / p.ne02);
43+
uint a_i3 = i3 / (p.ne23 / p.ne03);
44+
45+
uint b_i0 = i1 % p.ne10;
46+
uint b_i2 = i2;
47+
uint b_i3 = i3;
48+
49+
float sum = 0.0f;
50+
uint K = p.ne01;
51+
for (uint k = 0; k < K; k++) {
52+
uint aoff = get_aoffset() + a_i3*p.nb03 + a_i2*p.nb02 + k*p.nb01 + a_i0*p.nb00;
53+
uint boff = get_boffset() + b_i3*p.nb13 + b_i2*p.nb12 + k*p.nb11 + b_i0*p.nb10;
54+
sum += data_a[aoff] * data_b[boff];
55+
}
56+
57+
uint doff = get_doffset() + i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20;
58+
data_d[doff] = sum;
59+
}

‎ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1036,6 +1036,8 @@ void process_shaders() {
10361036
}
10371037
}
10381038

1039+
string_to_spv("out_prod_f32", "out_prod.comp", {});
1040+
10391041
string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
10401042

10411043
string_to_spv("conv_transpose_1d_f32", "conv_transpose_1d.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});

0 commit comments

Comments
 (0)