| {% macro flat_tail_open() %} |
| @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 |
| // dispatch's per-axis workgroup fold width (the dispatch caps x and spills the rest into y). |
| let invocation = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u; |
| // Tail-safe scalar x4 keeps vector-like dispatch density without requiring |
| // the logical tensor length (or its storage binding) to be vec4 aligned. |
| {% set itemsPerInvocation = itemsPerInvocation if itemsPerInvocation is defined else 4 %} |
| let begin = invocation * {{ itemsPerInvocation }}u; |
| let end = min(begin + {{ itemsPerInvocation }}u, params.count); |
| for (var i = begin; i < end; i = i + 1u) { |
| {%- endmacro %} |
| {% macro flat_tail_close() %} |
| } |
| {% endmacro %} |
| |
| {% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %} |
| fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 { |
| {% if out_numel == 0 %} |
| return 0u; |
| {% elif op_numel == 1 %} |
| return 0u; |
| {% elif op_same %} |
| return out_index; |
| {% else %} |
| var offset = 0u; |
| {% for axis in range(outRank) %} |
| {% set op_axis = axis - (outRank - opRank) %} |
| {% if op_axis >= 0 and opShape[op_axis] != 1 %} |
| {% set c_stride = namespace(value=1) %} |
| {% for j in range(axis + 1, outRank) %} |
| {% set c_stride.value = c_stride.value * outShape[j] %} |
| {% endfor %} |
| {% set op_stride = namespace(value=1) %} |
| {% for j in range(op_axis + 1, opRank) %} |
| {% set op_stride.value = op_stride.value * opShape[j] %} |
| {% endfor %} |
| {% if c_stride.value == 1 %} |
| let coord{{ axis }} = out_index % {{ outShape[axis] }}u; |
| {% else %} |
| let coord{{ axis }} = (out_index / {{ c_stride.value }}u) % {{ outShape[axis] }}u; |
| {% endif %} |
| {% if op_stride.value == 1 %} |
| offset = offset + coord{{ axis }}; |
| {% else %} |
| offset = offset + coord{{ axis }} * {{ op_stride.value }}u; |
| {% endif %} |
| {% endif %} |
| {% endfor %} |
| return offset; |
| {% endif %} |
| } |
| {%- endmacro %}{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %} |
| {% set op_numel = namespace(value=1) %} |
| {% for d in opShape %}{% set op_numel.value = op_numel.value * d %}{% endfor %} |
| {% set out_numel = namespace(value=1) %} |
| {% for d in outShape %}{% set out_numel.value = out_numel.value * d %}{% endfor %} |
| {{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %}) |
| {%- endmacro %}{% macro broadcast_offset_fn(fn_name, opShape, opRank, outShape, outRank) %} |
| {% set op_numel = namespace(value=1) %} |
| {% for d in opShape %} |
| {% set op_numel.value = op_numel.value * d %} |
| {% endfor %} |
| {% set out_numel = namespace(value=1) %} |
| {% for d in outShape %} |
| {% set out_numel.value = out_numel.value * d %} |
| {% endfor %} |
| {% set op_same = namespace(value=(opRank == outRank)) %} |
| {% if op_same.value %} |
| {% for axis in range(outRank) %} |
| {% if opShape[axis] != outShape[axis] %} |
| {% set op_same.value = false %} |
| {% endif %} |
| {% endfor %} |
| {% endif %} |
| {{ offset_fn(fn_name, opShape, opRank, op_same.value, op_numel.value, outShape, outRank, out_numel.value) }} |
| {%- endmacro %}{% macro binary_broadcast_offsets() %} |
| {% set aShape = aShape | default([]) %} |
| {% set aRank = aRank | default(0) %} |
| {% set bShape = bShape | default([]) %} |
| {% set bRank = bRank | default(0) %} |
| {% set cShape = cShape | default([]) %} |
| {% set cRank = cRank | default(0) %} |
| {{ broadcast_offset_fn("a_offset", aShape, aRank, cShape, cRank) }} |
| |
| {{ broadcast_offset_fn("b_offset", bShape, bRank, cShape, cRank) }} |
| {%- endmacro %} |
| |
| {{ env.wgsl.resourceDeclarations }} |
| |
| |
| {{ binary_broadcast_offsets() }} |
| |
| {{ flat_tail_open() }} |
| var value = a[{{ broadcast_offset_call("a_offset", aShape, cShape, "i") }}] & b[{{ broadcast_offset_call("b_offset", bShape, cShape, "i") }}]; |
| {% if logicalDtype == "uint8" %} |
| value = value & 0xffu; |
| {% endif %} |
| c[i] = value; |
| {{ flat_tail_close() -}} |
| } |
| |
