File size: 1,928 Bytes
9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f 959e638 9663d5f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | {{ env.wgsl.resourceDeclarations }}
{% set vec4PerThread = vec4PerThread %}
{% if vec4PerThread > 1 %}
const ITEMS: u32 = {{ vec4PerThread }}u;
{% endif %}
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
// 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
// per-axis dispatch fold width (the dispatch caps x and spills the rest into y).
{% if vec4PerThread > 1 %}
// Each invocation walks ITEMS vec4 groups a span apart. Consecutive lanes
// access consecutive words on every step, while each lane can keep several
// independent loads in flight.
let tid = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
let span = (params.count + ITEMS - 1u) / ITEMS;
for (var j = 0u; j < ITEMS; j = j + 1u) {
let i = tid + j * span;
if (i >= params.count) {
break;
}
{% else %}
let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
if (i >= params.count) {
return;
}
{% endif %}
{% set scalarOperand = scalarOperand if scalarOperand is defined else "" %}
{% if scalarOperand == "a" %}
// One-element operand: read once and splat across the vector.
let av = {{ vectorScalar }}(a[0]);
{% else %}
let av = a[i];
{% endif %}
{% if scalarOperand == "b" %}
let bv = {{ vectorScalar }}(b[0]);
{% else %}
let bv = b[i];
{% endif %}
{% if scalar == "i32" or scalar == "u32" %}
let r = av + bv;
// Narrow integer operations wrap modulo the logical dtype width; int8/uint8
// use i32/u32 storage.
{% if cDtype == "int8" %}
c[i] = (r << vec4<u32>(24u)) >> vec4<u32>(24u);
{% elif cDtype == "uint8" %}
c[i] = r & vec4<u32>(0xFFu);
{% else %}
c[i] = r;
{% endif %}
{% elif scalar == "f16" %}
c[i] = vec4<f16>(vec4<f32>(av) + vec4<f32>(bv));
{% else %}
c[i] = av + bv;
{% endif %}
{% if vec4PerThread > 1 %}
}
{% endif %}
}
|