{% if foldKind == "subgroup_full" or foldKind == "subgroup_butterfly" %}
enable subgroups;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

// The Engram gate, fused per (batch, sequence, hc_mult) row:
//   dot  = sum(RMSNorm(key) * RMSNorm(query)) / sqrt(HIDDEN)
//   out  = sigmoid(sign(dot) * sqrt(max(abs(dot), 1e-6))) * value
// Both inverse-RMS factors are row scalars, so they come out of the sum and the
// two sums of squares and the scaled cross term ride ONE pass as a single vec3
// accumulator. Normalizing per element first would need two passes and a
// different rounding.
//
// Each row owns a fixed LANE GROUP rather than a whole workgroup. LANES_PER_ROW
// is the smallest power of two that covers the row, capped by the workgroup
// (and, on the subgroup routes, by the subgroup) so several rows are resident at
// once, and widened toward the workgroup when there are too few rows to fill the
// dispatch. The value words stay in registers across the fold.
//
// Every accumulation is float32 whatever the tensor type; only the store narrows.
const ROW_WORDS: u32 = {{ rowWords }}u;
const HIDDEN: f32 = {{ hidden }}.0;
const HC_MULT: u32 = {{ hcMult }}u;
const EPS: f32 = {{ epsilon }};
const GATE_FLOOR: f32 = 0.000001;
const WG: u32 = {{ workgroupSize }}u;
const LANES_PER_ROW: u32 = {{ lanesPerRow }}u;
const ROWS_PER_GROUP: u32 = {{ rowsPerGroup }}u;
{% if foldKind == "workgroup_tree" %}

// One contiguous LANES_PER_ROW slice per row, carrying all three accumulators as
// one vector. The fold never crosses a slice boundary, so the rows sharing this
// array do not interfere.
var<workgroup> partials: array<vec3<f32>, WG>;
{% endif %}

fn sigmoid_safe(x: f32) -> f32 {
  if (x >= 0.0) {
    let z = exp(-x);
    return 1.0 / (1.0 + z);
  }
  let z = exp(x);
  return z / (1.0 + z);
}

// sign(0.0) is 0.0 in WGSL, so a zero dot product gives an argument of exactly
// zero and therefore a gate of exactly 0.5. A copysign spelling would map it to
// +sqrt(GATE_FLOOR) instead, which is the divergence upstream calls out.
fn engram_gate(dotValue: f32) -> f32 {
  return sigmoid_safe(sign(dotValue) * sqrt(max(abs(dotValue), GATE_FLOOR)));
}

@compute @workgroup_size(WG)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let lane = lid.x % LANES_PER_ROW;
  let localRow = lid.x / LANES_PER_ROW;
  // 2D-folded row index: wg.y carries the high bits past the dispatch's per-axis
  // workgroup fold width.
  let row = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * ROWS_PER_GROUP + localRow;
  let rowValid = row < params.rows;
  // A lane group past the last row folds an empty row and stores nothing, so the
  // collectives and barriers below stay in uniform control flow.
  let limit = select(0u, ROW_WORDS, rowValid);
  let base = row * ROW_WORDS;
  // The hyper-connection stream selects the scale rows; the token selects the one
  // value row that every stream of that token shares.
  let scaleBase = (row % HC_MULT) * ROW_WORDS;
  let valueBase = (row / HC_MULT) * ROW_WORDS;

  // x: sum(key^2), y: sum(query^2), z: sum(key * key_scale * query * query_scale).
  var acc = vec3<f32>(0.0);
{% if stageRow %}
{% for i in range(wordsPerLane) %}
{% if useVec4 %}
  var vv{{ i }} = vec4<f32>(0.0);
{% else %}
  var vv{{ i }} = 0.0;
{% endif %}
  let c{{ i }} = lane + {{ i }}u * LANES_PER_ROW;
  if (c{{ i }} < limit) {
{% if useVec4 %}
    let kv{{ i }} = vec4<f32>(key[base + c{{ i }}]);
    let qv{{ i }} = vec4<f32>(query[base + c{{ i }}]);
    let ks{{ i }} = vec4<f32>(kscale[scaleBase + c{{ i }}]);
    let qs{{ i }} = vec4<f32>(qscale[scaleBase + c{{ i }}]);
    vv{{ i }} = vec4<f32>(value[valueBase + c{{ i }}]);
    acc = acc + vec3<f32>(dot(kv{{ i }}, kv{{ i }}), dot(qv{{ i }}, qv{{ i }}), dot(kv{{ i }} * ks{{ i }}, qv{{ i }} * qs{{ i }}));
{% else %}
    let kv{{ i }} = f32(key[base + c{{ i }}]);
    let qv{{ i }} = f32(query[base + c{{ i }}]);
    let ks{{ i }} = f32(kscale[scaleBase + c{{ i }}]);
    let qs{{ i }} = f32(qscale[scaleBase + c{{ i }}]);
    vv{{ i }} = f32(value[valueBase + c{{ i }}]);
    acc = acc + vec3<f32>(kv{{ i }} * kv{{ i }}, qv{{ i }} * qv{{ i }}, (kv{{ i }} * ks{{ i }}) * (qv{{ i }} * qs{{ i }}));
{% endif %}
  }
{% endfor %}
{% else %}
  // The row's value words do not fit in registers at this width, so the statistic
  // pass and the store pass each walk it.
  for (var c = lane; c < limit; c = c + LANES_PER_ROW) {
{% if useVec4 %}
    let kv = vec4<f32>(key[base + c]);
    let qv = vec4<f32>(query[base + c]);
    let ks = vec4<f32>(kscale[scaleBase + c]);
    let qs = vec4<f32>(qscale[scaleBase + c]);
    acc = acc + vec3<f32>(dot(kv, kv), dot(qv, qv), dot(kv * ks, qv * qs));
{% else %}
    let kv = f32(key[base + c]);
    let qv = f32(query[base + c]);
    let ks = f32(kscale[scaleBase + c]);
    let qs = f32(qscale[scaleBase + c]);
    acc = acc + vec3<f32>(kv * kv, qv * qv, (kv * ks) * (qv * qs));
{% endif %}
  }
{% endif %}
{% if foldKind == "subgroup_full" %}
  let sums = subgroupAdd(acc);
{% elif foldKind == "subgroup_butterfly" %}
  // Segmented butterfly: only the low bits of the lane index are exchanged, so each
  // aligned group of LANES_PER_ROW lanes folds its own row and one subgroup carries
  // several rows without a barrier or any workgroup memory. All three accumulators
  // ride one shuffle.
  var folded = acc;
{% for step in range(foldSteps) %}
  folded = folded + subgroupShuffleXor(folded, 1u << {{ step }}u);
{% endfor %}
  let sums = folded;
{% elif foldKind == "workgroup_tree" %}
  partials[lid.x] = acc;
  workgroupBarrier();
  // LANES_PER_ROW is a power of two, so the halving stride is uniform across the
  // whole workgroup and every lane reaches every barrier.
  var stride = LANES_PER_ROW / 2u;
  loop {
    if (stride == 0u) {
      break;
    }
    if (lane < stride) {
      partials[lid.x] = partials[lid.x] + partials[lid.x + stride];
    }
    workgroupBarrier();
    stride = stride / 2u;
  }
  let sums = partials[localRow * LANES_PER_ROW];
{% else %}
  // One lane owns the whole row, so the accumulator already holds the row sums.
  let sums = acc;
{% endif %}
  let keyInv = inverseSqrt(sums.x / HIDDEN + EPS);
  let queryInv = inverseSqrt(sums.y / HIDDEN + EPS);
  let gate = engram_gate(sums.z * keyInv * queryInv / sqrt(HIDDEN));
{% if stageRow %}
{% for i in range(wordsPerLane) %}
  if (c{{ i }} < limit) {
{% if useVec4 %}
    out[base + c{{ i }}] = {{ vectorScalar }}(vv{{ i }} * gate);
{% else %}
    out[base + c{{ i }}] = {{ scalar }}(vv{{ i }} * gate);
{% endif %}
  }
{% endfor %}
{% else %}
  for (var c = lane; c < limit; c = c + LANES_PER_ROW) {
{% if useVec4 %}
    out[base + c] = {{ vectorScalar }}(vec4<f32>(value[valueBase + c]) * gate);
{% else %}
    out[base + c] = {{ scalar }}(f32(value[valueBase + c]) * gate);
{% endif %}
  }
{% endif %}
}