{{ env.wgsl.resourceDeclarations }}
// Rotary positional embedding. The dispatch pins one slice on z and folds that slice's own
// rotation groups across x and y, so every row of a slice shares one base offset and all the
// head geometry below is a compile-time constant: every division in the index math is by a
// literal rather than by a uniform.
//
// A slice is one batch at rank 3, where a token's heads are adjacent and a workgroup covers a
// whole token, so it reads each cos/sin row once. At rank 4 a token's heads are a sequence
// apart, so a slice is instead one batch and ITEMS consecutive heads: the invocation gathers
// its cos/sin row once and rotates the same rotation group in each of those heads.
const WG: u32 = {{ workgroupSize }}u;
const ITEMS: u32 = {{ itemsPerLane }}u;
// Rotation groups per head row, the row's stride, and the number of groups that rotate, all
// counted in the storage element the bindings use: a vec4 on the four-pair route and a scalar
// on the single-pair route.
const HEAD_UNITS: u32 = {{ headUnits }}u;
const HEAD_STRIDE: u32 = {{ headStride }}u;
const ROT_UNITS: u32 = {{ rotUnits }}u;
{% if hasTail %}
// First element past the rotary window.
const ROT_END: u32 = {{ rotStart }}u;
{% endif %}
const NUM_HEADS: u32 = {{ numHeads }}u;
{% if rank == 4 %}
const HEAD_BLOCKS: u32 = {{ headBlocks }}u;
{% endif %}
{% if hasTail %}

{% macro copy_tail(base) %}
      let tail = ROT_END + (group - ROT_UNITS) * 2u;
      y[{{ base }} + tail] = x[{{ base }} + tail];
{% if tailGuard %}
      // An odd head size leaves the last group with a single element to copy.
      if (tail + 1u < HEAD_STRIDE) {
        y[{{ base }} + tail + 1u] = x[{{ base }} + tail + 1u];
      }
{% else %}
      y[{{ base }} + tail + 1u] = x[{{ base }} + tail + 1u];
{% endif %}
{% endmacro %}
{% endif %}

{% macro rotate(base) %}
{% if interleaved and useVec4 %}
      // Four interleaved pairs span two vec4s: (w0.x, w0.y), (w0.z, w0.w), (w1.x, w1.y),
      // (w1.z, w1.w). Deinterleave, rotate, and reinterleave on the way back out.
      let pairBase = {{ base }} + group * 2u;
      let w0 = vec4<f32>(x[pairBase]);
      let w1 = vec4<f32>(x[pairBase + 1u]);
      let a = vec4<f32>(w0.x, w0.z, w1.x, w1.z);
      let b = vec4<f32>(w0.y, w0.w, w1.y, w1.w);
      let ra = a * cf - b * sf;
      let rb = a * sf + b * cf;
      y[pairBase] = {{ castOut }}(vec4<f32>(ra.x, rb.x, ra.y, rb.y));
      y[pairBase + 1u] = {{ castOut }}(vec4<f32>(ra.z, rb.z, ra.w, rb.w));
{% elif interleaved %}
      // Adjacent elements form the rotated pair.
      let pairBase = {{ base }} + group * 2u;
      let a = {{ castIn }}(x[pairBase]);
      let b = {{ castIn }}(x[pairBase + 1u]);
      y[pairBase] = {{ castOut }}(a * cf - b * sf);
      y[pairBase + 1u] = {{ castOut }}(a * sf + b * cf);
{% else %}
      // Split halves: ROT_UNITS is both the rotating group count and the offset from an
      // element of the first half to its partner in the second.
      let a = {{ castIn }}(x[{{ base }} + group]);
      let b = {{ castIn }}(x[{{ base }} + ROT_UNITS + group]);
      y[{{ base }} + group] = {{ castOut }}(a * cf - b * sf);
      y[{{ base }} + ROT_UNITS + group] = {{ castOut }}(a * sf + b * cf);
{% endif %}
{% endmacro %}

@compute @workgroup_size(WG)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
{% if rank == 3 %}
  // wg.y carries the high bits of the group index past the per-axis dispatch fold width.
  let slice = wg.z;
  let sliceUnits = params.sliceRows * HEAD_UNITS;
  let sliceBase = slice * params.sliceRows * HEAD_STRIDE;
  let laneUnit = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * (WG * ITEMS) + lid.x;
{% for item in range(itemsPerLane) %}
  {
    let unit = laneUnit + {{ item }}u * WG;
    if (unit < sliceUnits) {
      let group = unit % HEAD_UNITS;
      let row = unit / HEAD_UNITS;
      let base = sliceBase + row * HEAD_STRIDE;
      let token = row / NUM_HEADS;
{% if hasTail %}
      if (group >= ROT_UNITS) {
{{ copy_tail("base") }}
      } else {
{% endif %}
{% if posOffset %}
      // One base offset for the whole request: token s reads cache row p[0] + s.
      let position = position_ids[0] + token;
{% else %}
      // The (batch_size, sequence_length) table. params.sliceRows / NUM_HEADS is the sequence
      // length, so this slice's table row is slice * sequence_length + token.
      let position = position_ids[slice * (params.sliceRows / NUM_HEADS) + token];
{% endif %}
      let cache = position * ROT_UNITS + group;
      let cf = {{ castIn }}(cos_cache[cache]);
      let sf = {{ castIn }}(sin_cache[cache]);
{{ rotate("base") }}
{% if hasTail %}
      }
{% endif %}
    }
  }
{% endfor %}
{% else %}
  // A rank-4 slice is one batch and ITEMS consecutive heads; wg.y carries the high bits of
  // the group index past the per-axis dispatch fold width.
  let slice = wg.z;
  let headBase = slice % HEAD_BLOCKS * ITEMS;
  let batch = slice / HEAD_BLOCKS;
  let sliceUnits = params.sliceRows * HEAD_UNITS;
  let unit = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG + lid.x;
  if (unit >= sliceUnits) {
    return;
  }
  let group = unit % HEAD_UNITS;
  let token = unit / HEAD_UNITS;
  let rowBase = (batch * NUM_HEADS + headBase) * params.sliceRows + token;
{% if hasTail %}
  if (group >= ROT_UNITS) {
{% for item in range(itemsPerLane) %}
    {
      let base = (rowBase + {{ item }}u * params.sliceRows) * HEAD_STRIDE;
{{ copy_tail("base") }}
    }
{% endfor %}
    return;
  }
{% endif %}
{% if posOffset %}
  // One base offset for the whole request: token s reads cache row p[0] + s.
  let position = position_ids[0] + token;
{% else %}
  // The (batch_size, sequence_length) table; params.sliceRows is the sequence length.
  let position = position_ids[batch * params.sliceRows + token];
{% endif %}
  // Gathered once for the whole head block: cos and sin depend on the token and the rotation
  // group, never on the head.
  let cache = position * ROT_UNITS + group;
  let cf = {{ castIn }}(cos_cache[cache]);
  let sf = {{ castIn }}(sin_cache[cache]);
{% for item in range(itemsPerLane) %}
  {
    let base = (rowBase + {{ item }}u * params.sliceRows) * HEAD_STRIDE;
{{ rotate("base") }}
  }
{% endfor %}
{% endif %}
}