# ggml/src/ggml-sycl/* -> src/ggml-sycl/*
# ggml/src/ggml-sycl.cpp -> src/ggml-sycl.cpp
# ggml/src/ggml-vulkan.cpp -> src/ggml-vulkan.cpp
+ # ggml/src/vulkan-shaders/* -> src/vulkan-shaders/*
#
# ggml/include/ggml.h -> include/ggml.h
# ggml/include/ggml-alloc.h -> include/ggml-alloc.h
-e 's/\/ggml\/src\/ggml-sycl\//\/src\/ggml-sycl\//g' \
-e 's/\/ggml\/src\/ggml-sycl\.cpp/\/src\/ggml-sycl.cpp/g' \
-e 's/\/ggml\/src\/ggml-vulkan\.cpp/\/src\/ggml-vulkan.cpp/g' \
+ -e 's/\/ggml\/src\/vulkan-shaders\//\/src\/vulkan-shaders\//g' \
-e 's/\/ggml\/include\/ggml\.h/\/include\/ggml.h/g' \
-e 's/\/ggml\/include\/ggml-alloc\.h/\/include\/ggml-alloc.h/g' \
-e 's/\/ggml\/include\/ggml-backend\.h/\/include\/ggml-backend.h/g' \
cp -rpv ../llama.cpp/ggml/src/ggml-sycl/* src/ggml-sycl/
cp -rpv ../llama.cpp/ggml/src/ggml-sycl.cpp src/ggml-sycl.cpp
cp -rpv ../llama.cpp/ggml/src/ggml-vulkan.cpp src/ggml-vulkan.cpp
+cp -rpv ../llama.cpp/ggml/src/ggml-vulkan.cpp src/ggml-vulkan.cpp
+cp -rpv ../llama.cpp/ggml/src/vulkan-shaders/* src/vulkan-shaders/
cp -rpv ../llama.cpp/ggml/include/ggml.h include/ggml.h
cp -rpv ../llama.cpp/ggml/include/ggml-alloc.h include/ggml-alloc.h
--- /dev/null
+
+set(TARGET vulkan-shaders-gen)
+add_executable(${TARGET} vulkan-shaders-gen.cpp)
+install(TARGETS ${TARGET} RUNTIME)
+target_compile_features(${TARGET} PRIVATE cxx_std_11)
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_binary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(FLOAT_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]) + FLOAT_TYPE(data_b[src1_idx(gl_GlobalInvocationID.x)]));
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+
+#define BLOCK_SIZE 1024
+#define ASC 0
+
+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) buffer D {int data_d[];};
+
+layout (push_constant) uniform parameter {
+ uint ncols;
+ uint ncols_pad;
+ uint order;
+} p;
+
+shared int dst_row[BLOCK_SIZE];
+
+void swap(uint idx0, uint idx1) {
+ int tmp = dst_row[idx0];
+ dst_row[idx0] = dst_row[idx1];
+ dst_row[idx1] = tmp;
+}
+
+void main() {
+ // bitonic sort
+ const int col = int(gl_LocalInvocationID.x);
+ const uint row = gl_WorkGroupID.y;
+
+ if (col >= p.ncols_pad) {
+ return;
+ }
+
+ const uint row_offset = row * p.ncols;
+
+ // initialize indices
+ dst_row[col] = col;
+ barrier();
+
+ for (uint k = 2; k <= p.ncols_pad; k *= 2) {
+ for (uint j = k / 2; j > 0; j /= 2) {
+ const uint ixj = col ^ j;
+ if (ixj > col) {
+ if ((col & k) == 0) {
+ if (dst_row[col] >= p.ncols ||
+ (dst_row[ixj] < p.ncols && (p.order == ASC ?
+ data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]] :
+ data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]]))
+ ) {
+ swap(col, ixj);
+ }
+ } else {
+ if (dst_row[ixj] >= p.ncols ||
+ (dst_row[col] < p.ncols && (p.order == ASC ?
+ data_a[row_offset + dst_row[col]] < data_a[row_offset + dst_row[ixj]] :
+ data_a[row_offset + dst_row[col]] > data_a[row_offset + dst_row[ixj]]))
+ ) {
+ swap(col, ixj);
+ }
+ }
+ }
+ barrier();
+ }
+ }
+
+ if (col < p.ncols) {
+ data_d[row_offset + col] = dst_row[col];
+ }
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_unary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+ const FLOAT_TYPE val = FLOAT_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]);
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(val < p.param1 ? p.param1 : (val > p.param2 ? p.param2 : val));
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_unary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+#ifndef OPTIMIZATION_ERROR_WORKAROUND
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]);
+#else
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = data_a[src0_idx(gl_GlobalInvocationID.x)];
+#endif
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {float data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_GlobalInvocationID.x * 16;
+
+ if (i >= p.nel) {
+ return;
+ }
+
+ [[unroll]] for (uint l = 0; l < 16; l++) {
+ data_b[i + l] = D_TYPE(data_a[i + l]);
+ }
+}
--- /dev/null
+#if !defined(DATA_A_F32) && !defined(DATA_A_F16)
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#endif
+
+#if defined(DATA_A_F32)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]);
+}
+#endif
+
+#if defined(DATA_A_F16)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ return vec2(data_a[a_offset + ib], data_a[a_offset + ib + 1]);
+}
+#endif
+
+#if defined(DATA_A_Q4_0)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const float d = float(data_a[a_offset + ib].d);
+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+ return (vec2(vui & 0xF, vui >> 4) - 8.0f) * d;
+}
+#endif
+
+#if defined(DATA_A_Q4_1)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const float d = float(data_a[a_offset + ib].d);
+ const float m = float(data_a[a_offset + ib].m);
+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+ return vec2(vui & 0xF, vui >> 4) * d + m;
+}
+#endif
+
+#if defined(DATA_A_Q5_0)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const float d = float(data_a[a_offset + ib].d);
+ const uint uint_qh = uint(data_a[a_offset + ib].qh[1]) << 16 | data_a[a_offset + ib].qh[0];
+ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+ return (vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) - 16.0f) * d;
+}
+#endif
+
+#if defined(DATA_A_Q5_1)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const float d = float(data_a[a_offset + ib].d);
+ const float m = float(data_a[a_offset + ib].m);
+ const uint uint_qh = data_a[a_offset + ib].qh;
+ const ivec2 qh = ivec2(((uint_qh >> iqs) << 4) & 0x10, (uint_qh >> (iqs + 12)) & 0x10);
+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+ return vec2((vui & 0xF) | qh.x, (vui >> 4) | qh.y) * d + m;
+}
+#endif
+
+#if defined(DATA_A_Q8_0)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const float d = float(data_a[a_offset + ib].d);
+ return vec2(int(data_a[a_offset + ib].qs[iqs]), int(data_a[a_offset + ib].qs[iqs + 1])) * d;
+}
+#endif
+
+#if defined(DATA_A_IQ4_NL)
+vec2 dequantize(uint ib, uint iqs, uint a_offset) {
+ const float d = float(data_a[a_offset + ib].d);
+ const uint vui = uint(data_a[a_offset + ib].qs[iqs]);
+ return vec2(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]) * d;
+}
+#endif
--- /dev/null
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+ uint M;
+ uint K;
+ uint stride_a;
+ uint stride_b;
+ uint nel;
+} p;
+
+#include "types.comp"
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_iq4_nl data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid/32;
+ const uint ir = tid%32;
+ const uint ib = 32*i + ir;
+ if (ib >= p.nel / 32) {
+ return;
+ }
+
+ const uint q_idx = 8*il;
+ const uint b_idx = 1024*i + 32*ir + q_idx;
+
+ const float d = float(data_a[ib].d);
+
+ [[unroll]] for (uint l = 0; l < 8; ++l) {
+ data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]);
+ data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]);
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
+ const uint i = gl_WorkGroupID.x * 256 + wgy;
+ if (i >= p.M * p.K / QUANT_K) {
+ return;
+ }
+
+ const uint tid = gl_LocalInvocationID.x;
+ const uint ip = tid / 32;
+ const uint il = tid - 32 * ip;
+ const uint is = 8 * ip + il / 16;
+
+ const uint y_idx = i * QUANT_K + 128 * ip + il;
+
+ const uint ql_idx = 32 * ip + il;
+ const uint8_t qs = data_a[i].qs[32 * ip + il];
+
+ FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x);
+ FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y);
+ data_b[y_idx + 0] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+0] & 0xF) * ((qs >> 0) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+0] >> 4));
+ data_b[y_idx + 32] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+2] & 0xF) * ((qs >> 2) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+2] >> 4));
+ data_b[y_idx + 64] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+4] & 0xF) * ((qs >> 4) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+4] >> 4));
+ data_b[y_idx + 96] = D_TYPE(dall * FLOAT_TYPE((data_a[i].scales[is+6] & 0xF) * ((qs >> 6) & 3)) - dmin * FLOAT_TYPE(data_a[i].scales[is+6] >> 4));
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
+ const uint i = uint(gl_WorkGroupID.x * 256 + wgy);
+ if (i >= p.M * p.K / QUANT_K) {
+ return;
+ }
+
+ const uint r = gl_LocalInvocationID.x / 4;
+ const uint tid = r / 2;
+ const uint is0 = r % 2;
+ const uint l0 = 16 * is0 + 4 * (gl_LocalInvocationID.x % 4);
+ const uint n = tid / 4;
+ const uint j = tid - 4*n;
+
+ const uint8_t m = uint8_t(1 << (4*n + j));
+ const uint is = 8*n + 2*j + is0;
+ const uint shift = 2*j;
+
+ const int8_t us = int8_t(is < 4 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+8] >> 0) & 3) << 4) :
+ is < 8 ? (data_a[i].scales[is-0] & 0xF) | (((data_a[i].scales[is+4] >> 2) & 3) << 4) :
+ is < 12 ? (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is+0] >> 4) & 3) << 4) :
+ (data_a[i].scales[is-8] >> 4) | (((data_a[i].scales[is-4] >> 6) & 3) << 4));
+ const FLOAT_TYPE d_all = FLOAT_TYPE(data_a[i].d);
+ const FLOAT_TYPE dl = d_all * FLOAT_TYPE(us - 32);
+
+ const uint y_idx = i * QUANT_K + 128 * n + 32 * j;
+ const uint qs_idx = 32*n;
+
+ for (uint l = l0; l < l0 + 4; ++l) {
+ data_b[y_idx + l] = D_TYPE(dl * FLOAT_TYPE(int8_t((data_a[i].qs[qs_idx + l] >> shift) & 3) - (((data_a[i].hmask[l] & m) != 0) ? 0 : 4)));
+ }
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q4_0 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid/32;
+ const uint ir = tid%32;
+ const uint ib = 32*i + ir;
+ if (ib >= p.nel / 32) {
+ return;
+ }
+
+ const uint q_idx = 8*il;
+ const uint b_idx = 1024*i + 32*ir + q_idx;
+
+ const float d = float(data_a[ib].d);
+
+ [[unroll]] for (uint l = 0; l < 8; ++l) {
+ data_b[b_idx + l + 0] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] & 0xF) - 8.0f));
+ data_b[b_idx + l + 16] = D_TYPE(d * ((data_a[ib].qs[q_idx + l] >> 4) - 8.0f));
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q4_1 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid/32;
+ const uint ir = tid%32;
+ const uint ib = 32*i + ir;
+ if (ib >= p.nel / 32) {
+ return;
+ }
+
+ const uint b_idx = 1024*i + 32*ir + 8*il;
+
+ const float d = float(data_a[ib].d);
+ const float m = float(data_a[ib].m);
+
+ const uint q_idx = 8*il;
+
+ [[unroll]] for (uint l = 0; l < 8; ++l) {
+ data_b[b_idx + l + 0] = D_TYPE(d * (data_a[ib].qs[q_idx + l] & 0xF) + m);
+ data_b[b_idx + l + 16] = D_TYPE(d * (data_a[ib].qs[q_idx + l] >> 4) + m);
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
+ const uint i = gl_WorkGroupID.x * 256 + wgy;
+ if (i >= p.M * p.K / QUANT_K) {
+ return;
+ }
+
+ const uint tid = gl_LocalInvocationID.x;
+ const uint il = tid / 8;
+ const uint ir = tid % 8;
+ const uint is = 2 * il;
+ const uint n = 4;
+
+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x);
+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y);
+
+ const uint y_idx = i * QUANT_K + 64 * il + n * ir;
+ const uint qs_idx = 32*il + n * ir;
+
+ uint8_t sc;
+ uint8_t m;
+ if (is < 4) {
+ sc = uint8_t(data_a[i].scales[is] & 63);
+ m = uint8_t(data_a[i].scales[is + 4] & 63);
+ } else {
+ sc = uint8_t((data_a[i].scales[is + 4] & 0xF) | ((data_a[i].scales[is - 4] >> 6) << 4));
+ m = uint8_t((data_a[i].scales[is + 4] >> 4) | ((data_a[i].scales[is ] >> 6) << 4));
+ }
+ const FLOAT_TYPE d1 = dall * sc;
+ const FLOAT_TYPE m1 = dmin * m;
+
+ if (is < 4) {
+ sc = uint8_t(data_a[i].scales[is + 1] & 63);
+ m = uint8_t(data_a[i].scales[is + 5] & 63);
+ } else {
+ sc = uint8_t((data_a[i].scales[is + 5] & 0xF) | ((data_a[i].scales[is - 3] >> 6) << 4));
+ m = uint8_t((data_a[i].scales[is + 5] >> 4) | ((data_a[i].scales[is + 1] >> 6) << 4));
+ }
+ const FLOAT_TYPE d2 = dall * sc;
+ const FLOAT_TYPE m2 = dmin * m;
+
+ [[unroll]] for (uint l = 0; l < n; ++l) {
+ data_b[y_idx + l ] = D_TYPE(d1 * FLOAT_TYPE(data_a[i].qs[qs_idx + l] & 0xF) - m1);
+ data_b[y_idx + l + 32] = D_TYPE(d2 * FLOAT_TYPE(data_a[i].qs[qs_idx + l] >> 4) - m2);
+ }
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q5_0 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid/32;
+ const uint ir = tid%32;
+ const uint ib = 32*i + ir;
+ if (ib >= p.nel / 32) {
+ return;
+ }
+
+ const uint b_idx = 1024*i + 32*ir + 8*il;
+
+ const float d = float(data_a[ib].d);
+ const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0];
+
+ const uint q_idx = 8*il;
+
+ [[unroll]] for (uint l = 0; l < 8; ++l) {
+ const uint iqs = q_idx + l;
+ const uint vui = uint(data_a[ib].qs[iqs]);
+ data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10)) - 16.0f));
+ data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10)) - 16.0f));
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q5_1 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid/32;
+ const uint ir = tid%32;
+ const uint ib = 32*i + ir;
+ if (ib >= p.nel / 32) {
+ return;
+ }
+
+ const uint b_idx = 1024*i + 32*ir + 8*il;
+
+ const float d = float(data_a[ib].d);
+ const float m = float(data_a[ib].m);
+ const uint qh = data_a[ib].qh;
+
+ const uint q_idx = 8*il;
+
+ [[unroll]] for (uint l = 0; l < 8; ++l) {
+ const uint iqs = q_idx + l;
+ const uint vui = uint(data_a[ib].qs[iqs]);
+ data_b[b_idx + l + 0] = D_TYPE(d * (((vui & 0xF) | (((qh >> iqs) << 4) & 0x10))) + m);
+ data_b[b_idx + l + 16] = D_TYPE(d * (((vui >> 4) | ((qh >> (iqs + 12)) & 0x10))) + m);
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
+ const uint i = gl_WorkGroupID.x * 256 + wgy;
+ if (i >= p.M * p.K / QUANT_K) {
+ return;
+ }
+
+ const uint tid = gl_LocalInvocationID.x;
+ const uint il = tid / 16;
+ const uint ir = tid % 16;
+ const uint is = 2 * il;
+
+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[i].d.x);
+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[i].d.y);
+
+ const uint y_idx = i * QUANT_K + 64 * il + 2 * ir;
+ const uint qs_idx = 32*il + 2 * ir;
+ const uint qh_idx = 2 * ir;
+
+ uint8_t sc;
+ uint8_t m;
+ if (is < 4) {
+ sc = uint8_t(data_a[i].scales[is] & 63);
+ m = uint8_t(data_a[i].scales[is + 4] & 63);
+ } else {
+ sc = uint8_t((data_a[i].scales[is + 4] & 0xF) | ((data_a[i].scales[is - 4] >> 6) << 4));
+ m = uint8_t((data_a[i].scales[is + 4] >> 4) | ((data_a[i].scales[is ] >> 6) << 4));
+ }
+ const FLOAT_TYPE d1 = dall * sc;
+ const FLOAT_TYPE m1 = dmin * m;
+
+ if (is < 4) {
+ sc = uint8_t(data_a[i].scales[is + 1] & 63);
+ m = uint8_t(data_a[i].scales[is + 5] & 63);
+ } else {
+ sc = uint8_t((data_a[i].scales[is + 5] & 0xF) | ((data_a[i].scales[is - 3] >> 6) << 4));
+ m = uint8_t((data_a[i].scales[is + 5] >> 4) | ((data_a[i].scales[is + 1] >> 6) << 4));
+ }
+ const FLOAT_TYPE d2 = dall * sc;
+ const FLOAT_TYPE m2 = dmin * m;
+
+ const uint8_t hm1 = uint8_t(1 << (2 * il ));
+ const uint8_t hm2 = uint8_t(1 << (2 * il + 1));
+ data_b[y_idx ] = D_TYPE(d1 * FLOAT_TYPE((data_a[i].qs[qs_idx ] & 0xF) + (((data_a[i].qh[qh_idx ] & hm1) != 0) ? 16 : 0)) - m1);
+ data_b[y_idx + 1] = D_TYPE(d1 * FLOAT_TYPE((data_a[i].qs[qs_idx + 1] & 0xF) + (((data_a[i].qh[qh_idx + 1] & hm1) != 0) ? 16 : 0)) - m1);
+ data_b[y_idx + 32] = D_TYPE(d2 * FLOAT_TYPE((data_a[i].qs[qs_idx ] >> 4) + (((data_a[i].qh[qh_idx ] & hm2) != 0) ? 16 : 0)) - m2);
+ data_b[y_idx + 33] = D_TYPE(d2 * FLOAT_TYPE((data_a[i].qs[qs_idx + 1] >> 4) + (((data_a[i].qh[qh_idx + 1] & hm2) != 0) ? 16 : 0)) - m2);
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ [[unroll]] for (uint wgy = 0; wgy < 256; wgy++) {
+ const uint i = gl_WorkGroupID.x * 256 + wgy;
+ if (i >= p.M * p.K / QUANT_K) {
+ return;
+ }
+ const uint tid = gl_LocalInvocationID.x;
+ const uint ip = tid / 32;
+ const uint il = tid - 32 * ip;
+ const uint is = 8 * ip + il / 16;
+
+ const uint y_idx = i * QUANT_K + 128 * ip + il;
+
+ const uint ql_idx = 64 * ip + il;
+ const uint8_t qh = data_a[i].qh[32 * ip + il];
+
+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[i].d);
+
+ data_b[y_idx + 0] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 0] * (int8_t((data_a[i].ql[ql_idx + 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)));
+ data_b[y_idx + 32] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 2] * (int8_t((data_a[i].ql[ql_idx + 32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)));
+ data_b[y_idx + 64] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 4] * (int8_t((data_a[i].ql[ql_idx + 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)));
+ data_b[y_idx + 96] = D_TYPE(d * FLOAT_TYPE(data_a[i].scales[is + 6] * (int8_t((data_a[i].ql[ql_idx + 32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)));
+ }
+}
--- /dev/null
+#version 450
+
+#include "dequant_head.comp"
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {block_q8_0 data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_b[];};
+
+void main() {
+ const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64;
+
+ const uint tid = gl_LocalInvocationID.x % 64;
+ const uint il = tid/32;
+ const uint ir = tid%32;
+ const uint ib = 32*i + ir;
+ if (ib >= p.nel / 32) {
+ return;
+ }
+
+ const uint b_idx = 1024*i + 32*ir + 16*il;
+
+ const float d = float(data_a[ib].d);
+
+ const uint q_idx = 16*il;
+
+ [[unroll]] for (uint l = 0; l < 16; l += 2) {
+ data_b[b_idx + l ] = D_TYPE(d * data_a[ib].qs[q_idx + l ]);
+ data_b[b_idx + l + 1] = D_TYPE(d * data_a[ib].qs[q_idx + l + 1]);
+ }
+}
--- /dev/null
+#version 450
+
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_control_flow_attributes : enable
+
+layout (push_constant) uniform parameter
+{
+ uint ncols;
+ uint rows_per_channel;
+ uint n_past;
+} p;
+
+#include "types.comp"
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+void main() {
+ const uint col = gl_GlobalInvocationID.y;
+ const uint row = gl_GlobalInvocationID.x;
+
+ if (col >= p.ncols) {
+ return;
+ }
+
+ const uint i = row*p.ncols + col;
+ if (col > p.n_past + row % p.rows_per_channel) {
+ data_d[i] = D_TYPE(uintBitsToFloat(0xFF800000));
+ } else {
+ data_d[i] = D_TYPE(data_a[i]);
+ }
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_binary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(FLOAT_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]) / FLOAT_TYPE(data_b[src1_idx(gl_GlobalInvocationID.x)]));
+}
--- /dev/null
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+void main() {
+ const float GELU_COEF_A = 0.044715f;
+ const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f;
+ const uint i = gl_GlobalInvocationID.x;
+
+ if (i >= p.KX) {
+ return;
+ }
+
+ const float xi = float(data_a[i]);
+ const float val = SQRT_2_OVER_PI*xi*(1.0f + GELU_COEF_A*xi*xi);
+ data_d[i] = D_TYPE(0.5f*xi*(2.0f - 2.0f / (exp(2 * val) + 1)));
+}
--- /dev/null
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+ uint ne;
+ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
+ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
+ uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23;
+ uint d_offset;
+ float param1; float param2;
+} p;
+
+layout(local_size_x = 512, 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[];};
+
+uint src0_idx(uint idx) {
+ const uint i03 = idx / (p.ne02*p.ne01*p.ne00);
+ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;
+ const uint i02 = (idx - i03_offset) / (p.ne01*p.ne00);
+ const uint i02_offset = i02*p.ne01*p.ne00;
+ const uint i01 = (idx - i03_offset - i02_offset) / p.ne00;
+ const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00;
+ return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00;
+}
+
+uint src1_idx(uint idx) {
+ const uint i03 = idx / (p.ne02*p.ne01*p.ne00);
+ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;
+ const uint i02 = (idx - i03_offset) / (p.ne01*p.ne00);
+ const uint i02_offset = i02*p.ne01*p.ne00;
+ const uint i01 = (idx - i03_offset - i02_offset) / p.ne00;
+ const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00;
+
+ return (i03 % p.ne13)*p.nb13 + (i02 % p.ne12)*p.nb12 + (i01 % p.ne11)*p.nb11 + (i00 % p.ne10)*p.nb10;
+}
+
+uint dst_idx(uint idx) {
+ const uint i23 = idx / (p.ne22*p.ne21*p.ne20);
+ const uint i23_offset = i23 * p.ne22*p.ne21*p.ne20;
+ const uint i22 = (idx - i23_offset) / (p.ne21*p.ne20);
+ const uint i22_offset = i22*p.ne21*p.ne20;
+ const uint i21 = (idx - i23_offset - i22_offset) / p.ne20;
+ const uint i20 = idx - i23_offset - i22_offset - i21*p.ne20;
+ return i23*p.nb23 + i22*p.nb22 + i21*p.nb21 + i20*p.nb20;
+}
--- /dev/null
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+ uint KX;
+ uint KY;
+ float param1;
+ float param2;
+} p;
--- /dev/null
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+ uint ne;
+ uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
+ uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
+ uint d_offset;
+ float param1; float param2;
+} p;
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+uint src0_idx(uint idx) {
+ const uint i03 = idx / (p.ne02*p.ne01*p.ne00);
+ const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00;
+ const uint i02 = (idx - i03_offset) / (p.ne01*p.ne00);
+ const uint i02_offset = i02*p.ne01*p.ne00;
+ const uint i01 = (idx - i03_offset - i02_offset) / p.ne00;
+ const uint i00 = idx - i03_offset - i02_offset - i01*p.ne00;
+ return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00*p.nb00;
+}
+
+uint dst_idx(uint idx) {
+ const uint i13 = idx / (p.ne12*p.ne11*p.ne10);
+ const uint i13_offset = i13 * p.ne12*p.ne11*p.ne10;
+ const uint i12 = (idx - i13_offset) / (p.ne11*p.ne10);
+ const uint i12_offset = i12*p.ne11*p.ne10;
+ const uint i11 = (idx - i13_offset - i12_offset) / p.ne10;
+ const uint i10 = idx - i13_offset - i12_offset - i11*p.ne10;
+ return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10*p.nb10;
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_binary_head.comp"
+
+void main() {
+ const uint i00 = gl_GlobalInvocationID.x;
+ const uint i10 = gl_GlobalInvocationID.y;
+ const uint i11 = (gl_GlobalInvocationID.z)/p.ne12;
+ const uint i12 = (gl_GlobalInvocationID.z)%p.ne12;
+
+ if (i00 >= p.ne00) {
+ return;
+ }
+
+ const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12];
+
+ const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03;
+ const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23;
+
+#ifndef OPTIMIZATION_ERROR_WORKAROUND
+ data_d[d_offset + i00] = D_TYPE(data_a[a_offset + i00]);
+#else
+ data_d[d_offset + i00] = data_a[a_offset + i00];
+#endif
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_binary_head.comp"
+#include "dequant_funcs.comp"
+
+void main() {
+ const uint i00 = (gl_GlobalInvocationID.x)*2;
+ const uint i10 = gl_GlobalInvocationID.y;
+ const uint i11 = (gl_GlobalInvocationID.z)/p.ne12;
+ const uint i12 = (gl_GlobalInvocationID.z)%p.ne12;
+
+ if (i00 >= p.ne00) {
+ return;
+ }
+
+ const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12];
+
+ const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03;
+ const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23;
+
+ const uint ib = a_offset + i00/QUANT_K; // block index
+ const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index
+ const uint iybs = i00 - i00%QUANT_K; // dst block start index
+ const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
+
+ vec2 v = dequantize(ib, iqs, 0);
+
+ data_d[d_offset + iybs + iqs ] = D_TYPE(v.x);
+ data_d[d_offset + iybs + iqs + y_offset] = D_TYPE(v.y);
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_binary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(FLOAT_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]) * FLOAT_TYPE(data_b[src1_idx(gl_GlobalInvocationID.x)]));
+}
--- /dev/null
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer A {float data_a[];};
+layout (binding = 1) writeonly buffer D {float data_d[];};
+
+layout (push_constant) uniform parameter {
+ uint ne;
+ uint k_num;
+} p;
+
+void main() {
+ const uint idx = gl_GlobalInvocationID.x;
+
+ if (idx >= p.ne) {
+ return;
+ }
+
+ float result = 0.0f;
+
+ [[unroll]] for (uint i = 0; i < p.k_num; i++) {
+ result += data_a[i * p.ne + idx];
+ }
+
+ data_d[idx] = result;
+}
--- /dev/null
+#version 450
+
+#ifdef FLOAT16
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#endif
+
+#include "mul_mat_vec_base.comp"
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (constant_id = 0) const uint BLOCK_SIZE = 32;
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+void main() {
+ const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
+ const uint tid = gl_LocalInvocationID.x;
+
+ uint a_offset, b_offset, d_offset;
+ get_offsets(a_offset, b_offset, d_offset);
+
+ const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
+
+ tmp[tid] = FLOAT_TYPE(0.0f);
+
+ [[unroll]] for (uint i = 0; i < p.ncols/BLOCK_SIZE; i += 2) {
+ const uint col = i*BLOCK_SIZE + 2*tid;
+ const uint ib = (row*p.ncols + col)/QUANT_K; // block index
+ const uint iqs = (col%QUANT_K)/QUANT_R; // quant index
+ const uint iybs = col - col%QUANT_K; // y block start index
+
+ vec2 v = dequantize(ib, iqs, a_offset / QUANT_K);
+
+ // matrix multiplication
+ tmp[tid] += FLOAT_TYPE(v.x) * FLOAT_TYPE(data_b[b_offset + iybs + iqs]) +
+ FLOAT_TYPE(v.y) * FLOAT_TYPE(data_b[b_offset + iybs + iqs + y_offset]);
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (uint s = BLOCK_SIZE/2; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+ if (tid == 0) {
+ data_d[d_offset + row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_8bit_storage : require
+
+#define K_QUANTS_PER_ITERATION 2
+
+#ifdef MUL_MAT_ID
+#define EXPERT_COUNT 8
+#endif
+
+#include "types.comp"
+
+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
+
+#include "dequant_funcs.comp"
+
+layout (push_constant) uniform parameter
+{
+ uint ncols;
+ 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 ne11;
+#else
+ uint ne02;
+ uint ne12;
+ uint broadcast2;
+ uint broadcast3;
+#endif
+} p;
+
+void get_offsets(out uint a_offset, out uint b_offset, out uint d_offset) {
+#ifdef MUL_MAT_ID
+ const uint expert_idx = gl_GlobalInvocationID.y;
+#else
+ const uint batch_idx = gl_GlobalInvocationID.y;
+#endif
+
+#ifndef MUL_MAT_ID
+ 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;
+#else
+ const uint expert_id = data_ids[expert_idx];
+#endif
+
+ a_offset =
+#ifdef MUL_MAT_ID
+ expert_id * p.batch_stride_a;
+#else
+ batch_idx_a * p.batch_stride_a;
+#endif
+ b_offset =
+#ifdef MUL_MAT_ID
+ (expert_idx % p.ne11) * p.stride_b;
+#else
+ batch_idx * p.batch_stride_b;
+#endif
+ d_offset =
+#ifdef MUL_MAT_ID
+ expert_idx * p.stride_d;
+#else
+ batch_idx * p.batch_stride_d;
+#endif
+}
--- /dev/null
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+
+#define BLOCK_SIZE 32
+#define FLOAT_TYPE float
+
+layout(local_size_x = BLOCK_SIZE, 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 dst[];};
+
+layout (push_constant) uniform parameter
+{
+ uint ncols_x;
+ uint nrows_x;
+ uint row_stride_x;
+ uint channel_stride_x;
+ uint channel_x_divisor;
+ uint b_offset;
+ uint d_offset;
+} p;
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+void main() {
+ const uint tid = gl_LocalInvocationID.x;
+ const uint row_x = gl_GlobalInvocationID.y;
+ const uint channel = gl_GlobalInvocationID.z;
+ const uint channel_x = channel / p.channel_x_divisor;
+
+ const uint nrows_y = p.ncols_x;
+ const uint nrows_dst = p.nrows_x;
+ const uint row_dst = row_x;
+
+ const uint idst = channel*nrows_dst + row_dst;
+
+ tmp[tid] = 0.0f;
+
+ for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) {
+ const uint col_x = col_x0 + tid;
+
+ if (col_x >= p.ncols_x) {
+ break;
+ }
+
+ const uint row_y = col_x;
+
+ const uint ix = channel_x*p.channel_stride_x + row_x*p.row_stride_x + col_x;
+ const uint iy = channel*nrows_y + row_y;
+
+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]);
+
+ tmp[tid] += xi * FLOAT_TYPE(data_b[iy]);
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+
+ if (tid == 0) {
+ dst[idst] = tmp[0];
+ }
+}
--- /dev/null
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+
+#define BLOCK_SIZE 32
+#define FLOAT_TYPE float
+
+layout(local_size_x = BLOCK_SIZE, 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 dst[];};
+
+layout (push_constant) uniform parameter
+{
+ uint ncols_x;
+ uint nrows_x;
+ uint nchannels_x;
+ uint nchannels_y;
+ uint b_offset;
+ uint d_offset;
+} p;
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+void main() {
+ const uint tid = gl_LocalInvocationID.x;
+ const uint row_x = gl_GlobalInvocationID.y;
+ const uint channel = gl_GlobalInvocationID.z;
+ const uint channel_x = channel / (p.nchannels_y / p.nchannels_x);
+
+ const uint nrows_y = p.ncols_x;
+ const uint nrows_dst = p.nrows_x;
+ const uint row_dst = row_x;
+
+ tmp[tid] = FLOAT_TYPE(0.0f);
+
+ for (uint col_x0 = 0; col_x0 < p.ncols_x; col_x0 += BLOCK_SIZE) {
+ const uint col_x = col_x0 + tid;
+
+ if (col_x >= p.ncols_x) {
+ break;
+ }
+
+ // x is transposed and permuted
+ const uint ix = row_x*p.nchannels_x*p.ncols_x + channel_x*p.ncols_x + col_x;
+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[ix]);
+
+ const uint row_y = col_x;
+
+ // y is not transposed but permuted
+ const uint iy = channel*nrows_y + row_y;
+
+ tmp[tid] += xi * FLOAT_TYPE(data_b[iy]);
+ }
+
+ // dst is not transposed and not permuted
+ const uint idst = channel*nrows_dst + row_dst;
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+
+ if (tid == 0) {
+ dst[idst] = tmp[0];
+ }
+}
--- /dev/null
+#version 450
+
+#include "mul_mat_vec_base.comp"
+
+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
+
+shared FLOAT_TYPE tmp[32];
+
+void main() {
+ const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
+
+ uint a_offset, b_offset, d_offset;
+ get_offsets(a_offset, b_offset, d_offset);
+
+ const uint num_blocks_per_row = p.ncols / QUANT_K;
+ const uint ib0 = a_offset / QUANT_K + row*num_blocks_per_row;
+
+ const uint tid = gl_LocalInvocationID.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16
+ const uint ix = gl_LocalInvocationID.x%K_QUANTS_PER_ITERATION; // 0 or 0, 1
+
+ const uint step = 16/K_QUANTS_PER_ITERATION; // 16 or 8
+
+ const uint v_im = tid/step; // 0 or 1. 0 computes 0..., 1 computes 128...
+ const uint v_in = tid - step*v_im; // 0...15 or 0...7
+
+ const uint l0 = K_QUANTS_PER_ITERATION*v_in; // 0...15
+ const uint q_offset = 32*v_im + l0;
+ const uint s_offset = 8*v_im;
+ const uint y_offset = 128*v_im + l0;
+
+ tmp[16 * ix + tid] = FLOAT_TYPE(0.0); // partial sum for thread in warp
+
+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) {
+ const uint y_idx = i * QUANT_K + y_offset;
+
+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib0 + i].d.x);
+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib0 + i].d.y);
+
+ FLOAT_TYPE sum1 = FLOAT_TYPE(0.0);
+ FLOAT_TYPE sum2 = FLOAT_TYPE(0.0);
+ for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) {
+ sum1 += FLOAT_TYPE(data_b[b_offset + y_idx + l + 0]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l + 0] >> 0) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 16]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 1] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l +16] >> 0) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 32]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l + 0] >> 2) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 48]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 3] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l +16] >> 2) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 64]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l + 0] >> 4) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 80]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 5] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l +16] >> 4) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 96]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l + 0] >> 6) & 3)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l +112]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 7] & 0xF) * FLOAT_TYPE((data_a[ib0 + i].qs[q_offset + l +16] >> 6) & 3);
+ sum2 += FLOAT_TYPE(data_b[b_offset + y_idx + l + 0]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 0] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 16]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 1] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 32]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 2] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 48]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 3] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 64]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 4] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 80]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 5] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 96]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 6] >> 4) & 0xF)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l +112]) * FLOAT_TYPE((data_a[ib0 + i].scales[s_offset + 7] >> 4) & 0xF);
+ }
+ tmp[16 * ix + tid] += dall * sum1 - dmin * sum2;
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (uint s = 16; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+ if (tid == 0) {
+ data_d[d_offset + row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#version 450
+
+#include "mul_mat_vec_base.comp"
+
+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
+
+shared FLOAT_TYPE tmp[32];
+
+void main() {
+ const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
+
+ uint a_offset, b_offset, d_offset;
+ get_offsets(a_offset, b_offset, d_offset);
+
+ const uint num_blocks_per_row = p.ncols / QUANT_K;
+ const uint ib0 = a_offset / QUANT_K + row*num_blocks_per_row;
+
+ const uint tid = gl_LocalInvocationID.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16
+ const uint ix = gl_LocalInvocationID.x%K_QUANTS_PER_ITERATION; // 0 or 0, 1
+
+ const uint step = 16/K_QUANTS_PER_ITERATION; // 16 or 8
+
+ const uint v_im = tid/step; // 0 or 1. 0 computes 0..., 1 computes 128...
+ const uint v_in = tid - step*v_im; // 0...15 or 0...7
+
+ const uint8_t m = uint8_t(1 << (4 * v_im));
+
+ const uint l0 = K_QUANTS_PER_ITERATION*v_in; // 0...15
+ const uint q_offset = 32*v_im + l0;
+ const uint y_offset = 128*v_im + l0;
+
+ tmp[16 * ix + tid] = FLOAT_TYPE(0.0); // partial sum for thread in warp
+
+ const uint s_shift = 4 * v_im;
+
+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) {
+ const uint y_idx = i * QUANT_K + y_offset;
+
+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d);
+
+ FLOAT_TYPE sum = FLOAT_TYPE(0.0);
+ for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) {
+ sum += FLOAT_TYPE(data_b[b_offset + y_idx + l + 0]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[0] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[ 8] >> (s_shift + 0) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 0)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 32]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[2] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[10] >> (s_shift + 0) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 1)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 64]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[4] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[ 8] >> (s_shift + 2) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 2)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 96]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[6] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[10] >> (s_shift + 2) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l ] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l ] & (m << 3)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 16]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[1] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[ 9] >> (s_shift + 0) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] ) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 0)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 48]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[3] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[11] >> (s_shift + 0) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 2) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 1)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l + 80]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[5] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[ 9] >> (s_shift + 2) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 4) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 2)) != 0) ? 0 : 4))
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l +112]) * FLOAT_TYPE(int8_t(((data_a[ib0 + i].scales[7] >> s_shift) & 0xF) | ((data_a[ib0 + i].scales[11] >> (s_shift + 2) & 0x3) << 4)) - 32) * FLOAT_TYPE(((data_a[ib0 + i].qs[q_offset + l+16] >> 6) & 3) - (((data_a[ib0 + i].hmask[l0 + l+16] & (m << 3)) != 0) ? 0 : 4));
+ }
+ tmp[16 * ix + tid] += d * sum;
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (uint s = 16; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+ if (tid == 0) {
+ data_d[d_offset + row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#version 450
+
+#include "mul_mat_vec_base.comp"
+
+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
+
+shared FLOAT_TYPE tmp[32];
+
+void main() {
+ const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
+
+ uint a_offset, b_offset, d_offset;
+ get_offsets(a_offset, b_offset, d_offset);
+
+ const uint num_blocks_per_row = p.ncols / QUANT_K;
+ const uint ib0 = a_offset / QUANT_K + row*num_blocks_per_row;
+
+ const uint tid = gl_LocalInvocationID.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16
+ const uint ix = gl_LocalInvocationID.x%K_QUANTS_PER_ITERATION; // 0 or 0, 1
+
+ const uint step = 8/K_QUANTS_PER_ITERATION; // 8 or 4
+
+ const uint il = tid/step; // 0...3
+ const uint ir = tid - step*il; // 0...7 or 0...3
+ const uint n = 2 * K_QUANTS_PER_ITERATION; // 2 or 4
+
+ const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224
+ const uint v_in = il % 2;
+
+ const uint l0 = n * (2 * ir + v_in); // 0...15
+ const uint q_offset = 32*v_im + l0;
+ const uint y_offset = 64*v_im + l0;
+
+ tmp[16 * ix + tid] = FLOAT_TYPE(0.0); // partial sum for thread in warp
+
+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) {
+ const uint y1_idx = i * QUANT_K + y_offset;
+ const uint y2_idx = y1_idx + 128;
+
+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib0 + i].d.x);
+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib0 + i].d.y);
+
+ const uint8_t sc0 = uint8_t( data_a[ib0 + i].scales[v_im * 2 ] & 0x3f);
+ const uint8_t sc1 = uint8_t( data_a[ib0 + i].scales[v_im * 2 + 1] & 0x3f);
+ const uint8_t sc2 = uint8_t( data_a[ib0 + i].scales[v_im * 2 + 4] & 0x3f);
+ const uint8_t sc3 = uint8_t( data_a[ib0 + i].scales[v_im * 2 + 5] & 0x3f);
+ const uint8_t sc4 = uint8_t(( data_a[ib0 + i].scales[v_im * 2 + 8] & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 ] & 0xc0) >> 2));
+ const uint8_t sc5 = uint8_t(( data_a[ib0 + i].scales[v_im * 2 + 9] & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 + 1] & 0xc0) >> 2));
+ const uint8_t sc6 = uint8_t(((data_a[ib0 + i].scales[v_im * 2 + 8] >> 4) & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 + 4] & 0xc0) >> 2));
+ const uint8_t sc7 = uint8_t(((data_a[ib0 + i].scales[v_im * 2 + 9] >> 4) & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 + 5] & 0xc0) >> 2));
+
+#if K_QUANTS_PER_ITERATION == 2
+ const uint8_t q4_0 = uint8_t(data_a[ib0 + i].qs[q_offset ] & 0xf);
+ const uint8_t q4_1 = uint8_t(data_a[ib0 + i].qs[q_offset + 1] & 0xf);
+ const uint8_t q4_2 = uint8_t(data_a[ib0 + i].qs[q_offset + 2] & 0xf);
+ const uint8_t q4_3 = uint8_t(data_a[ib0 + i].qs[q_offset + 3] & 0xf);
+ const uint8_t q4_4 = uint8_t(data_a[ib0 + i].qs[q_offset ] >> 4);
+ const uint8_t q4_5 = uint8_t(data_a[ib0 + i].qs[q_offset + 1] >> 4);
+ const uint8_t q4_6 = uint8_t(data_a[ib0 + i].qs[q_offset + 2] >> 4);
+ const uint8_t q4_7 = uint8_t(data_a[ib0 + i].qs[q_offset + 3] >> 4);
+ const uint8_t q4_8 = uint8_t(data_a[ib0 + i].qs[q_offset + 64] & 0xf);
+ const uint8_t q4_9 = uint8_t(data_a[ib0 + i].qs[q_offset + 65] & 0xf);
+ const uint8_t q4_10 = uint8_t(data_a[ib0 + i].qs[q_offset + 66] & 0xf);
+ const uint8_t q4_11 = uint8_t(data_a[ib0 + i].qs[q_offset + 67] & 0xf);
+ const uint8_t q4_12 = uint8_t(data_a[ib0 + i].qs[q_offset + 64] >> 4);
+ const uint8_t q4_13 = uint8_t(data_a[ib0 + i].qs[q_offset + 65] >> 4);
+ const uint8_t q4_14 = uint8_t(data_a[ib0 + i].qs[q_offset + 66] >> 4);
+ const uint8_t q4_15 = uint8_t(data_a[ib0 + i].qs[q_offset + 67] >> 4);
+
+ const FLOAT_TYPE sx = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y1_idx]) * q4_0 + FLOAT_TYPE(data_b[b_offset + y1_idx + 1]) * q4_1 + FLOAT_TYPE(data_b[b_offset + y1_idx + 2]) * q4_2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 3]) * q4_3);
+ const FLOAT_TYPE sy = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y1_idx + 32]) * q4_4 + FLOAT_TYPE(data_b[b_offset + y1_idx + 33]) * q4_5 + FLOAT_TYPE(data_b[b_offset + y1_idx + 34]) * q4_6 + FLOAT_TYPE(data_b[b_offset + y1_idx + 35]) * q4_7);
+ const FLOAT_TYPE sz = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y2_idx]) * q4_8 + FLOAT_TYPE(data_b[b_offset + y2_idx + 1]) * q4_9 + FLOAT_TYPE(data_b[b_offset + y2_idx + 2]) * q4_10 + FLOAT_TYPE(data_b[b_offset + y2_idx + 3]) * q4_11);
+ const FLOAT_TYPE sw = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y2_idx + 32]) * q4_12 + FLOAT_TYPE(data_b[b_offset + y2_idx + 33]) * q4_13 + FLOAT_TYPE(data_b[b_offset + y2_idx + 34]) * q4_14 + FLOAT_TYPE(data_b[b_offset + y2_idx + 35]) * q4_15);
+ const FLOAT_TYPE smin = FLOAT_TYPE(
+ FLOAT_TYPE(data_b[b_offset + y1_idx ]) * sc2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 32]) * sc3 + FLOAT_TYPE(data_b[b_offset + y2_idx ]) * sc6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 32]) * sc7
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 1]) * sc2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 33]) * sc3 + FLOAT_TYPE(data_b[b_offset + y2_idx + 1]) * sc6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 33]) * sc7
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 2]) * sc2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 34]) * sc3 + FLOAT_TYPE(data_b[b_offset + y2_idx + 2]) * sc6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 34]) * sc7
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 3]) * sc2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 35]) * sc3 + FLOAT_TYPE(data_b[b_offset + y2_idx + 3]) * sc6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 35]) * sc7
+ );
+ tmp[16 * ix + tid] += FLOAT_TYPE(dall * (sx * sc0 + sy * sc1 + sz * sc4 + sw * sc5) - dmin * smin);
+#else
+ const uint8_t q4_0 = uint8_t(data_a[ib0 + i].qs[q_offset ] & 0xf);
+ const uint8_t q4_1 = uint8_t(data_a[ib0 + i].qs[q_offset + 1] & 0xf);
+ const uint8_t q4_2 = uint8_t(data_a[ib0 + i].qs[q_offset ] >> 4);
+ const uint8_t q4_3 = uint8_t(data_a[ib0 + i].qs[q_offset + 1] >> 4);
+ const uint8_t q4_4 = uint8_t(data_a[ib0 + i].qs[q_offset + 64] & 0xf);
+ const uint8_t q4_5 = uint8_t(data_a[ib0 + i].qs[q_offset + 65] & 0xf);
+ const uint8_t q4_6 = uint8_t(data_a[ib0 + i].qs[q_offset + 64] >> 4);
+ const uint8_t q4_7 = uint8_t(data_a[ib0 + i].qs[q_offset + 65] >> 4);
+
+ const FLOAT_TYPE sx = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y1_idx ]) * q4_0 + FLOAT_TYPE(data_b[b_offset + y1_idx + 1]) * q4_1);
+ const FLOAT_TYPE sy = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y1_idx + 32]) * q4_2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 33]) * q4_3);
+ const FLOAT_TYPE sz = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y2_idx ]) * q4_4 + FLOAT_TYPE(data_b[b_offset + y2_idx + 1]) * q4_5);
+ const FLOAT_TYPE sw = FLOAT_TYPE(FLOAT_TYPE(data_b[b_offset + y2_idx + 32]) * q4_6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 33]) * q4_7);
+ const FLOAT_TYPE smin = FLOAT_TYPE(
+ FLOAT_TYPE(data_b[b_offset + y1_idx]) * sc2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 32]) * sc3 + FLOAT_TYPE(data_b[b_offset + y2_idx]) * sc6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 32]) * sc7
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 1]) * sc2 + FLOAT_TYPE(data_b[b_offset + y1_idx + 33]) * sc3 + FLOAT_TYPE(data_b[b_offset + y2_idx + 1]) * sc6 + FLOAT_TYPE(data_b[b_offset + y2_idx + 33]) * sc7
+ );
+
+ tmp[16 * ix + tid] += FLOAT_TYPE(dall * (sx * FLOAT_TYPE(data_a[ib0 + i].scales[v_im] & 0x3f) + sy * FLOAT_TYPE(data_a[ib0 + i].scales[v_im + 1] & 0x3f) + sz * FLOAT_TYPE((data_a[ib0 + i].scales[v_im + 4] & 0x0f) | ((data_a[ib0 + i].scales[v_im] & 0xc0) >> 2)) + sw * FLOAT_TYPE((data_a[ib0 + i].scales[v_im + 5] & 0x0f) | ((data_a[ib0 + i].scales[v_im + 1] & 0xc0) >> 2))) - dmin * smin);
+#endif
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (uint s = 16; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+ if (tid == 0) {
+ data_d[d_offset + row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#version 450
+
+#include "mul_mat_vec_base.comp"
+
+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
+
+shared FLOAT_TYPE tmp[32];
+
+void main() {
+ const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
+
+ uint a_offset, b_offset, d_offset;
+ get_offsets(a_offset, b_offset, d_offset);
+
+ const uint num_blocks_per_row = p.ncols / QUANT_K;
+ const uint ib0 = a_offset / QUANT_K + row*num_blocks_per_row;
+
+ const uint tid = gl_LocalInvocationID.x/2; // 0...31 or 0...16
+ const uint ix = gl_LocalInvocationID.x%2; // 0 or 0, 1
+
+ const uint il = tid/4; // 0...3
+ const uint ir = tid - 4*il; // 0...7 or 0...3
+
+ const uint v_im = il / 2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224
+ const uint v_in = il % 2;
+
+ const uint l0 = 4*ir + 2*v_in; // 0...15
+ const uint q_offset = 32*v_im + l0;
+ const uint y_offset = 64*v_im + l0;
+
+ const uint8_t hm1 = uint8_t(1 << (2*v_im));
+ const uint8_t hm2 = uint8_t(hm1 << 4);
+
+ tmp[16 * ix + tid] = FLOAT_TYPE(0.0); // partial sum for thread in warp
+
+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += 2) {
+ const uint y1_idx = i * QUANT_K + y_offset;
+ const uint y2_idx = y1_idx + 128;
+
+ const FLOAT_TYPE dall = FLOAT_TYPE(data_a[ib0 + i].d.x);
+ const FLOAT_TYPE dmin = FLOAT_TYPE(data_a[ib0 + i].d.y);
+
+ const uint8_t sc0 = uint8_t( data_a[ib0 + i].scales[v_im * 2 ] & 0x3f);
+ const uint8_t sc1 = uint8_t( data_a[ib0 + i].scales[v_im * 2 + 1] & 0x3f);
+ const uint8_t sc2 = uint8_t( data_a[ib0 + i].scales[v_im * 2 + 4] & 0x3f);
+ const uint8_t sc3 = uint8_t( data_a[ib0 + i].scales[v_im * 2 + 5] & 0x3f);
+ const uint8_t sc4 = uint8_t(( data_a[ib0 + i].scales[v_im * 2 + 8] & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 ] & 0xc0) >> 2));
+ const uint8_t sc5 = uint8_t(( data_a[ib0 + i].scales[v_im * 2 + 9] & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 + 1] & 0xc0) >> 2));
+ const uint8_t sc6 = uint8_t(((data_a[ib0 + i].scales[v_im * 2 + 8] >> 4) & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 + 4] & 0xc0) >> 2));
+ const uint8_t sc7 = uint8_t(((data_a[ib0 + i].scales[v_im * 2 + 9] >> 4) & 0x0f) | ((data_a[ib0 + i].scales[v_im * 2 + 5] & 0xc0) >> 2));
+
+ const uint8_t q4_0 = uint8_t(data_a[ib0 + i].qs[q_offset ] & 0xf);
+ const uint8_t q4_1 = uint8_t(data_a[ib0 + i].qs[q_offset + 1] & 0xf);
+ const uint8_t q4_2 = uint8_t(data_a[ib0 + i].qs[q_offset + 16] & 0xf);
+ const uint8_t q4_3 = uint8_t(data_a[ib0 + i].qs[q_offset + 17] & 0xf);
+ const uint8_t q4_4 = uint8_t(data_a[ib0 + i].qs[q_offset ] >> 4);
+ const uint8_t q4_5 = uint8_t(data_a[ib0 + i].qs[q_offset + 1] >> 4);
+ const uint8_t q4_6 = uint8_t(data_a[ib0 + i].qs[q_offset + 16] >> 4);
+ const uint8_t q4_7 = uint8_t(data_a[ib0 + i].qs[q_offset + 17] >> 4);
+ const uint8_t q4_8 = uint8_t(data_a[ib0 + i].qs[q_offset + 64] & 0xf);
+ const uint8_t q4_9 = uint8_t(data_a[ib0 + i].qs[q_offset + 65] & 0xf);
+ const uint8_t q4_10 = uint8_t(data_a[ib0 + i].qs[q_offset + 80] & 0xf);
+ const uint8_t q4_11 = uint8_t(data_a[ib0 + i].qs[q_offset + 81] & 0xf);
+ const uint8_t q4_12 = uint8_t(data_a[ib0 + i].qs[q_offset + 64] >> 4);
+ const uint8_t q4_13 = uint8_t(data_a[ib0 + i].qs[q_offset + 65] >> 4);
+ const uint8_t q4_14 = uint8_t(data_a[ib0 + i].qs[q_offset + 80] >> 4);
+ const uint8_t q4_15 = uint8_t(data_a[ib0 + i].qs[q_offset + 81] >> 4);
+
+ const FLOAT_TYPE sx = FLOAT_TYPE(
+ FLOAT_TYPE(data_b[b_offset + y1_idx ]) * (q4_0 + (((data_a[ib0 + i].qh[l0 ] & hm1) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 1]) * (q4_1 + (((data_a[ib0 + i].qh[l0 + 1] & hm1) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 16]) * (q4_2 + (((data_a[ib0 + i].qh[l0 + 16] & hm1) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 17]) * (q4_3 + (((data_a[ib0 + i].qh[l0 + 17] & hm1) != 0) ? 16 : 0))
+ );
+ const FLOAT_TYPE sy = FLOAT_TYPE(
+ FLOAT_TYPE(data_b[b_offset + y1_idx + 32]) * (q4_4 + (((data_a[ib0 + i].qh[l0 ] & (hm1 << 1)) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 33]) * (q4_5 + (((data_a[ib0 + i].qh[l0 + 1] & (hm1 << 1)) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 48]) * (q4_6 + (((data_a[ib0 + i].qh[l0 + 16] & (hm1 << 1)) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y1_idx + 49]) * (q4_7 + (((data_a[ib0 + i].qh[l0 + 17] & (hm1 << 1)) != 0) ? 16 : 0))
+ );
+ const FLOAT_TYPE sz = FLOAT_TYPE(
+ FLOAT_TYPE(data_b[b_offset + y2_idx ]) * (q4_8 + (((data_a[ib0 + i].qh[l0 ] & hm2) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y2_idx + 1]) * (q4_9 + (((data_a[ib0 + i].qh[l0 + 1] & hm2) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y2_idx + 16]) * (q4_10 + (((data_a[ib0 + i].qh[l0 + 16] & hm2) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y2_idx + 17]) * (q4_11 + (((data_a[ib0 + i].qh[l0 + 17] & hm2) != 0) ? 16 : 0))
+ );
+ const FLOAT_TYPE sw = FLOAT_TYPE(
+ FLOAT_TYPE(data_b[b_offset + y2_idx + 32]) * (q4_12 + (((data_a[ib0 + i].qh[l0 ] & (hm2 << 1)) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y2_idx + 33]) * (q4_13 + (((data_a[ib0 + i].qh[l0 + 1] & (hm2 << 1)) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y2_idx + 48]) * (q4_14 + (((data_a[ib0 + i].qh[l0 + 16] & (hm2 << 1)) != 0) ? 16 : 0))
+ + FLOAT_TYPE(data_b[b_offset + y2_idx + 49]) * (q4_15 + (((data_a[ib0 + i].qh[l0 + 17] & (hm2 << 1)) != 0) ? 16 : 0))
+ );
+ const FLOAT_TYPE smin = FLOAT_TYPE(
+ (FLOAT_TYPE(data_b[b_offset + y1_idx]) + FLOAT_TYPE(data_b[b_offset + y1_idx + 1]) + FLOAT_TYPE(data_b[b_offset + y1_idx + 16]) + FLOAT_TYPE(data_b[b_offset + y1_idx + 17])) * sc2 + (FLOAT_TYPE(data_b[b_offset + y1_idx + 32]) + FLOAT_TYPE(data_b[b_offset + y1_idx + 33]) + FLOAT_TYPE(data_b[b_offset + y1_idx + 48]) + FLOAT_TYPE(data_b[b_offset + y1_idx + 49])) * sc3
+ + (FLOAT_TYPE(data_b[b_offset + y2_idx]) + FLOAT_TYPE(data_b[b_offset + y2_idx + 1]) + FLOAT_TYPE(data_b[b_offset + y2_idx + 16]) + FLOAT_TYPE(data_b[b_offset + y2_idx + 17])) * sc6 + (FLOAT_TYPE(data_b[b_offset + y2_idx + 32]) + FLOAT_TYPE(data_b[b_offset + y2_idx + 33]) + FLOAT_TYPE(data_b[b_offset + y2_idx + 48]) + FLOAT_TYPE(data_b[b_offset + y2_idx + 49])) * sc7
+ );
+ tmp[16 * ix + tid] += FLOAT_TYPE(dall * (sx * sc0 + sy * sc1 + sz * sc4 + sw * sc5) - dmin * smin);
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (uint s = 16; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+ if (tid == 0) {
+ data_d[d_offset + row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#version 450
+
+#include "mul_mat_vec_base.comp"
+
+layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
+
+shared FLOAT_TYPE tmp[32];
+
+void main() {
+ const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
+
+ uint a_offset, b_offset, d_offset;
+ get_offsets(a_offset, b_offset, d_offset);
+
+ const uint num_blocks_per_row = p.ncols / QUANT_K;
+ const uint ib0 = a_offset / QUANT_K + row*num_blocks_per_row;
+
+ const uint tid = gl_LocalInvocationID.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16
+ const uint ix = gl_LocalInvocationID.x%K_QUANTS_PER_ITERATION; // 0 or 0, 1
+
+ const uint step = 16/K_QUANTS_PER_ITERATION; // 16 or 8
+
+ const uint v_im = tid/step; // 0 or 1. 0 computes 0..., 1 computes 128...
+ const uint v_in = tid - step*v_im; // 0...15 or 0...7
+
+#if K_QUANTS_PER_ITERATION == 1
+ const uint l0 = v_in; // 0...15
+ const uint is = 0;
+#else
+ const uint l0 = 4 * v_in; // 0, 4, 8, ..., 28
+ const uint is = v_in / 4;
+#endif
+
+ const uint ql_offset = 64*v_im + l0;
+ const uint qh_offset = 32*v_im + l0;
+ const uint s_offset = 8*v_im + is;
+ const uint y_offset = 128*v_im + l0;
+
+ tmp[16 * ix + tid] = FLOAT_TYPE(0.0); // partial sum for thread in warp
+
+ [[unroll]] for (uint i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) {
+ const uint y_idx = i * QUANT_K + y_offset;
+
+ const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib0 + i].d);
+
+#if K_QUANTS_PER_ITERATION == 1
+ FLOAT_TYPE sum = FLOAT_TYPE(data_b[b_offset + y_idx + 0]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 0] & 0xF) | ((data_a[ib0 + i].qh[qh_offset + 0] & 0x03) << 4)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + 16]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 1]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 16] & 0xF) | ((data_a[ib0 + i].qh[qh_offset + 16] & 0x03) << 4)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + 32]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 32] & 0xF) | ((data_a[ib0 + i].qh[qh_offset + 0] & 0x0c) << 2)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + 48]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 3]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 48] & 0xF) | ((data_a[ib0 + i].qh[qh_offset + 16] & 0x0c) << 2)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + 64]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 0] >> 4) | ((data_a[ib0 + i].qh[qh_offset + 0] & 0x30) >> 0)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + 80]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 5]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 16] >> 4) | ((data_a[ib0 + i].qh[qh_offset + 16] & 0x30) >> 0)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + 96]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 32] >> 4) | ((data_a[ib0 + i].qh[qh_offset + 0] & 0xc0) >> 2)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx +112]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 7]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + 48] >> 4) | ((data_a[ib0 + i].qh[qh_offset + 16] & 0xc0) >> 2)) - 32);
+ tmp[16 * ix + tid] += sum;
+#else
+ FLOAT_TYPE sum = FLOAT_TYPE(0.0);
+ [[unroll]] for (int l = 0; l < 4; ++l) {
+ sum += FLOAT_TYPE(data_b[b_offset + y_idx + l+ 0]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 0]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + l+ 0] & 0xF) | (((data_a[ib0 + i].qh[qh_offset + l] >> 0) & 3) << 4)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l+32]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 2]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + l+32] & 0xF) | (((data_a[ib0 + i].qh[qh_offset + l] >> 2) & 3) << 4)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l+64]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 4]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + l+ 0] >> 4) | (((data_a[ib0 + i].qh[qh_offset + l] >> 4) & 3) << 4)) - 32)
+ + FLOAT_TYPE(data_b[b_offset + y_idx + l+96]) * FLOAT_TYPE(data_a[ib0 + i].scales[s_offset + 6]) * d * FLOAT_TYPE(int8_t((data_a[ib0 + i].ql[ql_offset + l+32] >> 4) | (((data_a[ib0 + i].qh[qh_offset + l] >> 6) & 3) << 4)) - 32);
+ }
+ tmp[16 * ix + tid] += sum;
+#endif
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (uint s = 16; s > 0; s >>= 1) {
+ if (tid < s) {
+ tmp[tid] += tmp[tid + s];
+ }
+ barrier();
+ }
+ if (tid == 0) {
+ data_d[d_offset + row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#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[3072];
+#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));
+#elif defined(DATA_A_IQ4_NL)
+ 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(kvalues_iq4nl[vui & 0xF], kvalues_iq4nl[vui >> 4]) * d;
+
+ buf_a[buf_idx ] = FLOAT_TYPE(v.x);
+ buf_a[buf_idx + 16] = FLOAT_TYPE(v.y);
+#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
+ }
+ }
+ }
+ }
+}
--- /dev/null
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+#define BLOCK_SIZE 512
+
+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+shared vec2 sum[BLOCK_SIZE];
+
+void main() {
+ const uint row = gl_WorkGroupID.x;
+ const uint tid = gl_LocalInvocationID.x;
+
+ sum[tid] = vec2(0.0f, 0.0f);
+
+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) {
+ const float xi = float(data_a[row*p.KX + col]);
+ sum[tid].x += xi;
+ sum[tid].y += xi * xi;
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ sum[tid] += sum[tid + s];
+ }
+ barrier();
+ }
+
+ const float mean = sum[0].x / p.KX;
+ const float var = sum[0].y / p.KX - mean * mean;
+ const float inv_std = inversesqrt(var + p.param1);
+
+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) {
+ data_d[row*p.KX + col] = D_TYPE((float(data_a[row*p.KX + col]) - mean) * inv_std);
+ }
+}
--- /dev/null
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+void main() {
+ const uint i = gl_GlobalInvocationID.x;
+
+ if (i >= p.KX) {
+ return;
+ }
+
+ data_d[i] = max(float(data_a[i]), 0);
+}
--- /dev/null
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+#define BLOCK_SIZE 512
+
+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE sum[BLOCK_SIZE];
+
+void main() {
+ const uint row = gl_WorkGroupID.x;
+ const uint tid = gl_LocalInvocationID.x;
+
+ sum[tid] = FLOAT_TYPE(0.0f); // partial sum for thread in warp
+
+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) {
+ const FLOAT_TYPE xi = FLOAT_TYPE(data_a[row*p.KX + col]);
+ sum[tid] += xi * xi;
+ }
+
+ // sum up partial sums and write back result
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ sum[tid] += sum[tid + s];
+ }
+ barrier();
+ }
+
+ const FLOAT_TYPE mean = sum[0] / FLOAT_TYPE(p.KX);
+ const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.param1));
+
+ [[unroll]] for (uint col = tid; col < p.KX; col += BLOCK_SIZE) {
+ data_d[row*p.KX + col] = D_TYPE(scale * FLOAT_TYPE(data_a[row*p.KX + col]));
+ }
+}
--- /dev/null
+#include "types.comp"
+
+#extension GL_EXT_shader_16bit_storage : require
+
+layout(local_size_x = 1, local_size_y = 256, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) readonly buffer Y {int data_pos[];};
+layout (binding = 2) readonly buffer Z {float data_ff[];};
+layout (binding = 3) writeonly buffer D {D_TYPE data_d[];};
+
+layout (push_constant) uniform parameter {
+ uint ncols;
+ uint n_dims;
+ float freq_scale;
+ uint p_delta_rows;
+ float freq_base;
+ float ext_factor;
+ float attn_factor;
+ float corr_dims[2];
+ float theta_scale;
+ uint has_ff;
+} p;
+
+float rope_yarn_ramp(const float low, const float high, const uint i0) {
+ const float y = (i0 / 2 - low) / max(0.001f, high - low);
+ return 1.0f - min(1.0f, max(0.0f, y));
+}
+
+void rope_yarn(const float theta_extrap, const uint i0, out float cos_theta, out float sin_theta) {
+ float mscale = p.attn_factor;
+ // Get n-d rotational scaling corrected for extrapolation
+ float theta_interp = p.freq_scale * theta_extrap;
+ float theta = theta_interp;
+ if (p.ext_factor != 0.0f) {
+ float ramp_mix = rope_yarn_ramp(p.corr_dims[0], p.corr_dims[1], i0) * p.ext_factor;
+ theta = theta_interp * (1 - ramp_mix) + theta_extrap * ramp_mix;
+
+ // Get n-d magnitude scaling corrected for interpolation
+ mscale *= 1.0f + 0.1f * log(1.0f / p.freq_scale);
+ }
+ cos_theta = cos(theta) * mscale;
+ sin_theta = sin(theta) * mscale;
+}
--- /dev/null
+#version 450
+
+#include "rope_head.comp"
+
+void main() {
+ const uint col = gl_GlobalInvocationID.y * 2;
+ const uint row = gl_GlobalInvocationID.x;
+
+ if (col >= p.ncols) {
+ return;
+ }
+
+ if (col >= p.n_dims) {
+ const uint i = row*p.ncols + col;
+
+ data_d[i + 0] = data_a[i + 0];
+ data_d[i + 1] = data_a[i + 1];
+
+ return;
+ }
+
+ const uint i = row*p.ncols + col/2;
+ const uint i2 = row/p.p_delta_rows;
+
+ const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f);
+
+ const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f;
+
+ float cos_theta, sin_theta;
+ rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta);
+
+ const float x0 = float(data_a[i + 0]);
+ const float x1 = float(data_a[i + p.n_dims/2]);
+
+ data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta);
+ data_d[i + p.n_dims/2] = D_TYPE(x0*sin_theta + x1*cos_theta);
+}
--- /dev/null
+#version 450
+
+#include "rope_head.comp"
+
+void main() {
+ const uint col = gl_GlobalInvocationID.y * 2;
+ const uint row = gl_GlobalInvocationID.x;
+
+ if (col >= p.ncols) {
+ return;
+ }
+
+ if (col >= p.n_dims) {
+ const uint i = row*p.ncols + col;
+
+ data_d[i + 0] = data_a[i + 0];
+ data_d[i + 1] = data_a[i + 1];
+
+ return;
+ }
+
+ const uint i = row*p.ncols + col;
+ const uint i2 = row/p.p_delta_rows;
+
+ const float theta_base = data_pos[i2] * pow(p.theta_scale, col/2.0f);
+
+ const float freq_factor = p.has_ff != 0 ? data_ff[col/2] : 1.0f;
+
+ float cos_theta, sin_theta;
+ rope_yarn(theta_base / freq_factor, col, cos_theta, sin_theta);
+
+ const float x0 = float(data_a[i + 0]);
+ const float x1 = float(data_a[i + 1]);
+
+ data_d[i + 0] = D_TYPE(x0*cos_theta - x1*sin_theta);
+ data_d[i + 1] = D_TYPE(x0*sin_theta + x1*cos_theta);
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_unary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(FLOAT_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]) * FLOAT_TYPE(p.param1));
+}
--- /dev/null
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
+
+void main() {
+ const uint i = gl_GlobalInvocationID.x;
+
+ if (i >= p.KX) {
+ return;
+ }
+
+ const float xi = float(data_a[i]);
+ data_d[i] = D_TYPE(xi / (1.0f + exp(-xi)));
+}
--- /dev/null
+#version 450
+
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+ uint KX;
+ uint KY;
+ float scale;
+ float max_bias;
+ float m0;
+ float m1;
+ uint n_head_log2;
+} p;
+
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+#define BLOCK_SIZE 512
+
+layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
+layout (binding = 1) readonly buffer Y {B_TYPE data_b[];};
+layout (binding = 2) buffer D {D_TYPE data_d[];};
+
+shared FLOAT_TYPE vals[BLOCK_SIZE];
+
+void main() {
+ const uint tid = gl_LocalInvocationID.x;
+ const uint rowx = gl_WorkGroupID.x;
+ const uint rowy = rowx % p.KY;
+
+ float slope = 1.0f;
+
+ // ALiBi
+ if (p.max_bias > 0.0f) {
+ const uint h = rowx/p.KY; // head index
+
+ const float base = h < p.n_head_log2 ? p.m0 : p.m1;
+ const uint exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1;
+
+ slope = pow(base, exp);
+ }
+
+ // Find max
+ FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000);
+
+ [[unroll]] for (uint col0 = 0; col0 < p.KX; col0 += BLOCK_SIZE) {
+ const uint col = col0 + tid;
+
+ if (col >= p.KX) {
+ break;
+ }
+
+ max_val = max(max_val, FLOAT_TYPE(data_a[rowx * p.KX + col]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)));
+ }
+ vals[tid] = max_val;
+
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ vals[tid] = max(vals[tid], vals[tid + s]);
+ }
+ barrier();
+ }
+
+ max_val = vals[0];
+ barrier();
+
+ // Sum up values
+ vals[tid] = FLOAT_TYPE(0.0f);
+
+ [[unroll]] for (uint col0 = 0; col0 < p.KX; col0 += BLOCK_SIZE) {
+ const uint col = col0 + tid;
+
+ if (col >= p.KX) {
+ break;
+ }
+
+ const uint i = rowx * p.KX + col;
+ const FLOAT_TYPE val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val);
+ vals[tid] += val;
+ data_d[i] = D_TYPE(val);
+ }
+
+ barrier();
+ [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) {
+ vals[tid] += vals[tid + s];
+ }
+ barrier();
+ }
+
+ const D_TYPE divisor = D_TYPE(vals[0]);
+
+ [[unroll]] for (uint col0 = 0; col0 < p.KX; col0 += BLOCK_SIZE) {
+ const uint col = col0 + tid;
+
+ if (col >= p.KX) {
+ break;
+ }
+
+ data_d[rowx*p.KX + col] /= divisor;
+ }
+}
--- /dev/null
+#version 450
+
+#include "types.comp"
+#include "generic_unary_head.comp"
+
+void main() {
+ if (gl_GlobalInvocationID.x >= p.ne) {
+ return;
+ }
+
+ const FLOAT_TYPE val = FLOAT_TYPE(data_a[src0_idx(gl_GlobalInvocationID.x)]);
+ data_d[p.d_offset + dst_idx(gl_GlobalInvocationID.x)] = D_TYPE(val * val);
+}
--- /dev/null
+#version 450
+
+#include "generic_head.comp"
+#include "types.comp"
+
+#extension GL_EXT_control_flow_attributes : enable
+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) writeonly buffer D {D_TYPE data_d[];};
+
+layout (constant_id = 0) const uint BLOCK_SIZE = 32;
+
+shared FLOAT_TYPE tmp[BLOCK_SIZE];
+
+void main() {
+ const uint row = gl_WorkGroupID.x;
+ const uint col = gl_LocalInvocationID.x;
+
+ tmp[col] = FLOAT_TYPE(0.0f);
+
+ for (uint i = col; i < p.KX; i += BLOCK_SIZE) {
+ tmp[col] += FLOAT_TYPE(data_a[row*p.KX + i]);
+ }
+
+ barrier();
+ [[unroll]] for (int s = int(BLOCK_SIZE) / 2; s > 0; s >>= 1) {
+ if (col < s) {
+ tmp[col] += tmp[col + s];
+ }
+ barrier();
+ }
+
+ if (col == 0) {
+ data_d[row] = D_TYPE(tmp[0]);
+ }
+}
--- /dev/null
+#if !defined(DATA_A_F32) && !defined(DATA_A_F16)
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#endif
+
+#if defined(DATA_A_F32)
+#define QUANT_K 1
+#define QUANT_R 1
+
+#ifndef LOAD_VEC_A
+#define A_TYPE float
+#elif LOAD_VEC_A == 4
+#define A_TYPE vec4
+#elif LOAD_VEC_A == 8
+#define A_TYPE mat2x4
+#endif
+#endif
+
+#if defined(DATA_A_F16)
+#define QUANT_K 1
+#define QUANT_R 1
+
+#ifndef LOAD_VEC_A
+#define A_TYPE float16_t
+#elif LOAD_VEC_A == 4
+#define A_TYPE f16vec4
+#elif LOAD_VEC_A == 8
+#define A_TYPE f16mat2x4
+#endif
+#endif
+
+#if defined(DATA_A_Q4_0)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 32
+#define QUANT_R 2
+
+struct block_q4_0
+{
+ float16_t d;
+ uint8_t qs[16];
+};
+
+#define A_TYPE block_q4_0
+#endif
+
+#if defined(DATA_A_Q4_1)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 32
+#define QUANT_R 2
+
+struct block_q4_1
+{
+ float16_t d;
+ float16_t m;
+ uint8_t qs[16];
+};
+
+#define A_TYPE block_q4_1
+#endif
+
+#if defined(DATA_A_Q5_0)
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+#define QUANT_K 32
+#define QUANT_R 2
+
+struct block_q5_0
+{
+ float16_t d;
+ uint16_t qh[2];
+ uint8_t qs[16];
+};
+
+#define A_TYPE block_q5_0
+#endif
+
+#if defined(DATA_A_Q5_1)
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
+#define QUANT_K 32
+#define QUANT_R 2
+
+struct block_q5_1
+{
+ float16_t d;
+ float16_t m;
+ uint qh;
+ uint8_t qs[16];
+};
+
+#define A_TYPE block_q5_1
+#endif
+
+#if defined(DATA_A_Q8_0)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 32
+#define QUANT_R 1
+
+struct block_q8_0
+{
+ float16_t d;
+ int8_t qs[32];
+};
+
+#define A_TYPE block_q8_0
+#endif
+
+// K-quants
+#if defined(DATA_A_Q2_K)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 256
+
+struct block_q2_K
+{
+ uint8_t scales[QUANT_K/16];
+ uint8_t qs[QUANT_K/4];
+ f16vec2 d;
+};
+
+#define A_TYPE block_q2_K
+#endif
+
+#if defined(DATA_A_Q3_K)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 256
+
+struct block_q3_K
+{
+ uint8_t hmask[QUANT_K/8];
+ uint8_t qs[QUANT_K/4];
+ uint8_t scales[12];
+ float16_t d;
+};
+
+#define A_TYPE block_q3_K
+#endif
+
+#if defined(DATA_A_Q4_K)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 256
+
+struct block_q4_K
+{
+ f16vec2 d;
+ uint8_t scales[3*QUANT_K/64];
+ uint8_t qs[QUANT_K/2];
+};
+
+#define A_TYPE block_q4_K
+#endif
+
+#if defined(DATA_A_Q5_K)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 256
+
+struct block_q5_K
+{
+ f16vec2 d;
+ uint8_t scales[12];
+ uint8_t qh[QUANT_K/8];
+ uint8_t qs[QUANT_K/2];
+};
+
+#define A_TYPE block_q5_K
+#endif
+
+#if defined(DATA_A_Q6_K)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 256
+
+struct block_q6_K
+{
+ uint8_t ql[QUANT_K/2];
+ uint8_t qh[QUANT_K/4];
+ int8_t scales[QUANT_K/16];
+ float16_t d;
+};
+
+#define A_TYPE block_q6_K
+#endif
+
+// IQuants
+
+#if defined(DATA_A_IQ4_NL)
+#extension GL_EXT_shader_16bit_storage : require
+#define QUANT_K 32
+#define QUANT_R 2
+
+struct block_iq4_nl
+{
+ float16_t d;
+ uint8_t qs[QUANT_K/2];
+};
+
+#define A_TYPE block_iq4_nl
+
+const int8_t kvalues_iq4nl[16] = {
+ int8_t(-127), int8_t(-104), int8_t(-83), int8_t(-65), int8_t(-49), int8_t(-35), int8_t(-22), int8_t(-10),
+ int8_t(1), int8_t(13), int8_t(25), int8_t(38), int8_t(53), int8_t(69), int8_t(89), int8_t(113)
+};
+#endif
--- /dev/null
+
+
+#include <iostream>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <stdexcept>
+#include <array>
+#include <vector>
+#include <map>
+#include <thread>
+#include <mutex>
+#include <future>
+#include <queue>
+#include <condition_variable>
+#include <cstdio>
+#include <cstring>
+#include <cstdlib>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#ifdef _WIN32
+ #include <windows.h>
+ #include <direct.h> // For _mkdir on Windows
+#else
+ #include <unistd.h>
+ #include <sys/wait.h>
+ #include <fcntl.h>
+#endif
+
+#define ASYNCIO_CONCURRENCY 64
+
+std::mutex lock;
+std::vector<std::pair<std::string, std::string>> shader_fnames;
+
+std::string GLSLC = "glslc";
+std::string input_dir = "vulkan-shaders";
+std::string output_dir = "/tmp";
+std::string target_hpp = "ggml-vulkan-shaders.hpp";
+std::string target_cpp = "ggml-vulkan-shaders.cpp";
+bool no_clean = false;
+
+const std::vector<std::string> type_names = {
+ "f32",
+ "f16",
+ "q4_0",
+ "q4_1",
+ "q5_0",
+ "q5_1",
+ "q8_0",
+ "q2_k",
+ "q3_k",
+ "q4_k",
+ "q5_k",
+ "q6_k",
+ "iq4_nl"
+};
+
+void execute_command(const std::string& command, std::string& stdout_str, std::string& stderr_str) {
+#ifdef _WIN32
+ HANDLE stdout_read, stdout_write;
+ HANDLE stderr_read, stderr_write;
+ SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE };
+
+ if (!CreatePipe(&stdout_read, &stdout_write, &sa, 0) ||
+ !SetHandleInformation(stdout_read, HANDLE_FLAG_INHERIT, 0)) {
+ throw std::runtime_error("Failed to create stdout pipe");
+ }
+
+ if (!CreatePipe(&stderr_read, &stderr_write, &sa, 0) ||
+ !SetHandleInformation(stderr_read, HANDLE_FLAG_INHERIT, 0)) {
+ throw std::runtime_error("Failed to create stderr pipe");
+ }
+
+ PROCESS_INFORMATION pi;
+ STARTUPINFOA si = { sizeof(STARTUPINFOA) };
+ si.dwFlags = STARTF_USESTDHANDLES;
+ si.hStdOutput = stdout_write;
+ si.hStdError = stderr_write;
+
+ std::vector<char> cmd(command.begin(), command.end());
+ cmd.push_back('\0');
+
+ if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) {
+ throw std::runtime_error("Failed to create process");
+ }
+
+ CloseHandle(stdout_write);
+ CloseHandle(stderr_write);
+
+ std::array<char, 128> buffer;
+ DWORD bytes_read;
+
+ while (ReadFile(stdout_read, buffer.data(), buffer.size(), &bytes_read, NULL) && bytes_read > 0) {
+ stdout_str.append(buffer.data(), bytes_read);
+ }
+
+ while (ReadFile(stderr_read, buffer.data(), buffer.size(), &bytes_read, NULL) && bytes_read > 0) {
+ stderr_str.append(buffer.data(), bytes_read);
+ }
+
+ CloseHandle(stdout_read);
+ CloseHandle(stderr_read);
+ WaitForSingleObject(pi.hProcess, INFINITE);
+ CloseHandle(pi.hProcess);
+ CloseHandle(pi.hThread);
+#else
+int stdout_pipe[2];
+ int stderr_pipe[2];
+
+ if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) {
+ throw std::runtime_error("Failed to create pipes");
+ }
+
+ pid_t pid = fork();
+ if (pid < 0) {
+ throw std::runtime_error("Failed to fork process");
+ }
+
+ if (pid == 0) {
+ close(stdout_pipe[0]);
+ close(stderr_pipe[0]);
+ dup2(stdout_pipe[1], STDOUT_FILENO);
+ dup2(stderr_pipe[1], STDERR_FILENO);
+ close(stdout_pipe[1]);
+ close(stderr_pipe[1]);
+ execl("/bin/sh", "sh", "-c", command.c_str(), (char*) nullptr);
+ _exit(EXIT_FAILURE);
+ } else {
+ close(stdout_pipe[1]);
+ close(stderr_pipe[1]);
+
+ std::array<char, 128> buffer;
+ ssize_t bytes_read;
+
+ while ((bytes_read = read(stdout_pipe[0], buffer.data(), buffer.size())) > 0) {
+ stdout_str.append(buffer.data(), bytes_read);
+ }
+
+ while ((bytes_read = read(stderr_pipe[0], buffer.data(), buffer.size())) > 0) {
+ stderr_str.append(buffer.data(), bytes_read);
+ }
+
+ close(stdout_pipe[0]);
+ close(stderr_pipe[0]);
+ waitpid(pid, nullptr, 0);
+ }
+#endif
+}
+
+bool directory_exists(const std::string& path) {
+ struct stat info;
+ if (stat(path.c_str(), &info) != 0) {
+ return false; // Path doesn't exist or can't be accessed
+ }
+ return (info.st_mode & S_IFDIR) != 0; // Check if it is a directory
+}
+
+bool create_directory(const std::string& path) {
+#ifdef _WIN32
+ return _mkdir(path.c_str()) == 0 || errno == EEXIST; // EEXIST means the directory already exists
+#else
+ return mkdir(path.c_str(), 0755) == 0 || errno == EEXIST; // 0755 is the directory permissions
+#endif
+}
+
+std::string to_uppercase(const std::string& input) {
+ std::string result = input;
+ for (char& c : result) {
+ c = std::toupper(c);
+ }
+ return result;
+}
+
+bool string_ends_with(const std::string& str, const std::string& suffix) {
+ if (suffix.size() > str.size()) {
+ return false;
+ }
+ return std::equal(suffix.rbegin(), suffix.rend(), str.rbegin());
+}
+
+#ifdef _WIN32
+ static const char path_separator = '\\';
+#else
+ static const char path_separator = '/';
+#endif
+
+std::string join_paths(const std::string& path1, const std::string& path2) {
+ return path1 + path_separator + path2;
+}
+
+std::string basename(const std::string &path) {
+ return path.substr(path.find_last_of("/\\") + 1);
+}
+
+void string_to_spv(const std::string& _name, const std::string& in_fname, const std::map<std::string, std::string>& defines, bool fp16 = true) {
+ std::string name = _name + (fp16 ? "" : "_fp32");
+ std::string out_fname = join_paths(output_dir, name + ".spv");
+ std::string in_path = join_paths(input_dir, in_fname);
+
+ std::vector<std::string> cmd = {GLSLC, "-fshader-stage=compute", "--target-env=vulkan1.2", "-O", in_path, "-o", out_fname};
+ for (const auto& define : defines) {
+ cmd.push_back("-D" + define.first + "=" + define.second);
+ }
+
+ std::string command;
+ for (const auto& part : cmd) {
+ command += part + " ";
+ }
+
+ std::string stdout_str, stderr_str;
+ try {
+ // std::cout << "Executing command: ";
+ // for (const auto& part : cmd) {
+ // std::cout << part << " ";
+ // }
+ // std::cout << std::endl;
+
+ execute_command(command, stdout_str, stderr_str);
+ if (!stderr_str.empty()) {
+ std::cerr << "cannot compile " << name << "\n\n" << command << "\n\n" << stderr_str << std::endl;
+ return;
+ }
+
+ std::lock_guard<std::mutex> guard(lock);
+ shader_fnames.push_back(std::make_pair(name, out_fname));
+ } catch (const std::exception& e) {
+ std::cerr << "Error executing command for " << name << ": " << e.what() << std::endl;
+ }
+}
+
+std::map<std::string, std::string> merge_maps(const std::map<std::string, std::string>& a, const std::map<std::string, std::string>& b) {
+ std::map<std::string, std::string> result = a;
+ result.insert(b.begin(), b.end());
+ return result;
+}
+
+void matmul_shaders(std::vector<std::future<void>>& tasks, bool fp16, bool matmul_id) {
+ std::string load_vec = fp16 ? "8" : "4";
+ std::string aligned_b_type_f32 = fp16 ? "mat2x4" : "vec4";
+ std::string aligned_b_type_f16 = fp16 ? "f16mat2x4" : "f16vec4";
+
+ std::map<std::string, std::string> base_dict = {{"FLOAT_TYPE", fp16 ? "float16_t" : "float"}};
+ std::string shader_name = "matmul";
+
+ if (matmul_id) {
+ base_dict["MUL_MAT_ID"] = "1";
+ shader_name = "matmul_id";
+ }
+
+ if (fp16) {
+ base_dict["FLOAT16"] = "1";
+ }
+
+ // Shaders with f16 B_TYPE
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv(shader_name + "_f32_f16", "mul_mm.comp", merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16);
+ }));
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv(shader_name + "_f32_f16_aligned", "mul_mm.comp", merge_maps(base_dict, {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}}), fp16);
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv(shader_name + "_f16", "mul_mm.comp", merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}), fp16);
+ }));
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv(shader_name + "_f16_aligned", "mul_mm.comp", merge_maps(base_dict, {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"D_TYPE", "float"}}), fp16);
+ }));
+
+ for (const auto& tname : type_names) {
+ std::string data_a_key = "DATA_A_" + to_uppercase(tname);
+ std::string load_vec_a = (tname == "f32" || tname == "f16") ? load_vec : "2";
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv(shader_name + "_" + tname + "_f32", "mul_mm.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}), fp16);
+ }));
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv(shader_name + "_" + tname + "_f32_aligned", "mul_mm.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"D_TYPE", "float"}}), fp16);
+ }));
+ }
+}
+
+void process_shaders(std::vector<std::future<void>>& tasks) {
+ std::cout << "ggml_vulkan: Generating and compiling shaders to SPIR-V" << std::endl;
+ std::map<std::string, std::string> base_dict = {{"FLOAT_TYPE", "float"}};
+
+ for (const auto& fp16 : {false, true}) {
+ matmul_shaders(tasks, fp16, false);
+ matmul_shaders(tasks, fp16, true);
+ }
+
+ for (const auto& tname : type_names) {
+ // mul mat vec
+ std::string data_a_key = "DATA_A_" + to_uppercase(tname);
+ std::string shader = (string_ends_with(tname, "_k")) ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp";
+
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
+ }));
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}));
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("mul_mat_vec_id_" + tname + "_f32", shader, merge_maps(base_dict, {{"MUL_MAT_ID", "1"}, {data_a_key, "1"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
+ }));
+
+ // Dequant shaders
+ if (tname != "f16") {
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}}));
+ }));
+ }
+
+ if (!string_ends_with(tname, "_k")) {
+ shader = (tname == "f32" || tname == "f16") ? "get_rows.comp" : "get_rows_quant.comp";
+
+ if (tname == "f16") {
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("get_rows_" + tname, shader, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});
+ }));
+ } else {
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("get_rows_" + tname, shader, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float16_t"}});
+ }));
+ }
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("get_rows_" + tname + "_f32", shader, {{data_a_key, "1"}, {"B_TYPE", "int"}, {"D_TYPE", "float"}});
+ }));
+ }
+ }
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+
+ // Norms
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
+ }));
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("add_f32", "add.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("split_k_reduce", "mul_mat_split_k_reduce.comp", {});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("div_f32", "div.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("scale_f32", "scale.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("sqr_f32", "square.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("clamp_f32", "clamp.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("gelu_f32", "gelu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("silu_f32", "silu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("relu_f32", "relu.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("diag_mask_inf_f32", "diag_mask_inf.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("soft_max_f32", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
+ }));
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("soft_max_f32_f16", "soft_max.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float16_t"}, {"D_TYPE", "float"}}));
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("rope_norm_f32", "rope_norm.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("rope_norm_f16", "rope_norm.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("rope_neox_f32", "rope_neox.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ }));
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("rope_neox_f16", "rope_neox.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [] {
+ string_to_spv("argsort_f32", "argsort.comp", {{"A_TYPE", "float"}});
+ }));
+
+ tasks.push_back(std::async(std::launch::async, [=] {
+ string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
+ }));
+}
+
+void write_output_files() {
+ FILE* hdr = fopen(target_hpp.c_str(), "w");
+ FILE* src = fopen(target_cpp.c_str(), "w");
+
+ fprintf(hdr, "#include <cstdint>\n\n");
+ fprintf(src, "#include \"%s\"\n\n", basename(target_hpp).c_str());
+
+ for (const auto& pair : shader_fnames) {
+ const std::string& name = pair.first;
+ const std::string& path = pair.second;
+ FILE* spv = fopen(path.c_str(), "rb");
+ if (!spv) {
+ std::cerr << "Error opening SPIR-V file: " << path << "\n";
+ continue;
+ }
+
+ fseek(spv, 0, SEEK_END);
+ size_t size = ftell(spv);
+ fseek(spv, 0, SEEK_SET);
+
+ std::vector<unsigned char> data(size);
+ size_t read_size = fread(data.data(), 1, size, spv);
+ fclose(spv);
+ if (read_size != size) {
+ std::cerr << "Error reading SPIR-V file: " << path << "\n";
+ continue;
+ }
+
+ fprintf(hdr, "extern unsigned char %s_data[%zu];\n", name.c_str(), size);
+ fprintf(hdr, "const uint64_t %s_len = %zu;\n\n", name.c_str(), size);
+
+ fprintf(src, "unsigned char %s_data[%zu] = {\n", name.c_str(), size);
+ for (size_t i = 0; i < size; ++i) {
+ fprintf(src, "0x%02x,", data[i]);
+ if ((i + 1) % 12 == 0) fprintf(src, "\n");
+ }
+ fprintf(src, "\n};\n\n");
+
+ if (!no_clean) {
+ std::remove(path.c_str());
+ }
+ }
+
+ fclose(hdr);
+ fclose(src);
+}
+
+int main(int argc, char** argv) {
+ std::map<std::string, std::string> args;
+ for (int i = 1; i < argc; i += 2) {
+ if (i + 1 < argc) {
+ args[argv[i]] = argv[i + 1];
+ }
+ }
+
+ if (args.find("--glslc") != args.end()) {
+ GLSLC = args["--glslc"]; // Path to glslc
+ }
+ if (args.find("--input-dir") != args.end()) {
+ input_dir = args["--input-dir"]; // Directory containing shader sources
+ }
+ if (args.find("--output-dir") != args.end()) {
+ output_dir = args["--output-dir"]; // Directory for containing SPIR-V output
+ }
+ if (args.find("--target-hpp") != args.end()) {
+ target_hpp = args["--target-hpp"]; // Path to generated header file
+ }
+ if (args.find("--target-cpp") != args.end()) {
+ target_cpp = args["--target-cpp"]; // Path to generated cpp file
+ }
+ if (args.find("--no-clean") != args.end()) {
+ no_clean = true; // Keep temporary SPIR-V files in output-dir after build
+ }
+
+ if (!directory_exists(input_dir)) {
+ std::cerr << "\"" << input_dir << "\" must be a valid directory containing shader sources" << std::endl;
+ return EXIT_FAILURE;
+ }
+
+ if (!directory_exists(output_dir)) {
+ if (!create_directory(output_dir)) {
+ std::cerr << "Error creating output directory: " << output_dir << "\n";
+ return EXIT_FAILURE;
+ }
+ }
+
+ std::vector<std::future<void>> tasks;
+ process_shaders(tasks);
+
+ for (auto& task : tasks) {
+ task.get();
+ }
+
+ write_output_files();
+
+ return EXIT_SUCCESS;
+}
overview / pkg / ggml / sources / ggml / commitdiff