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