Skip to content

Commit 7a6b4d0

Browse files
jeffbolznvjeffbolznv
committed
vulkan: Use coopmat2 for conv2d
1 parent 7d3553f commit 7a6b4d0

File tree

3 files changed

+77
-17
lines changed

3 files changed

+77
-17
lines changed

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

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3093,7 +3093,7 @@ static void ggml_vk_load_shaders(vk_device& device) {
30933093
uint32_t use_collectives = 0; // Enables subgroup ops for preventing the re-calculation of indices.
30943094
uint32_t conv2d_BS_NPQ = 128;
30953095
uint32_t conv2d_TS_K = 8;
3096-
uint32_t conv2d_SHMEM_PAD = 4;
3096+
uint32_t conv2d_SHMEM_PAD = device->coopmat2 ? 8 : 4; // 16B == 8 float16_t or 4 float
30973097

30983098
switch (s) {
30993099
default:
@@ -3137,14 +3137,25 @@ static void ggml_vk_load_shaders(vk_device& device) {
31373137
}
31383138
}
31393139

3140-
ggml_vk_create_pipeline(
3141-
device, device->pipeline_conv2d_f32[s], "conv2d_f32", conv2d_f32_len, conv2d_f32_data, "main", 3,
3142-
sizeof(vk_op_conv2d_push_constants), { conv2d_BS_K, conv2d_BS_NPQ, 1 },
3143-
{ conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD }, 1, true, use_collectives);
3144-
ggml_vk_create_pipeline(
3145-
device, device->pipeline_conv2d_f16_f32[s], "conv2d_f16_f32", conv2d_f16_f32_len, conv2d_f16_f32_data, "main", 3,
3146-
sizeof(vk_op_conv2d_push_constants), { conv2d_BS_K, conv2d_BS_NPQ, 1 },
3147-
{ conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD }, 1, true, use_collectives);
3140+
if (device->coopmat2) {
3141+
ggml_vk_create_pipeline(
3142+
device, device->pipeline_conv2d_f32[s], "conv2d_f32", conv2d_f32_cm2_len, conv2d_f32_cm2_data, "main", 3,
3143+
sizeof(vk_op_conv2d_push_constants), { conv2d_BS_K, conv2d_BS_NPQ, 1 },
3144+
{ conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD }, 1, true, use_collectives);
3145+
ggml_vk_create_pipeline(
3146+
device, device->pipeline_conv2d_f16_f32[s], "conv2d_f16_f32", conv2d_f16_f32_cm2_len, conv2d_f16_f32_cm2_data, "main", 3,
3147+
sizeof(vk_op_conv2d_push_constants), { conv2d_BS_K, conv2d_BS_NPQ, 1 },
3148+
{ conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD }, 1, true, use_collectives);
3149+
} else {
3150+
ggml_vk_create_pipeline(
3151+
device, device->pipeline_conv2d_f32[s], "conv2d_f32", conv2d_f32_len, conv2d_f32_data, "main", 3,
3152+
sizeof(vk_op_conv2d_push_constants), { conv2d_BS_K, conv2d_BS_NPQ, 1 },
3153+
{ conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD }, 1, true, use_collectives);
3154+
ggml_vk_create_pipeline(
3155+
device, device->pipeline_conv2d_f16_f32[s], "conv2d_f16_f32", conv2d_f16_f32_len, conv2d_f16_f32_data, "main", 3,
3156+
sizeof(vk_op_conv2d_push_constants), { conv2d_BS_K, conv2d_BS_NPQ, 1 },
3157+
{ conv2d_WG_SIZE, conv2d_BS_K, conv2d_BS_CRS, conv2d_BS_NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD }, 1, true, use_collectives);
3158+
}
31483159
}
31493160

31503161
ggml_vk_create_pipeline(device, device->pipeline_conv2d_dw_whcn_f32, "conv2d_dw_whcn_f32", conv2d_dw_whcn_f32_len, conv2d_dw_whcn_f32_data, "main", 3, sizeof(vk_op_conv2d_dw_push_constants), {512, 1, 1}, {}, 1);

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

Lines changed: 54 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
#version 450
22

33
#extension GL_EXT_control_flow_attributes : enable
4+
#ifdef COOPMAT2
5+
#extension GL_NV_cooperative_matrix2 : enable
6+
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
7+
#extension GL_KHR_memory_scope_semantics : enable
8+
#endif
49

510
#ifdef USE_COLLECTIVES
611
# extension GL_KHR_shader_subgroup_shuffle : enable
@@ -91,6 +96,12 @@ uint32_t n_elems_out = K * NPQ;
9196
// Number of blocktiles per input
9297
uint32_t NB_CRS = splitWork(CRS, BS_CRS);
9398

99+
#ifdef COOPMAT2
100+
#define SHMEM_TYPE float16_t
101+
#else
102+
#define SHMEM_TYPE float
103+
#endif
104+
94105
const uint32_t Ash_stride = BS_CRS + SHMEM_PAD;
95106
const uint32_t Bsh_stride = BS_NPQ + SHMEM_PAD;
96107

@@ -100,8 +111,8 @@ const uint32_t Bsh_numel = BS_CRS * BS_NPQ;
100111
const uint32_t Ash_len = BS_K * Ash_stride;
101112
const uint32_t Bsh_len = BS_CRS * Bsh_stride;
102113

103-
shared float Ash[Ash_len]; // K x CRS
104-
shared float Bsh[Bsh_len]; // CRS x NPQ
114+
shared SHMEM_TYPE Ash[Ash_len]; // K x CRS
115+
shared SHMEM_TYPE Bsh[Bsh_len]; // CRS x NPQ
105116

106117
// Threadtile sizes
107118
const uint32_t TS_NPQ = BS_K * BS_NPQ / WG_SIZE / TS_K;
@@ -110,10 +121,6 @@ const uint32_t TS_NPQ = BS_K * BS_NPQ / WG_SIZE / TS_K;
110121
const uint32_t NT_K = BS_K / TS_K;
111122
const uint32_t NT_NPQ = BS_NPQ / TS_NPQ;
112123

113-
float regA[TS_K];
114-
float regB[TS_NPQ];
115-
float regC[TS_K][TS_NPQ];
116-
117124
/*
118125
Compute
119126
KxCRS @ CRSxNPQ = K x NPQ
@@ -145,12 +152,36 @@ uint fastdiv(uint n, uint mp, uint L) {
145152
return (msbs + n) >> L;
146153
}
147154

155+
#ifdef COOPMAT2
156+
#define ACC_TYPE float16_t
157+
158+
ACC_TYPE perElemOpStore(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem)
159+
{
160+
uint32_t K_idx = B_idx_K * BS_K + r;
161+
uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + c;
162+
uint32_t N_idx = fastdiv(NPQ_idx, p.OWOHmp, p.OWOHL); // divide by p.OH * p.OW;
163+
uint32_t OH_idx = fastdiv(NPQ_idx - N_idx * p.OH * p.OW, p.OWmp, p.OWL); // divide by p.OW;
164+
uint32_t OW_idx = NPQ_idx - N_idx * p.OH * p.OW - OH_idx * p.OW;
165+
uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + K_idx * p.nb2 + N_idx * p.nb3;
166+
if (K_idx < K && NPQ_idx < NPQ) {
167+
dst_data[dst_idx] = D_TYPE(elem);
168+
}
169+
return elem;
170+
}
171+
#endif
172+
148173
void main() {
174+
#ifdef COOPMAT2
175+
coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator> matC;
176+
matC = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator>(0.0);
177+
#else
178+
float regC[TS_K][TS_NPQ];
149179
for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
150180
for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
151181
regC[T_ly][T_lx] = 0.0;
152182
}
153183
}
184+
#endif
154185
/* Advance block in CRS dim */
155186
for (uint32_t B_idx_CRS = 0; B_idx_CRS < NB_CRS; B_idx_CRS++) {
156187
uint32_t CRS_idx_a;
@@ -199,7 +230,7 @@ void main() {
199230
if (K_idx >= K || CRS_idx_a >= CRS) {
200231
val = 0.0;
201232
}
202-
Ash[B_ly * Ash_stride + B_lx] = val;
233+
Ash[B_ly * Ash_stride + B_lx] = SHMEM_TYPE(val);
203234
}
204235
/* Load input to B_block: (BS_CRS x BS_NPQ) */
205236
[[unroll]] for (uint32_t r_offset = 0; r_offset < BS_CRS; r_offset += BrpWg) {
@@ -244,11 +275,21 @@ void main() {
244275
if (CRS_idx_b >= CRS || NPQ_idx >= NPQ || H_idx < 0 || H_idx >= p.H || W_idx < 0 || W_idx >= p.W) {
245276
val = 0.0;
246277
}
247-
Bsh[B_ly * Bsh_stride + B_lx] = val;
278+
Bsh[B_ly * Bsh_stride + B_lx] = SHMEM_TYPE(val);
248279
}
249280
barrier();
281+
#ifdef COOPMAT2
282+
coopmat<float16_t, gl_ScopeWorkgroup, BS_K, BS_CRS, gl_MatrixUseA> matA;
283+
coopmat<float16_t, gl_ScopeWorkgroup, BS_CRS, BS_NPQ, gl_MatrixUseB> matB;
284+
285+
coopMatLoad(matA, Ash, 0, Ash_stride, gl_CooperativeMatrixLayoutRowMajor);
286+
coopMatLoad(matB, Bsh, 0, Bsh_stride, gl_CooperativeMatrixLayoutRowMajor);
287+
matC = coopMatMulAdd(matA, matB, matC);
288+
#else
250289
if (T_y * TS_K < K) {
251290
[[unroll]] for (uint32_t CRS_lidx = 0; CRS_lidx < BS_CRS; CRS_lidx++) {
291+
float regA[TS_K];
292+
float regB[TS_NPQ];
252293
for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
253294
regA[T_ly] = Ash[(T_y * TS_K + T_ly) * Ash_stride + CRS_lidx];
254295
}
@@ -262,9 +303,13 @@ void main() {
262303
}
263304
}
264305
}
306+
#endif
265307
barrier();
266308
}
267309
/* Save C* */
310+
#ifdef COOPMAT2
311+
coopMatPerElementNV(matC, matC, perElemOpStore);
312+
#else
268313
if (T_y * TS_K < K) {
269314
for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
270315
for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
@@ -280,4 +325,5 @@ void main() {
280325
}
281326
}
282327
}
328+
#endif
283329
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -658,6 +658,9 @@ void process_shaders() {
658658
string_to_spv("conv2d_f32", "conv2d_mm.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}});
659659
string_to_spv("conv2d_f16_f32", "conv2d_mm.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}});
660660

661+
string_to_spv("conv2d_f32", "conv2d_mm.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}, {"COOPMAT2", "1"}}, true, false, true);
662+
string_to_spv("conv2d_f16_f32", "conv2d_mm.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"USE_COLLECTIVES", "1"}, {"COOPMAT2", "1"}}, true, false, true);
663+
661664
string_to_spv("conv2d_dw_whcn_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"WHCN", "1"}}));
662665
string_to_spv("conv2d_dw_cwhn_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"CWHN", "1"}}));
663666

0 commit comments

Comments
 (0)