{% set precastOperands = precastOperands | default(false) %}
{% set precastScalar = precastScalar | default("f32") %}
enable subgroups;
{% if pinSubgroupSize32 %}
enable subgroup_size_control;
{% endif %}
enable chromium_experimental_subgroup_matrix;
diagnostic(off, chromium.subgroup_matrix_uniformity);
{{ env.wgsl.resourceDeclarations }}
// Rare-boundary fallback: reproduce the reference's binary64 scale division
// and accumulator product using integer significands; no f64 feature required.
fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
  let low = a.x + b.x;
  let middle = a.y + b.y;
  let with_carry = middle + u32(low < a.x);
  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
}
fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
}
fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
  var q = value;
  var sticky = false;
  var guard = false;
  for (var bit = 0u; bit < shift; bit++) {
    sticky = sticky || guard;
    guard = (q.x & 1u) != 0u;
    q = quant_u96_shr(q);
  }
  if (guard && (sticky || (q.x & 1u) != 0u)) {
    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
  }
  return q;
}
fn quant_float_significand(bits: u32) -> u32 {
  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
}
fn quant_float_exponent(bits: u32) -> i32 {
  return max(1, i32((bits >> 23u) & 255u)) - 150;
}
fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
  if (acc == 0) { return 0; }
  let a_bits = bitcast<u32>(a_s);
  let b_bits = bitcast<u32>(b_s);
  let y_bits = bitcast<u32>(y_s);
  let a = quant_float_significand(a_bits);
  let b = quant_float_significand(b_bits);
  let denominator = quant_float_significand(y_bits);
  if (a == 0u || b == 0u) { return 0; }
  // The exact product of two 24-bit significands occupies at most 48 bits.
  let p0 = (a & 65535u) * (b & 65535u);
  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
  let product_low = p0 + (cross << 16u);
  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
  var remainder = 0u;
  var significant_bits = 0u;
  var leading_exponent = 0;
  var significand = vec2<u32>(0u);
  // Long division generates the 53 bits of the binary64 multiplier. The
  // remainder fits in 25 bits, including its left shift. At most 24 leading
  // zero bits precede the first significant bit.
  loop {
    var incoming = 0u;
    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
    remainder = (remainder << 1u) | incoming;
    let digit = u32(remainder >= denominator);
    if (digit != 0u) { remainder -= denominator; }
    if (significant_bits != 0u || digit != 0u) {
      if (significant_bits == 0u) { leading_exponent = input_bit; }
      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
      significant_bits++;
    }
    input_bit--;
    if (significant_bits == 53u) { break; }
  }
  let twice_remainder = remainder << 1u;
  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
    let low = significand.x + 1u;
    significand = vec2<u32>(low, significand.y + u32(low == 0u));
  }
  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
  // significand by the exact accumulator into three little-endian u32 limbs.
  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
  var operand = vec3<u32>(significand, 0u);
  var product = vec3<u32>(0u);
  loop {
    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
    magnitude >>= 1u;
    if (magnitude == 0u) { break; }
    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
  }
  var product_top = 31 - i32(countLeadingZeros(product.x));
  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
  if (product_top + scale_exponent < -2) { return 0; }
  let product_shift = u32(max(0, product_top - 52));
  let rounded_product = quant_u96_round_shift(product, product_shift);
  let result_exponent = scale_exponent + i32(product_shift);
  var rounded = rounded_product;
  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
  else { rounded.x <<= u32(result_exponent); }
  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
}

// A power-of-two division and an exactly represented <=48-bit product need
// no binary64 boundary resolver. This includes ordinary dyadic quantization.
fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
  let significand = quant_float_significand(bitcast<u32>(scale));
  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
  return integer_bits + scale_bits <= 48u;
}
fn quant_normal_expansion_range(v: f32) -> bool {
  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
  // Reserve two binary32 significands plus four carry/error bits above the
  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
}

fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
  let estimate = numerator_hi / denominator;
  // Anything outside this range saturates after adding an 8-bit zero point;
  // returning early also keeps floor/i32 conversion away from Inf and overflow.
  if (estimate <= -1024.0) { return -1024; }
  if (estimate >= 1024.0) { return 1024; }

  let lo = floor(estimate);
  let hi = lo + 1.0;
  let midpoint = lo + 0.5;
  let midpoint_hi = fma(midpoint, denominator, 0.0);
  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
  // The high-word subtraction is exact for these similarly sized products
  // (Sterbenz); adding the residual difference therefore preserves its sign
  // even when the rounded high words straddle adjacent f32 values.
  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
  // Six FMA/product tails and their additions fit within this conservative
  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
  }
  if (midpoint_delta < 0.0) {
    return i32(lo);
  }
  if (midpoint_delta > 0.0) {
    return i32(hi);
  }
  let half_lo = floor(lo * 0.5);
  let lo_is_even = lo - half_lo * 2.0 == 0.0;
  return i32(select(hi, lo, lo_is_even));
}

// Multiplication by a power of two keeps every i32 bit in binary64. Round the
// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
  var rounded = 0u;
  if (exponent >= 0) {
    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
  } else {
    let shift = u32(-exponent);
    if (shift >= 32u) { return 0; }
    let whole = magnitude >> shift;
    let remainder = magnitude & ((1u << shift) - 1u);
    let half = 1u << (shift - 1u);
    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
  }
  return select(i32(rounded), -i32(rounded), acc < 0);
}

fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
  if (acc == 0) { return 0; }
  let a_bits = bitcast<u32>(a_s);
  let b_bits = bitcast<u32>(b_s);
  let y_bits = bitcast<u32>(y_s);
  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
    return quant_round_power_two(acc, exponent);
  }
  let acc_f = f32(acc);
  // Every i32 high part aligned to 256 is exactly representable in f32.
  // Recover the conversion residual without an out-of-range i32(acc_f)
  // conversion when INT_MAX rounds up to 2^31.
  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
  // Equal b/y scales cancel exactly. Bypass their product and division so a
  // rounded intermediate cannot move a near-half value onto the tie.
  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
    let direct_hi = fma(acc_f, a_s, 0.0);
    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
  }

  // Keep the scale product as a two-float expansion. Its FMA residual preserves
  // which side of a rounding midpoint a near-half value occupies.
  let ab_hi = fma(a_s, b_s, 0.0);
  let ab_lo = fma(a_s, b_s, -ab_hi);
  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
  let product_hi = fma(acc_f, ab_hi, 0.0);
  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
  // Compare the expansion directly. Renormalizing it adds a cancelling
  // subtraction/addition chain that can erase the residual during lowering.
  let denominator_bits = bitcast<u32>(y_s);
  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
}

// A conservative binary32 error filter accepts only a result safely inside one
// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
// path; this is an exact-result shortcut, not an approximate quantization mode.
fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
  if (acc == 0) { return 0; }
  let a_bits = bitcast<u32>(a_s);
  let b_bits = bitcast<u32>(b_s);
  let y_bits = bitcast<u32>(y_s);
  let ea = (a_bits >> 23u) & 255u;
  let eb = (b_bits >> 23u) & 255u;
  let ey = (y_bits >> 23u) & 255u;
  // Certify the range before arithmetic: overflow need not produce infinity.
  // Any nonzero i32 converts to a float with biased exponent <=158. These
  // scale-only bounds cover every accumulator and can be shared across lanes.
  // Reserve two significands plus carry bits above FTZ, and a factor of two
  // below overflow for products, division and the expansion midpoint product.
  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
      ea + eb < 180u || ea + eb > 347u ||
      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
    return quant_round_binary64(acc, a_s, b_s, y_s);
  }
  let scale = fma(a_s, b_s, 0.0);
  let product = fma(f32(acc), scale, 0.0);
  let estimate = product / y_s;
  if (estimate <= -1024.0) { return -1024; }
  if (estimate >= 1024.0) { return 1024; }
  let nearest = round(estimate);
  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
  // products, quotient, and distance arithmetic, plus binary64 rounding.
  let margin = max(1.0, abs(estimate)) * 0x1p-19;
  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
{% if device.adapterInfo.vendor != "apple" %}
  // A rounded half-integer estimate can hide which side the exact f32-scale
  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
  // the compact fast path; exact power-of-two ratios already have a cheap
  // and correct path in round_requantized_precise.
  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
      min(a_bits, min(b_bits, y_bits)) != 0u &&
      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
  }
{% endif %}
  return round_requantized_precise(acc, a_s, b_s, y_s);
}

// Adjusted byte integers fit exactly in f32. Each bounded f32 dot product
// has absolute magnitude at most 2^24; int32 adds then preserve wraparound.
const TILE_M: u32 = {{ sgmatTileM }}u;
const TILE_N: u32 = {{ sgmatTileN }}u;
const TILE_K: u32 = {{ sgmatTileK }}u;

{% if precastOperands %}
const PADDED_K: u32 = {{ preparedK }}u;
const PADDED_N: u32 = {{ preparedN }}u;
var<workgroup> matrix_out: array<f32, {{ sgmatSubgroups * 8 * 64 }}>;
{% else %}
var<workgroup> tile_A: array<f32, {{ sgmatTileM * sgmatTileK }}>;
var<workgroup> tile_B: array<f32, {{ sgmatTileK * sgmatTileN }}>;
{% if not sgmatReuseStaging %}
var<workgroup> scratch: array<array<array<f32, 64>, 4>, {{ sgmatSubgroups }}>;
{% endif %}
{% endif %}

@compute @workgroup_size({{ sgmatWorkgroup }}){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
fn main(@builtin(workgroup_id) wid: vec3<u32>,
{% if not precastOperands or not matrixSubgroupId %}
        @builtin(local_invocation_index) tid: u32,
{% endif %}
        @builtin(subgroup_invocation_id) lane: u32,
@builtin(subgroup_id) subgroup_index: u32) {
  let subgroup = subgroup_index;
  let base_A = (subgroup % {{ sgmatRowGroups }}u) * 16u;
  let base_B = (subgroup / {{ sgmatRowGroups }}u) * 32u;
  let row = lane / 4u;
  let col = (lane % 4u) * 2u;
  // Bound adjusted byte magnitudes from scalar zero points. A column-wise
  // B zero point keeps the full byte-difference bound for the whole tile.
  let max_a = {% if matrixZeroPerRow %}255u{% else %}u32(max(abs({{ 0 if aUnsigned else (0 - 128) }} - i32(a_zero_point[0])), abs({{ 255 if aUnsigned else 127 }} - i32(a_zero_point[0])))){% endif %};
  let max_b = {% if hasBZero and not bZeroPerColumn %}u32(max(abs({{ 0 if bUnsigned else (0 - 128) }} - i32(b_zero_point[0])), abs({{ 255 if bUnsigned else 127 }} - i32(b_zero_point[0])))){% else %}{{ 255 if hasBZero or bUnsigned else 128 }}u{% endif %};
  let partial_k = (16777216u / (max_a * max_b) / TILE_K) * TILE_K;
  var totals: array<i32, 16>;
  for (var partial = 0u; partial < params.K; partial += partial_k) {
{% for r in range(2) %}{% for c in range(4) %}
    var matC{{ r }}{{ c }}: subgroup_matrix_result<f32, 8, 8>;
{% endfor %}{% endfor %}
    let partial_end = min(params.K, partial + partial_k);
    for (var kb = partial; kb < partial_end; kb += TILE_K) {
{% if not precastOperands %}
      for (var index = tid; index < TILE_M * TILE_K; index += {{ sgmatWorkgroup }}u) {
        let m = wid.y * TILE_M + index / TILE_K;
        let k = kb + index % TILE_K;
        var value = f32(0.0);
        if (m < params.M && k < params.K) {
          value = f32(i32(a[m * params.K + k]) - i32(a_zero_point[{{ "m" if matrixZeroPerRow else "0" }}]));
        }
        tile_A[index] = value;
      }
      for (var index = tid; index < TILE_K * TILE_N; index += {{ sgmatWorkgroup }}u) {
        let k = kb + index / TILE_N;
        let n = wid.x * TILE_N + index % TILE_N;
        var value = f32(0.0);
        if (k < params.K && n < params.N) {
          value = f32(i32(b[{% if matrixBatchB %}wid.z * params.bBatchStride + {% endif %}k * params.N + n]) - i32(b_zero_point[{{ "n" if bZeroPerColumn else "0" }}]));
        }
        tile_B[index] = value;
      }
      workgroupBarrier();
{% endif %}
      for (var step = 0u; step < TILE_K; step += 8u) {
{% for r in range(2) %}
        let matA{{ r }} = subgroupMatrixLoad<subgroup_matrix_left<{{ precastScalar if precastOperands else "f32" }}, 8, 8>, row_major>(&{{ "a" if precastOperands else "tile_A" }}, {% if precastOperands %}(wid.y * TILE_M + base_A + {{ r * 8 }}u) * PADDED_K + kb + step{% else %}(base_A + {{ r * 8 }}u) * TILE_K + step{% endif %}, {{ "PADDED_K" if precastOperands else "TILE_K" }});
{% endfor %}
{% for c in range(4) %}
        let matB{{ c }} = subgroupMatrixLoad<subgroup_matrix_right<{{ precastScalar if precastOperands else "f32" }}, 8, 8>, row_major>(&{{ "b" if precastOperands else "tile_B" }}, {% if precastOperands %}{% if matrixBatchB %}wid.z * PADDED_K * PADDED_N + {% endif %}(kb + step) * PADDED_N + wid.x * TILE_N + base_B + {{ c * 8 }}u{% else %}step * TILE_N + base_B + {{ c * 8 }}u{% endif %}, {{ "PADDED_N" if precastOperands else "TILE_N" }});
{% endfor %}
{% for r in range(2) %}{% for c in range(4) %}
        matC{{ r }}{{ c }} = subgroupMatrixMultiplyAccumulate(matA{{ r }}, matB{{ c }}, matC{{ r }}{{ c }});
{% endfor %}{% endfor %}
      }
{% if not precastOperands %}
      workgroupBarrier();
{% endif %}
    }
{% if precastOperands %}
    // The prepared operands need no staging barriers. Publish all eight
    // matrices together before converting the bounded partial to int32.
{% for r in range(2) %}{% for c in range(4) %}
    subgroupMatrixStore<row_major>(&matrix_out, subgroup * 512u + {{ r * 256 + c * 64 }}u, matC{{ r }}{{ c }}, 8u);
{% endfor %}{% endfor %}
    workgroupBarrier();
{% for r in range(2) %}{% for c in range(4) %}{% for h in range(2) %}
    totals[{{ r * 8 + c * 2 + h }}] += i32(matrix_out[subgroup * 512u + {{ r * 256 + c * 64 }}u + row * 8u + col + {{ h }}u]);
{% endfor %}{% endfor %}{% endfor %}
    workgroupBarrier();
{% else %}
{% if sgmatReuseStaging %}
    // All staged A reads have completed. Reuse its storage for one result band
    // at a time: each subgroup needs 4 * 64 floats, exactly its A-tile share.
{% endif %}
{% for r in range(2) %}
{% for c in range(4) %}
    subgroupMatrixStore<row_major>({% if sgmatReuseStaging %}&tile_A{% else %}&scratch[subgroup][{{ c }}]{% endif %}, {% if sgmatReuseStaging %}subgroup * 256u + {{ c * 64 }}u{% else %}0u{% endif %}, matC{{ r }}{{ c }}, 8u);
{% endfor %}
    workgroupBarrier();
{% for c in range(4) %}{% for h in range(2) %}
    totals[{{ r * 8 + c * 2 + h }}] += i32({% if sgmatReuseStaging %}tile_A[subgroup * 256u + {{ c * 64 }}u + row * 8u + col + {{ h }}u]{% else %}scratch[subgroup][{{ c }}][row * 8u + col + {{ h }}u]{% endif %});
{% endfor %}{% endfor %}
    workgroupBarrier();
{% endfor %}
{% endif %}
  }
{% for r in range(2) %}{% for c in range(4) %}{% for h in range(2) %}
  {
    let m = wid.y * TILE_M + base_A + {{ r * 8 }}u + row;
    let n = wid.x * TILE_N + base_B + {{ c * 8 + h }}u + col;
    if (m < params.M && n < params.N) {
      let acc = totals[{{ r * 8 + c * 2 + h }}]{% if matrixHasBias %} + bias[{{ "m" if matrixBiasPerRow else "n" }}]{% endif %};
      let rounded = round_requantized(acc, f32(a_scale[0]), f32(b_scale[{{ "m" if matrixScalePerRow else ("n" if bScalePerColumn else "0") }}]), f32(y_scale[0]));
      let q = clamp(rounded + i32(y_zero_point[0]), {{ qMin }}, {{ qMax }});
      y[{% if matrixBatchB %}wid.z * params.yBatchStride + {% endif %}m * params.N + n] = {{ "u32(q)" if yUnsigned else "q" }};
    }
  }
{% endfor %}{% endfor %}{% endfor %}
}