| {% if usesF16 %} |
| enable f16; |
| {% endif %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }}) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) { |
| // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the |
| // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y). |
| let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u; |
| if (i >= params.count) { |
| return; |
| } |
| let av = a[i]; |
| let bv = b[i]; |
| {% 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 source.cDtype == "int8" %} |
| c[i] = (r << vec4<u32>(24u)) >> vec4<u32>(24u); |
| {% elif source.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 %} |
| } |
| |
