char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ruben Ortlamvulkan: use spec constant for matrix matrix multiplication A-type (#25773)91f6a6cf3

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
20.1 KiB544 linesraw
1#version 450
2
3#extension GL_EXT_control_flow_attributes : enable
4#extension GL_EXT_shader_16bit_storage : require
5
6#ifdef FLOAT16
7#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
8#endif
9#if defined(DATA_A_IQ1_M)
10#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
11#endif
12#if !defined(DATA_A_F32) && !defined(DATA_A_F16) && !defined(DATA_A_BF16)
13#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
14#endif
15
16#if defined(DATA_A_BF16) && defined(COOPMAT)
17#extension GL_EXT_bfloat16 : enable
18#endif
19
20#ifdef COOPMAT
21#extension GL_KHR_cooperative_matrix : enable
22#extension GL_KHR_memory_scope_semantics : enable
23#endif
24
25#if defined(COOPMAT) || defined(MUL_MAT_ID_USE_SUBGROUPS)
26#extension GL_KHR_shader_subgroup_basic : enable
27#extension GL_KHR_shader_subgroup_ballot : enable
28#endif
29
30#ifdef MUL_MAT_ID
31#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
32#endif
33
34#ifdef MULMAT_QUANT
35#include "ggml_type_ids.glsl"
36layout (constant_id = 12) const uint MmTypeA = 0;
37#endif
38
39#include "types.glsl"
40#include "dot_product_funcs.glsl"
41
42#ifndef MULMAT_QUANT
43#ifndef LOAD_VEC_A
44#define LOAD_VEC_A 1
45#endif
46#endif
47#ifndef LOAD_VEC_B
48#define LOAD_VEC_B 1
49#endif
50
51layout (constant_id = 11) const uint ALIGNED = 0;
52
53#ifdef MULMAT_QUANT
54
55uint mm_load_vec_a() {
56    switch (MmTypeA) {
57    case GGML_TYPE_Q1_0:
58    case GGML_TYPE_Q4_0:
59    case GGML_TYPE_Q4_1:
60    case GGML_TYPE_Q5_1:
61        return 8u;
62    case GGML_TYPE_Q2_0:
63    case GGML_TYPE_Q5_0:
64    case GGML_TYPE_Q8_0:
65    case GGML_TYPE_Q2_K:
66    case GGML_TYPE_Q4_K:
67    case GGML_TYPE_Q5_K:
68        return 4u;
69    default:
70        return 2u;
71    }
72}
73#endif
74
75#if !defined(TO_FLOAT_TYPE)
76#define TO_FLOAT_TYPE FLOAT_TYPE
77#endif
78
79layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
80
81#ifndef MULMAT_QUANT
82layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
83#if defined(DATA_A_F32)
84layout (binding = 0) readonly buffer A_SCALAR {float data_a_scalar[];};
85#elif defined(DATA_A_F16)
86layout (binding = 0) readonly buffer A_SCALAR {float16_t data_a_scalar[];};
87#elif defined(DATA_A_BF16)
88layout (binding = 0) readonly buffer A_SCALAR {uint16_t data_a_scalar[];};
89#endif
90#if defined(A_TYPE_PACKED16)
91layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16[];};
92#endif
93#if defined(A_TYPE_PACKED32)
94layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];};
95#endif
96#else
97// Unpacked struct aliases
98layout (binding = 0) readonly buffer BUF_Q1_0     { block_q1_0 data[];     } a_q1_0;
99layout (binding = 0) readonly buffer BUF_Q2_0     { block_q2_0 data[];     } a_q2_0;
100layout (binding = 0) readonly buffer BUF_Q2_K     { block_q2_K data[];     } a_q2_k;
101layout (binding = 0) readonly buffer BUF_Q3_K     { block_q3_K data[];     } a_q3_k;
102layout (binding = 0) readonly buffer BUF_Q4_K     { block_q4_K data[];     } a_q4_k;
103layout (binding = 0) readonly buffer BUF_Q5_K     { block_q5_K data[];     } a_q5_k;
104layout (binding = 0) readonly buffer BUF_Q6_K     { block_q6_K data[];     } a_q6_k;
105layout (binding = 0) readonly buffer BUF_TQ1_0    { block_tq1_0 data[];    } a_tq1_0;
106layout (binding = 0) readonly buffer BUF_TQ2_0    { block_tq2_0 data[];    } a_tq2_0;
107// Packed16 aliases
108layout (binding = 0) readonly buffer BUF_Q4_0_P16   { block_q4_0_packed16 data[];   } a_q4_0_p16;
109layout (binding = 0) readonly buffer BUF_Q5_0_P16   { block_q5_0_packed16 data[];   } a_q5_0_p16;
110layout (binding = 0) readonly buffer BUF_Q8_0_P16   { block_q8_0_packed16 data[];   } a_q8_0_p16;
111layout (binding = 0) readonly buffer BUF_Q3_K_P16   { block_q3_K_packed16 data[];   } a_q3_k_p16;
112layout (binding = 0) readonly buffer BUF_Q6_K_P16   { block_q6_K_packed16 data[];   } a_q6_k_p16;
113// Packed32 aliases
114layout (binding = 0) readonly buffer BUF_Q4_1_P32   { block_q4_1_packed32 data[];   } a_q4_1_p32;
115layout (binding = 0) readonly buffer BUF_Q5_1_P32   { block_q5_1_packed32 data[];   } a_q5_1_p32;
116layout (binding = 0) readonly buffer BUF_Q2_K_P32   { block_q2_K_packed32 data[];   } a_q2_k_p32;
117layout (binding = 0) readonly buffer BUF_Q4_K_P32   { block_q4_K_packed32 data[];   } a_q4_k_p32;
118layout (binding = 0) readonly buffer BUF_Q5_K_P32   { block_q5_K_packed32 data[];   } a_q5_k_p32;
119#endif
120
121layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
122layout (binding = 1) readonly buffer B_SCALAR {B_TYPE_SCALAR data_b_scalar[];};
123layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
124
125#ifdef MUL_MAT_ID
126layout (binding = 3) readonly buffer IDS {int data_ids[];};
127layout (binding = 4) readonly buffer Counts {int data_expert_count[];};
128#endif
129
130layout (push_constant) uniform parameter
131{
132    uint M;
133    uint N;
134    uint K;
135    uint stride_a;
136    uint stride_b;
137    uint stride_d;
138
139    uint batch_stride_a;
140    uint batch_stride_b;
141    uint batch_stride_d;
142
143#ifdef MUL_MAT_ID
144    uint nei0;
145    uint nei1;
146    uint nbi1;
147    uint ne11;
148    uint n_experts;
149    uint hoist_row_ids;
150#else
151    uint base_work_group_z;
152    uint num_batches;
153    uint k_split;
154    uint ne02;
155    uint ne12;
156    uint broadcast2;
157    uint broadcast3;
158#endif
159} p;
160
161layout (constant_id = 0) const uint BLOCK_SIZE = 64;
162layout (constant_id = 1) const uint BM = 64;
163layout (constant_id = 2) const uint BN = 64;
164layout (constant_id = 4) const uint WM = 32;
165layout (constant_id = 5) const uint WN = 32;
166layout (constant_id = 6) const uint WMITER = 2;
167layout (constant_id = 7) const uint TM = 4;
168layout (constant_id = 8) const uint TN = 2;
169layout (constant_id = 9) const uint TK = 1;  // Only needed for coopmat
170layout (constant_id = 10) const uint WARP = 32;
171
172#if defined(DATA_A_F32) || defined(DATA_A_F16)
173#define BK 32
174#define BK_STEP 4
175#else
176layout (constant_id = 3) const uint BK = 16;  // Assumed to be 32 if working with a quant
177#define BK_STEP 2
178#endif
179
180#ifdef COOPMAT
181#ifdef MULMAT_QUANT
182layout(constant_id = 13) const uint SHMEM_STRIDE_PAD = 4;
183layout(constant_id = 14) const bool APPLY_SLM_A_RESHAPE = false;
184#else
185layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4;
186layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false;
187#endif
188#else
189const uint SHMEM_STRIDE_PAD = 1;
190const bool APPLY_SLM_A_RESHAPE = false;
191#endif
192#define SHMEM_STRIDE (BK / 2 + SHMEM_STRIDE_PAD)
193
194shared FLOAT_TYPEV2 buf_a[BM * SHMEM_STRIDE];
195shared FLOAT_TYPEV2 buf_b[BN * SHMEM_STRIDE];
196
197#define NUM_WARPS (BLOCK_SIZE / WARP)
198
199#ifdef COOPMAT
200shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS];
201#endif
202
203#include "mul_mm_id_funcs.glsl"
204#include "mul_mm_funcs.glsl"
205
206#ifdef MULMAT_QUANT
207#include "iq_shmem_init.glsl"
208#endif
209
210void main() {
211    const uint ic = gl_WorkGroupID.y;
212
213#ifdef MUL_MAT_ID
214    const uint expert_idx = gl_WorkGroupID.z;
215    if (ic * BN >= data_expert_count[expert_idx]) {
216        return;
217    }
218#endif
219#if defined(NEEDS_INIT_IQ_SHMEM) || defined(MULMAT_QUANT)
220    init_iq_shmem(gl_WorkGroupSize);
221#endif
222
223#ifndef MUL_MAT_ID
224    const uint batch_idx = gl_WorkGroupID.z + p.base_work_group_z;
225
226    const uint i13 = batch_idx / p.ne12;
227    const uint i12 = batch_idx % p.ne12;
228
229    const uint i03 = i13 / p.broadcast3;
230    const uint i02 = i12 / p.broadcast2;
231
232    const uint batch_idx_a = i03 * p.ne02 + i02;
233#endif
234
235    const uint blocks_m = (p.M + BM - 1) / BM;
236    const uint ir = gl_WorkGroupID.x % blocks_m;
237    const uint ik = gl_WorkGroupID.x / blocks_m;
238
239    const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER);
240    const uint WSUBM = WM / WMITER;
241    const uint WSUBN = WN / WNITER;
242
243#ifdef COOPMAT
244    const uint warp_i = gl_SubgroupID;
245
246    const uint tiw = gl_SubgroupInvocationID;
247
248    const uint cms_per_row = WM / TM;
249    const uint cms_per_col = WN / TN;
250
251    const uint storestride = WARP / TM;
252    const uint store_r = tiw % TM;
253    const uint store_c = tiw / TM;
254#else
255    const uint warp_i = gl_LocalInvocationID.x / WARP;
256
257    const uint tiw = gl_LocalInvocationID.x % WARP;
258
259    const uint tiwr = tiw % (WSUBM / TM);
260    const uint tiwc = tiw / (WSUBM / TM);
261#endif
262
263    const uint warp_r = warp_i % (BM / WM);
264    const uint warp_c = warp_i / (BM / WM);
265
266#if defined(DATA_A_F32) || defined(DATA_A_F16) || defined(DATA_A_BF16)
267    const uint LOAD_VEC_A_EFF = (ALIGNED != 0) ? LOAD_VEC_A : 1;
268    const uint LOAD_VEC_BATCH_A = (ALIGNED != 0) ? 1 : 2;
269#elif !defined(MULMAT_QUANT)
270    const uint LOAD_VEC_A_EFF = LOAD_VEC_A;
271    const uint LOAD_VEC_BATCH_A = 1;
272#else
273    const uint LOAD_VEC_A_EFF = mm_load_vec_a();
274    const uint LOAD_VEC_BATCH_A = 1;
275#endif
276    const uint LOAD_VEC_B_EFF = (ALIGNED != 0) ? LOAD_VEC_B : 1;
277    const uint LOAD_VEC_BATCH_B = (ALIGNED != 0) ? 1 : 2;
278
279    const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A_EFF / LOAD_VEC_BATCH_A);
280    const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A_EFF / LOAD_VEC_BATCH_A);
281    const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B_EFF / LOAD_VEC_BATCH_B);
282    const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B_EFF / LOAD_VEC_BATCH_B);
283
284    const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A_EFF * LOAD_VEC_BATCH_A / BK;
285    const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK;
286
287#ifdef MUL_MAT_ID
288    if (p.hoist_row_ids != 0) {
289        load_row_ids_hoisted(expert_idx, ic);
290    } else {
291#ifdef MUL_MAT_ID_USE_SUBGROUPS
292        if (bitCount(p.nei0) == 1) {
293            load_row_ids(expert_idx, true, ic);
294        } else {
295            load_row_ids(expert_idx, false, ic);
296        }
297#else
298        _ne1 = 0;
299        for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
300            for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
301                if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
302                    if (_ne1 >= ic * BN) {
303                        row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
304                    }
305                    _ne1++;
306                }
307            }
308        }
309
310        barrier();
311#endif
312    }
313
314    // Workgroup has no work
315    if (ic * BN >= _ne1) return;
316#endif
317
318#ifdef MUL_MAT_ID
319    const uint start_k = 0;
320    const uint end_k = p.K;
321#else
322    const uint start_k = ik * p.k_split;
323    const uint end_k = min(p.K, (ik + 1) * p.k_split);
324#endif
325
326    uint pos_a =
327#ifdef MUL_MAT_ID
328        expert_idx * (p.batch_stride_a / LOAD_VEC_A_EFF) +
329#else
330        batch_idx_a * (p.batch_stride_a / LOAD_VEC_A_EFF) +
331#endif
332        (ir * BM * p.stride_a + start_k) / LOAD_VEC_A_EFF;
333#ifdef MUL_MAT_ID
334    uint pos_b = 0;
335#else
336    uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B_EFF;
337#endif
338
339#ifdef COOPMAT
340    coopmat<FLOAT_TYPE, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> cache_a;
341    coopmat<FLOAT_TYPE, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> cache_b;
342    coopmat<ACC_TYPE, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> sums[cms_per_row * cms_per_col];
343
344    [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) {
345        sums[i] = coopmat<ACC_TYPE, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
346    }
347#else
348    ACC_TYPEV2 sums[WMITER * TM * WNITER * TN/2];
349#if defined(DATA_A_F32) || defined(DATA_A_F16)
350    FLOAT_TYPEV4 cache_a[WMITER * TM];
351    FLOAT_TYPEV4 cache_b;
352#else
353    FLOAT_TYPEV2 cache_a[WMITER * TM];
354    FLOAT_TYPEV2 cache_b;
355#endif
356
357    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN/2; i++) {
358        sums[i] = ACC_TYPEV2(0.0f, 0.0f);
359    }
360#endif
361
362    for (uint block = start_k; block < end_k; block += BK) {
363        [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) {
364            load_a_to_shmem(pos_a, loadr_a, loadc_a + l, ir * BM + loadc_a + l, block, end_k);
365        }
366        [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) {
367#if !defined(MUL_MAT_ID)
368            load_b_to_shmem(pos_b, loadr_b, loadc_b + l, ic * BN + loadc_b + l, block, end_k);
369#else
370            load_b_to_shmem(pos_b, loadr_b, loadc_b + l, ic, _ne1, block, end_k);
371#endif
372        }
373
374        barrier();
375
376        pos_a += BK / LOAD_VEC_A_EFF;
377        pos_b += BK / LOAD_VEC_B_EFF;
378
379#ifdef COOPMAT
380        [[unroll]] for (uint i = 0; i < BK; i += TK) {
381            [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
382                // Load from shared into cache
383                coopMatLoad(cache_a, buf_a, a_shmem_index(warp_r * WM + cm_row * TM, i / 2), a_shmem_stride(), gl_CooperativeMatrixLayoutRowMajor);
384
385                [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
386                    coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
387
388                    sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, sums[cm_col * cms_per_row + cm_row]);
389                }
390            }
391        }
392#else
393        [[unroll]] for (uint i = 0; i < BK / BK_STEP; i++) {
394            // Load from shared into cache
395            [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
396                [[unroll]] for (uint j = 0; j < TM; j++) {
397                #if defined(DATA_A_F32) || defined(DATA_A_F16)
398                    cache_a[wsir * TM + j].xy = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + 2 * i    ];
399                    cache_a[wsir * TM + j].zw = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + 2 * i + 1];
400                #else
401                    cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * SHMEM_STRIDE + i];
402                #endif
403                }
404            }
405
406            [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
407                [[unroll]] for (uint cc = 0; cc < TN; cc++) {
408                #if defined(DATA_A_F32) || defined(DATA_A_F16)
409                    cache_b.xy = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + cc) * SHMEM_STRIDE + 2 * i    ];
410                    cache_b.zw = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + cc) * SHMEM_STRIDE + 2 * i + 1];
411                #else
412                    cache_b = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + cc) * SHMEM_STRIDE + i];
413                #endif
414
415                    [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
416                        [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) {
417                            // [WNITER][TN][WMITER][TM / 2] -> [wsic][cc][wsir][cr]
418                            const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr;
419                            sums[sums_idx].x = dot_product(cache_a[wsir * TM + 2 * cr    ], cache_b, sums[sums_idx].x);
420                            sums[sums_idx].y = dot_product(cache_a[wsir * TM + 2 * cr + 1], cache_b, sums[sums_idx].y);
421                        }
422                    }
423                }
424            }
425
426        }
427#endif
428
429        barrier();
430    }
431
432#if defined(ACC_TYPE_MAX)
433#ifdef COOPMAT
434    [[unroll]] for (uint j = 0; j < cms_per_row * cms_per_col; j++) {
435        [[unroll]] for (uint i = 0; i < sums[j].length(); ++i) {
436            sums[j][i] = clamp(sums[j][i], -ACC_TYPE_MAX, ACC_TYPE_MAX);
437        }
438    }
439#else
440    [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN/2; i++) {
441        sums[i].x = clamp(sums[i].x, -ACC_TYPE_MAX, ACC_TYPE_MAX);
442        sums[i].y = clamp(sums[i].y, -ACC_TYPE_MAX, ACC_TYPE_MAX);
443    }
444#endif
445#endif
446
447    const uint dr = ir * BM + warp_r * WM;
448    const uint dc = ic * BN + warp_c * WN;
449
450#ifndef MUL_MAT_ID
451    const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches;
452#endif
453
454#ifdef COOPMAT
455#ifdef MUL_MAT_ID
456    [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
457        [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
458            coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor);
459
460            barrier();
461            [[unroll]] for (uint col = 0; col < TN; col += storestride) {
462                const uint row_i = dc + cm_col * TN + col + store_c;
463                if (row_i >= _ne1) break;
464
465                const u16vec2 row_idx = row_ids[row_i - ic * BN];
466
467                if (dr + cm_row * TM + store_r < p.M) {
468                    data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]);
469                }
470            }
471            barrier();
472        }
473    }
474#else
475    const bool is_aligned = p.stride_d % 4 == 0;  // Assumption: D_TYPE == float
476
477    [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
478        [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
479            const bool is_in_bounds = dr + (cm_row + 1) * TM <= p.M && dc + (cm_col + 1) * TN <= p.N;
480
481            if (is_aligned && is_in_bounds) {
482                // Full coopMat is within bounds and stride_d is aligned with 16B
483                coopmat<D_TYPE, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> cm_dtype = coopmat<D_TYPE, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(sums[cm_col * cms_per_row + cm_row]);
484                coopMatStore(cm_dtype, data_d, offsets + (dc + cm_col * TN) * p.stride_d + dr + cm_row * TM, p.stride_d, gl_CooperativeMatrixLayoutColumnMajor);
485            } else if (is_in_bounds) {
486                // Full coopMat is within bounds, but stride_d is not aligned
487                coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor);
488
489                controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
490                [[unroll]] for (uint col = 0; col < TN; col += storestride) {
491                    data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]);
492                }
493                controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
494            } else if (dr + cm_row * TM < p.M && dc + cm_col * TN < p.N) {
495                // Partial coopMat is within bounds
496                coopMatStore(sums[cm_col * cms_per_row + cm_row], coopmat_stage, warp_i * TM * TN, TM, gl_CooperativeMatrixLayoutColumnMajor);
497
498                controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
499                [[unroll]] for (uint col = 0; col < TN; col += storestride) {
500                    if (dr + cm_row * TM + store_r < p.M && dc + cm_col * TN + col + store_c < p.N) {
501                        data_d[offsets + (dc + cm_col * TN + col + store_c) * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]);
502                    }
503                }
504                controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
505            }
506        }
507    }
508#endif // MUL_MAT_ID
509#else
510    [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
511        [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
512
513            const uint dr_warp = dr + wsir * WSUBM + tiwr * TM;
514            const uint dc_warp = dc + wsic * WSUBN + tiwc * TN;
515            [[unroll]] for (uint cc = 0; cc < TN; cc++) {
516#ifdef MUL_MAT_ID
517                const uint row_i = dc_warp + cc;
518                if (row_i >= _ne1) break;
519
520                const u16vec2 row_idx = row_ids[row_i - ic * BN];
521#endif // MUL_MAT_ID
522                [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) {
523                    const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr;
524#ifdef MUL_MAT_ID
525                    if (dr_warp + 2 * cr < p.M) {
526                        data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr] = D_TYPE(sums[sums_idx].x);
527                    }
528                    if (dr_warp + 2 * cr + 1 < p.M) {
529                        data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr + 1] = D_TYPE(sums[sums_idx].y);
530                    }
531#else
532                    if (dr_warp + 2 * cr < p.M && dc_warp + cc < p.N) {
533                        data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + 2 * cr] = D_TYPE(sums[sums_idx].x);
534                    }
535                    if (dr_warp + 2 * cr + 1 < p.M && dc_warp + cc < p.N) {
536                        data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + 2 * cr + 1] = D_TYPE(sums[sums_idx].y);
537                    }
538#endif // MUL_MAT_ID
539                }
540            }
541        }
542    }
543#endif // COOPMAT
544}