{{ env.wgsl.resourceDeclarations }}

// Engram n-gram hash ids. For every order n in [2, MAX_NGRAM] the kernel mixes the n causal
// shifts of the compressed tokenizer ids,
//   mix_n = ids[t] * m[0] ^ ids[t-1] * m[1] ^ ... ^ ids[t-n+1] * m[n-1],
// and emits mix_n modulo each of that order's head vocabulary sizes. mix_n extends mix_{n-1}
// by a single term, so one walk over the window produces every order.
//
// Each product is formed in u32 so that overflow wraps in two's complement, which is the
// arithmetic the operator specifies. A signed multiply that overflows has an indeterminate
// result in WGSL, so the bitcasts here are load-bearing rather than cosmetic.
//
// One dispatch covers both outputs. Workgroup row wg.y owns batch row b, and the tile index
// folded across wg.x/wg.z walks a span of SEQ + STATE_LENGTH slots: the first SEQ slots hash a
// token position, and the trailing STATE_LENGTH slots emit the carry state that the next call
// feeds back as past_ids. A slot past the end of the span exits, so a short sequence costs one
// workgroup rather than a second dispatch.
const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
const ITEMS: u32 = {{ tunables.ITEMS_PER_LANE }}u;
const STATE_LENGTH: u32 = {{ stateLength }}u;
const PAD_ID: i32 = {{ padId }};

// Euclidean remainder for a positive divisor: WGSL `%` truncates toward zero exactly as C
// does, so a negative dividend needs the one correction the operator specifies. A head whose
// vocabulary size is not strictly positive has no meaningful modulo, and clamping the divisor
// to 1 makes its remainder 0 -- the value every GPU implementation of this operator emits for
// that head, and the reason it stays a device-side division rather than a rejected input.
fn positive_mod(value: i32, divisor: i32) -> i32 {
  let r = value % max(divisor, 1);
  return select(r, r + divisor, r < 0);
}
{% macro hash_position(interior) %}
      let idBase = inputBase + t;
      var mix = bitcast<u32>(input_ids[idBase]) * mult0;
{% for k in range(1, maxNgramSize) %}
{% if interior %}
      mix = mix ^ (bitcast<u32>(input_ids[idBase - {{ k }}u]) * mult{{ k }});
{% else %}
      var token{{ k }} = PAD_ID;
      if (t >= {{ k }}u) {
        token{{ k }} = input_ids[idBase - {{ k }}u];
      }
{% if hasPast %}
      else {
        // past_ids is right-aligned, so position -1 is its last slot. t < k bounds the index
        // inside the window, which is why no second range test is needed here.
        token{{ k }} = past_ids[pastBase + STATE_LENGTH + t - {{ k }}u];
      }
{% endif %}
      mix = mix ^ (bitcast<u32>(token{{ k }}) * mult{{ k }});
{% endif %}
      let mixOrder{{ k }} = bitcast<i32>(mix);
{% endfor %}
{% for n in range(1, maxNgramSize) %}
{% for h in range(headsPerNgram) %}
      let head{{ (n - 1) * headsPerNgram + h }} = positive_mod(mixOrder{{ n }}, vocab{{ (n - 1) * headsPerNgram + h }});
{% endfor %}
{% endfor %}
      let outBase = idBase * {{ numHeads }}u;
{% for h in range(numHeads) %}
      hash_ids[outBase + {{ h }}u] = head{{ h }};
{% endfor %}
{% endmacro %}

@compute @workgroup_size(WG)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  // wg.y is a real axis -- one workgroup row per batch row -- so only the span tiles are
  // folded, across wg.x and wg.z at the per-axis workgroup fold width.
  let b = wg.y;
  let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u;
  let seq = params.seqLength;
  let span = seq + STATE_LENGTH;
  let inputBase = b * seq;
  let tileBase = tile * (WG * ITEMS);
{% if hasPast %}
  let pastBase = b * STATE_LENGTH;
{% endif %}

  // The multiplier and vocabulary tables are attribute-sized and every position reads all of
  // them, so an invocation stages them once instead of reloading them inside the mix walk.
{% for k in range(maxNgramSize) %}
  let mult{{ k }} = bitcast<u32>(multipliers[{{ k }}u]);
{% endfor %}
{% for h in range(numHeads) %}
  let vocab{{ h }} = vocab_sizes[{{ h }}u];
{% endfor %}

  // An interior tile is one whose every slot hashes a token position and whose every causal
  // shift lands inside input_ids. Both of the per-slot tests below, and the per-shift history
  // fallback, are then statically false, which is worth one workgroup-uniform branch.
  if (tileBase >= STATE_LENGTH && tileBase + WG * ITEMS <= seq) {
    for (var item = 0u; item < ITEMS; item = item + 1u) {
      let t = tileBase + lid.x + item * WG;
{{ hash_position(true) }}
    }
    return;
  }

  for (var item = 0u; item < ITEMS; item = item + 1u) {
    let slot = tileBase + lid.x + item * WG;
    if (slot >= span) {
      break;
    }
    if (slot < seq) {
      let t = slot;
{{ hash_position(false) }}
    } else {
      // Carry slot j holds the id at position seq - STATE_LENGTH + j of this call, and slot j
      // of past_ids when the call is shorter than the window. Writing it from this dispatch is
      // safe because present_ids never shares its buffer with past_ids here.
      var carry = PAD_ID;
      if (slot >= STATE_LENGTH) {
        carry = input_ids[inputBase + slot - STATE_LENGTH];
      }
{% if hasPast %}
      else {
        carry = past_ids[pastBase + slot];
      }
{% endif %}
      present_ids[b * STATE_LENGTH + slot - seq] = carry;
    }
  }
}