hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

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

Jeff Bolzvulkan: combine duplicated fastdiv functions, rename the one optimizing small divs (#27526)c9ca51c1f2026-08-29

1#version 450
2
3#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
9
10#ifdef COOPMAT
11#extension GL_KHR_cooperative_matrix : enable
12#extension GL_KHR_shader_subgroup_basic : enable
13#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
14#extension GL_KHR_memory_scope_semantics : enable
15#endif
16
17#include "types.glsl"
18#include "utils.glsl"
19
20// shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j
21layout(binding = 0) readonly buffer A {
22    A_TYPE knl_data[];
23};  // src0 - kernel:   [KW, KH, KD, IC*OC]
24
25layout(binding = 1) readonly buffer B {
26    B_TYPE src_data[];
27};  // src1 - input:    [IW, IH, ID, IC*N] -- channel_first format
28
29layout(binding = 2) writeonly buffer D {
30    D_TYPE dst_data[];
31};  // dst - result:    [OW, OH, OD, OC*N]
32
33layout(push_constant) uniform parameter {
34    // I/O channels, batch size
35    uint32_t OC;
36    uint32_t IC;
37    uint32_t N;
38
39    // Tensor spatial sizes: input, output
40    uint32_t IW;
41    uint32_t IH;
42    uint32_t ID;
43    uint32_t OW;
44    uint32_t OH;
45    uint32_t OD;
46
47    // Strides in elements
48    uint32_t nb01;
49    uint32_t nb02;
50    uint32_t nb03;
51
52    uint32_t nb11;
53    uint32_t nb12;
54    uint32_t nb13;
55
56    uint32_t nb1;
57    uint32_t nb2;
58    uint32_t nb3;
59
60    // fastdiv helper values
61    uint32_t OWmp;   uint32_t OWL;
62    uint32_t OWOHmp; uint32_t OWOHL;
63    uint32_t OWOHODmp; uint32_t OWOHODL;
64}
65
66p;
67
68layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
69// Blocktile sizes
70layout(constant_id = 1) const uint BS_K            = 128;
71layout(constant_id = 2) const uint BS_CRS          = 16;
72layout(constant_id = 3) const uint BS_NPQ          = 128;
73// Thread-tile sizes
74layout(constant_id = 4) const uint TS_K            = 8;
75layout(constant_id = 5) const uint SHMEM_PAD       = 4;
76// Stride, padding, dilation
77layout(constant_id = 6)  const uint s0             = 1;
78layout(constant_id = 7)  const uint s1             = 1;
79layout(constant_id = 8)  const uint s2             = 1;
80layout(constant_id = 9)  const uint p0             = 0;
81layout(constant_id = 10) const uint p1             = 0;
82layout(constant_id = 11) const uint p2             = 0;
83layout(constant_id = 12) const uint d0             = 1;
84layout(constant_id = 13) const uint d1             = 1;
85layout(constant_id = 14) const uint d2             = 1;
86// Kernel spatial sizes
87layout(constant_id = 15) const uint KW             = 1;
88layout(constant_id = 16) const uint KH             = 1;
89layout(constant_id = 17) const uint KD             = 1;
90// when set, skip bounds checks and address clamps (K/CRS/NPQ are tile-aligned)
91layout(constant_id = 18) const uint aligned        = 0;
92// stage cm2 result through shmem (Csh) for coalesced stores. cm1 always does this.
93layout(constant_id = 19) const uint csh_store      = 0;
94
95#ifdef COOPMAT
96// cm1 subgroup tile: each subgroup computes a WM x WN region as a grid of
97// TM x TN x TK fragments. Requires WM%TM == WN%TN == BS_K%WM == BS_NPQ%WN ==
98// BS_CRS%TK == 0, and WG_SIZE == (BS_K/WM) * (BS_NPQ/WN) * subgroup_size.
99layout(constant_id = 20) const uint WM             = 32;
100layout(constant_id = 21) const uint WN             = 32;
101const uint TM = 16;
102const uint TN = 16;
103const uint TK = 16;
104const uint cms_per_row = WM / TM;
105const uint cms_per_col = WN / TN;
106const uint warps_M     = BS_K / WM;
107const uint warps_N     = BS_NPQ / WN;
108#endif
109
110// without padding, ID_idx/IH_idx/IW_idx are in bounds by construction
111const bool dhw_in_bounds = (p0 == 0) && (p1 == 0) && (p2 == 0);
112
113uint32_t       tid     = gl_LocalInvocationID.x;
114const uint32_t WG_SIZE = gl_WorkGroupSize.x;
115
116uint splitWork(uint work_size, uint block_size) {
117    return (block_size + work_size - 1) / block_size;
118}
119
120uint32_t K   = p.OC;
121uint32_t CRS = p.IC * KD * KH * KW;
122uint32_t NPQ = p.N * p.OD * p.OH * p.OW;
123
124// Number of blocktiles per input
125uint32_t NB_CRS = splitWork(CRS, BS_CRS);
126
127#if defined(COOPMAT2) || defined(COOPMAT)
128#define SHMEM_TYPE float16_t
129#else
130#define SHMEM_TYPE float
131#endif
132
133const uint32_t Ash_stride = BS_CRS + SHMEM_PAD;
134const uint32_t Bsh_stride = BS_NPQ + SHMEM_PAD;
135
136const uint32_t Ash_len = BS_K * Ash_stride;
137const uint32_t Bsh_len = BS_CRS * Bsh_stride;
138
139shared SHMEM_TYPE Ash[Ash_len];  // K x CRS
140shared SHMEM_TYPE Bsh[Bsh_len];  // CRS x NPQ
141
142#if defined(COOPMAT2) || defined(COOPMAT)
143// stage matC through shmem so global stores are row-major (NPQ-contiguous)
144const uint32_t Csh_stride = BS_NPQ;
145#ifdef COOPMAT
146const uint32_t Csh_len    = BS_K * Csh_stride;
147#else
148const uint32_t Csh_len    = csh_store != 0 ? BS_K * Csh_stride : 8; // 8 to workaround compiler bug
149#endif
150shared SHMEM_TYPE Csh[Csh_len];  // K x NPQ
151#endif
152
153// Threadtile sizes
154const uint32_t TS_NPQ = BS_K * BS_NPQ / WG_SIZE / TS_K;
155
156// Number of threadtiles per blocktile
157const uint32_t NT_NPQ = BS_NPQ / TS_NPQ;
158
159/*
160Compute
161KxCRS @ CRSxNPQ = K x NPQ
162K=OC
163C=IC
164D,R,S=KD,KH,KW
165Z,P,Q=OD,OH,OW
166*/
167
168uint32_t B_idx_K   = gl_WorkGroupID.x;
169uint32_t B_idx_NPQ = gl_WorkGroupID.y + gl_WorkGroupID.z * 512;
170
171uint32_t T_y = tid / NT_NPQ;
172uint32_t T_x = tid % NT_NPQ;
173
174uint32_t       Ar    = tid / BS_CRS;
175uint32_t       Ac    = tid % BS_CRS;
176const uint32_t ArpWg = WG_SIZE / BS_CRS;
177
178uint32_t       Br    = tid / BS_NPQ;
179uint32_t       Bc    = tid % BS_NPQ;
180const uint32_t BrpWg = WG_SIZE / BS_NPQ;
181
182void split_crs(uint32_t crs_idx, out uint32_t ic, out uint32_t kd, out uint32_t kh, out uint32_t kw) {
183    const uint32_t KHKW = KH * KW;
184    const uint32_t KDKHKW = KD * KHKW;
185    ic = crs_idx / KDKHKW;
186    uint32_t rem = crs_idx - ic * KDKHKW;
187    kd = rem / KHKW;
188    rem = rem - kd * KHKW;
189    kh = rem / KW;
190    kw = rem - kh * KW;
191}
192
193void split_npq(uint32_t npq_idx, out uint32_t n, out uint32_t od, out uint32_t oh, out uint32_t ow) {
194    const uint32_t OWOH = p.OW * p.OH;
195    n = fastdiv(npq_idx, p.OWOHODmp, p.OWOHODL);
196    uint32_t rem = npq_idx - n * p.OD * OWOH;
197    od = fastdiv(rem, p.OWOHmp, p.OWOHL);
198    rem = rem - od * OWOH;
199    oh = fastdiv(rem, p.OWmp, p.OWL);
200    ow = rem - oh * p.OW;
201}
202
203#ifdef COOPMAT2
204#define ACC_TYPE float16_t
205
206ACC_TYPE perElemOpStore(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem)
207{
208    uint32_t K_idx   = B_idx_K * BS_K + r;
209    uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + c;
210    uint32_t N_idx;
211    uint32_t OD_idx;
212    uint32_t OH_idx;
213    uint32_t OW_idx;
214    split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
215    uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + OD_idx * p.nb2 + (N_idx * p.OC + K_idx) * p.nb3;
216    if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
217        dst_data[dst_idx] = D_TYPE(elem);
218    }
219    return elem;
220}
221#endif
222
223void main() {
224    if (B_idx_NPQ * BS_NPQ >= NPQ) {
225        return;
226    }
227
228#ifdef COOPMAT2
229    coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator> matC;
230    matC = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator>(0.0);
231#elif defined(COOPMAT)
232    coopmat<float16_t, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> sums[cms_per_row * cms_per_col];
233    [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) {
234        sums[i] = coopmat<float16_t, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0);
235    }
236    const uint warp_r = gl_SubgroupID / warps_N;
237    const uint warp_c = gl_SubgroupID % warps_N;
238#else
239    float regC[TS_K][TS_NPQ];
240    for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
241        for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
242            regC[T_ly][T_lx] = 0.0;
243        }
244    }
245#endif
246    /* Advance block in CRS dim */
247    [[dont_unroll]] for (uint32_t B_idx_CRS = 0; B_idx_CRS < NB_CRS; B_idx_CRS++) {
248        uint32_t CRS_idx_a = B_idx_CRS * BS_CRS + Ac;
249        uint32_t IC_idx_a;
250        uint32_t KD_idx_a;
251        uint32_t KH_idx_a;
252        uint32_t KW_idx_a;
253        split_crs(CRS_idx_a, IC_idx_a, KD_idx_a, KH_idx_a, KW_idx_a);
254
255        /* Load kernel to A_block: (BS_K x BS_CRS)*/
256        UNROLL for (uint32_t r_offset = 0; r_offset < BS_K; r_offset += ArpWg) {
257            uint32_t B_ly    = r_offset + Ar;
258            uint32_t B_lx    = Ac;
259            uint32_t K_idx   = B_idx_K * BS_K + B_ly; /* Global K_idx (row index of A)*/
260            uint32_t knl_idx = KW_idx_a + KH_idx_a * p.nb01 + KD_idx_a * p.nb02 + (K_idx * p.IC + IC_idx_a) * p.nb03;
261            if (aligned == 0) {
262                knl_idx = min(knl_idx, K * CRS - 1);
263            }
264            float    val     = knl_data[knl_idx];
265            if (aligned == 0 && (K_idx >= K || CRS_idx_a >= CRS)) {
266                val = 0.0;
267            }
268            Ash[B_ly * Ash_stride + B_lx] = SHMEM_TYPE(val);
269        }
270        /* Load input to B_block: (BS_CRS x BS_NPQ) */
271        UNROLL for (uint32_t r_offset = 0; r_offset < BS_CRS; r_offset += BrpWg) {
272            uint32_t B_ly          = r_offset + Br;             /* Row index of B block */
273            uint32_t B_lx          = Bc;
274            uint32_t NPQ_idx       = B_idx_NPQ * BS_NPQ + B_lx; /* Global NPQ index (column index of B) */
275            uint32_t N_idx;
276            uint32_t OD_idx;
277            uint32_t OH_idx;
278            uint32_t OW_idx;
279            split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
280
281            uint32_t CRS_idx_b = B_idx_CRS * BS_CRS + B_ly;
282            uint32_t IC_idx_b;
283            uint32_t KD_idx_b;
284            uint32_t KH_idx_b;
285            uint32_t KW_idx_b;
286            split_crs(CRS_idx_b, IC_idx_b, KD_idx_b, KH_idx_b, KW_idx_b);
287
288            uint32_t ID_idx = OD_idx * s2 + KD_idx_b * d2 - p2;
289            uint32_t IH_idx = OH_idx * s1 + KH_idx_b * d1 - p1;
290            uint32_t IW_idx = OW_idx * s0 + KW_idx_b * d0 - p0;
291
292            uint32_t src_idx = IW_idx + IH_idx * p.nb11 + ID_idx * p.nb12 + (N_idx * p.IC + IC_idx_b) * p.nb13;
293            // skip clamp when address can't go OOB
294            if (aligned == 0 || !dhw_in_bounds) {
295                src_idx = min(src_idx, p.IC * p.N * p.IW * p.IH * p.ID - 1);
296            }
297            float val = src_data[src_idx];
298            bool oob = false;
299            if (aligned == 0 && (CRS_idx_b >= CRS || NPQ_idx >= NPQ)) {
300                oob = true;
301            }
302            // also catches lower-bound underflow (idx wraps to 0x80000000+)
303            if (!dhw_in_bounds && (ID_idx >= p.ID || IH_idx >= p.IH || IW_idx >= p.IW)) {
304                oob = true;
305            }
306            if (oob) {
307                val = 0.0;
308            }
309            Bsh[B_ly * Bsh_stride + B_lx] = SHMEM_TYPE(val);
310        }
311        barrier();
312#ifdef COOPMAT2
313        coopmat<float16_t, gl_ScopeWorkgroup, BS_K, BS_CRS, gl_MatrixUseA> matA;
314        coopmat<float16_t, gl_ScopeWorkgroup, BS_CRS, BS_NPQ, gl_MatrixUseB> matB;
315
316        coopMatLoad(matA, Ash, 0, Ash_stride, gl_CooperativeMatrixLayoutRowMajor);
317        coopMatLoad(matB, Bsh, 0, Bsh_stride, gl_CooperativeMatrixLayoutRowMajor);
318        matC = coopMatMulAdd(matA, matB, matC);
319#elif defined(COOPMAT)
320        // each subgroup multiplies its grid of fragments per TK-sized CRS chunk
321        [[unroll]] for (uint k_step = 0; k_step < BS_CRS / TK; k_step++) {
322            coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> cache_a[cms_per_row];
323            [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
324                const uint a_off = (warp_r * WM + cm_row * TM) * Ash_stride + k_step * TK;
325                coopMatLoad(cache_a[cm_row], Ash, a_off, Ash_stride, gl_CooperativeMatrixLayoutRowMajor);
326            }
327            [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
328                coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> cache_b;
329                const uint b_off = k_step * TK * Bsh_stride + warp_c * WN + cm_col * TN;
330                coopMatLoad(cache_b, Bsh, b_off, Bsh_stride, gl_CooperativeMatrixLayoutRowMajor);
331                [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
332                    sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a[cm_row], cache_b, sums[cm_col * cms_per_row + cm_row]);
333                }
334            }
335        }
336#else
337        if (T_y * TS_K < K) {
338            UNROLL for (uint32_t CRS_lidx = 0; CRS_lidx < BS_CRS; CRS_lidx++) {
339                float regA[TS_K];
340                float regB[TS_NPQ];
341                for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
342                    regA[T_ly] = Ash[(T_y * TS_K + T_ly) * Ash_stride + CRS_lidx];
343                }
344                for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
345                    regB[T_lx] = Bsh[CRS_lidx * Bsh_stride + T_x * TS_NPQ + T_lx];
346                }
347                for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
348                    for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
349                        regC[T_ly][T_lx] = fma(regA[T_ly], regB[T_lx], regC[T_ly][T_lx]);
350                    }
351                }
352            }
353        }
354#endif
355        barrier();
356    }
357    /* Save C* */
358#if defined(COOPMAT2) || defined(COOPMAT)
359    // stage matC into Csh, then write to dst with coalesced NPQ-contiguous stores
360#ifdef COOPMAT
361    const bool use_staged_store = true;
362#else
363    const bool use_staged_store = (csh_store != 0);
364#endif
365    if (use_staged_store) {
366#ifdef COOPMAT
367        // cm1: each subgroup stores its fragment grid into its Csh slot
368        [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
369            [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
370                const uint csh_off = (warp_r * WM + cm_row * TM) * Csh_stride + warp_c * WN + cm_col * TN;
371                coopMatStore(sums[cm_col * cms_per_row + cm_row], Csh, csh_off, Csh_stride, gl_CooperativeMatrixLayoutRowMajor);
372            }
373        }
374#else
375        coopMatStore(matC, Csh, 0, Csh_stride, gl_CooperativeMatrixLayoutRowMajor);
376#endif
377        barrier();
378
379        // cooperative shmem->global: WG threads spread across BS_NPQ (the
380        // contiguous direction of dst), each iter covers store_rows_per_iter K-rows
381        const uint32_t store_rows_per_iter = WG_SIZE / BS_NPQ;
382        const uint32_t store_iters         = BS_K / store_rows_per_iter;
383        const uint32_t k_thread_offset     = tid / BS_NPQ;
384        const uint32_t npq_thread          = tid % BS_NPQ;
385        [[unroll]] for (uint32_t i = 0; i < store_iters; i++) {
386            uint32_t k_local = i * store_rows_per_iter + k_thread_offset;
387            uint32_t K_idx   = B_idx_K * BS_K + k_local;
388            uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + npq_thread;
389            uint32_t N_idx;
390            uint32_t OD_idx;
391            uint32_t OH_idx;
392            uint32_t OW_idx;
393            split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
394            uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + OD_idx * p.nb2 + (N_idx * p.OC + K_idx) * p.nb3;
395            if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
396                dst_data[dst_idx] = D_TYPE(Csh[k_local * Csh_stride + npq_thread]);
397            }
398        }
399    }
400#ifdef COOPMAT2
401    else {
402        coopMatPerElementNV(matC, matC, perElemOpStore);
403    }
404#endif
405#else
406    if (T_y * TS_K < K) {
407        for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
408            for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
409                uint32_t K_idx   = B_idx_K * BS_K + T_y * TS_K + T_ly;
410                uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + T_x * TS_NPQ + T_lx;
411                uint32_t N_idx;
412                uint32_t OD_idx;
413                uint32_t OH_idx;
414                uint32_t OW_idx;
415                split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
416                uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + OD_idx * p.nb2 + (N_idx * p.OC + K_idx) * p.nb3;
417                if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
418                    dst_data[dst_idx] = D_TYPE(regC[T_ly][T_lx]);
419                }
420            }
421        }
422    }
423#endif
424}