/ vulkan-shaders / mul_mm.comp
mul_mm.comp
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 10 #ifdef MUL_MAT_ID 11 #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require 12 #endif 13 14 #include "types.comp" 15 16 #ifndef LOAD_VEC_A 17 #define LOAD_VEC_A 1 18 #endif 19 #ifndef LOAD_VEC_B 20 #define LOAD_VEC_B 1 21 #endif 22 23 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; 24 25 layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; 26 layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; 27 layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; 28 29 #ifdef MUL_MAT_ID 30 layout (binding = 3) readonly buffer IDS {int data_ids[];}; 31 #endif 32 33 layout (push_constant) uniform parameter 34 { 35 uint M; 36 uint N; 37 uint K; 38 uint stride_a; 39 uint stride_b; 40 uint stride_d; 41 42 uint batch_stride_a; 43 uint batch_stride_b; 44 uint batch_stride_d; 45 46 #ifdef MUL_MAT_ID 47 uint nei0; 48 uint nei1; 49 uint nbi1; 50 uint ne11; 51 #else 52 uint k_split; 53 uint ne02; 54 uint ne12; 55 uint broadcast2; 56 uint broadcast3; 57 #endif 58 } p; 59 60 layout (constant_id = 1) const uint BM = 64; 61 layout (constant_id = 2) const uint BN = 64; 62 layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant 63 layout (constant_id = 4) const uint WM = 32; 64 layout (constant_id = 5) const uint WN = 32; 65 layout (constant_id = 6) const uint WMITER = 2; 66 layout (constant_id = 7) const uint TM = 4; 67 layout (constant_id = 8) const uint TN = 2; 68 layout (constant_id = 9) const uint WARP = 32; 69 70 shared FLOAT_TYPE buf_a[BM * (BK+1)]; 71 shared FLOAT_TYPE buf_b[BN * (BK+1)]; 72 73 #ifdef MUL_MAT_ID 74 shared u16vec2 row_ids[2048]; 75 #endif 76 77 void main() { 78 #ifdef MUL_MAT_ID 79 const uint expert_idx = gl_GlobalInvocationID.z; 80 #else 81 const uint batch_idx = gl_GlobalInvocationID.z; 82 83 const uint i13 = batch_idx / p.ne12; 84 const uint i12 = batch_idx % p.ne12; 85 86 const uint i03 = i13 / p.broadcast3; 87 const uint i02 = i12 / p.broadcast2; 88 89 const uint batch_idx_a = i03 * p.ne02 + i02; 90 #endif 91 92 const uint blocks_m = (p.M + BM - 1) / BM; 93 const uint ir = gl_WorkGroupID.x % blocks_m; 94 const uint ik = gl_WorkGroupID.x / blocks_m; 95 const uint ic = gl_WorkGroupID.y; 96 97 const uint warp_i = gl_LocalInvocationID.x / WARP; 98 const uint warp_r = warp_i % (BM / WM); 99 const uint warp_c = warp_i / (BM / WM); 100 101 const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER); 102 const uint WSUBM = WM / WMITER; 103 const uint WSUBN = WN / WNITER; 104 105 const uint tiw = gl_LocalInvocationID.x % WARP; 106 const uint tiwr = tiw % (WSUBM / TM); 107 const uint tiwc = tiw / (WSUBM / TM); 108 109 const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A); 110 const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A); 111 const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B); 112 const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B); 113 114 const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A / BK; 115 const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B / BK; 116 117 #ifdef MUL_MAT_ID 118 uint _ne1 = 0; 119 for (uint ii1 = 0; ii1 < p.nei1; ii1++) { 120 for (uint ii0 = 0; ii0 < p.nei0; ii0++) { 121 if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { 122 row_ids[_ne1] = u16vec2(ii0, ii1); 123 _ne1++; 124 } 125 } 126 } 127 128 barrier(); 129 130 // Workgroup has no work 131 if (ic * BN >= _ne1) return; 132 #endif 133 134 #ifdef MUL_MAT_ID 135 const uint start_k = 0; 136 const uint end_k = p.K; 137 #else 138 const uint start_k = ik * p.k_split; 139 const uint end_k = min(p.K, (ik + 1) * p.k_split); 140 #endif 141 142 uint pos_a = ( 143 #ifdef MUL_MAT_ID 144 expert_idx * p.batch_stride_a + 145 #else 146 batch_idx_a * p.batch_stride_a + 147 #endif 148 ir * BM * p.stride_a + start_k) / LOAD_VEC_A; 149 #ifdef MUL_MAT_ID 150 uint pos_b = 0; 151 #else 152 uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B; 153 #endif 154 155 float sums[WMITER * TM * WNITER * TN]; 156 FLOAT_TYPE cache_a[WMITER * TM]; 157 FLOAT_TYPE cache_b[WNITER * TN]; 158 159 [[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) { 160 sums[i] = 0.0f; 161 } 162 163 [[unroll]] for (uint block = start_k; block < end_k; block += BK) { 164 [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) { 165 166 #if defined(DATA_A_F32) || defined(DATA_A_F16) 167 #if LOAD_VEC_A == 8 168 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 169 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 170 buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x); 171 buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y); 172 buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z); 173 buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w); 174 buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x); 175 buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y); 176 buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z); 177 buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w); 178 #elif LOAD_VEC_A == 4 179 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 180 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 181 buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x); 182 buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y); 183 buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z); 184 buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w); 185 #else 186 if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) { 187 buf_a[(loadc_a + l) * (BK+1) + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]); 188 } else { 189 buf_a[(loadc_a + l) * (BK+1) + loadr_a] = FLOAT_TYPE(0.0f); 190 } 191 #endif 192 #elif defined(DATA_A_Q4_0) 193 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 194 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a; 195 196 const uint ib = idx / 16; 197 const uint iqs = idx & 0xF; 198 199 const float d = float(data_a[ib].d); 200 const uint vui = uint(data_a[ib].qs[iqs]); 201 const vec2 v = (vec2(vui & 0xF, vui >> 4) - 8.0f) * d; 202 203 buf_a[buf_idx ] = FLOAT_TYPE(v.x); 204 buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); 205 #elif defined(DATA_A_Q4_1) 206 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 207 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a; 208 209 const uint ib = idx / 16; 210 const uint iqs = idx & 0xF; 211 212 const float d = float(data_a[ib].d); 213 const float m = float(data_a[ib].m); 214 const uint vui = uint(data_a[ib].qs[iqs]); 215 const vec2 v = vec2(vui & 0xF, vui >> 4) * d + m; 216 217 buf_a[buf_idx ] = FLOAT_TYPE(v.x); 218 buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); 219 #elif defined(DATA_A_Q5_0) 220 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 221 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a; 222 223 const uint ib = idx / 16; 224 const uint iqs = idx & 0xF; 225 226 const float d = float(data_a[ib].d); 227 const uint uint_qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; 228 const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); 229 const uint vui = uint(data_a[ib].qs[iqs]); 230 const vec2 v = (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d; 231 232 buf_a[buf_idx ] = FLOAT_TYPE(v.x); 233 buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); 234 #elif defined(DATA_A_Q5_1) 235 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 236 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a; 237 238 const uint ib = idx / 16; 239 const uint iqs = idx & 0xF; 240 241 const float d = float(data_a[ib].d); 242 const float m = float(data_a[ib].m); 243 const uint uint_qh = data_a[ib].qh; 244 const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10); 245 const uint vui = uint(data_a[ib].qs[iqs]); 246 const vec2 v = vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m; 247 248 buf_a[buf_idx ] = FLOAT_TYPE(v.x); 249 buf_a[buf_idx + 16] = FLOAT_TYPE(v.y); 250 #elif defined(DATA_A_Q8_0) 251 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 252 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 253 254 const uint ib = idx / 16; 255 const uint iqs = (idx & 0xF) * 2; 256 257 const float d = float(data_a[ib].d); 258 const vec2 v = vec2(int(data_a[ib].qs[iqs]), int(data_a[ib].qs[iqs + 1])) * d; 259 260 buf_a[buf_idx ] = FLOAT_TYPE(v.x); 261 buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); 262 #elif defined(DATA_A_Q2_K) 263 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 264 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 265 266 const uint ib = idx / 128; // 2 values per idx 267 const uint iqs = idx % 128; // 0..127 268 269 const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30 270 const uint scalesi = iqs / 8; // 0..15 271 const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 272 273 const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); 274 const uint scales = data_a[ib].scales[scalesi]; 275 const vec2 d = vec2(data_a[ib].d); 276 277 const vec2 v = d.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - d.y * float(scales >> 4); 278 279 buf_a[buf_idx ] = FLOAT_TYPE(v.x); 280 buf_a[buf_idx + 1] = FLOAT_TYPE(v.y); 281 #elif defined(DATA_A_Q3_K) 282 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 283 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 284 285 const uint ib = idx / 128; // 2 values per idx 286 const uint iqs = idx % 128; // 0..127 287 288 const uint n = iqs / 64; // 0,1 289 const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 290 const uint hmi = (iqs % 16) * 2; // 0,2,4..30 291 const uint j = (iqs % 64) / 4; // 0..3 292 const uint is = iqs / 8; // 0..15 293 const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 294 const uint qsshift = halfsplit * 2; // 0,2,4,6 295 const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128 296 297 const int8_t us = int8_t(is < 4 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+8] >> 0) & 3) << 4) : 298 is < 8 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+4] >> 2) & 3) << 4) : 299 is < 12 ? (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is+0] >> 4) & 3) << 4) : 300 (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is-4] >> 6) & 3) << 4)); 301 const float dl = float(data_a[ib].d) * float(us - 32); 302 303 buf_a[buf_idx ] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi ] >> qsshift) & 3) - (((data_a[ib].hmask[hmi ] & m) != 0) ? 0 : 4))); 304 buf_a[buf_idx + 1] = FLOAT_TYPE(dl * float(int8_t((data_a[ib].qs[qsi + 1] >> qsshift) & 3) - (((data_a[ib].hmask[hmi + 1] & m) != 0) ? 0 : 4))); 305 #elif defined(DATA_A_Q4_K) 306 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 307 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 308 309 const uint ib = idx / 128; // 2 values per idx 310 const uint iqs = idx % 128; // 0..127 311 312 const uint n = iqs / 32; // 0,1,2,3 313 const uint b = (iqs % 32) / 16; // 0,1 314 const uint is = 2 * n + b; // 0..7 315 const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 316 317 const vec2 loadd = vec2(data_a[ib].d); 318 319 uint8_t sc; 320 uint8_t mbyte; 321 if (is < 4) { 322 sc = uint8_t(data_a[ib].scales[is ] & 63); 323 mbyte = uint8_t(data_a[ib].scales[is + 4] & 63); 324 } else { 325 sc = uint8_t((data_a[ib].scales[is + 4] & 0xF) | ((data_a[ib].scales[is - 4] >> 6) << 4)); 326 mbyte = uint8_t((data_a[ib].scales[is + 4] >> 4) | ((data_a[ib].scales[is ] >> 6) << 4)); 327 } 328 const float d = loadd.x * sc; 329 const float m = loadd.y * mbyte; 330 331 buf_a[buf_idx ] = FLOAT_TYPE(d * float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) - m); 332 buf_a[buf_idx + 1] = FLOAT_TYPE(d * float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) - m); 333 #elif defined(DATA_A_Q5_K) 334 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 335 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 336 337 const uint ib = idx / 128; // 2 values per idx 338 const uint iqs = idx % 128; // 0..127 339 340 const uint n = iqs / 32; // 0,1,2,3 341 const uint b = (iqs % 32) / 16; // 0,1 342 const uint is = 2 * n + b; // 0..7 343 const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 344 const uint qhi = (iqs % 16) * 2; // 0,2,4..30 345 346 const uint8_t hm = uint8_t(1 << (iqs / 16)); 347 348 const vec2 loadd = vec2(data_a[ib].d); 349 350 uint8_t sc; 351 uint8_t mbyte; 352 if (is < 4) { 353 sc = uint8_t(data_a[ib].scales[is ] & 63); 354 mbyte = uint8_t(data_a[ib].scales[is + 4] & 63); 355 } else { 356 sc = uint8_t((data_a[ib].scales[is + 4] & 0xF) | ((data_a[ib].scales[is - 4] >> 6) << 4)); 357 mbyte = uint8_t((data_a[ib].scales[is + 4] >> 4) | ((data_a[ib].scales[is ] >> 6) << 4)); 358 } 359 const float d = loadd.x * sc; 360 const float m = loadd.y * mbyte; 361 362 buf_a[buf_idx ] = FLOAT_TYPE(d * (float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi ] & hm) != 0 ? 16 : 0)) - m); 363 buf_a[buf_idx + 1] = FLOAT_TYPE(d * (float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) + float((data_a[ib].qh[qhi + 1] & hm) != 0 ? 16 : 0)) - m); 364 #elif defined(DATA_A_Q6_K) 365 const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a; 366 const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A; 367 368 const uint ib = idx / 128; // 2 values per idx 369 const uint iqs = idx % 128; // 0..127 370 371 const uint n = iqs / 64; // 0,1 372 const uint b = (iqs % 64) / 32; // 0,1 373 const uint is_b = (iqs % 16) / 8; // 0,1 374 const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 375 const uint is = 8 * n + qhshift + is_b; // 0..15 376 const uint qsi = n * 64 + (iqs % 32) * 2; // 0,2,4..126 377 const uint qhi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 378 379 const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]); 380 381 buf_a[buf_idx ] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi ] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi ] >> qhshift) & 3) << 4)) - 32)); 382 buf_a[buf_idx + 1] = FLOAT_TYPE(dscale * float(int8_t(((data_a[ib].ql[qsi + 1] >> (b * 4)) & 0xF) | (((data_a[ib].qh[qhi + 1] >> qhshift) & 3) << 4)) - 32)); 383 #endif 384 } 385 [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) { 386 #if LOAD_VEC_B == 8 387 #ifdef MUL_MAT_ID 388 const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; 389 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 + loadr_b; 390 #else 391 const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; 392 #endif 393 const uint buf_idx = (loadc_b + l) * (BK+1) + loadr_b * LOAD_VEC_B; 394 buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x); 395 buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y); 396 buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z); 397 buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w); 398 buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x); 399 buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y); 400 buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z); 401 buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w); 402 #elif LOAD_VEC_B == 4 403 #ifdef MUL_MAT_ID 404 const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l]; 405 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 + loadr_b; 406 #else 407 const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b; 408 #endif 409 const uint buf_idx = (loadc_b + l) * (BK+1) + loadr_b * LOAD_VEC_B; 410 buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx].x); 411 buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx].y); 412 buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx].z); 413 buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx].w); 414 #elif !MUL_MAT_ID 415 if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) { 416 buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]); 417 } else { 418 buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(0.0f); 419 } 420 #else 421 const uint row_i = ic * BN + loadc_b + l; 422 if (row_i < _ne1) { 423 const u16vec2 row_idx = row_ids[row_i]; 424 buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(data_b[pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + loadr_b]); 425 } else { 426 buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(0.0f); 427 } 428 #endif 429 } 430 431 barrier(); 432 433 pos_a += BK / LOAD_VEC_A; 434 pos_b += BK / LOAD_VEC_B; 435 436 for (uint i = 0; i < BK; i++) { 437 // Load from shared into cache 438 [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { 439 [[unroll]] for (uint j = 0; j < TM; j++) { 440 cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * (BK+1) + i]; 441 } 442 } 443 [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { 444 [[unroll]] for (uint j = 0; j < TN; j++) { 445 cache_b[wsic * TN + j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * (BK+1) + i]; 446 } 447 } 448 449 [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { 450 [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { 451 [[unroll]] for (uint cc = 0; cc < TN; cc++) { 452 [[unroll]] for (uint cr = 0; cr < TM; cr++) { 453 sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr] += float(cache_a[wsir * TM + cr]) * float(cache_b[wsic * TN + cc]); 454 } 455 } 456 } 457 } 458 } 459 460 barrier(); 461 } 462 463 const uint dr = ir * BM + warp_r * WM; 464 const uint dc = ic * BN + warp_c * WN; 465 466 #ifndef MUL_MAT_ID 467 const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z; 468 #endif 469 470 [[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) { 471 [[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) { 472 473 const uint dr_warp = dr + wsir * WSUBM + tiwr * TM; 474 const uint dc_warp = dc + wsic * WSUBN + tiwc * TN; 475 [[unroll]] for (uint cc = 0; cc < TN; cc++) { 476 #ifdef MUL_MAT_ID 477 const uint row_i = dc_warp + cc; 478 if (row_i >= _ne1) break; 479 480 const u16vec2 row_idx = row_ids[row_i]; 481 #endif 482 [[unroll]] for (uint cr = 0; cr < TM; cr++) { 483 #ifdef MUL_MAT_ID 484 data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); 485 #else 486 if (dr_warp + cr < p.M && dc_warp + cc < p.N) { 487 data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]); 488 } 489 #endif 490 } 491 } 492 } 493 } 494 }