| {% 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 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. |
| let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }}; |
| if ({{ name }} >= {{ bound }}) { |
| return; |
| }{% endmacro %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| const BIT_WIDTH: u32 = {{ bitWidth }}u; |
| |
| {{ broadcast_offset_fn("x_offset", xShape, xRank, zShape, zRank) }} |
| |
| {{ broadcast_offset_fn("y_offset", yShape, yRank, zShape, zRank) }} |
| |
| @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }}) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| {{ flat_index_2d(tunables.WORKGROUP_SIZE) }} |
| let amount = shift[{{ broadcast_offset_call("y_offset", yShape, zShape, "i") }}]; |
| var value = 0u; |
| if (amount < BIT_WIDTH) { |
| {% if leftShift %} |
| value = x[{{ broadcast_offset_call("x_offset", xShape, zShape, "i") }}] << amount; |
| {% else %} |
| value = x[{{ broadcast_offset_call("x_offset", xShape, zShape, "i") }}] >> amount; |
| {% endif %} |
| {% if logicalDtype == "uint8" %} |
| value = value & 0xffu; |
| {% endif %} |
| } |
| z[i] = value; |
| } |
| |
