test
This commit is contained in:
parent
3c3543ab56
commit
be5295bd51
1 changed files with 7 additions and 14 deletions
|
@ -1,5 +1,7 @@
|
||||||
#version 450
|
#version 450
|
||||||
|
|
||||||
|
#extension GL_KHR_shader_subgroup_arithmetic: enable
|
||||||
|
|
||||||
#ifdef FLOAT16
|
#ifdef FLOAT16
|
||||||
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
|
||||||
#endif
|
#endif
|
||||||
|
@ -10,8 +12,6 @@ layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
||||||
|
|
||||||
layout (constant_id = 0) const uint BLOCK_SIZE = 32;
|
layout (constant_id = 0) const uint BLOCK_SIZE = 32;
|
||||||
|
|
||||||
shared FLOAT_TYPE tmp[BLOCK_SIZE];
|
|
||||||
|
|
||||||
void main() {
|
void main() {
|
||||||
const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
|
const uint row = gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z;
|
||||||
const uint tid = gl_LocalInvocationID.x;
|
const uint tid = gl_LocalInvocationID.x;
|
||||||
|
@ -28,7 +28,7 @@ void main() {
|
||||||
|
|
||||||
const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
|
const uint y_offset = QUANT_R == 1 ? 1 : QUANT_K/2;
|
||||||
|
|
||||||
tmp[tid] = FLOAT_TYPE(0.0f);
|
FLOAT_TYPE tmp = FLOAT_TYPE(0.0f);
|
||||||
|
|
||||||
[[unroll]] for (uint i = 0; i < p.ncols/block_size; i += 2) {
|
[[unroll]] for (uint i = 0; i < p.ncols/block_size; i += 2) {
|
||||||
const uint col = i*block_size + 2*tid;
|
const uint col = i*block_size + 2*tid;
|
||||||
|
@ -39,18 +39,11 @@ void main() {
|
||||||
vec2 v = dequantize(ib, iqs, a_offset / QUANT_K);
|
vec2 v = dequantize(ib, iqs, a_offset / QUANT_K);
|
||||||
|
|
||||||
// matrix multiplication
|
// matrix multiplication
|
||||||
tmp[tid] = fma(FLOAT_TYPE(v.x), FLOAT_TYPE(data_b[b_offset + iybs + iqs]), fma(FLOAT_TYPE(v.y), FLOAT_TYPE(data_b[b_offset + iybs + iqs + y_offset]), tmp[tid]));
|
tmp = fma(FLOAT_TYPE(v.x), FLOAT_TYPE(data_b[b_offset + iybs + iqs]), fma(FLOAT_TYPE(v.y), FLOAT_TYPE(data_b[b_offset + iybs + iqs + y_offset]), tmp));
|
||||||
}
|
}
|
||||||
|
|
||||||
// sum up partial sums and write back result
|
// sum up partial sums and write back result
|
||||||
barrier();
|
tmp = subgroupAdd(tmp);
|
||||||
[[unroll]] for (uint s = block_size/2; s > 0; s >>= 1) {
|
if (tid == 0)
|
||||||
if (tid < s) {
|
data_d[d_offset + row] = D_TYPE(tmp);
|
||||||
tmp[tid] += tmp[tid + s];
|
|
||||||
}
|
|
||||||
barrier();
|
|
||||||
}
|
|
||||||
if (tid == 0) {
|
|
||||||
data_d[d_offset + row] = D_TYPE(tmp[0]);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue