Skip to content

Commit 8dc1d9f

Browse files
0cc4m0cc4m
committed
Vulkan: Add GLU ops and shaders
1 parent 98a5019 commit 8dc1d9f

File tree

5 files changed

+215
-2
lines changed

5 files changed

+215
-2
lines changed

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

Lines changed: 91 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -431,6 +431,10 @@ struct vk_device_struct {
431431
vk_pipeline pipeline_tanh[2];
432432
vk_pipeline pipeline_sigmoid[2];
433433

434+
vk_pipeline pipeline_geglu[2];
435+
vk_pipeline pipeline_reglu[2];
436+
vk_pipeline pipeline_swiglu[2];
437+
434438
vk_pipeline pipeline_leaky_relu_f32;
435439
vk_pipeline pipeline_silu_back_f32;
436440
vk_pipeline pipeline_diag_mask_inf_f32;
@@ -2728,6 +2732,15 @@ static void ggml_vk_load_shaders(vk_device& device) {
27282732
CREATE_UNARY(sigmoid)
27292733
#undef CREATE_UNARY
27302734

2735+
#define CREATE_GLU(name) \
2736+
ggml_vk_create_pipeline(device, device->pipeline_ ## name [0], #name "_f32", name ## _f32_len, name ## _f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); \
2737+
ggml_vk_create_pipeline(device, device->pipeline_ ## name [1], #name "_f16", name ## _f16_len, name ## _f16_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
2738+
2739+
CREATE_GLU(geglu)
2740+
CREATE_GLU(reglu)
2741+
CREATE_GLU(swiglu)
2742+
#undef CREATE_GLU
2743+
27312744
ggml_vk_create_pipeline(device, device->pipeline_leaky_relu_f32, "leaky_relu_f32", leaky_relu_f32_len, leaky_relu_f32_data, "main", 2, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
27322745
ggml_vk_create_pipeline(device, device->pipeline_silu_back_f32, "silu_back_f32", silu_back_f32_len, silu_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
27332746

@@ -6415,6 +6428,24 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
64156428
break;
64166429
}
64176430
return nullptr;
6431+
case GGML_OP_GLU:
6432+
if ((src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) ||
6433+
(dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16) ||
6434+
(src0->type != dst->type)) {
6435+
return nullptr;
6436+
}
6437+
6438+
switch (ggml_get_glu_op(dst)) {
6439+
case GGML_GLU_OP_GEGLU:
6440+
return ctx->device->pipeline_geglu[dst->type == GGML_TYPE_F16];
6441+
case GGML_GLU_OP_REGLU:
6442+
return ctx->device->pipeline_reglu[dst->type == GGML_TYPE_F16];
6443+
case GGML_GLU_OP_SWIGLU:
6444+
return ctx->device->pipeline_swiglu[dst->type == GGML_TYPE_F16];
6445+
default:
6446+
break;
6447+
}
6448+
return nullptr;
64186449
case GGML_OP_DIAG_MASK_INF:
64196450
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
64206451
return ctx->device->pipeline_diag_mask_inf_f32;
@@ -6791,6 +6822,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
67916822
case GGML_OP_SOFT_MAX_BACK:
67926823
case GGML_OP_SUM_ROWS:
67936824
case GGML_OP_ARGMAX:
6825+
case GGML_OP_GLU:
67946826
{
67956827
const uint32_t nr = ggml_nrows(src0);
67966828
if (nr > 262144) {
@@ -7507,6 +7539,14 @@ static void ggml_vk_unary(ggml_backend_vk_context * ctx, vk_context& subctx, con
75077539
ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_UNARY, { (uint32_t)ggml_nelements(src0), 0, 0.0f, 0.0f }, dryrun);
75087540
}
75097541

7542+
static void ggml_vk_glu(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) {
7543+
GGML_ASSERT(src0->ne[0] / 2 == dst->ne[0]);
7544+
7545+
const uint32_t swapped = (uint32_t)dst->op_params[1];
7546+
7547+
ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_GLU, { (uint32_t)src0->ne[0], swapped, 0.0f, 0.0f }, dryrun);
7548+
}
7549+
75107550
static void ggml_vk_diag_mask_inf(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst, bool dryrun = false) {
75117551
int32_t * op_params = (int32_t *)dst->op_params;
75127552
ggml_vk_op_f32<vk_op_diag_mask_push_constants>(ctx, subctx, src0, nullptr, nullptr, dst, GGML_OP_DIAG_MASK_INF, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0] }, dryrun);
@@ -8718,6 +8758,16 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * nod
87188758
return false;
87198759
}
87208760
break;
8761+
case GGML_OP_GLU:
8762+
switch (ggml_get_glu_op(node)) {
8763+
case GGML_GLU_OP_GEGLU:
8764+
case GGML_GLU_OP_REGLU:
8765+
case GGML_GLU_OP_SWIGLU:
8766+
break;
8767+
default:
8768+
return false;
8769+
}
8770+
break;
87218771
case GGML_OP_REPEAT:
87228772
case GGML_OP_REPEAT_BACK:
87238773
case GGML_OP_GET_ROWS:
@@ -8810,6 +8860,7 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * nod
88108860
case GGML_OP_RMS_NORM_BACK:
88118861
case GGML_OP_L2_NORM:
88128862
case GGML_OP_UNARY:
8863+
case GGML_OP_GLU:
88138864
case GGML_OP_DIAG_MASK_INF:
88148865
case GGML_OP_SOFT_MAX:
88158866
case GGML_OP_SOFT_MAX_BACK:
@@ -8947,6 +8998,17 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * nod
89478998
return false;
89488999
}
89499000
break;
9001+
case GGML_OP_GLU:
9002+
switch (ggml_get_glu_op(node)) {
9003+
case GGML_GLU_OP_GEGLU:
9004+
case GGML_GLU_OP_REGLU:
9005+
case GGML_GLU_OP_SWIGLU:
9006+
ggml_vk_glu(ctx, compute_ctx, src0, node, dryrun);
9007+
break;
9008+
default:
9009+
return false;
9010+
}
9011+
break;
89509012
case GGML_OP_DIAG_MASK_INF:
89519013
ggml_vk_diag_mask_inf(ctx, compute_ctx, src0, node, dryrun);
89529014

@@ -9072,8 +9134,9 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_tensor * nod
90729134
if (!ok) {
90739135
if (node->op == GGML_OP_UNARY) {
90749136
std::cerr << __func__ << ": error: op not supported UNARY " << node->name << " (" << ggml_unary_op_name(static_cast<ggml_unary_op>(node->op_params[0])) << ")" << std::endl;
9075-
}
9076-
else {
9137+
} else if (node->op == GGML_OP_GLU) {
9138+
std::cerr << __func__ << ": error: op not supported GLU " << node->name << " (" << ggml_glu_op_name(static_cast<ggml_glu_op>(node->op_params[0])) << ")" << std::endl;
9139+
} else {
90779140
std::cerr << __func__ << ": error: op not supported " << node->name << " (" << ggml_op_name(node->op) << ")" << std::endl;
90789141
}
90799142
}
@@ -9152,6 +9215,17 @@ static bool ggml_vk_compute_forward(ggml_backend_vk_context * ctx, ggml_tensor *
91529215
return false;
91539216
}
91549217
break;
9218+
case GGML_OP_GLU:
9219+
switch (ggml_get_glu_op(tensor)) {
9220+
case GGML_GLU_OP_GEGLU:
9221+
case GGML_GLU_OP_REGLU:
9222+
case GGML_GLU_OP_SWIGLU:
9223+
buf = tensor->buffer;
9224+
break;
9225+
default:
9226+
return false;
9227+
}
9228+
break;
91559229
case GGML_OP_MUL_MAT:
91569230
case GGML_OP_MUL_MAT_ID:
91579231
case GGML_OP_FLASH_ATTN_EXT:
@@ -9923,6 +9997,19 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
99239997
return false;
99249998
}
99259999
break;
10000+
case GGML_OP_GLU:
10001+
switch (ggml_get_glu_op(op)) {
10002+
case GGML_GLU_OP_GEGLU:
10003+
case GGML_GLU_OP_REGLU:
10004+
case GGML_GLU_OP_SWIGLU:
10005+
return ggml_is_contiguous(op->src[0]) &&
10006+
(op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
10007+
(op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
10008+
(op->src[0]->type == op->type);
10009+
default:
10010+
return false;
10011+
}
10012+
break;
992610013
case GGML_OP_MUL_MAT:
992710014
case GGML_OP_MUL_MAT_ID:
992810015
{
@@ -10637,6 +10724,8 @@ static void ggml_vk_check_results_0(ggml_tensor * tensor) {
1063710724
std::cerr << "Missing vk_check_results OP: " << ggml_op_name(tensor->op) << std::endl;
1063810725
GGML_ABORT("fatal error");
1063910726
}
10727+
} else if (tensor->op == GGML_OP_GLU) {
10728+
tensor_clone = ggml_glu(ggml_ctx, src_clone[0], (ggml_glu_op) tensor->op_params[0], tensor->op_params[1]);
1064010729
} else if (tensor->op == GGML_OP_CPY || tensor->op == GGML_OP_DUP) {
1064110730
if (src1 == nullptr) {
1064210731
tensor_clone = ggml_dup(ggml_ctx, src_clone[0]);

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
#version 450
2+
3+
#include "generic_head.comp"
4+
#include "types.comp"
5+
6+
#extension GL_EXT_control_flow_attributes : enable
7+
8+
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
9+
10+
layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
11+
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
12+
13+
layout (constant_id = 0) const uint BLOCK_SIZE = 32;
14+
15+
void main() {
16+
const float GELU_COEF_A = 0.044715f;
17+
const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
18+
19+
const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
20+
const uint col = gl_LocalInvocationID.x;
21+
22+
const uint offset = p.KX / 2;
23+
24+
const bool swapped = p.KY > 0;
25+
26+
if (!swapped) {
27+
for (uint i = col; i < offset; i += BLOCK_SIZE) {
28+
const uint idx = row * p.KX + i;
29+
30+
const float xi = float(data_a[idx]);
31+
const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi);
32+
data_d[row * offset + i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1)) * float(data_a[idx + offset]));
33+
}
34+
} else {
35+
for (uint i = col; i < offset; i += BLOCK_SIZE) {
36+
const uint idx = row * p.KX + i;
37+
38+
const float xi = float(data_a[idx + offset]);
39+
const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi);
40+
data_d[row * offset + i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1)) * float(data_a[idx]));
41+
}
42+
}
43+
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
#version 450
2+
3+
#include "generic_head.comp"
4+
#include "types.comp"
5+
6+
#extension GL_EXT_control_flow_attributes : enable
7+
8+
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
9+
10+
layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
11+
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
12+
13+
layout (constant_id = 0) const uint BLOCK_SIZE = 32;
14+
15+
void main() {
16+
const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
17+
const uint col = gl_LocalInvocationID.x;
18+
19+
const uint offset = p.KX / 2;
20+
21+
const bool swapped = p.KY > 0;
22+
23+
if (!swapped) {
24+
for (uint i = col; i < offset; i += BLOCK_SIZE) {
25+
const uint idx = row * p.KX + i;
26+
27+
data_d[row * offset + i] = D_TYPE(max(float(data_a[idx]), 0.0f) * float(data_a[idx + offset]));
28+
}
29+
} else {
30+
for (uint i = col; i < offset; i += BLOCK_SIZE) {
31+
const uint idx = row * p.KX + i;
32+
33+
data_d[row * offset + i] = D_TYPE(max(float(data_a[idx + offset]), 0.0f) * float(data_a[idx]));
34+
}
35+
}
36+
}

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
#version 450
2+
3+
#include "generic_head.comp"
4+
#include "types.comp"
5+
6+
#extension GL_EXT_control_flow_attributes : enable
7+
8+
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
9+
10+
layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
11+
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
12+
13+
layout (constant_id = 0) const uint BLOCK_SIZE = 32;
14+
15+
void main() {
16+
const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
17+
const uint col = gl_LocalInvocationID.x;
18+
19+
const uint offset = p.KX / 2;
20+
21+
const bool swapped = p.KY > 0;
22+
23+
if (!swapped) {
24+
for (uint i = col; i < offset; i += BLOCK_SIZE) {
25+
const uint idx = row * p.KX + i;
26+
27+
const float xi = float(data_a[idx]);
28+
data_d[row * offset + i] = D_TYPE(xi / (1.0f + exp(-xi)) * float(data_a[idx + offset]));
29+
}
30+
} else {
31+
for (uint i = col; i < offset; i += BLOCK_SIZE) {
32+
const uint idx = row * p.KX + i;
33+
34+
const float xi = float(data_a[idx + offset]);
35+
data_d[row * offset + i] = D_TYPE(xi / (1.0f + exp(-xi)) * float(data_a[idx]));
36+
}
37+
}
38+
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,13 @@ void process_shaders() {
585585
string_to_spv("sigmoid_f16", "sigmoid.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
586586
string_to_spv("sigmoid_f32", "sigmoid.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
587587

588+
string_to_spv("geglu_f16", "geglu.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
589+
string_to_spv("geglu_f32", "geglu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
590+
string_to_spv("reglu_f16", "reglu.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
591+
string_to_spv("reglu_f32", "reglu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
592+
string_to_spv("swiglu_f16", "swiglu.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
593+
string_to_spv("swiglu_f32", "swiglu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
594+
588595
string_to_spv("leaky_relu_f32", "leaky_relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
589596
string_to_spv("silu_back_f32", "silu_back.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
590597

0 commit comments

Comments
 (0)