| {{ 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" %} |
| let avf = vec4<f32>(av); |
| let bvf = vec4<f32>(bv); |
| let quotient = avf / bvf; |
| let rounded = round(quotient); |
| let snapProduct = rounded * bvf == avf; |
| let snapClose = abs(quotient - rounded) < vec4<f32>(0.5); |
| c[i] = vec4<f16>(vec4<f32>( |
| select(quotient.x, rounded.x, snapProduct.x && snapClose.x), |
| select(quotient.y, rounded.y, snapProduct.y && snapClose.y), |
| select(quotient.z, rounded.z, snapProduct.z && snapClose.z), |
| select(quotient.w, rounded.w, snapProduct.w && snapClose.w) |
| )); |
| {% else %} |
| let quotient = av / bv; |
| let rounded = round(quotient); |
| let snapProduct = rounded * bv == av; |
| let snapClose = abs(quotient - rounded) < vec4<f32>(0.5); |
| c[i] = vec4<f32>( |
| select(quotient.x, rounded.x, snapProduct.x && snapClose.x), |
| select(quotient.y, rounded.y, snapProduct.y && snapClose.y), |
| select(quotient.z, rounded.z, snapProduct.z && snapClose.z), |
| select(quotient.w, rounded.w, snapProduct.w && snapClose.w) |
| ); |
| {% endif %} |
| {% if vec4PerThread > 1 %} |
| } |
| {% endif %} |
| } |
| |
