mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2025-11-04 09:32:00 +00:00
7c7836d9d4
* Refactor shaders, extract GLSL code from ggml_vk_generate_shaders.py into vulkan-shaders directory * Improve debug log code * Add memory debug output option * Fix flake8 * Fix unnecessary high llama-3 VRAM use
495 lines
20 KiB
Plaintext
495 lines
20 KiB
Plaintext
#version 450
|
|
|
|
#extension GL_EXT_control_flow_attributes : enable
|
|
#extension GL_EXT_shader_16bit_storage : require
|
|
|
|
#ifdef FLOAT16
|
|
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
|
#endif
|
|
|
|
#ifdef MUL_MAT_ID
|
|
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
|
|
#endif
|
|
|
|
#include "types.comp"
|
|
|
|
#ifndef LOAD_VEC_A
|
|
#define LOAD_VEC_A 1
|
|
#endif
|
|
#ifndef LOAD_VEC_B
|
|
#define LOAD_VEC_B 1
|
|
#endif
|
|
|
|
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
|
|
layout (binding = 1) readonly buffer B {B_TYPE data_b[];};
|
|
layout (binding = 2) writeonly buffer D {D_TYPE data_d[];};
|
|
|
|
#ifdef MUL_MAT_ID
|
|
layout (binding = 3) readonly buffer IDS {int data_ids[];};
|
|
#endif
|
|
|
|
layout (push_constant) uniform parameter
|
|
{
|
|
uint M;
|
|
uint N;
|
|
uint K;
|
|
uint stride_a;
|
|
uint stride_b;
|
|
uint stride_d;
|
|
|
|
uint batch_stride_a;
|
|
uint batch_stride_b;
|
|
uint batch_stride_d;
|
|
|
|
#ifdef MUL_MAT_ID
|
|
uint nei0;
|
|
uint nei1;
|
|
uint nbi1;
|
|
uint ne11;
|
|
#else
|
|
uint k_split;
|
|
uint ne02;
|
|
uint ne12;
|
|
uint broadcast2;
|
|
uint broadcast3;
|
|
#endif
|
|
} p;
|
|
|
|
layout (constant_id = 1) const uint BM = 64;
|
|
layout (constant_id = 2) const uint BN = 64;
|
|
layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant
|
|
layout (constant_id = 4) const uint WM = 32;
|
|
layout (constant_id = 5) const uint WN = 32;
|
|
layout (constant_id = 6) const uint WMITER = 2;
|
|
layout (constant_id = 7) const uint TM = 4;
|
|
layout (constant_id = 8) const uint TN = 2;
|
|
layout (constant_id = 9) const uint WARP = 32;
|
|
|
|
shared FLOAT_TYPE buf_a[BM * (BK+1)];
|
|
shared FLOAT_TYPE buf_b[BN * (BK+1)];
|
|
|
|
#ifdef MUL_MAT_ID
|
|
shared u16vec2 row_ids[2048];
|
|
#endif
|
|
|
|
void main() {
|
|
#ifdef MUL_MAT_ID
|
|
const uint expert_idx = gl_GlobalInvocationID.z;
|
|
#else
|
|
const uint batch_idx = gl_GlobalInvocationID.z;
|
|
|
|
const uint i13 = batch_idx / p.ne12;
|
|
const uint i12 = batch_idx % p.ne12;
|
|
|
|
const uint i03 = i13 / p.broadcast3;
|
|
const uint i02 = i12 / p.broadcast2;
|
|
|
|
const uint batch_idx_a = i03 * p.ne02 + i02;
|
|
#endif
|
|
|
|
const uint blocks_m = (p.M + BM - 1) / BM;
|
|
const uint ir = gl_WorkGroupID.x % blocks_m;
|
|
const uint ik = gl_WorkGroupID.x / blocks_m;
|
|
const uint ic = gl_WorkGroupID.y;
|
|
|
|
const uint warp_i = gl_LocalInvocationID.x / WARP;
|
|
const uint warp_r = warp_i % (BM / WM);
|
|
const uint warp_c = warp_i / (BM / WM);
|
|
|
|
const uint WNITER = (WM * WN) / (WARP * TM * TN * WMITER);
|
|
const uint WSUBM = WM / WMITER;
|
|
const uint WSUBN = WN / WNITER;
|
|
|
|
const uint tiw = gl_LocalInvocationID.x % WARP;
|
|
const uint tiwr = tiw % (WSUBM / TM);
|
|
const uint tiwc = tiw / (WSUBM / TM);
|
|
|
|
const uint loadr_a = gl_LocalInvocationID.x % (BK / LOAD_VEC_A);
|
|
const uint loadc_a = gl_LocalInvocationID.x / (BK / LOAD_VEC_A);
|
|
const uint loadr_b = gl_LocalInvocationID.x % (BK / LOAD_VEC_B);
|
|
const uint loadc_b = gl_LocalInvocationID.x / (BK / LOAD_VEC_B);
|
|
|
|
const uint loadstride_a = gl_WorkGroupSize.x * LOAD_VEC_A / BK;
|
|
const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B / BK;
|
|
|
|
#ifdef MUL_MAT_ID
|
|
uint _ne1 = 0;
|
|
for (uint ii1 = 0; ii1 < p.nei1; ii1++) {
|
|
for (uint ii0 = 0; ii0 < p.nei0; ii0++) {
|
|
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
|
|
row_ids[_ne1] = u16vec2(ii0, ii1);
|
|
_ne1++;
|
|
}
|
|
}
|
|
}
|
|
|
|
barrier();
|
|
|
|
// Workgroup has no work
|
|
if (ic * BN >= _ne1) return;
|
|
#endif
|
|
|
|
#ifdef MUL_MAT_ID
|
|
const uint start_k = 0;
|
|
const uint end_k = p.K;
|
|
#else
|
|
const uint start_k = ik * p.k_split;
|
|
const uint end_k = min(p.K, (ik + 1) * p.k_split);
|
|
#endif
|
|
|
|
uint pos_a = (
|
|
#ifdef MUL_MAT_ID
|
|
expert_idx * p.batch_stride_a +
|
|
#else
|
|
batch_idx_a * p.batch_stride_a +
|
|
#endif
|
|
ir * BM * p.stride_a + start_k) / LOAD_VEC_A;
|
|
#ifdef MUL_MAT_ID
|
|
uint pos_b = 0;
|
|
#else
|
|
uint pos_b = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / LOAD_VEC_B;
|
|
#endif
|
|
|
|
float sums[WMITER * TM * WNITER * TN];
|
|
FLOAT_TYPE cache_a[WMITER * TM];
|
|
FLOAT_TYPE cache_b[WNITER * TN];
|
|
|
|
[[unroll]] for (uint i = 0; i < WMITER*TM*WNITER*TN; i++) {
|
|
sums[i] = 0.0f;
|
|
}
|
|
|
|
[[unroll]] for (uint block = start_k; block < end_k; block += BK) {
|
|
[[unroll]] for (uint l = 0; l < BM; l += loadstride_a) {
|
|
|
|
#if defined(DATA_A_F32) || defined(DATA_A_F16)
|
|
#if LOAD_VEC_A == 8
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x);
|
|
buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y);
|
|
buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z);
|
|
buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w);
|
|
buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x);
|
|
buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y);
|
|
buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z);
|
|
buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w);
|
|
#elif LOAD_VEC_A == 4
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x);
|
|
buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y);
|
|
buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z);
|
|
buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w);
|
|
#else
|
|
if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) {
|
|
buf_a[(loadc_a + l) * (BK+1) + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]);
|
|
} else {
|
|
buf_a[(loadc_a + l) * (BK+1) + loadr_a] = FLOAT_TYPE(0.0f);
|
|
}
|
|
#endif
|
|
#elif defined(DATA_A_Q4_0)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a;
|
|
|
|
const uint ib = idx / 16;
|
|
const uint iqs = idx & 0xF;
|
|
|
|
const float d = float(data_a[ib].d);
|
|
const uint vui = uint(data_a[ib].qs[iqs]);
|
|
const vec2 v = (vec2(vui & 0xF, vui >> 4) - 8.0f) * d;
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(v.x);
|
|
buf_a[buf_idx + 16] = FLOAT_TYPE(v.y);
|
|
#elif defined(DATA_A_Q4_1)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a;
|
|
|
|
const uint ib = idx / 16;
|
|
const uint iqs = idx & 0xF;
|
|
|
|
const float d = float(data_a[ib].d);
|
|
const float m = float(data_a[ib].m);
|
|
const uint vui = uint(data_a[ib].qs[iqs]);
|
|
const vec2 v = vec2(vui & 0xF, vui >> 4) * d + m;
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(v.x);
|
|
buf_a[buf_idx + 16] = FLOAT_TYPE(v.y);
|
|
#elif defined(DATA_A_Q5_0)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a;
|
|
|
|
const uint ib = idx / 16;
|
|
const uint iqs = idx & 0xF;
|
|
|
|
const float d = float(data_a[ib].d);
|
|
const uint uint_qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0];
|
|
const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
|
|
const uint vui = uint(data_a[ib].qs[iqs]);
|
|
const vec2 v = (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d;
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(v.x);
|
|
buf_a[buf_idx + 16] = FLOAT_TYPE(v.y);
|
|
#elif defined(DATA_A_Q5_1)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a;
|
|
|
|
const uint ib = idx / 16;
|
|
const uint iqs = idx & 0xF;
|
|
|
|
const float d = float(data_a[ib].d);
|
|
const float m = float(data_a[ib].m);
|
|
const uint uint_qh = data_a[ib].qh;
|
|
const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
|
|
const uint vui = uint(data_a[ib].qs[iqs]);
|
|
const vec2 v = vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m;
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(v.x);
|
|
buf_a[buf_idx + 16] = FLOAT_TYPE(v.y);
|
|
#elif defined(DATA_A_Q8_0)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
|
|
const uint ib = idx / 16;
|
|
const uint iqs = (idx & 0xF) * 2;
|
|
|
|
const float d = float(data_a[ib].d);
|
|
const vec2 v = vec2(int(data_a[ib].qs[iqs]), int(data_a[ib].qs[iqs + 1])) * d;
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(v.x);
|
|
buf_a[buf_idx + 1] = FLOAT_TYPE(v.y);
|
|
#elif defined(DATA_A_Q2_K)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
|
|
const uint ib = idx / 128; // 2 values per idx
|
|
const uint iqs = idx % 128; // 0..127
|
|
|
|
const uint qsi = (iqs / 64) * 32 + (iqs % 16) * 2; // 0,2,4..30
|
|
const uint scalesi = iqs / 8; // 0..15
|
|
const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6
|
|
|
|
const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]);
|
|
const uint scales = data_a[ib].scales[scalesi];
|
|
const vec2 d = vec2(data_a[ib].d);
|
|
|
|
const vec2 v = d.x * float(scales & 0xF) * vec2((qs >> qsshift) & 3) - d.y * float(scales >> 4);
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(v.x);
|
|
buf_a[buf_idx + 1] = FLOAT_TYPE(v.y);
|
|
#elif defined(DATA_A_Q3_K)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
|
|
const uint ib = idx / 128; // 2 values per idx
|
|
const uint iqs = idx % 128; // 0..127
|
|
|
|
const uint n = iqs / 64; // 0,1
|
|
const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62
|
|
const uint hmi = (iqs % 16) * 2; // 0,2,4..30
|
|
const uint j = (iqs % 64) / 4; // 0..3
|
|
const uint is = iqs / 8; // 0..15
|
|
const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3
|
|
const uint qsshift = halfsplit * 2; // 0,2,4,6
|
|
const uint m = 1 << (4 * n + halfsplit); // 1,2,4,8,16,32,64,128
|
|
|
|
const int8_t us = int8_t(is < 4 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+8] >> 0) & 3) << 4) :
|
|
is < 8 ? (data_a[ib].scales[is-0] & 0xF) | (((data_a[ib].scales[is+4] >> 2) & 3) << 4) :
|
|
is < 12 ? (data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is+0] >> 4) & 3) << 4) :
|
|
(data_a[ib].scales[is-8] >> 4) | (((data_a[ib].scales[is-4] >> 6) & 3) << 4));
|
|
const float dl = float(data_a[ib].d) * float(us - 32);
|
|
|
|
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)));
|
|
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)));
|
|
#elif defined(DATA_A_Q4_K)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
|
|
const uint ib = idx / 128; // 2 values per idx
|
|
const uint iqs = idx % 128; // 0..127
|
|
|
|
const uint n = iqs / 32; // 0,1,2,3
|
|
const uint b = (iqs % 32) / 16; // 0,1
|
|
const uint is = 2 * n + b; // 0..7
|
|
const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126
|
|
|
|
const vec2 loadd = vec2(data_a[ib].d);
|
|
|
|
uint8_t sc;
|
|
uint8_t mbyte;
|
|
if (is < 4) {
|
|
sc = uint8_t(data_a[ib].scales[is ] & 63);
|
|
mbyte = uint8_t(data_a[ib].scales[is + 4] & 63);
|
|
} else {
|
|
sc = uint8_t((data_a[ib].scales[is + 4] & 0xF) | ((data_a[ib].scales[is - 4] >> 6) << 4));
|
|
mbyte = uint8_t((data_a[ib].scales[is + 4] >> 4) | ((data_a[ib].scales[is ] >> 6) << 4));
|
|
}
|
|
const float d = loadd.x * sc;
|
|
const float m = loadd.y * mbyte;
|
|
|
|
buf_a[buf_idx ] = FLOAT_TYPE(d * float((data_a[ib].qs[qsi ] >> (b * 4)) & 0xF) - m);
|
|
buf_a[buf_idx + 1] = FLOAT_TYPE(d * float((data_a[ib].qs[qsi + 1] >> (b * 4)) & 0xF) - m);
|
|
#elif defined(DATA_A_Q5_K)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
|
|
const uint ib = idx / 128; // 2 values per idx
|
|
const uint iqs = idx % 128; // 0..127
|
|
|
|
const uint n = iqs / 32; // 0,1,2,3
|
|
const uint b = (iqs % 32) / 16; // 0,1
|
|
const uint is = 2 * n + b; // 0..7
|
|
const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126
|
|
const uint qhi = (iqs % 16) * 2; // 0,2,4..30
|
|
|
|
const uint8_t hm = uint8_t(1 << (iqs / 16));
|
|
|
|
const vec2 loadd = vec2(data_a[ib].d);
|
|
|
|
uint8_t sc;
|
|
uint8_t mbyte;
|
|
if (is < 4) {
|
|
sc = uint8_t(data_a[ib].scales[is ] & 63);
|
|
mbyte = uint8_t(data_a[ib].scales[is + 4] & 63);
|
|
} else {
|
|
sc = uint8_t((data_a[ib].scales[is + 4] & 0xF) | ((data_a[ib].scales[is - 4] >> 6) << 4));
|
|
mbyte = uint8_t((data_a[ib].scales[is + 4] >> 4) | ((data_a[ib].scales[is ] >> 6) << 4));
|
|
}
|
|
const float d = loadd.x * sc;
|
|
const float m = loadd.y * mbyte;
|
|
|
|
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);
|
|
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);
|
|
#elif defined(DATA_A_Q6_K)
|
|
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
|
|
const uint buf_idx = (loadc_a + l) * (BK+1) + loadr_a * LOAD_VEC_A;
|
|
|
|
const uint ib = idx / 128; // 2 values per idx
|
|
const uint iqs = idx % 128; // 0..127
|
|
|
|
const uint n = iqs / 64; // 0,1
|
|
const uint b = (iqs % 64) / 32; // 0,1
|
|
const uint is_b = (iqs % 16) / 8; // 0,1
|
|
const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6
|
|
const uint is = 8 * n + qhshift + is_b; // 0..15
|
|
const uint qsi = n * 64 + (iqs % 32) * 2; // 0,2,4..126
|
|
const uint qhi = n * 32 + (iqs % 16) * 2; // 0,2,4..62
|
|
|
|
const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]);
|
|
|
|
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));
|
|
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));
|
|
#endif
|
|
}
|
|
[[unroll]] for (uint l = 0; l < BN; l += loadstride_b) {
|
|
#if LOAD_VEC_B == 8
|
|
#ifdef MUL_MAT_ID
|
|
const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l];
|
|
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;
|
|
#else
|
|
const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b;
|
|
#endif
|
|
const uint buf_idx = (loadc_b + l) * (BK+1) + loadr_b * LOAD_VEC_B;
|
|
buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x);
|
|
buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y);
|
|
buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z);
|
|
buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w);
|
|
buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x);
|
|
buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y);
|
|
buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z);
|
|
buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w);
|
|
#elif LOAD_VEC_B == 4
|
|
#ifdef MUL_MAT_ID
|
|
const u16vec2 row_idx = row_ids[ic * BN + loadc_b + l];
|
|
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;
|
|
#else
|
|
const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b;
|
|
#endif
|
|
const uint buf_idx = (loadc_b + l) * (BK+1) + loadr_b * LOAD_VEC_B;
|
|
buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx].x);
|
|
buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx].y);
|
|
buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx].z);
|
|
buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx].w);
|
|
#elif !MUL_MAT_ID
|
|
if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) {
|
|
buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]);
|
|
} else {
|
|
buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(0.0f);
|
|
}
|
|
#else
|
|
const uint row_i = ic * BN + loadc_b + l;
|
|
if (row_i < _ne1) {
|
|
const u16vec2 row_idx = row_ids[row_i];
|
|
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]);
|
|
} else {
|
|
buf_b[(loadc_b + l) * (BK+1) + loadr_b] = FLOAT_TYPE(0.0f);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
barrier();
|
|
|
|
pos_a += BK / LOAD_VEC_A;
|
|
pos_b += BK / LOAD_VEC_B;
|
|
|
|
for (uint i = 0; i < BK; i++) {
|
|
// Load from shared into cache
|
|
[[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
|
|
[[unroll]] for (uint j = 0; j < TM; j++) {
|
|
cache_a[wsir * TM + j] = buf_a[(warp_r * WM + wsir * WSUBM + tiwr * TM + j) * (BK+1) + i];
|
|
}
|
|
}
|
|
[[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
|
|
[[unroll]] for (uint j = 0; j < TN; j++) {
|
|
cache_b[wsic * TN + j] = buf_b[(warp_c * WN + wsic * WSUBN + tiwc * TN + j) * (BK+1) + i];
|
|
}
|
|
}
|
|
|
|
[[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
|
|
[[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
|
|
[[unroll]] for (uint cc = 0; cc < TN; cc++) {
|
|
[[unroll]] for (uint cr = 0; cr < TM; cr++) {
|
|
sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr] += float(cache_a[wsir * TM + cr]) * float(cache_b[wsic * TN + cc]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
barrier();
|
|
}
|
|
|
|
const uint dr = ir * BM + warp_r * WM;
|
|
const uint dc = ic * BN + warp_c * WN;
|
|
|
|
#ifndef MUL_MAT_ID
|
|
const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * gl_NumWorkGroups.z;
|
|
#endif
|
|
|
|
[[unroll]] for (uint wsic = 0; wsic < WNITER; wsic++) {
|
|
[[unroll]] for (uint wsir = 0; wsir < WMITER; wsir++) {
|
|
|
|
const uint dr_warp = dr + wsir * WSUBM + tiwr * TM;
|
|
const uint dc_warp = dc + wsic * WSUBN + tiwc * TN;
|
|
[[unroll]] for (uint cc = 0; cc < TN; cc++) {
|
|
#ifdef MUL_MAT_ID
|
|
const uint row_i = dc_warp + cc;
|
|
if (row_i >= _ne1) break;
|
|
|
|
const u16vec2 row_idx = row_ids[row_i];
|
|
#endif
|
|
[[unroll]] for (uint cr = 0; cr < TM; cr++) {
|
|
#ifdef MUL_MAT_ID
|
|
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]);
|
|
#else
|
|
if (dr_warp + cr < p.M && dc_warp + cc < p.N) {
|
|
data_d[offsets + (dc_warp + cc) * p.stride_d + dr_warp + cr] = D_TYPE(sums[(wsic * TN + cc) * (WMITER * TM) + wsir * TM + cr]);
|
|
}
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|