wehub-resource-sync
84 lines
2.4 KiB
Plaintext
84 lines
2.4 KiB
Plaintext
layout(set=0, binding=0) writeonly buffer d0{
|
|
FLOAT data[];
|
|
} uOutput;
|
|
|
|
layout(set=0, binding=1) readonly buffer s0{
|
|
FLOAT data[];
|
|
} uInputA;
|
|
|
|
layout(set=0, binding=2) readonly buffer s1{
|
|
FLOAT data[];
|
|
} uInputB;
|
|
|
|
layout(set=0, binding=3) readonly buffer s2{
|
|
FLOAT data[];
|
|
} uInputC;
|
|
|
|
layout(set=0, binding=4) readonly buffer s3{
|
|
int data[];
|
|
} uOOffset;
|
|
|
|
layout(set=0, binding=5) readonly buffer s4{
|
|
int data[];
|
|
} uAOffset;
|
|
|
|
layout(set=0, binding=6) readonly buffer s5{
|
|
int data[];
|
|
} uBOffset;
|
|
|
|
layout(set=0, binding=7) readonly buffer s6{
|
|
int data[];
|
|
} uCOffset;
|
|
|
|
layout(set=0, binding=8) uniform constBuffer{
|
|
ivec4 size;
|
|
ivec4 stride_o;
|
|
ivec4 stride_a;
|
|
ivec4 stride_b;
|
|
ivec4 stride_c; // Bias
|
|
ivec4 step;
|
|
ivec4 iter;
|
|
} uConstant;
|
|
|
|
layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
void main()
|
|
{
|
|
ivec3 posTmp = ivec3(gl_GlobalInvocationID);
|
|
int e = uConstant.size.x;
|
|
int l = uConstant.size.y;
|
|
int h = uConstant.size.z;
|
|
int n = uConstant.size.w;
|
|
int eh = e * h;
|
|
if (posTmp.x < eh * n)
|
|
{
|
|
int regionInsideIndex = posTmp.x % eh;
|
|
int regionOutsideIndex = posTmp.x / eh;
|
|
int X = regionInsideIndex % e;
|
|
int Y = regionInsideIndex / e;
|
|
int srcBasicOffset, dstBasicOffset;
|
|
ivec4 index = ivec4(regionOutsideIndex, regionOutsideIndex, regionOutsideIndex, regionOutsideIndex);
|
|
if (uConstant.iter.x >= 0) {
|
|
index.x = int(uOOffset.data[regionOutsideIndex]);
|
|
}
|
|
if (uConstant.iter.y >= 0) {
|
|
index.y = int(uAOffset.data[regionOutsideIndex]);
|
|
}
|
|
if (uConstant.iter.z >= 0) {
|
|
index.z = int(uBOffset.data[regionOutsideIndex]);
|
|
}
|
|
if (uConstant.iter.w >= 0) {
|
|
index.w = int(uCOffset.data[regionOutsideIndex]);
|
|
}
|
|
ivec4 offset = index * uConstant.step;
|
|
FLOAT value = FLOAT(0.0);
|
|
int aOffset = offset.y + uConstant.stride_a.w + X * uConstant.stride_a.x;
|
|
int bOffset = offset.z + uConstant.stride_b.w + Y * uConstant.stride_b.z;
|
|
for (int i=0; i<l; ++i) {
|
|
value = value + uInputA.data[aOffset + i * uConstant.stride_a.y] * uInputB.data[bOffset + i * uConstant.stride_b.y];
|
|
}
|
|
value = value + uInputC.data[offset.w + Y * uConstant.stride_c.z + X * uConstant.stride_c.x];
|
|
uOutput.data[offset.x + uConstant.stride_o.w + X * uConstant.stride_o.x + Y * uConstant.stride_o.z] = value;
|
|
}
|
|
}
|