Kernels:

---

webgpu-kernels

/

ai.onnx.BitShift

Like 0

Follow

WebGPU Kernels Community 42

{% 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;
}