{{ env.wgsl.resourceDeclarations }}

const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
const TILE: u32 = {{ sharedTileSlots }}u;
const STATE_LENGTH: u32 = {{ stateLength }}u;
const HEADS: u32 = {{ numHeads }}u;
const HEADS_PER_ORDER: u32 = {{ headsPerNgram }}u;
{% if not hasPast %}
const PAD_ID: i32 = {{ padId }};
{% endif %}
// Order-major storage coalesces prefix writes and shares each value across heads.
var<workgroup> prefixes: array<i32, {{ sharedTileSlots * stateLength }}>;

fn positive_mod(value: i32, divisor: i32) -> i32 {
  let r = value % max(divisor, 1);
  return select(r, r + divisor, r < 0);
}
@compute @workgroup_size(WG)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let b = wg.y;
  let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u;
  let seq = params.seqLength;
  let tileBase = tile * TILE;
  let inputBase = b * seq;
{% if hasPast %}
  let pastBase = b * STATE_LENGTH;
{% endif %}
  if (lid.x < TILE) {
    let t = tileBase + lid.x;
    if (t < seq) {
      // Unsigned products preserve the specified two's-complement wrapping.
      var mix = bitcast<u32>(input_ids[inputBase + t]) * bitcast<u32>(multipliers[0u]);
{% for k in range(1, maxNgramSize) %}
{% if hasPast %}
      var token{{ k }}: i32;
      if (t >= {{ k }}u) { token{{ k }} = input_ids[inputBase + t - {{ k }}u]; }
      else { token{{ k }} = past_ids[pastBase + STATE_LENGTH + t - {{ k }}u]; }
{% else %}
      var token{{ k }} = PAD_ID;
      if (t >= {{ k }}u) { token{{ k }} = input_ids[inputBase + t - {{ k }}u]; }
{% endif %}
      mix = mix ^ (bitcast<u32>(token{{ k }}) * bitcast<u32>(multipliers[{{ k }}u]));
      prefixes[{{ k - 1 }}u * TILE + lid.x] = bitcast<i32>(mix);
{% endfor %}
    } else if (t < seq + STATE_LENGTH) {
{% if hasPast %}
      var carry: i32;
      if (t >= STATE_LENGTH) { carry = input_ids[inputBase + t - STATE_LENGTH]; }
      else { carry = past_ids[pastBase + t]; }
{% else %}
      var carry = PAD_ID;
      if (t >= STATE_LENGTH) { carry = input_ids[inputBase + t - STATE_LENGTH]; }
{% endif %}
      present_ids[b * STATE_LENGTH + t - seq] = carry;
    }
  }
  // Every lane reaches the barrier, including carry-only and partial tiles.
  workgroupBarrier();
  var tokens = 0u;
  if (tileBase < seq) { tokens = min(TILE, seq - tileBase); }
  for (var index = lid.x; index < tokens * HEADS; index = index + WG) {
    let token = index / HEADS;
    let head = index % HEADS;
    let mix = prefixes[(head / HEADS_PER_ORDER) * TILE + token];
    hash_ids[(inputBase + tileBase) * HEADS + index] = positive_mod(mix, vocab_sizes[head]);
  }
}