{{ env.wgsl.resourceDeclarations }} // Same-shape vec4 bitwise binary (and/or/xor): 4 contiguous elements per lane // (128-bit loads/stores). uint8 storage uses one u32 slot per element, so // the result is masked to the low byte per lane. Same semantics as the scalar // broadcast kernel when A, B, C share a shape. {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %} {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %} // 2D-folded flat index: gid.y carries the high bits past the dispatch's // per-axis workgroup fold width. {% if bound == "" %} let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %} let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }}; if ({{ name }} >= {{ bound }}) { return; }{% endif %}{% endmacro %} {% if vec4PerThread > 1 %} const ITEMS: u32 = {{ vec4PerThread }}u; {% endif %} @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }}) fn main(@builtin(global_invocation_id) gid: vec3) { {% 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. {{ flat_index_2d(tunables.WORKGROUP_SIZE, "tid", "") }} let span = (params.count + ITEMS - 1u) / ITEMS; // Lanes at or past span exist only because the dispatch rounds up to whole // workgroups. Lane span + k would start on lane k's second group and rewrite // up to ITEMS - 1 groups another lane already stored. if (tid >= span) { return; } for (var j = 0u; j < ITEMS; j = j + 1u) { let i = tid + j * span; if (i >= params.count) { break; } {% else %} {{ flat_index_2d(tunables.WORKGROUP_SIZE) }} {% endif %} var value = a[i] & b[i]; {% if cDtype == "uint8" %} value = value & vec4(0xffu); {% endif %} c[i] = value; {% if vec4PerThread > 1 %} } {% endif %} }