wehub-resource-sync
72 lines
2.1 KiB
Plaintext
72 lines
2.1 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 s3{
|
|
int data[];
|
|
} uOOffset;
|
|
|
|
layout(set=0, binding=4) readonly buffer s4{
|
|
int data[];
|
|
} uAOffset;
|
|
|
|
layout(set=0, binding=5) readonly buffer s5{
|
|
int data[];
|
|
} uBOffset;
|
|
|
|
layout(set=0, binding=6) 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]);
|
|
}
|
|
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];
|
|
}
|
|
uOutput.data[offset.x + uConstant.stride_o.w + X * uConstant.stride_o.x + Y * uConstant.stride_o.z] = value;
|
|
}
|
|
}
|