1void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uint idx_m, const uint block, const uint end_k) {
2#if defined(DATA_A_F32) || defined(DATA_A_F16)
3#if LOAD_VEC_A == 8
4 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
5 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
6 FLOAT_TYPE_VEC8 aa = FLOAT_TYPE_VEC8(data_a[idx]);
7 buf_a[buf_idx ] = aa[0].xy;
8 buf_a[buf_idx + 1] = aa[0].zw;
9 buf_a[buf_idx + 2] = aa[1].xy;
10 buf_a[buf_idx + 3] = aa[1].zw;
11#elif LOAD_VEC_A == 4
12 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
13 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
14 FLOAT_TYPE_VEC4 aa = FLOAT_TYPE_VEC4(data_a[idx]);
15 buf_a[buf_idx ] = aa.xy;
16 buf_a[buf_idx + 1] = aa.zw;
17#else // LOAD_VEC_BATCH_A == 2
18 const uint idx = pos_a + col * p.stride_a + row * 2;
19 const uint buf_idx = col * SHMEM_STRIDE + row;
20 if (idx_m < p.M && block + row * 2 + 1 < end_k) {
21 buf_a[buf_idx] = FLOAT_TYPE_VEC2(data_a[idx],
22 data_a[idx + 1]);
23 } else if (idx_m < p.M && block + row * 2 < end_k) {
24 buf_a[buf_idx] = FLOAT_TYPE_VEC2(data_a[idx], 0.0f);
25 } else {
26 buf_a[buf_idx] = FLOAT_TYPE_VEC2(0.0f);
27 }
28#endif
29#elif defined(DATA_A_BF16)
30#if LOAD_VEC_A == 4
31 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
32 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
33 FLOAT_TYPE_VEC4 aa = FLOAT_TYPE_VEC4(TO_FLOAT_TYPE(data_a[idx]));
34 buf_a[buf_idx ] = aa.xy;
35 buf_a[buf_idx + 1] = aa.zw;
36#else // LOAD_VEC_BATCH_A == 2
37 const uint idx = pos_a + col * p.stride_a + row * 2;
38 const uint buf_idx = col * SHMEM_STRIDE + row;
39 if (idx_m < p.M && block + row * 2 + 1 < end_k) {
40 buf_a[buf_idx] = FLOAT_TYPE_VEC2(TO_FLOAT_TYPE(data_a[idx]),
41 TO_FLOAT_TYPE(data_a[idx + 1]));
42 } else if (idx_m < p.M && block + row * 2 < end_k) {
43 buf_a[buf_idx] = FLOAT_TYPE_VEC2(TO_FLOAT_TYPE(data_a[idx]), 0.0f);
44 } else {
45 buf_a[buf_idx] = FLOAT_TYPE_VEC2(0.0f);
46 }
47#endif
48#elif defined(DATA_A_Q4_0)
49 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
50 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
51
52 const uint ib = idx / 4;
53 const uint iqs = idx & 0x03;
54
55 const float d = float(data_a_packed16[ib].d);
56 const uint vui = uint(data_a_packed16[ib].qs[2*iqs]) | (uint(data_a_packed16[ib].qs[2*iqs + 1]) << 16);
57 const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d;
58 const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d;
59
60 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v0.xy);
61 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2(v0.zw);
62 buf_a[buf_idx + 8] = FLOAT_TYPE_VEC2(v1.xy);
63 buf_a[buf_idx + 9] = FLOAT_TYPE_VEC2(v1.zw);
64#elif defined(DATA_A_Q4_1)
65 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
66 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
67
68 const uint ib = idx / 4;
69 const uint iqs = idx & 0x03;
70
71 const vec2 dm = vec2(data_a_packed32[ib].dm);
72 const uint vui = data_a_packed32[ib].qs[iqs];
73 const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y;
74 const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y;
75
76 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v0.xy);
77 buf_a[buf_idx + 1 ] = FLOAT_TYPE_VEC2(v0.zw);
78 buf_a[buf_idx + 8 ] = FLOAT_TYPE_VEC2(v1.xy);
79 buf_a[buf_idx + 9 ] = FLOAT_TYPE_VEC2(v1.zw);
80#elif defined(DATA_A_Q5_0)
81 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
82 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
83
84 const uint ib = idx / 8;
85 const uint iqs = idx & 0x07;
86
87 const float d = float(data_a_packed16[ib].d);
88 const uint uint_qh = uint(data_a_packed16[ib].qh[1]) << 16 | uint(data_a_packed16[ib].qh[0]);
89 const ivec2 qh0 = ivec2(((uint_qh >> 2*iqs) << 4) & 0x10, (uint_qh >> (2*iqs + 12)) & 0x10);
90 const ivec2 qh1 = ivec2(((uint_qh >> (2*iqs + 1)) << 4) & 0x10, (uint_qh >> (2*iqs + 13)) & 0x10);
91
92 const uint vui = uint(data_a_packed16[ib].qs[iqs]);
93 const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d;
94
95 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v.xz);
96 buf_a[buf_idx + 8] = FLOAT_TYPE_VEC2(v.yw);
97#elif defined(DATA_A_Q5_1)
98 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
99 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
100
101 const uint ib = idx / 4;
102 const uint iqs = idx & 0x03;
103
104 const vec2 dm = vec2(data_a_packed32[ib].dm);
105 const uint uint_qh = data_a_packed32[ib].qh;
106 const uvec2 qh0 = uvec2(((uint_qh >> 4*iqs) << 4) & 0x10, (uint_qh >> (4*iqs + 12)) & 0x10);
107 const uvec2 qh1 = uvec2(((uint_qh >> (4*iqs + 1)) << 4) & 0x10, (uint_qh >> (4*iqs + 13)) & 0x10);
108 const uvec2 qh2 = uvec2(((uint_qh >> (4*iqs + 2)) << 4) & 0x10, (uint_qh >> (4*iqs + 14)) & 0x10);
109 const uvec2 qh3 = uvec2(((uint_qh >> (4*iqs + 3)) << 4) & 0x10, (uint_qh >> (4*iqs + 15)) & 0x10);
110
111 const uint vui = data_a_packed32[ib].qs[iqs];
112 const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y;
113 const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y;
114
115 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v0.xz);
116 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2(v1.xz);
117 buf_a[buf_idx + 8] = FLOAT_TYPE_VEC2(v0.yw);
118 buf_a[buf_idx + 9] = FLOAT_TYPE_VEC2(v1.yw);
119#elif defined(DATA_A_Q8_0)
120 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
121 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
122
123 const uint ib = idx / 8;
124 const uint iqs = idx & 0x07;
125
126 const float d = float(data_a_packed16[ib].d);
127 const i8vec2 v0 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs])).xy; // vec4 used due to #12147
128 const i8vec2 v1 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs + 1])).xy;
129 const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d;
130
131 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v.xy);
132 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2(v.zw);
133#elif defined(DATA_A_Q2_K)
134 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
135 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
136
137 const uint ib = idx / 64; // 4 values per idx
138 const uint iqs = (idx % 64) * 2; // 0,2,4..126
139
140 const uint qsi = (iqs / 64) * 16 + (iqs % 16); // 0..15
141 const uint scalesi = iqs / 8; // 0..15
142 const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6
143
144 const vec4 qs = vec4(unpack8((data_a_packed32[ib].qs[qsi / 2] >> qsshift) & 0x03030303));
145 const uint scales = data_a[ib].scales[scalesi];
146 const vec2 dm = vec2(data_a[ib].dm);
147
148 const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4);
149
150 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v.xy);
151 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2(v.zw);
152#elif defined(DATA_A_Q3_K)
153 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
154 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
155
156 const uint ib = idx / 128; // 2 values per idx
157 const uint iqs = idx % 128; // 0..127
158
159 const uint n = iqs / 64; // 0,1
160 const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62
161 const uint hmi = (iqs % 16) * 2; // 0,2,4..30
162 const uint j = (iqs % 64) / 4; // 0..3
163 const uint is = iqs / 8; // 0..15
164 const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3
165 const uint qsshift = halfsplit * 2; // 0,2,4,6
166
167 const int8_t us = int8_t(((data_a[ib].scales[is % 8] >> (4 * int(is / 8))) & 0xF)
168 | (((data_a[ib].scales[8 + (is % 4)] >> (2 * int(is / 4))) & 3) << 4));
169 const float dl = float(data_a[ib].d) * float(us - 32);
170
171 const vec2 qs = vec2(unpack8((uint(data_a_packed16[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy);
172 const vec2 hm = vec2(unpack8(((uint(data_a_packed16[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy);
173
174 buf_a[buf_idx] = FLOAT_TYPE_VEC2(dl * (qs.x - hm.x),
175 dl * (qs.y - hm.y));
176#elif defined(DATA_A_Q4_K)
177 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
178 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
179
180 const uint ib = idx / 64; // 4 values per idx
181 const uint iqs = (idx % 64) * 2; // 0,2,4..126
182
183 const uint n = iqs / 32; // 0,1,2,3
184 const uint b = (iqs % 32) / 16; // 0,1
185 const uint is = 2 * n + b; // 0..7
186 const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126
187
188 const vec2 loadd = vec2(data_a[ib].dm);
189
190 const uint scidx0 = (is < 4) ? is : (is + 4);
191 const uint scidx1 = (is < 4) ? is : (is - 4);
192 const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0;
193 const uint scidxshift1 = (is < 4) ? 0 : 2;
194 const uint mbidx0 = is + 4;
195 const uint mbidx1 = (is < 4) ? is + 4 : is;
196 const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0;
197 const uint mbidxshift0 = (is < 4) ? 0 : 4;
198 const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0;
199 const uint mbidxshift1 = (is < 4) ? 0 : 2;
200
201 const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1));
202 const uint8_t mbyte = uint8_t((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0 | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1));
203
204 const float d = loadd.x * sc;
205 const float m = -loadd.y * mbyte;
206
207 const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F));
208
209 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(fma(d, q.x, m), fma(d, q.y, m));
210 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2(fma(d, q.z, m), fma(d, q.w, m));
211#elif defined(DATA_A_Q5_K)
212 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
213 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
214
215 const uint ib = idx / 64; // 4 values per idx
216 const uint iqs = (idx % 64) * 2; // 0,2,4..126
217
218 const uint n = iqs / 32; // 0,1,2,3
219 const uint b = (iqs % 32) / 16; // 0,1
220 const uint is = 2 * n + b; // 0..7
221 const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126
222 const uint qhi = (iqs % 16) * 2; // 0,2,4..30
223
224 const vec2 loadd = vec2(data_a[ib].dm);
225
226 const uint scidx0 = (is < 4) ? is : (is + 4);
227 const uint scidx1 = (is < 4) ? is : (is - 4);
228 const uint scidxmask1 = (is < 4) ? 0x30 : 0xC0;
229 const uint scidxshift1 = (is < 4) ? 0 : 2;
230 const uint mbidx0 = is + 4;
231 const uint mbidx1 = (is < 4) ? is + 4 : is;
232 const uint mbidxmask0 = (is < 4) ? 0xF : 0xF0;
233 const uint mbidxshift0 = (is < 4) ? 0 : 4;
234 const uint mbidxmask1 = (is < 4) ? 0x30 : 0xC0;
235 const uint mbidxshift1 = (is < 4) ? 0 : 2;
236
237 const uint8_t sc = uint8_t((data_a[ib].scales[scidx0] & 0xF) | ((data_a[ib].scales[scidx1] & scidxmask1) >> scidxshift1));
238 const uint8_t mbyte = uint8_t(((data_a[ib].scales[mbidx0] & mbidxmask0) >> mbidxshift0) | ((data_a[ib].scales[mbidx1] & mbidxmask1) >> mbidxshift1));
239
240 const float d = loadd.x * sc;
241 const float m = -loadd.y * mbyte;
242
243 const uint qs = (data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F;
244 const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4;
245 const vec4 q = vec4(unpack8(qs | qh));
246
247 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(fma(d, q.x, m), fma(d, q.y, m));
248 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2(fma(d, q.z, m), fma(d, q.w, m));
249#elif defined(DATA_A_Q6_K)
250 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
251 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
252
253 const uint ib = idx / 128; // 2 values per idx
254 const uint iqs = idx % 128; // 0..127
255
256 const uint n = iqs / 64; // 0,1
257 const uint b = ((iqs % 64) / 32) * 4; // 0,4
258 const uint is_b = (iqs % 16) / 8; // 0,1
259 const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6
260 const uint is = 8 * n + qhshift + is_b; // 0..15
261 const uint qsi = n * 32 + (iqs % 32); // 0..63
262 const uint qhi = n * 16 + (iqs % 16); // 0..31
263
264 const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]);
265
266 const uint ql = (uint(data_a_packed16[ib].ql[qsi]) >> b) & 0x0F0F;
267 const uint qh = (uint(data_a_packed16[ib].qh[qhi]) >> qhshift) & 0x0303;
268 const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale;
269
270 buf_a[buf_idx] = FLOAT_TYPE_VEC2(q.x, q.y);
271#elif defined(DATA_A_IQ1_S)
272 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
273 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
274
275 const uint ib = idx / 32; // 8 values per idx
276 const uint ib32 = (idx % 32) / 4; // 0..7
277 const uint ib8 = idx % 32;
278
279 const float d = float(data_a[ib].d);
280 const uint qh = data_a[ib].qh[ib32];
281 const uint qs = data_a[ib].qs[ib8];
282 const float dl = d * (2 * bitfieldExtract(qh, 12, 3) + 1);
283 const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
284 const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]);
285
286 [[unroll]] for (int k = 0; k < 4; ++k) {
287 buf_a[buf_idx + k] = FLOAT_TYPE_VEC2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta),
288 dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta));
289 }
290#elif defined(DATA_A_IQ1_M)
291 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
292 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
293
294 const uint ib = idx / 32; // 8 values per idx
295 const uint ib8 = idx % 32;
296 const uint ib16 = ib8 / 2;
297
298 const uint16_t[4] scales = data_a[ib].scales;
299 const u16vec4 s = u16vec4(scales[0], scales[1], scales[2], scales[3]) >> 12;
300 const float d = float(unpackHalf2x16(s.x | (s.y << 4) | (s.z << 8) | (s.w << 12)).x);
301 const uint sc = scales[ib8 / 8];
302 const uint qs = data_a[ib].qs[ib8];
303 const uint qh = data_a[ib].qh[ib16] >> (4 * (ib8 & 1));
304 const float dl = d * (2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1);
305 const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
306 const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
307
308 [[unroll]] for (int k = 0; k < 4; ++k) {
309 buf_a[buf_idx + k] = FLOAT_TYPE_VEC2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta),
310 dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta));
311 }
312#elif defined(DATA_A_IQ2_XXS)
313 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
314 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
315
316 const uint ib = idx / 32; // 8 values per idx
317 const uint ib32 = (idx % 32) / 4; // 0..7
318 const uint ib8 = idx % 4;
319
320 const float d = float(data_a[ib].d);
321 const uint qs = data_a[ib].qs[8 * ib32 + ib8];
322 const uint signs = pack32(u8vec4(
323 data_a[ib].qs[8*ib32 + 4],
324 data_a[ib].qs[8*ib32 + 5],
325 data_a[ib].qs[8*ib32 + 6],
326 data_a[ib].qs[8*ib32 + 7]
327 ));
328 const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + (signs >> 28)));
329 const uint32_t sign7 = bitfieldExtract(signs, 7 * int(ib8), 7);
330 const uint sign = sign7 | (bitCount(sign7) << 7);
331 const uvec2 grid = iq2xxs_grid[qs];
332 const vec4 grid0 = vec4(unpack8(grid.x));
333 const vec4 grid1 = vec4(unpack8(grid.y));
334
335 buf_a[buf_idx ] = db * FLOAT_TYPE_VEC2((sign & 1) != 0 ? -grid0.x : grid0.x,
336 (sign & 2) != 0 ? -grid0.y : grid0.y);
337 buf_a[buf_idx + 1] = db * FLOAT_TYPE_VEC2((sign & 4) != 0 ? -grid0.z : grid0.z,
338 (sign & 8) != 0 ? -grid0.w : grid0.w);
339 buf_a[buf_idx + 2] = db * FLOAT_TYPE_VEC2((sign & 16) != 0 ? -grid1.x : grid1.x,
340 (sign & 32) != 0 ? -grid1.y : grid1.y);
341 buf_a[buf_idx + 3] = db * FLOAT_TYPE_VEC2((sign & 64) != 0 ? -grid1.z : grid1.z,
342 (sign & 128) != 0 ? -grid1.w : grid1.w);
343#elif defined(DATA_A_IQ2_XS)
344 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
345 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
346
347 const uint ib = idx / 32; // 8 values per idx
348 const uint ib32 = (idx % 32) / 4; // 0..7
349 const uint ib8 = idx % 4; // 0..3
350
351 const float d = float(data_a[ib].d);
352 const uint scale = (data_a[ib].scales[ib32] >> (2 * (ib8 & 2))) & 0xf;
353 const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + scale));
354 const uint qs = data_a[ib].qs[4 * ib32 + ib8];
355 const uint sign7 = qs >> 9;
356 const uint sign = sign7 | (bitCount(sign7) << 7);
357 const uvec2 grid = iq2xs_grid[qs & 511];
358 const vec4 grid0 = vec4(unpack8(grid.x));
359 const vec4 grid1 = vec4(unpack8(grid.y));
360
361 buf_a[buf_idx ] = db * FLOAT_TYPE_VEC2((sign & 1) != 0 ? -grid0.x : grid0.x,
362 (sign & 2) != 0 ? -grid0.y : grid0.y);
363 buf_a[buf_idx + 1] = db * FLOAT_TYPE_VEC2((sign & 4) != 0 ? -grid0.z : grid0.z,
364 (sign & 8) != 0 ? -grid0.w : grid0.w);
365 buf_a[buf_idx + 2] = db * FLOAT_TYPE_VEC2((sign & 16) != 0 ? -grid1.x : grid1.x,
366 (sign & 32) != 0 ? -grid1.y : grid1.y);
367 buf_a[buf_idx + 3] = db * FLOAT_TYPE_VEC2((sign & 64) != 0 ? -grid1.z : grid1.z,
368 (sign & 128) != 0 ? -grid1.w : grid1.w);
369#elif defined(DATA_A_IQ2_S)
370 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
371 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
372
373 const uint ib = idx / 32; // 8 values per idx
374 const uint ib8 = idx % 32; // 0..31
375 const uint ib32 = ib8 / 4; // 0..7
376
377 const uint scale = (data_a[ib].scales[ib32] >> (2 * (ib8 & 2))) & 0xf;
378 const uint qs = data_a[ib].qs[ib8];
379 const uint qh = data_a[ib].qh[ib32];
380 const uint qhshift = 2 * (ib8 % 4);
381 const uint sign = data_a[ib].qs[QUANT_K / 8 + ib8];
382
383 const float d = float(data_a[ib].d);
384 const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + scale));
385 const uvec2 grid = iq2s_grid[qs | ((qh << (8 - qhshift)) & 0x300)];
386 const vec4 grid0 = vec4(unpack8(grid.x));
387 const vec4 grid1 = vec4(unpack8(grid.y));
388
389 buf_a[buf_idx ] = db * FLOAT_TYPE_VEC2((sign & 1) != 0 ? -grid0.x : grid0.x,
390 (sign & 2) != 0 ? -grid0.y : grid0.y);
391 buf_a[buf_idx + 1] = db * FLOAT_TYPE_VEC2((sign & 4) != 0 ? -grid0.z : grid0.z,
392 (sign & 8) != 0 ? -grid0.w : grid0.w);
393 buf_a[buf_idx + 2] = db * FLOAT_TYPE_VEC2((sign & 16) != 0 ? -grid1.x : grid1.x,
394 (sign & 32) != 0 ? -grid1.y : grid1.y);
395 buf_a[buf_idx + 3] = db * FLOAT_TYPE_VEC2((sign & 64) != 0 ? -grid1.z : grid1.z,
396 (sign & 128) != 0 ? -grid1.w : grid1.w);
397#elif defined(DATA_A_IQ3_XXS)
398 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
399 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
400
401 const uint ib = idx / 64; // 4 values per idx
402 const uint iqs = idx % 64; // 0..63
403 const uint is = QUANT_K / 4 + 4 * (iqs / 8); // 8 values
404
405 const float d = float(data_a[ib].d);
406 const uint qs = data_a[ib].qs[iqs];
407 const uint signs = pack32(u16vec2(
408 data_a_packed16[ib].qs[is/2],
409 data_a_packed16[ib].qs[is/2+1]
410 ));
411 const float db = d * 0.5 * (0.5 + (signs >> 28));
412 const uint32_t sign7 = bitfieldExtract(signs, 7 * (int(iqs / 2) % 4), 7);
413 const uint sign = (sign7 | (bitCount(sign7) << 7)) >> (4 * (idx % 2));
414 const uint grid = iq3xxs_grid[qs];
415 const vec4 v = db * vec4(unpack8(grid));
416
417 buf_a[buf_idx ] = FLOAT_TYPE_VEC2((sign & 1) != 0 ? -v.x : v.x,
418 (sign & 2) != 0 ? -v.y : v.y);
419 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2((sign & 4) != 0 ? -v.z : v.z,
420 (sign & 8) != 0 ? -v.w : v.w);
421#elif defined(DATA_A_IQ3_S)
422 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
423 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
424
425 const uint ib = idx / 64; // 4 values per idx
426 const uint iqs = idx % 64; // 0..63
427 const uint iqh = iqs / 8;
428
429 const float d = float(data_a[ib].d);
430 const uint qs = data_a[ib].qs[iqs];
431 const uint qh = data_a[ib].qh[iqh];
432 const int8_t sign = int8_t(data_a[ib].signs[iqs / 2] >> (4 * (idx % 2)));
433 const uint scale = data_a[ib].scales[iqs / 16];
434 const i8vec2 sign01 = i8vec2(1 - (2 & i8vec2(sign << 1, sign)));
435 const float db = d * (1 + 2 * ((scale >> (4 * (iqh & 1))) & 0xf));
436 const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)];
437 const vec4 v = db * vec4(unpack8(grid));
438
439 buf_a[buf_idx ] = FLOAT_TYPE_VEC2((sign & 1) != 0 ? -v.x : v.x,
440 (sign & 2) != 0 ? -v.y : v.y);
441 buf_a[buf_idx + 1] = FLOAT_TYPE_VEC2((sign & 4) != 0 ? -v.z : v.z,
442 (sign & 8) != 0 ? -v.w : v.w);
443#elif defined(DATA_A_IQ4_XS)
444 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
445 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
446
447 const uint ib = idx / 128; // 2 values per idx
448 const uint ib32 = (idx % 128) / 16; // 0..7
449 const uint iq = 16 * ib32 + 2 * (idx % 8);
450
451 const uint sl = (data_a[ib].scales_l[ib32/2] >> (4 * (ib32 & 1))) & 0xF;
452 const uint sh = ((data_a[ib].scales_h) >> (2 * ib32)) & 3;
453 const uint qshift = (idx & 8) >> 1;
454 u8vec2 qs = unpack8((uint(data_a_packed16[ib].qs[iq/2]) >> qshift) & 0x0F0F).xy;
455
456 const float d = float(data_a[ib].d);
457 const vec2 v = d * float(int(sl | (sh << 4)) - 32) * vec2(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y]);
458
459 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(v.xy);
460#elif defined(DATA_A_IQ4_NL)
461 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
462 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
463
464 const uint ib = idx / 8;
465 const uint iqs = idx & 0x07;
466
467 const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d);
468 const uint vui = uint(data_a_packed16[ib].qs[iqs]);
469
470 buf_a[buf_idx ] = d * FLOAT_TYPE_VEC2(kvalues_iq4nl[vui & 0xF],
471 kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]);
472 buf_a[buf_idx + 8] = d * FLOAT_TYPE_VEC2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)],
473 kvalues_iq4nl[vui >> 12]);
474#elif defined(DATA_A_MXFP4)
475 const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
476 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
477
478 const uint ib = idx / 8;
479 const uint iqs = (idx & 0x07) * 2;
480
481 const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
482 const uint vui = uint(data_a[ib].qs[iqs]);
483 const uint vui2 = uint(data_a[ib].qs[iqs+1]);
484
485 buf_a[buf_idx ] = FLOAT_TYPE_VEC2(kvalues_mxfp4[vui & 0xF] * d,
486 kvalues_mxfp4[vui2 & 0xF] * d);
487 buf_a[buf_idx + 8] = FLOAT_TYPE_VEC2(kvalues_mxfp4[vui >> 4] * d,
488 kvalues_mxfp4[vui2 >> 4] * d);
489#endif
490}
491
492#if !defined(MUL_MAT_ID)
493void load_b_to_shmem(const uint pos_b, const uint row, const uint col, const uint idx_n, const uint block, const uint end_k) {
494#if LOAD_VEC_B == 8
495 // Not supported for b_type bf16 because bf16mat2x4 does not exist
496 const uint idx = pos_b + col * p.stride_b / LOAD_VEC_B + row;
497 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2;
498 FLOAT_TYPE_VEC8 bb = FLOAT_TYPE_VEC8(data_b[idx]);
499 buf_b[buf_idx + 0] = bb[0].xy;
500 buf_b[buf_idx + 1] = bb[0].zw;
501 buf_b[buf_idx + 2] = bb[1].xy;
502 buf_b[buf_idx + 3] = bb[1].zw;
503#elif LOAD_VEC_B == 4
504 const uint idx = pos_b + col * p.stride_b / LOAD_VEC_B + row;
505 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2;
506#if defined(DATA_B_BF16)
507 FLOAT_TYPE_VEC4 bb = FLOAT_TYPE_VEC4(TO_FLOAT_TYPE(data_b[idx]));
508#else
509 FLOAT_TYPE_VEC4 bb = FLOAT_TYPE_VEC4(data_b[idx]);
510#endif
511 buf_b[buf_idx + 0] = bb.xy;
512 buf_b[buf_idx + 1] = bb.zw;
513#else // LOAD_VEC_BATCH_B == 2
514 const uint idx = pos_b + col * p.stride_b + row * 2;
515 const uint buf_idx = col * SHMEM_STRIDE + row;
516 if (idx_n < p.N && block + row * 2 + 1 < end_k) {
517 buf_b[buf_idx] = FLOAT_TYPE_VEC2(TO_FLOAT_TYPE(data_b[idx]),
518 TO_FLOAT_TYPE(data_b[idx + 1]));
519 } else if (idx_n < p.N && block + row * 2 < end_k) {
520 buf_b[buf_idx] = FLOAT_TYPE_VEC2(TO_FLOAT_TYPE(data_b[idx]), 0.0f);
521 } else {
522 buf_b[buf_idx] = FLOAT_TYPE_VEC2(0.0f);
523 }
524#endif
525}
526#else
527void load_b_to_shmem(const uint pos_b, const uint row, const uint col, const uint ic, const uint _ne1, const uint block, const uint end_k) {
528#if LOAD_VEC_B == 8
529 // Not supported for b_type bf16 because bf16mat2x4 does not exist
530 const u16vec2 row_idx = row_ids[col];
531 const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + row;
532 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2;
533 FLOAT_TYPE_VEC8 bb = FLOAT_TYPE_VEC8(data_b[idx]);
534 buf_b[buf_idx + 0] = bb[0].xy;
535 buf_b[buf_idx + 1] = bb[0].zw;
536 buf_b[buf_idx + 2] = bb[1].xy;
537 buf_b[buf_idx + 3] = bb[1].zw;
538#elif LOAD_VEC_B == 4
539 const u16vec2 row_idx = row_ids[col];
540 const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + row;
541 const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2;
542#if defined(DATA_B_BF16)
543 FLOAT_TYPE_VEC4 bb = FLOAT_TYPE_VEC4(TO_FLOAT_TYPE(data_b[idx]));
544#else
545 FLOAT_TYPE_VEC4 bb = FLOAT_TYPE_VEC4(data_b[idx]);
546#endif
547 buf_b[buf_idx + 0] = bb.xy;
548 buf_b[buf_idx + 1] = bb.zw;
549#else // LOAD_VEC_BATCH_B == 2
550 const uint row_i = ic * BN + col;
551 const uint buf_idx = col * SHMEM_STRIDE + row;
552 if (row_i < _ne1 && block + row * 2 + 1 < end_k) {
553 const u16vec2 row_idx = row_ids[col];
554 const uint idx = pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + row * 2;
555 buf_b[buf_idx] = FLOAT_TYPE_VEC2(TO_FLOAT_TYPE(data_b[idx]),
556 TO_FLOAT_TYPE(data_b[idx + 1]));
557 } else if (row_i < _ne1 && block + row * 2 < end_k) {
558 const u16vec2 row_idx = row_ids[col];
559 const uint idx = pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + row * 2;
560 buf_b[buf_idx] = FLOAT_TYPE_VEC2(TO_FLOAT_TYPE(data_b[idx]), 0.0f);
561 } else {
562 buf_b[buf_idx] = FLOAT_TYPE_VEC2(0.0f);
563 }
564#endif
565}
566#endif