// The gate kernel of wai.quantum.circuit and wai.quantum.circuit2 in 32-bit
// lanes: a line-for-line transcription of quantum_fabric::lanes (mul_wide,
// fxmul, cmul, cadd, apply) that also carries every result as an exact 64-bit
// value and reports any result that leaves the 32-bit range.
//
// Every value is a u32. A signed amplitude or matrix entry is the bit pattern
// of its i32 value. This file holds no signed type, no float, no division, no
// conversion or bitcast, no workgroup memory, no barrier, and no shift whose
// amount is not a literal; no constant expression overflows. Any edit to this
// file changes the implementation id, and every certificate must be earned again.
struct Gate {
row0: vec4<u32>, // m00.re, m00.im, m01.re, m01.im
row1: vec4<u32>, // m10.re, m10.im, m11.re, m11.im
tbit: u32, // 1 << target
ctrl: u32, // the controls' bits, OR-ed
pairs: u32, // 2^(n-1)
row_stride: u32, // invocations per grid row: 64 x workgroups along x
index: u32, // the gate's position in the op-log
pad0: u32,
pad1: u32,
pad2: u32,
}
@group(0) @binding(0) var<storage, read_write> re: array<u32>;
@group(0) @binding(1) var<storage, read_write> im: array<u32>;
// report[0]: bit 1, a result left the 32-bit range; bit 2, an index outside the
// state; bit 4, a malformed record. report[1]: the lowest op-log index that set
// bit 1. The host writes (0, 0xffffffff) before a run.
@group(0) @binding(2) var<storage, read_write> report: array<atomic<u32>, 2>;
@group(0) @binding(3) var<uniform> g: Gate;
const SIGN: u32 = 0x80000000u;
const HALF: u32 = 0x20000000u; // 2^29, the rounding constant of fx
const LOW16: u32 = 0x0000ffffu;
const EXT30: u32 = 0xfffffffcu; // the bits an arithmetic shift by 30 copies from the sign
const ONES: u32 = 0xffffffffu;
fn neg(x: u32) -> bool { return (x & SIGN) != 0u; }
fn bit(c: bool) -> u32 { return select(0u, 1u, c); }
// lanes::mul_wide: the 64-bit two's-complement product of two i32 patterns, as (hi, lo).
fn mul_wide(a: u32, b: u32) -> vec2<u32> {
let negative = neg(a) != neg(b);
let x = select(a, 0u - a, neg(a)); // i32::unsigned_abs
let y = select(b, 0u - b, neg(b));
let x1 = x >> 16u;
let x0 = x & LOW16;
let y1 = y >> 16u;
let y0 = y & LOW16;
let p00 = x0 * y0; // no partial product reaches 2^32
let p01 = x0 * y1;
let p10 = x1 * y0;
let p11 = x1 * y1;
let mid = p01 + p10;
let mid_carry = bit(mid < p01);
let lo = p00 + (mid << 16u);
let lo_carry = bit(lo < p00);
let hi = p11 + (mid >> 16u) + (mid_carry << 16u) + lo_carry;
if (negative) {
let lo_n = ~lo + 1u;
let hi_n = ~hi + bit(lo_n == 0u);
return vec2<u32>(hi_n, lo_n);
}
return vec2<u32>(hi, lo);
}
// The exact 64-bit value of fx(a, b) = (a*b + 2^29) >> 30, as (hi, lo).
// lo is lanes::fxmul(a, b).
fn fx64(a: u32, b: u32) -> vec2<u32> {
let w = mul_wide(a, b);
let lo2 = w.y + HALF;
let hi2 = w.x + bit(lo2 < w.y);
let lo = (lo2 >> 30u) | (hi2 << 2u);
let hi = (hi2 >> 30u) | select(0u, EXT30, neg(hi2));
return vec2<u32>(hi, lo);
}
fn add64(x: vec2<u32>, y: vec2<u32>) -> vec2<u32> {
let lo = x.y + y.y;
return vec2<u32>(x.x + y.x + bit(lo < x.y), lo);
}
fn sub64(x: vec2<u32>, y: vec2<u32>) -> vec2<u32> {
return vec2<u32>(x.x - y.x - bit(x.y < y.y), x.y - y.y);
}
// Whether a 64-bit value is an i32: its high word repeats its low word's sign.
fn fits(v: vec2<u32>) -> bool { return v.x == select(0u, ONES, neg(v.y)); }
// One output amplitude m_r0*a0 + m_r1*a1, m = (m_r0.re, m_r0.im, m_r1.re, m_r1.im),
// each part exact; its low words are lanes::cadd(cmul(m_r0, a0), cmul(m_r1, a1)).
// Returns (re.hi, re.lo, im.hi, im.lo).
fn row(m: vec4<u32>, a0: vec2<u32>, a1: vec2<u32>) -> vec4<u32> {
let r = add64(sub64(fx64(m.x, a0.x), fx64(m.y, a0.y)),
sub64(fx64(m.z, a1.x), fx64(m.w, a1.y)));
let i = add64(add64(fx64(m.x, a0.y), fx64(m.y, a0.x)),
add64(fx64(m.z, a1.y), fx64(m.w, a1.x)));
return vec4<u32>(r, i);
}
@compute @workgroup_size(64)
fn apply(@builtin(global_invocation_id) id: vec3<u32>) {
let tbit = g.tbit;
let span = g.pairs << 1u; // 2^n
// A malformed record is reported and never run.
if (tbit == 0u || (tbit & (tbit - 1u)) != 0u || tbit > g.pairs
|| (g.ctrl & tbit) != 0u || (g.ctrl & ~(span - 1u)) != 0u) {
atomicOr(&report[0], 4u);
return;
}
let p = id.y * g.row_stride + id.x;
if (p >= g.pairs) { return; }
let low = p & (tbit - 1u);
let i = ((p ^ low) << 1u) | low; // p with a 0 inserted at the target bit
let j = i | tbit;
if ((i & g.ctrl) != g.ctrl) { return; }
if (j >= arrayLength(&re) || j >= arrayLength(&im)) {
atomicOr(&report[0], 2u);
return;
}
let a0 = vec2<u32>(re[i], im[i]);
let a1 = vec2<u32>(re[j], im[j]);
let o0 = row(g.row0, a0, a1);
let o1 = row(g.row1, a0, a1);
re[i] = o0.y;
im[i] = o0.w;
re[j] = o1.y;
im[j] = o1.w;
if (!(fits(o0.xy) && fits(o0.zw) && fits(o1.xy) && fits(o1.zw))) {
atomicOr(&report[0], 1u);
atomicMin(&report[1], g.index);
}
}
// The arithmetic probe: mul_wide and fx64 of each case, for detection and for
// the full-range check. Its bindings are distinct from apply's.
struct Probe { a: u32, b: u32, hi: u32, lo: u32, fx_hi: u32, fx_lo: u32, pad0: u32, pad1: u32, }
@group(0) @binding(4) var<storage, read_write> cases: array<Probe>;
@group(0) @binding(5) var<uniform> probe_grid: vec4<u32>; // x: case count, y: invocations per grid row
@compute @workgroup_size(64)
fn probe(@builtin(global_invocation_id) id: vec3<u32>) {
let k = id.y * probe_grid.y + id.x;
if (k >= probe_grid.x || k >= arrayLength(&cases)) { return; }
let w = mul_wide(cases[k].a, cases[k].b);
let f = fx64(cases[k].a, cases[k].b);
cases[k].hi = w.x;
cases[k].lo = w.y;
cases[k].fx_hi = f.x;
cases[k].fx_lo = f.y;
}