diff options
Diffstat (limited to 'ggml/src/kompute-shaders')
34 files changed, 1611 insertions, 0 deletions
diff --git a/ggml/src/kompute-shaders/common.comp b/ggml/src/kompute-shaders/common.comp new file mode 100644 index 00000000..62d62b02 --- /dev/null +++ b/ggml/src/kompute-shaders/common.comp @@ -0,0 +1,102 @@ +#extension GL_EXT_shader_16bit_storage: require +#extension GL_EXT_shader_8bit_storage: require +#extension GL_EXT_shader_explicit_arithmetic_types_float16: require +#extension GL_EXT_shader_explicit_arithmetic_types_int8: require +#extension GL_EXT_shader_explicit_arithmetic_types_int16: require +#extension GL_EXT_control_flow_attributes: enable +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_EXT_debug_printf : enable + +#define QK4_0 32 +#define QK4_1 32 + +#define GELU_COEF_A 0.044715 +#define SQRT_2_OVER_PI 0.79788456080286535587989211986876 +#define TWOPI_F 6.283185307179586f + +#define QK_K 256 + +#define u8BufToU16(buf, idx) (((uint16_t(buf[idx + 1]) << 8)) | buf[idx]) +#define u8BufToFloat16(buf, idx) uint16BitsToHalf u8BufToU16(buf, idx) +#define u8BufToU32(buf, idx) (((uint32_t u8BufToU16(buf, idx + 2) << 8 | buf[idx + 1]) << 8) | buf[idx]) +#define u8BufToFloat(buf, idx) uintBitsToFloat u8BufToU32(buf, idx) + +#define sizeof_block_q4_0 0x12 +struct block_q4_0 { + float16_t d; + uint8_t qs[QK4_0 / 2]; +}; +mat4 dequantize_q4_0(const block_q4_0 xb, uint il) { + const float d1 = il != 0 ? (xb.d / 16.f) : xb.d; + const float d2 = d1 / 256.f; + const float md = -8.f * xb.d; + const uint16_t mask0 = il != 0 ? uint16_t(0x00F0) : uint16_t(0x000F); + const uint16_t mask1 = mask0 << 8; + + mat4 reg; + for (int i=0;i<8;i++) { + uint16_t b = (uint16_t(xb.qs[2 * i + 1]) << 8) | uint16_t(xb.qs[2 * i]); + reg[i/2][2*(i%2)+0] = d1 * (b & mask0) + md; + reg[i/2][2*(i%2)+1] = d2 * (b & mask1) + md; + } + return reg; +} + +#define sizeof_block_q4_1 0x14 +struct block_q4_1 { + float16_t d; + float16_t m; + uint8_t qs[QK4_1 / 2]; +}; +mat4 dequantize_q4_1(const block_q4_1 xb, uint il) { + const float d1 = il != 0 ? (xb.d / 16.f) : xb.d; + const float d2 = d1 / 256.f; + const float m = xb.m; + const uint16_t mask0 = il != 0 ? uint16_t(0x00F0) : uint16_t(0x000F); + const uint16_t mask1 = mask0 << 8; + + mat4 reg; + for (int i=0;i<8;i++) { + uint16_t b = (uint16_t(xb.qs[2 * i + 1]) << 8) | uint16_t(xb.qs[2 * i]); + reg[i/2][2*(i%2)+0] = ((b & mask0) * d1) + m; + reg[i/2][2*(i%2)+1] = ((b & mask1) * d2) + m; + } + return reg; +} + +#define sizeof_block_q6_k 210 +struct block_q6_k { + uint8_t ql[QK_K/2]; // quants, lower 4 bits + uint8_t qh[QK_K/4]; // quants, upper 2 bits + int8_t scales[QK_K/16]; // scales, quantized with 8 bits + float16_t d; // super-block scale +}; +mat4 dequantize_q6_k(const block_q6_k xb, uint il) { + const float16_t d_all = xb.d; + + const uint qlIndex = 64*(il/8) + 32*((il/2)&1) + 16*(il&1); + const uint qhIndex = 32*(il/8) + 16*(il&1); + float16_t sc = xb.scales[(il%2) + 2 * ((il/2))]; + il = (il/2) & 3; + + const uint16_t kmask1 = il>1 ? uint16_t(il>2 ? 192 : 48) : uint16_t(il>0 ? 12 : 3); + const uint16_t kmask2 = il>1 ? uint8_t(0xF0) : uint8_t(0x0F); + const float16_t coef = il>1 ? float16_t(1.f/16.f) : float16_t(1.f); + const float16_t ml = float16_t(d_all * sc * 32.f); + const float16_t dl = float16_t(d_all * sc * coef); + mat4 reg; + for (int i = 0; i < 16; ++i) { + const float16_t q = (il&1) != 0 ? ((xb.ql[qlIndex + i] & kmask2) | ((xb.qh[qhIndex + i] & kmask1) << 2)) + : ((xb.ql[qlIndex + i] & kmask2) | ((xb.qh[qhIndex + i] & kmask1) << 4)); + reg[i/4][i%4] = dl * q - ml; + } + return reg; +} + + +#define QK8_0 32 +// struct block_q8_0 { +// float16_t d; // delta +// int8_t qs[QK8_0]; // quants +// }; +#define sizeof_block_q8_0 34 diff --git a/ggml/src/kompute-shaders/op_add.comp b/ggml/src/kompute-shaders/op_add.comp new file mode 100644 index 00000000..b7b76a79 --- /dev/null +++ b/ggml/src/kompute-shaders/op_add.comp @@ -0,0 +1,58 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1024) in; + +layout(binding = 0) buffer restrict readonly tensorInA { float inA[]; }; +layout(binding = 1) buffer restrict readonly tensorInB { float inB[]; }; +layout(binding = 2) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb00; + int nb01; + int nb02; + int nb03; + int ne10; + int ne11; + int ne12; + int ne13; + int nb10; + int nb11; + int nb12; + int nb13; + int ne0; + int nb0; + int nb1; + int nb2; + int nb3; + //int offs; // TODO: needed for GGML_OP_ACC, see metal code +} pcs; + +// general-purpose kernel for addition of two tensors +// pros: works for non-contiguous tensors, supports broadcast across dims 1, 2 and 3 +// cons: not very efficient +void main() { + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const uint i13 = i03 % pcs.ne13; + const uint i12 = i02 % pcs.ne12; + const uint i11 = i01 % pcs.ne11; + + int offs = 0; // TMP (see above) + + uint src0_off = uint((i03*pcs.nb03 + i02*pcs.nb02 + i01*pcs.nb01 + offs) / 4); + uint src1_off = uint((i13*pcs.nb13 + i12*pcs.nb12 + i11*pcs.nb11 ) / 4); + uint dst_off = uint((i03*pcs.nb3 + i02*pcs.nb2 + i01*pcs.nb1 + offs) / 4); + + for (uint i0 = gl_LocalInvocationID.x; i0 < pcs.ne0; i0 += gl_WorkGroupSize.x) { + const uint i10 = i0 % pcs.ne10; + out_[pcs.outOff + dst_off + i0] = inA[pcs.inAOff + src0_off + i0] + inB[pcs.inBOff + src1_off + i10]; + } +} diff --git a/ggml/src/kompute-shaders/op_addrow.comp b/ggml/src/kompute-shaders/op_addrow.comp new file mode 100644 index 00000000..2376a6b8 --- /dev/null +++ b/ggml/src/kompute-shaders/op_addrow.comp @@ -0,0 +1,25 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorInA { float inA[]; }; +layout(binding = 1) buffer restrict readonly tensorInB { float inB[]; }; +layout(binding = 2) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inAOff; + uint inBOff; + uint outOff; + uint row; +} pcs; + +void main() { + const uint baseIndex = gl_WorkGroupID.x * 4; + + for (uint x = 0; x < 4; x++) { + const uint i = baseIndex + x; + out_[i + pcs.outOff] = inA[i + pcs.inAOff] + inB[(i % pcs.row) + pcs.inBOff]; + } +} diff --git a/ggml/src/kompute-shaders/op_cpy_f16_f16.comp b/ggml/src/kompute-shaders/op_cpy_f16_f16.comp new file mode 100644 index 00000000..d57247d2 --- /dev/null +++ b/ggml/src/kompute-shaders/op_cpy_f16_f16.comp @@ -0,0 +1,52 @@ +#version 450 + +#include "common.comp" + +#define IN_TYPE float16_t +#define IN_TYPE_SIZE 2 +#define OUT_TYPE float16_t +#define OUT_TYPE_SIZE 2 + +layout(local_size_x = 1024) in; + +layout (binding = 0) readonly buffer tensorIn { IN_TYPE in_[]; }; +layout (binding = 1) writeonly buffer tensorOut { OUT_TYPE out_[]; }; + +layout (push_constant) uniform parameter { + uint inOff; + uint outOff; + int ne00; + int ne01; + int ne02; + uint nb00; + uint nb01; + uint nb02; + uint nb03; + int ne0; + int ne1; + int ne2; + uint nb0; + uint nb1; + uint nb2; + uint nb3; +} pcs; + +void main() { + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const int n = int(i03)*pcs.ne02*pcs.ne01*pcs.ne00 + int(i02)*pcs.ne01*pcs.ne00 + int(i01)*pcs.ne00; + + const int i3 = n / (pcs.ne2*pcs.ne1*pcs.ne0); + const int i2 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0) / (pcs.ne1*pcs.ne0); + const int i1 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0) / pcs.ne0; + const int i0 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0 - i1*pcs.ne0); + + const uint dst_data = (i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / OUT_TYPE_SIZE + pcs.outOff; // Based from out_ + + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + const uint src = uint((i03*pcs.nb03 + i02*pcs.nb02 + i01*pcs.nb01 + i00*pcs.nb00) / IN_TYPE_SIZE) + pcs.inOff; // Based from in_ + out_[dst_data+i00] = OUT_TYPE(in_[src]); + } +} diff --git a/ggml/src/kompute-shaders/op_cpy_f16_f32.comp b/ggml/src/kompute-shaders/op_cpy_f16_f32.comp new file mode 100644 index 00000000..b568bcd7 --- /dev/null +++ b/ggml/src/kompute-shaders/op_cpy_f16_f32.comp @@ -0,0 +1,52 @@ +#version 450 + +#include "common.comp" + +#define IN_TYPE float16_t +#define IN_TYPE_SIZE 2 +#define OUT_TYPE float +#define OUT_TYPE_SIZE 4 + +layout(local_size_x = 1024) in; + +layout (binding = 0) readonly buffer tensorIn { IN_TYPE in_[]; }; +layout (binding = 1) writeonly buffer tensorOut { OUT_TYPE out_[]; }; + +layout (push_constant) uniform parameter { + uint inOff; + uint outOff; + int ne00; + int ne01; + int ne02; + uint nb00; + uint nb01; + uint nb02; + uint nb03; + int ne0; + int ne1; + int ne2; + uint nb0; + uint nb1; + uint nb2; + uint nb3; +} pcs; + +void main() { + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const int n = int(i03)*pcs.ne02*pcs.ne01*pcs.ne00 + int(i02)*pcs.ne01*pcs.ne00 + int(i01)*pcs.ne00; + + const int i3 = n / (pcs.ne2*pcs.ne1*pcs.ne0); + const int i2 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0) / (pcs.ne1*pcs.ne0); + const int i1 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0) / pcs.ne0; + const int i0 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0 - i1*pcs.ne0); + + const uint dst_data = (i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / OUT_TYPE_SIZE + pcs.outOff; // Based from out_ + + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + const uint src = uint((i03*pcs.nb03 + i02*pcs.nb02 + i01*pcs.nb01 + i00*pcs.nb00) / IN_TYPE_SIZE) + pcs.inOff; // Based from in_ + out_[dst_data+i00] = OUT_TYPE(in_[src]); + } +} diff --git a/ggml/src/kompute-shaders/op_cpy_f32_f16.comp b/ggml/src/kompute-shaders/op_cpy_f32_f16.comp new file mode 100644 index 00000000..99b22834 --- /dev/null +++ b/ggml/src/kompute-shaders/op_cpy_f32_f16.comp @@ -0,0 +1,52 @@ +#version 450 + +#include "common.comp" + +#define IN_TYPE float +#define IN_TYPE_SIZE 4 +#define OUT_TYPE float16_t +#define OUT_TYPE_SIZE 2 + +layout(local_size_x = 1024) in; + +layout (binding = 0) readonly buffer tensorIn { IN_TYPE in_[]; }; +layout (binding = 1) writeonly buffer tensorOut { OUT_TYPE out_[]; }; + +layout (push_constant) uniform parameter { + uint inOff; + uint outOff; + int ne00; + int ne01; + int ne02; + uint nb00; + uint nb01; + uint nb02; + uint nb03; + int ne0; + int ne1; + int ne2; + uint nb0; + uint nb1; + uint nb2; + uint nb3; +} pcs; + +void main() { + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const int n = int(i03)*pcs.ne02*pcs.ne01*pcs.ne00 + int(i02)*pcs.ne01*pcs.ne00 + int(i01)*pcs.ne00; + + const int i3 = n / (pcs.ne2*pcs.ne1*pcs.ne0); + const int i2 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0) / (pcs.ne1*pcs.ne0); + const int i1 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0) / pcs.ne0; + const int i0 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0 - i1*pcs.ne0); + + const uint dst_data = (i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / OUT_TYPE_SIZE + pcs.outOff; // Based from out_ + + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + const uint src = uint((i03*pcs.nb03 + i02*pcs.nb02 + i01*pcs.nb01 + i00*pcs.nb00) / IN_TYPE_SIZE) + pcs.inOff; // Based from in_ + out_[dst_data+i00] = OUT_TYPE(in_[src]); + } +} diff --git a/ggml/src/kompute-shaders/op_cpy_f32_f32.comp b/ggml/src/kompute-shaders/op_cpy_f32_f32.comp new file mode 100644 index 00000000..2fc99849 --- /dev/null +++ b/ggml/src/kompute-shaders/op_cpy_f32_f32.comp @@ -0,0 +1,52 @@ +#version 450 + +#include "common.comp" + +#define IN_TYPE float +#define IN_TYPE_SIZE 4 +#define OUT_TYPE float +#define OUT_TYPE_SIZE 4 + +layout(local_size_x = 1024) in; + +layout (binding = 0) readonly buffer tensorIn { IN_TYPE in_[]; }; +layout (binding = 1) writeonly buffer tensorOut { OUT_TYPE out_[]; }; + +layout (push_constant) uniform parameter { + uint inOff; + uint outOff; + int ne00; + int ne01; + int ne02; + uint nb00; + uint nb01; + uint nb02; + uint nb03; + int ne0; + int ne1; + int ne2; + uint nb0; + uint nb1; + uint nb2; + uint nb3; +} pcs; + +void main() { + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const int n = int(i03)*pcs.ne02*pcs.ne01*pcs.ne00 + int(i02)*pcs.ne01*pcs.ne00 + int(i01)*pcs.ne00; + + const int i3 = n / (pcs.ne2*pcs.ne1*pcs.ne0); + const int i2 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0) / (pcs.ne1*pcs.ne0); + const int i1 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0) / pcs.ne0; + const int i0 = (n - i3*pcs.ne2*pcs.ne1*pcs.ne0 - i2*pcs.ne1*pcs.ne0 - i1*pcs.ne0); + + const uint dst_data = (i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / OUT_TYPE_SIZE + pcs.outOff; // Based from out_ + + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + const uint src = uint((i03*pcs.nb03 + i02*pcs.nb02 + i01*pcs.nb01 + i00*pcs.nb00) / IN_TYPE_SIZE) + pcs.inOff; // Based from in_ + out_[dst_data+i00] = OUT_TYPE(in_[src]); + } +} diff --git a/ggml/src/kompute-shaders/op_diagmask.comp b/ggml/src/kompute-shaders/op_diagmask.comp new file mode 100644 index 00000000..291c3fc1 --- /dev/null +++ b/ggml/src/kompute-shaders/op_diagmask.comp @@ -0,0 +1,30 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; + uint n_past; + int ne00; + int ne01; +} pcs; + +void main() { + const uint i02 = gl_WorkGroupID.z; + const uint i01 = gl_WorkGroupID.y; + const uint i00 = gl_WorkGroupID.x; + + const uint index = i02*pcs.ne01*pcs.ne00 + i01*pcs.ne00 + i00; + + if (i00 > pcs.n_past + i01) { + out_[index + pcs.outOff] = uintBitsToFloat(0xFF800000); + } else { + out_[index + pcs.outOff] = in_[index + pcs.inOff]; + } +} diff --git a/ggml/src/kompute-shaders/op_gelu.comp b/ggml/src/kompute-shaders/op_gelu.comp new file mode 100644 index 00000000..9d8c5371 --- /dev/null +++ b/ggml/src/kompute-shaders/op_gelu.comp @@ -0,0 +1,22 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; }; +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; +} pcs; + +void main() { + const uint baseIndex = gl_WorkGroupID.x * 8; + + for (uint x = 0; x < 8; x++) { + const uint i = baseIndex + x; + const float y = in_[i + pcs.inOff]; + out_[i + pcs.outOff] = 0.5*y*(1.0 + tanh(clamp(SQRT_2_OVER_PI*y*(1.0 + GELU_COEF_A*y*y), -15.0, 15.0))); + } +} diff --git a/ggml/src/kompute-shaders/op_getrows.comp b/ggml/src/kompute-shaders/op_getrows.comp new file mode 100644 index 00000000..1a5581b2 --- /dev/null +++ b/ggml/src/kompute-shaders/op_getrows.comp @@ -0,0 +1,17 @@ +void main() { + const uint i = gl_WorkGroupID.x; + const int r = inB[i + pcs.inBOff]; + + int z = 0; + for (uint ind = gl_LocalInvocationID.x; ind < pcs.ne00/16; ind += gl_WorkGroupSize.x) { + const uint inIndex = (r * pcs.nb01 + pcs.inAOff) + ind/NL * SIZE_OF_BLOCK; + const mat4 result = dequantize_block(inIndex, ind%NL); + for (uint j = 0; j < 4; ++j) { + for (uint k = 0; k < 4; ++k) { + const uint outIndex = i * pcs.nb1/BYTES_FOR_TYPE + pcs.outOff + z; + out_[outIndex] = result[j][k]; + ++z; + } + } + } +} diff --git a/ggml/src/kompute-shaders/op_getrows_f16.comp b/ggml/src/kompute-shaders/op_getrows_f16.comp new file mode 100644 index 00000000..48c93610 --- /dev/null +++ b/ggml/src/kompute-shaders/op_getrows_f16.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout (binding = 0) readonly buffer tensorInA { float16_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { int inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb01; + int nb1; +} pcs; + +void dequantize_row_f16(uint x /*Based from inA unaligned*/, uint y /*Based from out_*/, int k) { + for (int j = 0; j < k; j++) { + out_[y + j] = inA[x + j]; + } +} + +void main() { + const uint i = gl_WorkGroupID.x; + const int r = inB[i + pcs.inBOff]; + + dequantize_row_f16(r*pcs.nb01/2/*bytes for float16*/ + pcs.inAOff, i*pcs.nb1/4 + pcs.outOff, pcs.ne00); +} diff --git a/ggml/src/kompute-shaders/op_getrows_f32.comp b/ggml/src/kompute-shaders/op_getrows_f32.comp new file mode 100644 index 00000000..9d7acdaf --- /dev/null +++ b/ggml/src/kompute-shaders/op_getrows_f32.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout (binding = 0) readonly buffer tensorInA { float inA[]; }; +layout (binding = 1) readonly buffer tensorInB { int inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb01; + int nb1; +} pcs; + +void dequantize_row_f32(uint x /*Based from inA unaligned*/, uint y /*Based from out_*/, int k) { + for (int j = 0; j < k; j++) { + out_[y + j] = inA[x + j]; + } +} + +void main() { + const uint i = gl_WorkGroupID.x; + const int r = inB[i + pcs.inBOff]; + + dequantize_row_f32(r*pcs.nb01/4 + pcs.inAOff, i*pcs.nb1/4 + pcs.outOff, pcs.ne00); +} diff --git a/ggml/src/kompute-shaders/op_getrows_q4_0.comp b/ggml/src/kompute-shaders/op_getrows_q4_0.comp new file mode 100644 index 00000000..32b2e891 --- /dev/null +++ b/ggml/src/kompute-shaders/op_getrows_q4_0.comp @@ -0,0 +1,38 @@ +#version 450 + +#include "common.comp" + +#define NL 2 +#define BYTES_FOR_TYPE 4 /*bytes for float*/ +#define SIZE_OF_BLOCK sizeof_block_q4_0 + +layout(local_size_x = 1) in; + +layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { int inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb01; + int nb1; +} pcs; + +block_q4_0 get_unaligned_block_q4_0(uint index) { + block_q4_0 fres; + fres.d = u8BufToFloat16(inA, index); + [[unroll]] for (uint it = 0; it != QK4_0 / 2; it++) { + fres.qs[it] = inA[index+2+it]; + } + return fres; +} + +mat4 dequantize_block(uint index, uint il) { + const block_q4_0 block = get_unaligned_block_q4_0(index); + return dequantize_q4_0(block, il); +} + +#include "op_getrows.comp" diff --git a/ggml/src/kompute-shaders/op_getrows_q4_1.comp b/ggml/src/kompute-shaders/op_getrows_q4_1.comp new file mode 100644 index 00000000..87f2fbe1 --- /dev/null +++ b/ggml/src/kompute-shaders/op_getrows_q4_1.comp @@ -0,0 +1,39 @@ +#version 450 + +#include "common.comp" + +#define NL 2 +#define BYTES_FOR_TYPE 4 /*bytes for float*/ +#define SIZE_OF_BLOCK sizeof_block_q4_1 + +layout(local_size_x = 1) in; + +layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { int inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb01; + int nb1; +} pcs; + +block_q4_1 get_unaligned_block_q4_1(uint index) { + block_q4_1 fres; + fres.d = u8BufToFloat16(inA, index); + fres.m = u8BufToFloat16(inA, index+2); + [[unroll]] for (uint it = 0; it != QK4_1 / 2; it++) { + fres.qs[it] = inA[index+4+it]; + } + return fres; +} + +mat4 dequantize_block(uint index, uint il) { + const block_q4_1 block = get_unaligned_block_q4_1(index); + return dequantize_q4_1(block, il); +} + +#include "op_getrows.comp" diff --git a/ggml/src/kompute-shaders/op_getrows_q6_k.comp b/ggml/src/kompute-shaders/op_getrows_q6_k.comp new file mode 100644 index 00000000..9ce3545d --- /dev/null +++ b/ggml/src/kompute-shaders/op_getrows_q6_k.comp @@ -0,0 +1,44 @@ +#version 450 + +#include "common.comp" + +#define NL 16 +#define BYTES_FOR_TYPE 4 /*bytes for float*/ +#define SIZE_OF_BLOCK sizeof_block_q6_k + +layout(local_size_x = 1) in; + +layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { int inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb01; + int nb1; +} pcs; + +block_q6_k get_unaligned_block_q6_k(uint index) { + block_q6_k fres; + [[unroll]] for (uint it = 0; it != QK_K / 2; it++) { + fres.ql[it] = inA[index + it]; + } + [[unroll]] for (uint it = 0; it != QK_K / 4; it++) { + fres.qh[it] = inA[index + QK_K/2 + it]; + } + [[unroll]] for (uint it = 0; it != QK_K / 16; it++) { + fres.scales[it] = int8_t(inA[index + QK_K/2 + QK_K/4 + it]); + } + fres.d = u8BufToFloat16(inA, index + QK_K/2 + QK_K/4 + QK_K/16); + return fres; +} + +mat4 dequantize_block(uint index, uint il) { + const block_q6_k block = get_unaligned_block_q6_k(index); + return dequantize_q6_k(block, il); +} + +#include "op_getrows.comp" diff --git a/ggml/src/kompute-shaders/op_mul.comp b/ggml/src/kompute-shaders/op_mul.comp new file mode 100644 index 00000000..c92647c4 --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul.comp @@ -0,0 +1,52 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1024) in; + +layout(binding = 0) buffer restrict readonly tensorInA { float inA[]; }; +layout(binding = 1) buffer restrict readonly tensorInB { float inB[]; }; +layout(binding = 2) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int nb00; + int nb01; + int nb02; + int nb03; + int ne10; + int ne11; + int ne12; + int ne13; + int nb10; + int nb11; + int nb12; + int nb13; + int ne0; + int nb0; + int nb1; + int nb2; + int nb3; +} pcs; + +void main() { + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const uint i13 = i03 % pcs.ne13; + const uint i12 = i02 % pcs.ne12; + const uint i11 = i01 % pcs.ne11; + + uint src0_off = uint((i03*pcs.nb03 + i02*pcs.nb02 + i01*pcs.nb01) / 4); + uint src1_off = uint((i13*pcs.nb13 + i12*pcs.nb12 + i11*pcs.nb11) / 4); + uint dst_off = uint((i03*pcs.nb3 + i02*pcs.nb2 + i01*pcs.nb1) / 4); + + for (uint i0 = gl_LocalInvocationID.x; i0 < pcs.ne0; i0 += gl_WorkGroupSize.x) { + const uint i10 = i0 % pcs.ne10; + out_[pcs.outOff + dst_off + i0] = inA[pcs.inAOff + src0_off + i0] * inB[pcs.inBOff + src1_off + i10]; + } +} diff --git a/ggml/src/kompute-shaders/op_mul_mat_f16.comp b/ggml/src/kompute-shaders/op_mul_mat_f16.comp new file mode 100644 index 00000000..8f0a9031 --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mat_f16.comp @@ -0,0 +1,67 @@ +#version 450 + +#include "common.comp" + +#extension GL_KHR_shader_subgroup_arithmetic : require + +layout(local_size_x_id = 0) in; + +layout (binding = 0) readonly buffer tensorInA { float16_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { float inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int ne01; + int ne02; + uint nb00; + uint nb01; + uint nb02; + int ne10; + int ne11; + int ne12; + uint nb10; + uint nb11; + uint nb12; + int ne0; + int ne1; + uint r2; + uint r3; +} pcs; + +#define N_F16_F32 4 + +void main() { + const uint r0 = gl_WorkGroupID.x; + const uint rb = gl_WorkGroupID.y*N_F16_F32; + const uint im = gl_WorkGroupID.z; + + const uint i12 = im%pcs.ne12; + const uint i13 = im/pcs.ne12; + + const uint offset0 = r0*pcs.nb01 + (i12/pcs.r2)*pcs.nb02 + (i13/pcs.r3)*pcs.nb02*pcs.ne02; + + const uint x = offset0 / 2 + pcs.inAOff; // Based from inA + + for (uint row = 0; row < N_F16_F32; ++row) { + uint r1 = rb + row; + if (r1 >= pcs.ne11) { + break; + } + + const uint y = (r1*pcs.nb11 + im*pcs.nb12) / 4 + pcs.inBOff; // Based from inB + + float sumf = 0; + for (uint i = gl_SubgroupInvocationID.x; i < pcs.ne00; i += gl_SubgroupSize) { + sumf += float(inA[x+i]) * float(inB[y+i]); + } + + const float all_sum = subgroupAdd(sumf); + if (subgroupElect()) { + out_[im*pcs.ne1*pcs.ne0 + r1*pcs.ne0 + r0 + pcs.outOff] = all_sum; + } + } +} diff --git a/ggml/src/kompute-shaders/op_mul_mat_mat_f32.comp b/ggml/src/kompute-shaders/op_mul_mat_mat_f32.comp new file mode 100644 index 00000000..d1ca4ad6 --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mat_mat_f32.comp @@ -0,0 +1,51 @@ +#version 450 + +#include "common.comp" + +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_EXT_debug_printf : enable + +// device subgroup size +layout (local_size_x_id = 0) in; + +layout(binding = 0) readonly buffer tensorInA { float inA[]; }; +layout(binding = 1) readonly buffer tensorInB { float inB[]; }; +layout(binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout(push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int ne01; + int ne02; + int ne11; + int ne12; + uint nb01; + uint nb02; + uint nb11; + uint nb12; + uint nb1; + uint nb2; +} +pcs; + + +void main() { + uvec3 gid = gl_WorkGroupID; + + uint bc_ab = pcs.ne12 > pcs.ne02 ? gid.z / (pcs.ne12 / pcs.ne02) : gid.z; + uint bc_ba = pcs.ne02 > pcs.ne12 ? gid.z / (pcs.ne02 / pcs.ne12) : gid.z; + + const uint x = (gid.x*pcs.nb01 + bc_ab*pcs.nb02) / 4 + pcs.inAOff; // Based from inA + const uint y = (gid.y*pcs.nb11 + bc_ba*pcs.nb12) / 4 + pcs.inBOff; // based from inB + float sum = 0.0f; + for (uint i = gl_SubgroupInvocationID.x; i < pcs.ne00; i += gl_SubgroupSize) { + sum += float(inA[x+i]) * float(inB[y+i]); + } + + const float all_sum = subgroupAdd(sum); + if (subgroupElect()) { + out_[gid.z*(pcs.nb2/4) + gid.y*(pcs.nb1/4) + gid.x + pcs.outOff] = all_sum; + } +} diff --git a/ggml/src/kompute-shaders/op_mul_mat_q4_0.comp b/ggml/src/kompute-shaders/op_mul_mat_q4_0.comp new file mode 100644 index 00000000..b0cea8bb --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mat_q4_0.comp @@ -0,0 +1,33 @@ +#version 450 + +#include "common.comp" + +#define BLOCKS_IN_QUANT QK4_0 +#define SIZE_OF_BLOCK sizeof_block_q4_0 +#define N_ROWS 4 + +#include "op_mul_mv_q_n_pre.comp" + +// The q4_0 version of this function +float block_q_n_dot_y(uint block_index, uint yb, uint il) { + vec2 acc = vec2(0.0, 0.0); + const uint index = (block_index) * SIZE_OF_BLOCK + pcs.inAOff; + float d = float(u8BufToFloat16(inA, index)); + float sumy = 0.0f; + for (int i = 0; i < BLOCKS_IN_QUANT/4; i+=2) { + const uint16_t b = u8BufToU16(inA, index + 2 + il + i); + + const float yl0 = inB[yb + i]; + const float yl1 = inB[yb + i + 1]; + const float yl8 = inB[yb + i + BLOCKS_IN_QUANT/2]; + const float yl9 = inB[yb + i + BLOCKS_IN_QUANT/2 + 1]; + + sumy += yl0 + yl1 + yl8 + yl9; + + acc[0] += yl0 * (b & 0x000F) + yl1 / 256.f * (b & 0x0F00); + acc[1] += yl8 / 16.f * (b & 0x00F0) + yl9 / 4096.f * (b & 0xF000); + } + return d * (sumy * -8.f + acc[0] + acc[1]); +} + +#include "op_mul_mv_q_n.comp" diff --git a/ggml/src/kompute-shaders/op_mul_mat_q4_1.comp b/ggml/src/kompute-shaders/op_mul_mat_q4_1.comp new file mode 100644 index 00000000..8582c61a --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mat_q4_1.comp @@ -0,0 +1,35 @@ +#version 450 + +#include "common.comp" + +#define BLOCKS_IN_QUANT QK4_1 +#define SIZE_OF_BLOCK sizeof_block_q4_1 +#define N_ROWS 4 + +#include "op_mul_mv_q_n_pre.comp" + +// The q4_1 version of this function +float block_q_n_dot_y(uint block_index, uint yb, uint il) { + vec2 acc = vec2(0.0, 0.0); + const uint index = (block_index) * SIZE_OF_BLOCK + pcs.inAOff; + float d = float(u8BufToFloat16(inA, index)); + float m = float(u8BufToFloat16(inA, index+2)); + + float sumy = 0.0f; + for (int i = 0; i < BLOCKS_IN_QUANT/4; i+=2) { + const uint16_t b = u8BufToU16(inA, index + 4 + il + i); + + const float yl0 = inB[yb + i]; + const float yl1 = inB[yb + i + 1]; + const float yl8 = inB[yb + i + BLOCKS_IN_QUANT/2]; + const float yl9 = inB[yb + i + BLOCKS_IN_QUANT/2 + 1]; + + sumy += yl0 + yl1 + yl8 + yl9; + + acc[0] += yl0 * (b & 0x000F) + yl1 / 256.f * (b & 0x0F00); + acc[1] += yl8 / 16.f * (b & 0x00F0) + yl9 / 4096.f * (b & 0xF000); + } + return d * (acc[0] + acc[1]) + sumy * m; +} + +#include "op_mul_mv_q_n.comp" diff --git a/ggml/src/kompute-shaders/op_mul_mat_q6_k.comp b/ggml/src/kompute-shaders/op_mul_mat_q6_k.comp new file mode 100644 index 00000000..c9baebdf --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mat_q6_k.comp @@ -0,0 +1,94 @@ +#version 450 + +#include "common.comp" + +#define SIZE_OF_BLOCK sizeof_block_q6_k + +layout(local_size_x_id = 0) in; +layout(local_size_y_id = 1) in; +layout(local_size_z = 1) in; + +layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { float inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int ne10; + int ne0; + int ne1; + int ne01; + int gqa; +} pcs; + +void main() { + const uint8_t kmask1 = uint8_t(0x03); + const uint8_t kmask2 = uint8_t(0x0C); + const uint8_t kmask3 = uint8_t(0x30); + const uint8_t kmask4 = uint8_t(0xC0); + + const uint nb = pcs.ne00/QK_K; + + const uint r0 = gl_WorkGroupID.x; + const uint r1 = gl_WorkGroupID.y; + const uint r2 = gl_WorkGroupID.z; + + const uint row = (r0 * gl_NumSubgroups + gl_SubgroupID); + const uint offset0 = r2/pcs.gqa*(nb*pcs.ne0); + const uint x = row * nb + offset0; // Based from inA without base offset + const uint yy = r1*pcs.ne10 + r2*pcs.ne00*pcs.ne1+pcs.inBOff; // Based from inB + + float sumf = 0; + + // bits of invocation ID for gl_SubgroupSize=32: + // x x x x x + // 4 3 2 1 0 + // ( tid ) ix + // ip ( il ) + + const uint block_stride = gl_SubgroupSize / 16; // number of blocks each subgroup processes + const uint tid = gl_SubgroupInvocationID/block_stride; // first block_stride groups have tid=0 + const uint ix = gl_SubgroupInvocationID%block_stride; // first block is 0..block_stride-1 + const uint ip = tid/8; // first or second half of block (0 or 1) + const uint il = tid%8; // each half has 8 parts, one per scale + const uint n = 4; // 4 scales at a time (and 4 sums) + const uint l0 = n*il; // offset into half-block, 0..28 + const uint is = 8*ip + l0/16; // 0, 1, 8, 9 + + const uint y_offset = 128*ip + l0; + const uint q_offset_l = 64*ip + l0; + const uint q_offset_h = 32*ip + l0; + + for (uint i = ix; i < nb; i += block_stride) { + + const uint baseIndex = (x + i) * SIZE_OF_BLOCK + pcs.inAOff; + + const uint qlIndex = q_offset_l; + const uint q2Index = qlIndex + QK_K/8; + const uint qhIndex = q_offset_h; + const uint y = yy + i * QK_K + y_offset; + + float sums[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (uint l = 0; l < n; ++l) { + const uint8_t currentQ1 = inA[baseIndex + qlIndex + l]; + const uint8_t currentQ2 = inA[baseIndex + q2Index + l]; + const uint8_t currentQh = inA[baseIndex + QK_K/2 + qhIndex + l]; + + sums[0] += inB[y+l+ 0] * (int8_t((currentQ1 & 0xF) | ((currentQh & kmask1) << 4)) - 32); + sums[1] += inB[y+l+32] * (int8_t((currentQ2 & 0xF) | ((currentQh & kmask2) << 2)) - 32); + sums[2] += inB[y+l+64] * (int8_t((currentQ1 >> 4) | ((currentQh & kmask3) << 0)) - 32); + sums[3] += inB[y+l+96] * (int8_t((currentQ2 >> 4) | ((currentQh & kmask4) >> 2)) - 32); + } + + float d = u8BufToFloat16(inA, baseIndex + QK_K/2 + QK_K/4 + QK_K/16); + sumf += d * (sums[0] * int8_t(inA[baseIndex + QK_K/2 + QK_K/4 + is]) + sums[1] * int8_t(inA[baseIndex + QK_K/2 + QK_K/4 + 2 + is]) + sums[2] * int8_t(inA[baseIndex + QK_K/2 + QK_K/4 + 4 + is]) + sums[3] * int8_t(inA[baseIndex + QK_K/2 + QK_K/4 + 6 + is])); + } + + const float tot = subgroupAdd(sumf); + if (subgroupElect()) { + out_[r1*pcs.ne0 + r2*pcs.ne0*pcs.ne1 + row + pcs.outOff] = tot; + } +} diff --git a/ggml/src/kompute-shaders/op_mul_mat_q8_0.comp b/ggml/src/kompute-shaders/op_mul_mat_q8_0.comp new file mode 100644 index 00000000..34d015e9 --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mat_q8_0.comp @@ -0,0 +1,73 @@ +#version 450 + +#include "common.comp" + +#include "op_mul_mv_q_n_pre.comp" + +#define SIZE_OF_D 2 + +#define N_DST 4 // each SIMD group works on 4 rows +#define N_SIMDGROUP 2 // number of SIMD groups in a thread group +#define N_SIMDWIDTH 32 // assuming SIMD group size is 32 + +#define NB_Q8_0 8 + +void main() { + // NB: hack to make compatible with AMD GPUs that have a subgroup size of 64 + if (gl_SubgroupInvocationID > 31) + return; + + const int nr = N_DST; + const int nsg = N_SIMDGROUP; + const int nw = N_SIMDWIDTH; + + const int nb = pcs.ne00/QK8_0; + const uint r0 = gl_WorkGroupID.x; + const uint r1 = gl_WorkGroupID.y; + const uint im = gl_WorkGroupID.z; + + const uint first_row = (r0 * nsg + gl_SubgroupID) * nr; + + const uint i12 = im%pcs.ne12; + const uint i13 = im/pcs.ne12; + + const uint offset0 = first_row * nb + (i12/pcs.r2)*(nb*pcs.ne01) + (i13/pcs.r3)*(nb*pcs.ne01*pcs.ne02); + + const uint x = offset0*sizeof_block_q8_0 + pcs.inAOff; // Based from inA + const uint y = r1*pcs.ne10 + im*pcs.ne00*pcs.ne1 + pcs.inBOff; // based from inB + + float yl[NB_Q8_0]; + float sumf[N_DST]={0.f, 0.f, 0.f, 0.f}; + + const uint ix = gl_SubgroupInvocationID.x/4; + const uint il = gl_SubgroupInvocationID.x%4; + + uint yb = y + ix * QK8_0 + NB_Q8_0*il; + + // each thread in a SIMD group deals with NB_Q8_0 quants at a time + for (uint ib = ix; ib < nb; ib += nw/4) { + for (int i = 0; i < NB_Q8_0; ++i) { + yl[i] = inB[yb + i]; + } + + for (int row = 0; row < nr; row++) { + const uint block_offset = (ib+row*nb) * sizeof_block_q8_0; + float sumq = 0.f; + for (int iq = 0; iq < NB_Q8_0; ++iq) { + const int8_t qs_iq = int8_t(inA[x + block_offset + SIZE_OF_D + NB_Q8_0*il + iq]); + sumq += qs_iq * yl[iq]; + } + const float16_t d = u8BufToFloat16(inA, x + block_offset); + sumf[row] += sumq*d; + } + + yb += NB_Q8_0 * nw; + } + + for (int row = 0; row < nr; ++row) { + const float tot = subgroupAdd(sumf[row]); + if (subgroupElect() && first_row + row < pcs.ne01) { + out_[r1*pcs.ne0 + im*pcs.ne0*pcs.ne1 + first_row + row] = tot; + } + } +} diff --git a/ggml/src/kompute-shaders/op_mul_mv_q_n.comp b/ggml/src/kompute-shaders/op_mul_mv_q_n.comp new file mode 100644 index 00000000..440b5ab2 --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mv_q_n.comp @@ -0,0 +1,48 @@ +void main() { + // NB: hack to make compatible with AMD GPUs that have a subgroup size of 64 + if (gl_SubgroupInvocationID > 31) + return; + + const uint nb = uint(pcs.ne00/BLOCKS_IN_QUANT); + + const uint r0 = gl_WorkGroupID.x; + const uint r1 = gl_WorkGroupID.y; + const uint im = gl_WorkGroupID.z; + + const uint first_row = (r0 * gl_NumSubgroups + gl_SubgroupID) * N_ROWS; + + const uint i12 = im%pcs.ne12; + const uint i13 = im/pcs.ne12; + + const uint offset0 = first_row * nb + (i12/pcs.r2)*(nb*pcs.ne01) + (i13/pcs.r3)*(nb*pcs.ne01*pcs.ne02); + + const uint x = offset0; // Based from inA without base offset + const uint y = r1*uint(pcs.ne10)+im*pcs.ne00*pcs.ne1+pcs.inBOff; // Based from inB + + float sumf[N_ROWS] = {0.0f, 0.0f, 0.0f, 0.0f}; + + const uint ix = gl_SubgroupInvocationID/2; + const uint il = (BLOCKS_IN_QUANT/4)*(gl_SubgroupInvocationID%2); + + uint yb = y + ix * BLOCKS_IN_QUANT + il; + + //debugPrintfEXT("gl_NumSubgroups=%d, gl_SubgroupID=%d, gl_SubgroupInvocationID=%d, glSubgroupSize=%d, gl_WorkGroupSize.x=%d, gl_WorkGroupSize.y=%d, gl_WorkGroupSize.z=%d\n", + // gl_NumSubgroups, gl_SubgroupID, gl_SubgroupInvocationID, gl_SubgroupSize, + // gl_WorkGroupSize.x, gl_WorkGroupSize.y, gl_WorkGroupSize.z); + + for (uint ib = ix; ib < nb; ib += 16) { + for (int row = 0; row < N_ROWS; row++) { + const uint block_index = x + ib + row * nb; + sumf[row] += block_q_n_dot_y(block_index, yb, il); + } + + yb += BLOCKS_IN_QUANT * 16; + } + + for (int row = 0; row < N_ROWS; ++row) { + const float tot = subgroupAdd(sumf[row]); + if (first_row + row < pcs.ne01 && subgroupElect()) { + out_[r1*pcs.ne0 + im*pcs.ne0*pcs.ne1 + first_row + row + pcs.outOff] = tot; + } + } +} diff --git a/ggml/src/kompute-shaders/op_mul_mv_q_n_pre.comp b/ggml/src/kompute-shaders/op_mul_mv_q_n_pre.comp new file mode 100644 index 00000000..7912b09a --- /dev/null +++ b/ggml/src/kompute-shaders/op_mul_mv_q_n_pre.comp @@ -0,0 +1,22 @@ +layout(local_size_x_id = 0) in; +layout(local_size_y = 1) in; +layout(local_size_z = 1) in; + +layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; }; +layout (binding = 1) readonly buffer tensorInB { float inB[]; }; +layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int ne01; + int ne02; + int ne10; + int ne12; + int ne0; + int ne1; + uint r2; + uint r3; +} pcs; diff --git a/ggml/src/kompute-shaders/op_norm.comp b/ggml/src/kompute-shaders/op_norm.comp new file mode 100644 index 00000000..ad0c3c01 --- /dev/null +++ b/ggml/src/kompute-shaders/op_norm.comp @@ -0,0 +1,84 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 256) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; + uint ne00; + uint nb01; + float eps; +} pcs; + +shared float sum[gl_WorkGroupSize.x]; + +void main() { + const uint x = (gl_WorkGroupID.x*pcs.nb01/4) + pcs.inOff; // Based from in_ + // MEAN + // parallel sum + sum[gl_LocalInvocationID.x] = 0.0; + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + sum[gl_LocalInvocationID.x] += in_[x+i00]; + } + + // reduce + barrier(); + memoryBarrierShared(); + [[unroll]] for (uint i = gl_WorkGroupSize.x/2; i > 0; i /= 2) { + if (gl_LocalInvocationID.x < i) { + sum[gl_LocalInvocationID.x] += sum[gl_LocalInvocationID.x + i]; + } + barrier(); + memoryBarrierShared(); + } + + // broadcast + if (gl_LocalInvocationID.x == 0) { + sum[0] /= float(pcs.ne00); + } + barrier(); + memoryBarrierShared(); + const float mean = sum[0]; + + // recenter + const uint y = (gl_WorkGroupID.x*pcs.ne00) + pcs.outOff; // Based from out_ + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + out_[y+i00] = in_[x+i00] - mean; + } + + // VARIANCE + // parallel sum + sum[gl_LocalInvocationID.x] = 0.0; + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + sum[gl_LocalInvocationID.x] += out_[y+i00] * out_[y+i00]; + } + + // reduce + barrier(); + memoryBarrierShared(); + [[unroll]] for (uint i = gl_WorkGroupSize.x/2; i > 0; i /= 2) { + if (gl_LocalInvocationID.x < i) { + sum[gl_LocalInvocationID.x] += sum[gl_LocalInvocationID.x + i]; + } + barrier(); + memoryBarrierShared(); + } + + // broadcast + if (gl_LocalInvocationID.x == 0) { + sum[0] /= float(pcs.ne00); + } + barrier(); + memoryBarrierShared(); + const float variance = sum[0]; + + const float scale = 1.0f/sqrt(variance + pcs.eps); + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + out_[y+i00] *= scale; + } +} diff --git a/ggml/src/kompute-shaders/op_relu.comp b/ggml/src/kompute-shaders/op_relu.comp new file mode 100644 index 00000000..52a601fe --- /dev/null +++ b/ggml/src/kompute-shaders/op_relu.comp @@ -0,0 +1,21 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; }; +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; +} pcs; + +void main() { + const uint baseIndex = gl_WorkGroupID.x * 4; + + for (uint x = 0; x < 4; x++) { + const uint i = baseIndex + x; + out_[i + pcs.outOff] = max(0.0, in_[i + pcs.inOff]); + } +} diff --git a/ggml/src/kompute-shaders/op_rmsnorm.comp b/ggml/src/kompute-shaders/op_rmsnorm.comp new file mode 100644 index 00000000..da658c16 --- /dev/null +++ b/ggml/src/kompute-shaders/op_rmsnorm.comp @@ -0,0 +1,53 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 512) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; + uint ne00; + uint nb01; + float eps; +} pcs; + +shared float sum[gl_WorkGroupSize.x]; + +void main() { + const uint x = (gl_WorkGroupID.x*pcs.nb01/4) + pcs.inOff; // Based from in_ + + // parallel sum + sum[gl_LocalInvocationID.x] = 0.0; + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + sum[gl_LocalInvocationID.x] += in_[x+i00] * in_[x+i00]; + } + + // reduce + barrier(); + memoryBarrierShared(); + [[unroll]] for (uint i = gl_WorkGroupSize.x/2; i > 0; i /= 2) { + if (gl_LocalInvocationID.x < i) { + sum[gl_LocalInvocationID.x] += sum[gl_LocalInvocationID.x + i]; + } + barrier(); + memoryBarrierShared(); + } + + // broadcast + if (gl_LocalInvocationID.x == 0) { + sum[0] /= float(pcs.ne00); + } + barrier(); + memoryBarrierShared(); + + const float scale = 1.0f/sqrt(sum[0] + pcs.eps); + + const uint y = (gl_WorkGroupID.x*pcs.ne00) + pcs.outOff; // Based from out_ + for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += gl_WorkGroupSize.x) { + out_[y+i00] = in_[x+i00] * scale; + } +} diff --git a/ggml/src/kompute-shaders/op_rope_f16.comp b/ggml/src/kompute-shaders/op_rope_f16.comp new file mode 100644 index 00000000..1a4058b3 --- /dev/null +++ b/ggml/src/kompute-shaders/op_rope_f16.comp @@ -0,0 +1,73 @@ +#version 450 + +#include "rope_common.comp" + +layout(binding = 0) buffer restrict readonly tensorInA { float16_t inA[]; }; +layout(binding = 1) buffer restrict readonly tensorInB { int inB[]; }; +layout(binding = 2) buffer restrict writeonly tensorOut { float16_t out_[]; }; + +void main() { + const uint i3 = gl_WorkGroupID.z; + const uint i2 = gl_WorkGroupID.y; + const uint i1 = gl_WorkGroupID.x; + + const bool is_neox = (pcs.mode & 2) != 0; + + float corr_dims[2]; + rope_yarn_corr_dims(pcs.n_dims, pcs.n_ctx_orig, pcs.freq_base, pcs.beta_fast, pcs.beta_slow, corr_dims); + + const float theta_scale = pow(pcs.freq_base, -2.0/pcs.n_dims); + + const int p = inB[pcs.inBOff + i2]; + + float theta = float(p); + + if (!is_neox) { + for (uint i0 = 0; i0 < pcs.ne0; i0 += 2) { + float cos_theta, sin_theta; + rope_yarn(theta, pcs.freq_scale, corr_dims, i0, pcs.ext_factor, pcs.attn_factor, cos_theta, sin_theta); + + theta *= theta_scale; + + const uint src = uint((i3*pcs.nb03 + i2*pcs.nb02 + i1*pcs.nb01 + i0*pcs.nb00) / 2) + pcs.inAOff; // Based from in + const uint dst_data = uint((i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / 2) + pcs.outOff; // Based from out_ + + const float x0 = float(inA[src]); + const float x1 = float(inA[src+1]); + + out_[dst_data] = float16_t(x0*cos_theta - x1*sin_theta); + out_[dst_data+1] = float16_t(x0*sin_theta + x1*cos_theta); + } + } else { + const float inv_ndims = -1.f/pcs.n_dims; + for (uint ic = 0; ic < pcs.n_dims; ic += 2) { + const uint cur_rot = ic; + + float cos_theta, sin_theta; + rope_yarn(theta, pcs.freq_scale, corr_dims, cur_rot, pcs.ext_factor, pcs.attn_factor, cos_theta, sin_theta); + + theta *= theta_scale; + + const uint i0 = ic/2; + + const uint src = uint((i3*pcs.nb03 + i2*pcs.nb02 + i1*pcs.nb01 + i0*pcs.nb00) / 2) + pcs.inAOff; // Based from in + const uint dst_data = uint((i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / 2) + pcs.outOff; // Based from out_ + + const float x0 = float(inA[src]); + const float x1 = float(inA[src+pcs.n_dims/2]); + + out_[dst_data] = float16_t(x0*cos_theta - x1*sin_theta); + out_[dst_data+pcs.n_dims/2] = float16_t(x0*sin_theta + x1*cos_theta); + } + + for (uint ic = pcs.n_dims; ic < pcs.ne0; ic += 2) { + const uint i0 = ic; + + const uint src = uint((i3*pcs.nb03 + i2*pcs.nb02 + i1*pcs.nb01 + i0*pcs.nb00) / 2) + pcs.inAOff; // Based from in + const uint dst_data = uint((i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / 2) + pcs.outOff; // Based from out_ + + out_[dst_data + 0] = inA[src + 0]; + out_[dst_data + 1] = inA[src + 1]; + } + } +} diff --git a/ggml/src/kompute-shaders/op_rope_f32.comp b/ggml/src/kompute-shaders/op_rope_f32.comp new file mode 100644 index 00000000..65e03827 --- /dev/null +++ b/ggml/src/kompute-shaders/op_rope_f32.comp @@ -0,0 +1,73 @@ +#version 450 + +#include "rope_common.comp" + +layout(binding = 0) buffer restrict readonly tensorInA { float inA[]; }; +layout(binding = 1) buffer restrict readonly tensorInB { int inB[]; }; +layout(binding = 2) buffer restrict writeonly tensorOut { float out_[]; }; + +void main() { + const uint i3 = gl_WorkGroupID.z; + const uint i2 = gl_WorkGroupID.y; + const uint i1 = gl_WorkGroupID.x; + + const bool is_neox = (pcs.mode & 2) != 0; + + float corr_dims[2]; + rope_yarn_corr_dims(pcs.n_dims, pcs.n_ctx_orig, pcs.freq_base, pcs.beta_fast, pcs.beta_slow, corr_dims); + + const float theta_scale = pow(pcs.freq_base, -2.0/pcs.n_dims); + + const int p = inB[pcs.inBOff + i2]; + + float theta = float(p); + + if (!is_neox) { + for (uint i0 = 0; i0 < pcs.ne0; i0 += 2) { + float cos_theta, sin_theta; + rope_yarn(theta, pcs.freq_scale, corr_dims, i0, pcs.ext_factor, pcs.attn_factor, cos_theta, sin_theta); + + theta *= theta_scale; + + const uint src = uint((i3*pcs.nb03 + i2*pcs.nb02 + i1*pcs.nb01 + i0*pcs.nb00) / 4) + pcs.inAOff; // Based from in + const uint dst_data = uint((i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / 4) + pcs.outOff; // Based from out_ + + const float x0 = inA[src]; + const float x1 = inA[src+1]; + + out_[dst_data] = x0*cos_theta - x1*sin_theta; + out_[dst_data+1] = x0*sin_theta + x1*cos_theta; + } + } else { + const float inv_ndims = -1.f/pcs.n_dims; + for (uint ic = 0; ic < pcs.n_dims; ic += 2) { + const uint cur_rot = ic; + + float cos_theta, sin_theta; + rope_yarn(theta, pcs.freq_scale, corr_dims, cur_rot, pcs.ext_factor, pcs.attn_factor, cos_theta, sin_theta); + + theta *= theta_scale; + + const uint i0 = ic/2; + + const uint src = uint((i3*pcs.nb03 + i2*pcs.nb02 + i1*pcs.nb01 + i0*pcs.nb00) / 4) + pcs.inAOff; // Based from in + const uint dst_data = uint((i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / 4) + pcs.outOff; // Based from out_ + + const float x0 = inA[src]; + const float x1 = inA[src+pcs.n_dims/2]; + + out_[dst_data] = x0*cos_theta - x1*sin_theta; + out_[dst_data+pcs.n_dims/2] = x0*sin_theta + x1*cos_theta; + } + + for (uint ic = pcs.n_dims; ic < pcs.ne0; ic += 2) { + const uint i0 = ic; + + const uint src = uint((i3*pcs.nb03 + i2*pcs.nb02 + i1*pcs.nb01 + i0*pcs.nb00) / 4) + pcs.inAOff; // Based from in + const uint dst_data = uint((i3*pcs.nb3 + i2*pcs.nb2 + i1*pcs.nb1 + i0*pcs.nb0) / 4) + pcs.outOff; // Based from out_ + + out_[dst_data + 0] = inA[src + 0]; + out_[dst_data + 1] = inA[src + 1]; + } + } +} diff --git a/ggml/src/kompute-shaders/op_scale.comp b/ggml/src/kompute-shaders/op_scale.comp new file mode 100644 index 00000000..bdae2673 --- /dev/null +++ b/ggml/src/kompute-shaders/op_scale.comp @@ -0,0 +1,19 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; + float scale; +} pcs; + +void main() { + const uint i = gl_WorkGroupID.x; + out_[i + pcs.outOff] = in_[i + pcs.inOff] * pcs.scale; +} diff --git a/ggml/src/kompute-shaders/op_scale_8.comp b/ggml/src/kompute-shaders/op_scale_8.comp new file mode 100644 index 00000000..ada69754 --- /dev/null +++ b/ggml/src/kompute-shaders/op_scale_8.comp @@ -0,0 +1,23 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; + float scale; +} pcs; + +void main() { + const uint baseIndex = gl_WorkGroupID.x * 8; + + for (uint x = 0; x < 8; x++) { + const uint i = baseIndex + x; + out_[i + pcs.outOff] = in_[i + pcs.inOff] * pcs.scale; + } +} diff --git a/ggml/src/kompute-shaders/op_silu.comp b/ggml/src/kompute-shaders/op_silu.comp new file mode 100644 index 00000000..0fb8e4b7 --- /dev/null +++ b/ggml/src/kompute-shaders/op_silu.comp @@ -0,0 +1,22 @@ +#version 450 + +#include "common.comp" + +layout(local_size_x = 1) in; + +layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; }; +layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; }; +layout(push_constant) uniform PushConstants { + uint inOff; + uint outOff; +} pcs; + +void main() { + const uint baseIndex = gl_WorkGroupID.x * 4; + + for (uint x = 0; x < 4; x++) { + const uint i = baseIndex + x; + const float y = in_[i + pcs.inOff]; + out_[i + pcs.outOff] = y / (1.0 + exp(-y)); + } +} diff --git a/ggml/src/kompute-shaders/op_softmax.comp b/ggml/src/kompute-shaders/op_softmax.comp new file mode 100644 index 00000000..7bc9176c --- /dev/null +++ b/ggml/src/kompute-shaders/op_softmax.comp @@ -0,0 +1,56 @@ +// TODO: implement multi-simd softmax (llama.cpp commit e16b9fa4) + +#version 450 + +#include "common.comp" + +layout(local_size_x_id = 0) in; + +layout(binding = 0) buffer restrict readonly tensorInA { float inA[]; }; +layout(binding = 1) buffer restrict readonly tensorInB { float inB[]; }; +layout(binding = 2) buffer restrict writeonly tensorOut { float out_[]; }; + +layout(push_constant) uniform PushConstants { + uint inAOff; + uint inBOff; + uint outOff; + int ne00; + int ne01; + int ne02; + float scale; + int mask; +} pcs; + +void main() { + if (gl_SubgroupInvocationID > 31) + return; + + const uint i03 = gl_WorkGroupID.z; + const uint i02 = gl_WorkGroupID.y; + const uint i01 = gl_WorkGroupID.x; + + const uint extra_off = i03*pcs.ne02*pcs.ne01*pcs.ne00 + i02*pcs.ne01*pcs.ne00 + i01*pcs.ne00; + const uint psrc0 = extra_off + pcs.inAOff; // Based from inA + const uint pmask = i01*pcs.ne00 + pcs.inBOff; // Based from inB + const uint pdst = extra_off + pcs.outOff; // Based from out_ + + // parallel max + float localMax = uintBitsToFloat(0xFF800000); + for (uint i00 = gl_SubgroupInvocationID.x; i00 < pcs.ne00; i00 += 32) { + localMax = max(localMax, inA[psrc0 + i00]*pcs.scale + (pcs.mask!=0 ? inB[pmask + i00] : 0.0f)); + } + float max_ = subgroupMax(localMax); + + // parallel sum + float localSum = 0.0f; + for (uint i00 = gl_SubgroupInvocationID.x; i00 < pcs.ne00; i00 += 32) { + const float exp_psrc0 = exp(inA[psrc0 + i00]*pcs.scale + (pcs.mask!=0 ? inB[pmask + i00] : 0.0f) - max_); + localSum += exp_psrc0; + out_[pdst + i00] = exp_psrc0; + } + + const float sum = subgroupAdd(localSum); + for (uint i00 = gl_SubgroupInvocationID.x; i00 < pcs.ne00; i00 += 32) { + out_[pdst + i00] /= sum; + } +} diff --git a/ggml/src/kompute-shaders/rope_common.comp b/ggml/src/kompute-shaders/rope_common.comp new file mode 100644 index 00000000..7b9394cb --- /dev/null +++ b/ggml/src/kompute-shaders/rope_common.comp @@ -0,0 +1,67 @@ +#include "common.comp" + +// TODO: use a local size of 32 or more (Metal uses 1024) +layout(local_size_x = 1) in; + +layout (push_constant) uniform parameter { + uint inAOff; + uint inBOff; + uint outOff; + int n_dims; + int mode; + int n_ctx_orig; + float freq_base; + float freq_scale; + float ext_factor; + float attn_factor; + float beta_fast; + float beta_slow; + uint nb00; + uint nb01; + uint nb02; + uint nb03; + int ne0; + uint nb0; + uint nb1; + uint nb2; + uint nb3; +} pcs; + +float rope_yarn_ramp(const float low, const float high, const float i0) { + const float y = (i0 / 2 - low) / max(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +// YaRN algorithm based on LlamaYaRNScaledRotaryEmbedding.py from https://github.com/jquesnelle/yarn +// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng. +void rope_yarn( + float theta_extrap, float freq_scale, float corr_dims[2], float i0, float ext_factor, float mscale, + out float cos_theta, out float sin_theta +) { + // Get n-d rotational scaling corrected for extrapolation + float theta_interp = freq_scale * theta_extrap; + float theta = theta_interp; + if (ext_factor != 0.0f) { + float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], i0) * 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 / freq_scale); + } + cos_theta = cos(theta) * mscale; + sin_theta = sin(theta) * mscale; +} + +// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get +// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` +float rope_yarn_corr_factor(int n_dims, int n_ctx_orig, float n_rot, float base) { + return n_dims * log(n_ctx_orig / (n_rot * TWOPI_F)) / (2 * log(base)); +} + +void rope_yarn_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, out float dims[2] +) { + // start and end correction dims + dims[0] = max(0.0f, floor(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_fast, freq_base))); + dims[1] = min(n_dims - 1.0f, ceil(rope_yarn_corr_factor(n_dims, n_ctx_orig, beta_slow, freq_base))); +} |