wehub-resource-sync
81 lines
2.7 KiB
Plaintext
81 lines
2.7 KiB
Plaintext
// Per-row minmax pass 1 (partial reduction along K).
|
|
// - vec4 accumulators (4 independent min/max chains, no serial dependency).
|
|
// - Hoist all per-element K-bounds checks out of the hot loop. Tail (only
|
|
// when last partialBlock and K not multiple of 4) handled by a single
|
|
// post-loop iteration with scalar checks.
|
|
// - Final horizontal reduction vec4 -> scalar.
|
|
|
|
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
layout(constant_id = 3) const uint M_TILE = 64;
|
|
// K-tile granularity (in uvec4 units along K). Hard-coded here to keep
|
|
// the tuning point in one place; host (VulkanConv1x1CoopA8.cpp) mirrors
|
|
// this literal when computing partialBlocks for grid dispatch.
|
|
const uint K4_BLOCK = 64u;
|
|
|
|
layout(binding = 0) readonly buffer SrcActivation {
|
|
FLOAT4 A4[];
|
|
};
|
|
|
|
layout(binding = 1) writeonly buffer PartialMinBuffer {
|
|
FLOAT PartialMin[];
|
|
};
|
|
|
|
layout(binding = 2) writeonly buffer PartialMaxBuffer {
|
|
FLOAT PartialMax[];
|
|
};
|
|
|
|
layout(push_constant) uniform PushConstants {
|
|
uint M;
|
|
uint K;
|
|
uint padM;
|
|
uint partialBlocks;
|
|
} pc;
|
|
|
|
void main() {
|
|
uint tileM = gl_WorkGroupID.x;
|
|
uint partialBlock = gl_WorkGroupID.y;
|
|
uint laneM = gl_LocalInvocationID.x;
|
|
uint m = tileM * M_TILE + laneM;
|
|
uint k4Start = partialBlock * K4_BLOCK;
|
|
uint k4Count = (pc.K + 3u) / 4u;
|
|
uint k4End = min(k4Start + K4_BLOCK, k4Count);
|
|
|
|
bool kAligned = (pc.K & 3u) == 0u;
|
|
bool isLastBlock = (k4End == k4Count);
|
|
bool needTail = !kAligned && isLastBlock;
|
|
uint k4FullEnd = needTail ? (k4End - 1u) : k4End;
|
|
|
|
vec4 minV4 = vec4(3.402823466e+38);
|
|
vec4 maxV4 = vec4(-3.402823466e+38);
|
|
|
|
if (m < pc.M) {
|
|
for (uint k4 = k4Start; k4 < k4FullEnd; ++k4) {
|
|
vec4 v = vec4(A4[k4 * pc.M + m]);
|
|
minV4 = min(minV4, v);
|
|
maxV4 = max(maxV4, v);
|
|
}
|
|
if (needTail) {
|
|
uint k4 = k4End - 1u;
|
|
uint kBase = k4 * 4u;
|
|
vec4 v = vec4(A4[k4 * pc.M + m]);
|
|
if (kBase + 0u < pc.K) { minV4.x = min(minV4.x, v.x); maxV4.x = max(maxV4.x, v.x); }
|
|
if (kBase + 1u < pc.K) { minV4.y = min(minV4.y, v.y); maxV4.y = max(maxV4.y, v.y); }
|
|
if (kBase + 2u < pc.K) { minV4.z = min(minV4.z, v.z); maxV4.z = max(maxV4.z, v.z); }
|
|
if (kBase + 3u < pc.K) { minV4.w = min(minV4.w, v.w); maxV4.w = max(maxV4.w, v.w); }
|
|
}
|
|
} else if (m < pc.padM) {
|
|
minV4 = vec4(0.0);
|
|
maxV4 = vec4(0.0);
|
|
}
|
|
|
|
float minV = min(min(minV4.x, minV4.y), min(minV4.z, minV4.w));
|
|
float maxV = max(max(maxV4.x, maxV4.y), max(maxV4.z, maxV4.w));
|
|
|
|
if (m < pc.padM) {
|
|
uint idx = partialBlock * pc.padM + m;
|
|
PartialMin[idx] = FLOAT(minV);
|
|
PartialMax[idx] = FLOAT(maxV);
|
|
}
|
|
}
|