Signed-off-by: Jared Van Bortel <jared@nomic.ai> Co-authored-by: niansa <anton-sa@web.de> Co-authored-by: Adam Treat <treat.adam@gmail.com> Co-authored-by: Aaron Miller <apage43@ninjawhale.com> Co-authored-by: ToKiNoBug <tokinobug@163.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com> Co-authored-by: slaren <slarengh@gmail.com>
68 lines
1.6 KiB
Plaintext
68 lines
1.6 KiB
Plaintext
#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;
|
|
}
|
|
}
|
|
}
|