File size: 1,989 Bytes
ea833c8 a79717a ea833c8 f402c03 a79717a f402c03 ea833c8 a79717a f402c03 a79717a f402c03 a79717a f402c03 ea833c8 a79717a ea833c8 a79717a ea833c8 | 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 | {{ 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<u32>) {
{% 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<u32>(0xffu);
{% endif %}
c[i] = value;
{% if vec4PerThread > 1 %}
}
{% endif %}
}
|