diff --git a/wai/wai-quantum-serve/Cargo.lock b/wai/wai-quantum-serve/Cargo.lock index 34b20486b..4164a5571 100644 --- a/wai/wai-quantum-serve/Cargo.lock +++ b/wai/wai-quantum-serve/Cargo.lock @@ -589,7 +589,7 @@ dependencies = [ [[package]] name = "wai-quantum" -version = "0.3.25" +version = "0.3.26" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-quantum/Cargo.lock b/wai/wai-quantum/Cargo.lock index c2286cbf0..7de852a9d 100644 --- a/wai/wai-quantum/Cargo.lock +++ b/wai/wai-quantum/Cargo.lock @@ -415,7 +415,7 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" [[package]] name = "wai-quantum" -version = "0.3.25" +version = "0.3.26" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-quantum/Cargo.toml b/wai/wai-quantum/Cargo.toml index cfc38cdda..3b2d1d6b7 100644 --- a/wai/wai-quantum/Cargo.toml +++ b/wai/wai-quantum/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "wai-quantum" -version = "0.3.25" +version = "0.3.26" edition = "2024" rust-version = "1.97" license = "Apache-2.0" @@ -56,6 +56,9 @@ quantum_resource_receipt = ["quantum_resource", "_signing"] # Noisy Clifford circuits at QEC scale: bit-packed Pauli-frame sampling, detector # error models by backward sensitivity, and exact minimum-weight matching. quantum_frame = [] +# Tensor-network contraction: hyper-optimised contraction trees, slicing, and +# amplitudes of circuits far past the state vector, reproducible bit for bit. +quantum_tn = [] # Mixed states and the observables that distinguish quantum theory from a classical # account: density matrices, the partial trace, exact Pauli expectations, CHSH. quantum_info = ["quantum"] @@ -95,13 +98,17 @@ quantum_phasor_meter = ["_signing"] # Everything. full = [ - "quantum", "quantum_toolchain", "quantum_mps", "quantum_stabilizer", "quantum_pauli", "quantum_spd", "quantum_spd_receipt", "quantum_bptn", "quantum_resource", "quantum_resource_receipt", "quantum_frame", "quantum_qir", "quantum_info", "quantum_qudit", "quantum_comm", "quantum_source", "quantum_channel", + "quantum", "quantum_toolchain", "quantum_mps", "quantum_stabilizer", "quantum_pauli", "quantum_spd", "quantum_spd_receipt", "quantum_bptn", "quantum_resource", "quantum_resource_receipt", "quantum_frame", "quantum_tn", "quantum_qir", "quantum_info", "quantum_qudit", "quantum_comm", "quantum_source", "quantum_channel", "quantum_receipt", "quantum_ops", "quantum_cal", "quantum_mitigate", "quantum_qec", "quantum_control", "quantum_noise", "quantum_compile", "quantum_atom", "quantum_qbom", "quantum_vml", "quantum_phasor", "quantum_qfhrr", "quantum_kernel", "quantum_qdata", "quantum_phasor_meter", ] +[[example]] +name = "tn_path" +required-features = ["quantum_tn"] + [[example]] name = "qec_threshold" required-features = ["quantum_frame"] diff --git a/wai/wai-quantum/README.md b/wai/wai-quantum/README.md index f1682622b..f0be5abe4 100644 --- a/wai/wai-quantum/README.md +++ b/wai/wai-quantum/README.md @@ -70,6 +70,7 @@ never creates information. | `quantum_mitigate` | zero-noise extrapolation, readout M3, classical shadows | | `quantum_qec` | CSS codes (distance verified by exhaustive search, not asserted) — **rotated surface** `[[d²,1,d]]`, toric, bivariate-bicycle, and the **`[[7,1,3]]` colour code**, whose self-duality buys a *transversal* logical Hadamard — one round, no ancillas, no surgery — verified on the simulator, with the surface code shown failing the same move — with **two decoders**: union-find (weighted matching, single-shot, space-time with faulty measurement, **and circuit-level with hook errors**; **imports and exports** Stim-syntax detector error models, so other tools' noise models decode here and ours decode there) and Relay-BP (for qLDPC) | | `quantum_frame` | **noisy Clifford circuits at QEC scale**: reads the ecosystem's circuit text; samples detection events by bit-packed Pauli-frame simulation (64 shots per word); derives the detector error model by backward sensitivity, refusing non-deterministic detectors; generates the rotated surface-code memory experiment byte for byte as the reference layout; and decodes by **exact minimum-weight perfect matching** (Edmonds' blossom on integer weights). Measures logical error rates and fits the code model resource estimates assume | +| `quantum_tn` | **tensor-network contraction**: circuits (any unitary on any qubits) become networks; a seeded hyper-optimiser finds contraction trees from randomised-greedy and multilevel-partition families, reconfigures them optimally in subtrees, and slices them to a memory cap with a second search around the slices; slices contract in a fixed order, threaded, to the same bits everywhere. Single amplitudes or open-output batches | | `quantum_compile` | Clifford routing + stabilizer-tableau equivalence proof | | `quantum_atom` | neutral-atom register preparation (Hungarian / LSAP) | | `quantum_qir` | **QIR export + ingest** — emit and parse QIR (the QIR Alliance's LLVM-based IR); emitted modules validated with `llvm-as`, and 9 of 10 catalog algorithms re-ingest to a bit-identical statevector | @@ -175,6 +176,29 @@ fitted code model comes out at `a ≈ 0.023`, `p* ≈ 0.0100`. That is a measure value for the constants a resource estimate (`quantum_resource`) otherwise takes on trust, and it can be fed straight into one. +## Tensor networks + +An amplitude of a circuit is a tensor network, and the order it is contracted +in decides whether it costs `2³⁰` operations or `2⁶⁰`. `quantum_tn` searches for +that order, slices it to fit memory, and contracts it: + +```sh +wai-quantum tn --file circuit.jsonl --bits 0110… # one amplitude +wai-quantum tn --file circuit.jsonl --open 0,5,9 # a batch over three open outputs +wai-quantum tn --file circuit.jsonl --plan-only --cap-log2 30 +``` + +Circuits are JSON lines, one gate per line: `{"q":[a,b],"m":[re,im, …]}`. + +On single amplitudes of random circuits on the 54-qubit advantage-experiment +grid, the search finds trees of `10¹⁰·⁵` operations at 10 cycles and `10¹⁸·⁸` at +20, in seconds. An independently developed hyper-optimiser, given the same +networks, lands within a factor of 2.5 of these at every depth, in either +direction. The 54-qubit, 10-cycle amplitude, contracted here in 16 slices, +agrees with an independent contraction to 12 significant figures. The module +documentation states what slicing costs on deep circuits and what the scalar +executor can and cannot contract. + ## Determinism Simulation is **byte-exact**: amplitudes are dyadic fixed-point, so a circuit diff --git a/wai/wai-quantum/examples/tn_path.rs b/wai/wai-quantum/examples/tn_path.rs new file mode 100644 index 000000000..8f83cb337 --- /dev/null +++ b/wai/wai-quantum/examples/tn_path.rs @@ -0,0 +1,63 @@ +//! Contraction trees for a circuit's amplitude, from the command line. +//! +//! Reads a circuit as JSON lines, one gate per line: +//! `{"q":[0,1],"m":[re,im, re,im, …]}` (row-major, rows = outputs, the first +//! qubit most significant). Writes the simplified network's shape (for any +//! other tool to search the same network) and the best plan found. +//! +//! ```text +//! cargo run --release --example tn_path --features quantum_tn -- CIRCUIT N TRIALS SEED [LOG2_CAP] [SHAPE_OUT] [--contract] +//! ``` +use wai_quantum::quantum_tn::{contract_plan, find_path, Families, Gate, Network, PathOptions}; + +fn main() { + let a: Vec = std::env::args().collect(); + let text = std::fs::read_to_string(&a[1]).expect("read circuit"); + let n: u32 = a[2].parse().unwrap(); + let trials: usize = a[3].parse().unwrap(); + let seed: u64 = a[4].parse().unwrap(); + let cap: Option = a.get(5).filter(|s| *s != "-").map(|s| 2f64.powf(s.parse().unwrap())); + let shape_out = a.get(6).filter(|s| *s != "-"); + let gates: Vec = text + .lines() + .filter(|l| !l.trim().is_empty()) + .map(|l| { + let v: serde_json::Value = serde_json::from_str(l).expect("json gate"); + let qubits: Vec = v["q"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect(); + let m: Vec = v["m"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect(); + let matrix = m.chunks(2).map(|c| wai_quantum::quantum_tn::complex(c[0], c[1])).collect(); + Gate { qubits, matrix } + }) + .collect(); + let mut net = Network::amplitude(n, &gates, &vec![0; n as usize], &[]).expect("network"); + let raw = net.tensors.len(); + net.simplify(); + let shape = net.shape(); + if let Some(path) = shape_out { + let inputs: Vec = shape.iter().map(|s| format!("{s:?}")).collect(); + let dims: Vec = net.dims.iter().map(|d| d.to_string()).collect(); + std::fs::write(path, format!("{{\"inputs\":[{}],\"output\":{:?},\"dims\":[{}]}}\n", inputs.join(","), net.open, dims.join(","))).unwrap(); + } + let t0 = std::time::Instant::now(); + let families = match std::env::var("TN_FAMILIES").as_deref() { + Ok("greedy") => Families::Greedy, + Ok("partition") => Families::Partition, + _ => Families::Both, + }; + let reconfigure = std::env::var("TN_RECONF").ok().and_then(|v| v.parse().ok()).unwrap_or(8); + let plan = find_path(&shape, &net.dims, &net.open, &PathOptions { trials, seed, max_size: cap, reconfigure, families }); + let dt = t0.elapsed().as_secs_f64(); + println!( + "tensors {raw} -> {} log10(flops) {:.3} log2(max size) {:.1} slices {} search {:.1}s", + shape.len(), + plan.total_flops().log10(), + plan.per_slice.max_size.log2(), + plan.slices, + dt + ); + if a.iter().any(|x| x == "--contract") { + let t0 = std::time::Instant::now(); + let t = contract_plan(&net, &plan).expect("contract"); + println!("amplitude {:.12e} {:+.12e}i contract {:.1}s", t.data[0].re, t.data[0].im, t0.elapsed().as_secs_f64()); + } +} diff --git a/wai/wai-quantum/src/lib.rs b/wai/wai-quantum/src/lib.rs index e2628802f..b413b264b 100644 --- a/wai/wai-quantum/src/lib.rs +++ b/wai/wai-quantum/src/lib.rs @@ -17,7 +17,7 @@ #[cfg(feature = "_signing")] mod merkle; -#[cfg(any(feature = "quantum_mps", feature = "quantum_bptn"))] +#[cfg(any(feature = "quantum_mps", feature = "quantum_bptn", feature = "quantum_tn"))] mod linalg; #[cfg(any( @@ -28,7 +28,8 @@ mod linalg; feature = "quantum_phasor", feature = "quantum_kernel", feature = "quantum_resource", - feature = "quantum_frame" + feature = "quantum_frame", + feature = "quantum_tn" ))] mod repro; @@ -98,6 +99,9 @@ pub mod quantum_frame; #[cfg(feature = "quantum_frame")] pub mod quantum_match; +#[cfg(feature = "quantum_tn")] +pub mod quantum_tn; + #[cfg(feature = "quantum_info")] pub mod quantum_info; diff --git a/wai/wai-quantum/src/linalg.rs b/wai/wai-quantum/src/linalg.rs index 48e2f95a7..235c3cc86 100644 --- a/wai/wai-quantum/src/linalg.rs +++ b/wai/wai-quantum/src/linalg.rs @@ -13,6 +13,9 @@ pub struct C { pub re: f64, pub im: f64, } +// The tensor-network layer uses only the type and its constants; the +// arithmetic helpers serve the decomposition backends. +#[cfg_attr(not(any(feature = "quantum_mps", feature = "quantum_bptn")), allow(dead_code))] impl C { pub(crate) const ZERO: C = C { re: 0.0, im: 0.0 }; pub(crate) const ONE: C = C { re: 1.0, im: 0.0 }; diff --git a/wai/wai-quantum/src/quantum_tn.rs b/wai/wai-quantum/src/quantum_tn.rs new file mode 100644 index 000000000..f751b3811 --- /dev/null +++ b/wai/wai-quantum/src/quantum_tn.rs @@ -0,0 +1,1676 @@ +//! Tensor-network contraction — `wai.quantum.tn`. +//! +//! A circuit's amplitude `⟨x|U|0⟩` is a tensor network: one tensor per gate, +//! one per input and output basis vector, wired along the qubits' timelines. +//! Contracting it in a good order costs exponentially less than the state +//! vector when the circuit's entanglement is structured. That is why sliced +//! tensor-network contraction is the method behind most classical +//! reproductions of quantum-advantage experiments. The cost of an order is +//! everything: the same network can need `2³⁰` operations or `2⁶⁰`, depending +//! only on the order chosen. +//! +//! This module is that pipeline, in pure Rust, reproducible bit for bit. +//! +//! - **Build.** [`Network::amplitude`] turns gates (any unitary on any qubits, +//! given as a matrix) and an output bitstring into a network. Any subset of +//! the outputs can be left open, giving a batch of amplitudes at once. +//! - **Simplify.** Vectors and one-qubit gates are absorbed into their +//! neighbours before any search. They never cost more than the tensor they +//! join. +//! - **Order.** [`find_path`] runs a seeded hyper-optimiser over two +//! families of contraction trees: +//! - randomised greedy, scoring a pair by +//! `size(out)/c − c·(size(a) + size(b))` with Gumbel noise; +//! - recursive bipartition of the tensor graph, multilevel. It coarsens by +//! heavy-edge matching, cuts the coarsest graph several ways, and refines +//! at every level on the way back, with greedy leaves. +//! +//! Trials alternate families until one leads, then favour it. In the second +//! half, half the trials sample near the best parameters found so far. The +//! best few trees are then reconfigured: subtrees of up to eight tensors are +//! re-contracted optimally, by dynamic programming over subsets, until +//! nothing improves. +//! - **Slice.** [`slice`] fixes, one at a time, the index that leaves the least +//! total work, until the largest intermediate fits the cap. Each candidate is +//! priced incrementally. A second search round then looks for a tree for the +//! network with the first round's slices fixed (all of them, or the more +//! valuable half), so the tree routes around the sliced bonds. Each slice is +//! an independent contraction of the same tree, and the results are summed in +//! order. +//! - **Contract.** Every pairwise contraction is a permutation plus a complex +//! matrix product in a fixed accumulation order, using `+ − ×` only. The +//! result is the same bits on every machine, threaded or not. +//! +//! Contraction trees here are scored the way the field scores them: complex +//! multiply-adds summed over every contraction, and the largest intermediate. +//! So a path found here is comparable with any other tool's on the same +//! network. +//! +//! # What it finds +//! +//! The networks are single amplitudes of random circuits on the 54-qubit +//! advantage-experiment grid: `m` cycles of random √X/√Y/√W gates and +//! `fSim(π/2, π/6)` on the staggered coupler pattern, simplified to 174–395 +//! tensors. Costs are log₁₀ of operations, unsliced, with 256 trials, each +//! search taking under ten seconds on one machine: +//! +//! | cycles `m` | 10 | 12 | 14 | 16 | 20 | +//! |---|---|---|---|---|---| +//! | log₁₀ operations | 10.5 | 13.9 | 14.6 | 17.6 | 18.8 | +//! +//! An independently developed hyper-optimiser, given the same networks, +//! finds trees whose costs differ from these by at most 0.4 in log₁₀ (a +//! factor of 2.5) at every depth, lower at some depths and higher at others. +//! +//! The contraction is checked end to end at full size. The 54-qubit, 10-cycle +//! amplitude computed here, sliced 16 ways, agrees with an independent +//! contraction of the same circuit to 12 significant figures. +//! +//! # Honest boundaries +//! +//! - **Tight caps cost extra work on deep circuits.** At `m = 10` slicing is +//! essentially free: the same cost at a cap 16 times smaller than the +//! unsliced intermediate. Far below the natural size it is not: at `m = 14` +//! with a `2³⁰` cap, a plan costs about twelve times its unsliced cost. +//! - **The executor is scalar.** It runs about a billion complex multiply-adds +//! per second per core, threaded over slices as memory allows, with no SIMD +//! kernels or GPU. +//! Plans worth `10¹³` operations or more are for planning and verification +//! arithmetic here, not for contraction. +//! - **Networks are simple graphs.** Every index joins at most two tensors. +//! Diagonal-gate hyperedges are not exploited. + +pub use crate::linalg::C; +use crate::repro::ln; +use std::collections::BTreeMap; + +/// A complex number for building gates. +pub fn complex(re: f64, im: f64) -> C { + C { re, im } +} + +/// A gate: a unitary on `qubits`, as a `2ᵏ × 2ᵏ` row-major matrix whose rows +/// are outputs and columns inputs, with `qubits[0]` the most significant bit. +#[derive(Clone, Debug)] +pub struct Gate { + pub qubits: Vec, + pub matrix: Vec, +} + +/// A dense tensor: row-major over `inds`, the first index most significant. +#[derive(Clone, Debug)] +pub struct Tensor { + pub inds: Vec, + pub data: Vec, +} + +/// A tensor network: tensors, the dimension of every index, and the indices +/// left open (in the order the result is laid out). +#[derive(Clone, Debug)] +pub struct Network { + pub tensors: Vec, + pub dims: Vec, + pub open: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum TnError { + /// A gate's matrix does not have `4^k` entries, repeats a qubit, or acts + /// outside the register. + BadGate(usize), + /// A bitstring or open-qubit list that does not fit the register. + BadOutput, + /// The plan's network does not match the network it is applied to. + PlanMismatch, +} + +impl Network { + /// The network for `⟨bits|U|0…0⟩` on `n` qubits, with the qubits in + /// `open_qubits` left open instead of projected: the result is then the + /// `2^|open|` amplitudes over those qubits (first listed most + /// significant), with the other outputs fixed to `bits`. + pub fn amplitude(n: u32, gates: &[Gate], bits: &[u8], open_qubits: &[u32]) -> Result { + if bits.len() != n as usize || bits.iter().any(|&b| b > 1) || open_qubits.iter().any(|&q| q >= n) { + return Err(TnError::BadOutput); + } + let mut next = 0u32; + let mut fresh = || { + next += 1; + next - 1 + }; + let mut tensors = Vec::new(); + let mut wire: Vec = (0..n).map(|_| fresh()).collect(); + for &w in &wire { + tensors.push(Tensor { inds: vec![w], data: vec![C::ONE, C::ZERO] }); + } + for (gi, g) in gates.iter().enumerate() { + let k = g.qubits.len(); + let mut seen = g.qubits.clone(); + seen.sort_unstable(); + seen.dedup(); + if k == 0 || seen.len() != k || g.qubits.iter().any(|&q| q >= n) || g.matrix.len() != 1 << (2 * k) { + return Err(TnError::BadGate(gi)); + } + let outs: Vec = (0..k).map(|_| fresh()).collect(); + let mut inds = outs.clone(); + inds.extend(g.qubits.iter().map(|&q| wire[q as usize])); + tensors.push(Tensor { inds, data: g.matrix.clone() }); + for (i, &q) in g.qubits.iter().enumerate() { + wire[q as usize] = outs[i]; + } + } + let mut open = Vec::new(); + for q in 0..n { + if !open_qubits.contains(&q) { + let v = if bits[q as usize] == 0 { vec![C::ONE, C::ZERO] } else { vec![C::ZERO, C::ONE] }; + tensors.push(Tensor { inds: vec![wire[q as usize]], data: v }); + } + } + for &q in open_qubits { + open.push(wire[q as usize]); + } + Ok(Network { tensors, dims: vec![2; next as usize], open }) + } + + /// Absorb every tensor of rank at most two into a neighbour, while that + /// never grows the neighbour, until none is left (or only one tensor + /// remains). Deterministic: always the lowest-numbered candidate first. + pub fn simplify(&mut self) { + loop { + let owners = self.owners(); + let mut done = true; + for t in 0..self.tensors.len() { + let rank = self.tensors[t].inds.len(); + if rank > 2 || self.tensors.len() == 1 { + continue; + } + // The neighbour sharing an index with t, whose rank does not grow. + let mut target = None; + for &i in &self.tensors[t].inds { + if let Some(&u) = owners.get(&i).and_then(|o| o.iter().find(|&&u| u != t)) { + let merged = sym_diff(&self.tensors[t].inds, &self.tensors[u].inds, &self.open).len(); + if merged <= self.tensors[u].inds.len() { + target = Some(u); + break; + } + } + } + if let Some(u) = target { + let (a, b) = if t < u { (t, u) } else { (u, t) }; + let tb = self.tensors.remove(b); + let ta = self.tensors.remove(a); + // Keep the larger tensor's index order in front. + let merged = if ta.inds.len() >= tb.inds.len() { + contract(&ta, &tb, &self.dims, &self.open) + } else { + contract(&tb, &ta, &self.dims, &self.open) + }; + self.tensors.insert(a, merged); + done = false; + break; + } + } + if done { + return; + } + } + } + + fn owners(&self) -> BTreeMap> { + let mut o: BTreeMap> = BTreeMap::new(); + for (t, x) in self.tensors.iter().enumerate() { + for &i in &x.inds { + o.entry(i).or_default().push(t); + } + } + o + } + + /// The structure a path is found for: each tensor's index set. + pub fn shape(&self) -> Vec> { + self.tensors.iter().map(|t| t.inds.clone()).collect() + } +} + +/// Indices in exactly one of `a`, `b`, plus open indices in either. +fn sym_diff(a: &[u32], b: &[u32], open: &[u32]) -> Vec { + let mut out: Vec = a.iter().copied().filter(|i| !b.contains(i) || open.contains(i)).collect(); + out.extend(b.iter().copied().filter(|i| !a.contains(i))); + out +} + +/// Permute a tensor's data so its indices come in `order`. +fn permute(t: &Tensor, order: &[u32], dims: &[usize]) -> Vec { + if t.inds == order { + return t.data.clone(); + } + let r = t.inds.len(); + // Strides of the source, by source position. + let mut src_stride = vec![1usize; r]; + for p in (0..r.saturating_sub(1)).rev() { + src_stride[p] = src_stride[p + 1] * dims[t.inds[p + 1] as usize]; + } + let pos: Vec = order.iter().map(|i| t.inds.iter().position(|j| j == i).unwrap()).collect(); + let out_dims: Vec = order.iter().map(|&i| dims[i as usize]).collect(); + let total = t.data.len(); + let mut out = Vec::with_capacity(total); + let mut idx = vec![0usize; r]; + let mut src = 0usize; + for _ in 0..total { + out.push(t.data[src]); + // Odometer over the output dimensions, tracking the source offset. + for p in (0..r).rev() { + idx[p] += 1; + src += src_stride[pos[p]]; + if idx[p] < out_dims[p] { + break; + } + src -= src_stride[pos[p]] * out_dims[p]; + idx[p] = 0; + } + } + out +} + +/// Contract two tensors over their shared (non-open) indices. The result's +/// indices are `a`'s remaining ones, then `b`'s. +pub fn contract(a: &Tensor, b: &Tensor, dims: &[usize], open: &[u32]) -> Tensor { + let shared: Vec = a.inds.iter().copied().filter(|i| b.inds.contains(i) && !open.contains(i)).collect(); + let a_rem: Vec = a.inds.iter().copied().filter(|i| !shared.contains(i)).collect(); + let b_rem: Vec = b.inds.iter().copied().filter(|i| !shared.contains(i)).collect(); + let size = |v: &[u32]| v.iter().map(|&i| dims[i as usize]).product::(); + let (m, k, n) = (size(&a_rem), size(&shared), size(&b_rem)); + let mut ao = a_rem.clone(); + ao.extend(&shared); + let mut bo = shared.clone(); + bo.extend(&b_rem); + let ad = permute(a, &ao, dims); + let bd = permute(b, &bo, dims); + let mut out = vec![C::ZERO; m * n]; + for i in 0..m { + let row = &mut out[i * n..(i + 1) * n]; + for kk in 0..k { + let x = ad[i * k + kk]; + if x.re == 0.0 && x.im == 0.0 { + continue; + } + let brow = &bd[kk * n..(kk + 1) * n]; + for (o, y) in row.iter_mut().zip(brow) { + o.re += x.re * y.re - x.im * y.im; + o.im += x.re * y.im + x.im * y.re; + } + } + } + let mut inds = a_rem; + inds.extend(b_rem); + Tensor { inds, data: out } +} + +// --------------------------------------------------------------------------- +// Contraction paths +// --------------------------------------------------------------------------- + +/// A contraction order in single-static-assignment form: step `s` contracts +/// two live tensors into tensor `n + s`, where `n` is the number of inputs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Path { + pub steps: Vec<(usize, usize)>, +} + +/// What a path costs on a network shape. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Cost { + /// Complex multiply-adds, summed over every contraction. + pub flops: f64, + /// Elements of the largest intermediate tensor. + pub max_size: f64, +} + +impl Cost { + pub fn log10_flops(&self) -> f64 { + ln(self.flops) / core::f64::consts::LN_10 + } + pub fn log2_size(&self) -> f64 { + ln(self.max_size) / core::f64::consts::LN_2 + } +} + +fn set_size(s: &[u32], dims: &[usize]) -> f64 { + s.iter().map(|&i| dims[i as usize] as f64).product() +} + +/// Union of two index sets (a's order, then b's extras). +fn union(a: &[u32], b: &[u32]) -> Vec { + let mut u = a.to_vec(); + u.extend(b.iter().copied().filter(|i| !a.contains(i))); + u +} + +/// The cost of `path` on a network of index sets `shape`. +pub fn path_cost(shape: &[Vec], dims: &[usize], open: &[u32], path: &Path) -> Cost { + let mut live: Vec> = shape.to_vec(); + let mut flops = 0.0; + let mut max_size: f64 = shape.iter().map(|s| set_size(s, dims)).fold(0.0, f64::max); + for &(a, b) in &path.steps { + flops += set_size(&union(&live[a], &live[b]), dims); + let out = sym_diff(&live[a], &live[b], open); + max_size = max_size.max(set_size(&out, dims)); + live.push(out); + } + Cost { flops, max_size } +} + +/// A deterministic random stream (splitmix64). +pub(crate) struct Rng(u64); + +impl Rng { + pub(crate) fn new(seed: u64) -> Rng { + Rng(seed) + } + fn next(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ (z >> 31) + } + /// Uniform on `(0, 1)`. + fn unit(&mut self) -> f64 { + ((self.next() >> 11) as f64 + 0.5) * (1.0 / 9_007_199_254_740_992.0) + } + fn below(&mut self, n: usize) -> usize { + (self.next() % n as u64) as usize + } +} + +/// Randomised greedy: repeatedly contract the pair sharing an index whose +/// result grows memory least, `size(out)/costmod − costmod·(size(a) + size(b))`, +/// perturbed by Gumbel noise of the given temperature (0 = plain greedy). A +/// large `costmod` favours contractions that retire big tensors, which is +/// where good paths for circuit networks tend to be. +/// Disconnected parts are joined smallest first at the end. +pub fn greedy_path(shape: &[Vec], dims: &[usize], open: &[u32], costmod: f64, temperature: f64, seed: u64) -> Path { + let mut rng = Rng::new(seed); + let mut live: Vec>> = shape.iter().cloned().map(Some).collect(); + let mut steps = Vec::new(); + let ids: Vec = (0..shape.len()).collect(); + greedy_steps(&ids, &mut live, &mut steps, dims, open, costmod, temperature, &mut rng); + Path { steps } +} + +/// Greedy over a subset of live tensors, appending its steps; returns the id +/// of the subset's result. Indices leaving the subset stay on the result. +#[allow(clippy::too_many_arguments)] +fn greedy_steps( + ids: &[usize], + live: &mut Vec>>, + steps: &mut Vec<(usize, usize)>, + dims: &[usize], + open: &[u32], + costmod: f64, + temperature: f64, + rng: &mut Rng, +) -> usize { + if ids.len() == 1 { + return ids[0]; + } + let mut owners: BTreeMap> = BTreeMap::new(); + for &t in ids { + for &i in live[t].as_ref().unwrap() { + owners.entry(i).or_default().push(t); + } + } + let score = |a: &[u32], b: &[u32], rng: &mut Rng| -> f64 { + let out = set_size(&sym_diff(a, b, open), dims); + let c = out / costmod - costmod * (set_size(a, dims) + set_size(b, dims)); + let c = if c >= 0.0 { ln(1.0 + c) } else { -ln(1.0 - c) }; + if temperature > 0.0 { c - temperature * -ln(-ln(rng.unit())) } else { c } + }; + // Candidates as (score, a, b), best = smallest; ties broken by ids. + let mut cands: Vec<(f64, usize, usize)> = Vec::new(); + let mut seen = std::collections::BTreeSet::new(); + for ts in owners.values() { + for x in 0..ts.len() { + for y in x + 1..ts.len() { + let (a, b) = (ts[x].min(ts[y]), ts[x].max(ts[y])); + if seen.insert((a, b)) { + let sc = score(live[a].as_ref().unwrap(), live[b].as_ref().unwrap(), rng); + cands.push((sc, a, b)); + } + } + } + } + let mut members: Vec = ids.to_vec(); + loop { + let mut best: Option = None; + for (ci, c) in cands.iter().enumerate() { + if best.is_none_or(|bi| { + let b = cands[bi]; + c.0 < b.0 || (c.0 == b.0 && (c.1, c.2) < (b.1, b.2)) + }) { + best = Some(ci); + } + } + let Some(bi) = best else { break }; + let (_, a, b) = cands.swap_remove(bi); + cands.retain(|c| c.1 != a && c.1 != b && c.2 != a && c.2 != b); + let sa = live[a].take().unwrap(); + let sb = live[b].take().unwrap(); + let out = sym_diff(&sa, &sb, open); + let id = live.len(); + steps.push((a, b)); + members.retain(|&t| t != a && t != b); + for &i in &out { + if let Some(o) = owners.get_mut(&i) { + o.retain(|&t| t != a && t != b); + for &t in o.iter() { + let sc = score(live[t].as_ref().unwrap(), &out, rng); + cands.push((sc, t.min(id), t.max(id))); + } + o.push(id); + } + } + live.push(Some(out)); + members.push(id); + } + // Join what is left (disconnected parts), smallest first. + while members.len() > 1 { + members.sort_by(|&x, &y| { + set_size(live[x].as_ref().unwrap(), dims) + .partial_cmp(&set_size(live[y].as_ref().unwrap(), dims)) + .unwrap() + .then(x.cmp(&y)) + }); + let (a, b) = (members[0], members[1]); + let sa = live[a].take().unwrap(); + let sb = live[b].take().unwrap(); + steps.push((a, b)); + let id = live.len(); + live.push(Some(sym_diff(&sa, &sb, open))); + members.drain(0..2); + members.push(id); + } + members[0] +} + +/// Parameters of one recursive-bipartition tree. +#[derive(Clone, Copy, Debug, PartialEq)] +struct PartitionParams { + /// Allowed deviation of a part from half, as a fraction. + imbalance: f64, + /// Parts at most this large are finished by greedy. + cutoff: usize, + /// Refinement passes per bisection. + passes: usize, + costmod: f64, + temperature: f64, +} + +/// A contraction tree by recursive balanced bipartition of the tensor graph. +fn partition_path(shape: &[Vec], dims: &[usize], open: &[u32], p: PartitionParams, seed: u64) -> Path { + let mut rng = Rng::new(seed); + let mut live: Vec>> = shape.iter().cloned().map(Some).collect(); + let mut steps = Vec::new(); + let ids: Vec = (0..shape.len()).collect(); + partition_rec(&ids, &mut live, &mut steps, dims, open, &p, &mut rng); + Path { steps } +} + +fn partition_rec( + ids: &[usize], + live: &mut Vec>>, + steps: &mut Vec<(usize, usize)>, + dims: &[usize], + open: &[u32], + p: &PartitionParams, + rng: &mut Rng, +) -> usize { + if ids.len() <= p.cutoff.max(2) { + return greedy_steps(ids, live, steps, dims, open, p.costmod, p.temperature, rng); + } + let (left, right) = bisect(ids, live, dims, p, rng); + let a = partition_rec(&left, live, steps, dims, open, p, rng); + let b = partition_rec(&right, live, steps, dims, open, p, rng); + let sa = live[a].take().unwrap(); + let sb = live[b].take().unwrap(); + steps.push((a, b)); + live.push(Some(sym_diff(&sa, &sb, open))); + live.len() - 1 +} + +/// A weighted graph for partitioning: node weights and merged edge weights. +struct WGraph { + nw: Vec, + adj: Vec>, +} + +impl WGraph { + fn n(&self) -> usize { + self.nw.len() + } +} + +/// Merge parallel edges by summing their weights; neighbours in ascending order. +fn merged(adj: Vec>) -> Vec> { + adj.into_iter() + .map(|mut v| { + v.sort_by_key(|a| a.0); + let mut out: Vec<(usize, f64)> = Vec::with_capacity(v.len()); + for (u, w) in v { + match out.last_mut() { + Some(last) if last.0 == u => last.1 += w, + _ => out.push((u, w)), + } + } + out + }) + .collect() +} + +/// One level of coarsening by heavy-edge matching: nodes in random order, +/// each unmatched node joins its unmatched neighbour with the heaviest edge. +/// Returns the coarse graph and each fine node's coarse node. +fn coarsen(g: &WGraph, rng: &mut Rng) -> (WGraph, Vec) { + let n = g.n(); + let mut order: Vec = (0..n).collect(); + for i in (1..n).rev() { + order.swap(i, rng.below(i + 1)); + } + let mut mate = vec![usize::MAX; n]; + for &v in &order { + if mate[v] != usize::MAX { + continue; + } + let mut best: Option<(f64, usize)> = None; + for &(u, w) in &g.adj[v] { + if u != v && mate[u] == usize::MAX && best.is_none_or(|(bw, _)| w > bw) { + best = Some((w, u)); + } + } + match best { + Some((_, u)) => { + mate[v] = u; + mate[u] = v; + } + None => mate[v] = v, + } + } + let mut map = vec![usize::MAX; n]; + let mut nw = Vec::new(); + for v in 0..n { + if map[v] == usize::MAX { + let c = nw.len(); + map[v] = c; + let u = mate[v]; + let mut w = g.nw[v]; + if u != v { + map[u] = c; + w += g.nw[u]; + } + nw.push(w); + } + } + let mut adj = vec![Vec::new(); nw.len()]; + for v in 0..n { + for &(u, w) in &g.adj[v] { + if map[v] != map[u] { + adj[map[v]].push((map[u], w)); + } + } + } + (WGraph { nw, adj: merged(adj) }, map) +} + +/// Gain-ordered single moves that keep the side's weight within `[lo, hi]`, +/// keeping the best cut seen in each pass. +fn refine(g: &WGraph, side: &mut Vec, lo: f64, hi: f64, passes: usize) { + let n = g.n(); + let cut = |s: &[bool]| -> f64 { + let mut c = 0.0; + for v in 0..n { + for &(u, w) in &g.adj[v] { + if v < u && s[v] != s[u] { + c += w; + } + } + } + c + }; + for _ in 0..passes { + let mut locked = vec![false; n]; + let mut cur = side.clone(); + let mut weight: f64 = (0..n).filter(|&v| cur[v]).map(|v| g.nw[v]).sum(); + let mut cur_cut = cut(&cur); + let mut best_cut = cur_cut; + let mut best = cur.clone(); + let mut gain: Vec = (0..n) + .map(|v| g.adj[v].iter().map(|&(u, w)| if cur[u] == cur[v] { -w } else { w }).sum()) + .collect(); + for _ in 0..n { + let mut pick: Option<(f64, usize)> = None; + for v in 0..n { + if locked[v] { + continue; + } + let nwgt = if cur[v] { weight - g.nw[v] } else { weight + g.nw[v] }; + if nwgt < lo || nwgt > hi { + continue; + } + if pick.is_none_or(|(gp, _)| gain[v] > gp) { + pick = Some((gain[v], v)); + } + } + let Some((gv, v)) = pick else { break }; + locked[v] = true; + weight = if cur[v] { weight - g.nw[v] } else { weight + g.nw[v] }; + cur[v] = !cur[v]; + cur_cut -= gv; + gain[v] = -gv; + for &(u, w) in &g.adj[v] { + // u's edge to v flipped between cut and uncut. + gain[u] += if cur[u] == cur[v] { -2.0 * w } else { 2.0 * w }; + } + if cur_cut < best_cut - 1e-9 { + best_cut = cur_cut; + best = cur.clone(); + } + } + if best == *side { + break; + } + *side = best; + } +} + +/// Grow one side breadth-first from a random seed until it weighs `target`. +fn grow(g: &WGraph, target: f64, rng: &mut Rng) -> Vec { + let n = g.n(); + let mut side = vec![false; n]; + let mut weight = 0.0; + let mut queue = std::collections::VecDeque::new(); + let mut start = rng.below(n); + while weight < target { + if queue.is_empty() { + let mut tries = 0; + while side[start] && tries < n { + start = (start + 1) % n; + tries += 1; + } + if side[start] { + break; + } + side[start] = true; + weight += g.nw[start]; + queue.push_back(start); + continue; + } + let v = queue.pop_front().unwrap(); + for &(u, _) in &g.adj[v] { + if weight < target && !side[u] { + side[u] = true; + weight += g.nw[u]; + queue.push_back(u); + } + } + } + side +} + +/// Split `ids` in two with a small weighted cut, multilevel: coarsen by +/// heavy-edge matching, partition the coarsest graph several ways (grown +/// breadth-first, refined), keep the best, then project it back up, +/// refining at every level. The two sides weigh within the imbalance of half. +fn bisect(ids: &[usize], live: &[Option>], dims: &[usize], p: &PartitionParams, rng: &mut Rng) -> (Vec, Vec) { + let n = ids.len(); + let mut by_index: BTreeMap> = BTreeMap::new(); + for (k, &t) in ids.iter().enumerate() { + for &i in live[t].as_ref().unwrap() { + by_index.entry(i).or_default().push(k); + } + } + let mut adj: Vec> = vec![Vec::new(); n]; + for (&i, ks) in &by_index { + if ks.len() == 2 { + let w = ln(dims[i as usize] as f64) / core::f64::consts::LN_2; + adj[ks[0]].push((ks[1], w)); + adj[ks[1]].push((ks[0], w)); + } + } + let mut levels = vec![WGraph { nw: vec![1.0; n], adj: merged(adj) }]; + let mut maps: Vec> = Vec::new(); + while levels.last().unwrap().n() > 24 { + let (coarse, map) = coarsen(levels.last().unwrap(), rng); + if coarse.n() as f64 > 0.95 * levels.last().unwrap().n() as f64 { + break; + } + levels.push(coarse); + maps.push(map); + } + let total = n as f64; + let lo = (total * (0.5 - p.imbalance / 2.0)).max(1.0); + let hi = (total * (0.5 + p.imbalance / 2.0)).min(total - 1.0); + let coarsest = levels.last().unwrap(); + let mut side: Vec = Vec::new(); + let mut best_cut = f64::INFINITY; + for _ in 0..6 { + let target = lo + (hi - lo) * rng.unit(); + let mut s = grow(coarsest, target, rng); + refine(coarsest, &mut s, lo, hi, p.passes); + let w: f64 = (0..coarsest.n()).filter(|&v| s[v]).map(|v| coarsest.nw[v]).sum(); + let c: f64 = (0..coarsest.n()) + .flat_map(|v| coarsest.adj[v].iter().filter(move |&&(u, _)| v < u).map(move |&(u, w)| (v, u, w))) + .filter(|&(v, u, _)| s[v] != s[u]) + .map(|t| t.2) + .sum(); + let fits = w >= lo && w <= hi; + if fits && c < best_cut { + best_cut = c; + side = s; + } else if side.is_empty() { + side = s; + } + } + for lvl in (0..maps.len()).rev() { + let map = &maps[lvl]; + let mut fine: Vec = map.iter().map(|&c| side[c]).collect(); + refine(&levels[lvl], &mut fine, lo, hi, p.passes); + side = fine; + } + // A degenerate split (everything on one side) falls back to halves. + if side.iter().all(|&x| x) || side.iter().all(|&x| !x) { + side = (0..n).map(|k| k < n / 2).collect(); + } + let left: Vec = (0..n).filter(|&k| side[k]).map(|k| ids[k]).collect(); + let right: Vec = (0..n).filter(|&k| !side[k]).map(|k| ids[k]).collect(); + (left, right) +} + +// --------------------------------------------------------------------------- +// Subtree reconfiguration +// --------------------------------------------------------------------------- + +/// Re-contract small subtrees optimally: for each internal node, expand its +/// subtree into a frontier of at most `k` tensors (always splitting the +/// costliest), then find the cheapest order of that frontier by dynamic +/// programming over subsets. Never makes the path costlier. +pub fn reconfigure(shape: &[Vec], dims: &[usize], open: &[u32], path: &Path, k: usize) -> Path { + let n = shape.len(); + // Tree: node -> children; node sets. + let mut sets: Vec> = shape.to_vec(); + let mut kids: Vec> = vec![None; n]; + for &(a, b) in &path.steps { + sets.push(sym_diff(&sets[a], &sets[b], open)); + kids.push(Some((a, b))); + } + if kids.len() <= n { + return path.clone(); + } + let flops = |a: &[u32], b: &[u32]| set_size(&union(a, b), dims); + let root = kids.len() - 1; + // Visit internal nodes top-down; the tree is rewritten in place. + let mut order = vec![root]; + let mut i = 0; + while i < order.len() { + if let Some((a, b)) = kids[order[i]] { + order.push(a); + order.push(b); + } + i += 1; + } + for &v in &order { + if kids[v].is_none() { + continue; + } + // Frontier: split the internal node with the costliest contraction. + let mut frontier = vec![v]; + loop { + let mut pick: Option<(f64, usize)> = None; + for (fi, &f) in frontier.iter().enumerate() { + if let Some((a, b)) = kids[f] { + let c = flops(&sets[a], &sets[b]); + if pick.is_none_or(|(pc, _)| c > pc) { + pick = Some((c, fi)); + } + } + } + let Some((_, fi)) = pick else { break }; + if frontier.len() + 1 > k { + break; + } + let f = frontier.swap_remove(fi); + let (a, b) = kids[f].unwrap(); + frontier.push(a); + frontier.push(b); + } + let m = frontier.len(); + if m < 3 { + continue; + } + // Current cost of the subtree above the frontier. + let mut current = 0.0; + let mut stack = vec![v]; + while let Some(x) = stack.pop() { + if frontier.contains(&x) { + continue; + } + if let Some((a, b)) = kids[x] { + current += flops(&sets[a], &sets[b]); + stack.push(a); + stack.push(b); + } + } + // DP over subsets of the frontier. + let full = (1usize << m) - 1; + let mut out: Vec> = vec![Vec::new(); full + 1]; + let mut cost = vec![f64::INFINITY; full + 1]; + let mut split = vec![0usize; full + 1]; + for j in 0..m { + out[1 << j] = sets[frontier[j]].clone(); + cost[1 << j] = 0.0; + } + for sub in 1..=full { + if sub.count_ones() < 2 { + continue; + } + // Its output: indices appearing once among its members, or open. + let low = sub.isolate_lowest_one(); + let rest = sub ^ low; + out[sub] = sym_diff(&out[low], &out[rest], open); + // Enumerate splits containing the lowest member on the left. + let mut a = (sub - 1) & sub; + while a > 0 { + if a & low != 0 { + let b = sub ^ a; + let c = cost[a] + cost[b] + flops(&out[a], &out[b]); + if c < cost[sub] { + cost[sub] = c; + split[sub] = a; + } + } + a = (a - 1) & sub; + } + } + if cost[full] >= current * (1.0 - 1e-12) { + continue; + } + // Rebuild the subtree from the DP splits, reusing v as its root. + fn build( + sub: usize, + root: Option, + frontier: &[usize], + split: &[usize], + out: &[Vec], + sets: &mut Vec>, + kids: &mut Vec>, + ) -> usize { + if sub.count_ones() == 1 { + return frontier[sub.trailing_zeros() as usize]; + } + let a = split[sub]; + let l = build(a, None, frontier, split, out, sets, kids); + let r = build(sub ^ a, None, frontier, split, out, sets, kids); + let id = match root { + Some(v) => v, + None => { + sets.push(Vec::new()); + kids.push(None); + sets.len() - 1 + } + }; + sets[id] = out[sub].clone(); + kids[id] = Some((l, r)); + id + } + build(full, Some(v), &frontier, &split, &out, &mut sets, &mut kids); + } + // Back to SSA order: post-order from the root, numbering as we go. + let mut ssa: Vec> = vec![None; kids.len()]; + for (t, slot) in ssa.iter_mut().enumerate().take(n) { + *slot = Some(t); + } + let mut steps = Vec::new(); + let mut next = n; + let mut stack = vec![(root, false)]; + while let Some((x, expanded)) = stack.pop() { + let Some((a, b)) = kids[x] else { continue }; + if expanded { + steps.push((ssa[a].unwrap(), ssa[b].unwrap())); + ssa[x] = Some(next); + next += 1; + } else { + stack.push((x, true)); + stack.push((b, false)); + stack.push((a, false)); + } + } + Path { steps } +} + +// --------------------------------------------------------------------------- +// Slicing and the hyper-optimiser +// --------------------------------------------------------------------------- + +/// Dimensions with the sliced indices fixed (dimension one). +fn sliced_dims(dims: &[usize], sliced: &[u32]) -> Vec { + let mut d = dims.to_vec(); + for &i in sliced { + d[i as usize] = 1; + } + d +} + +/// Fix indices, one at a time, until the largest intermediate holds at most +/// `max_size` elements: each time the index whose fixing leaves the least +/// total work (per-slice cost times the number of slices). Returns the sliced +/// indices. Each candidate is priced incrementally, touching only the +/// contractions that carry it. +pub fn slice(shape: &[Vec], dims: &[usize], open: &[u32], path: &Path, max_size: f64) -> Vec { + // Every step's flop set and output set, and every input's set. + let mut sets: Vec> = shape.to_vec(); + let mut flop_sets: Vec> = Vec::with_capacity(path.steps.len()); + for &(a, b) in &path.steps { + flop_sets.push(union(&sets[a], &sets[b])); + sets.push(sym_diff(&sets[a], &sets[b], open)); + } + let nd = dims.len(); + // Where each index appears: in flop sets, and in tensor sets (inputs and outputs). + let mut in_flops: Vec> = vec![Vec::new(); nd]; + for (s, f) in flop_sets.iter().enumerate() { + for &i in f { + in_flops[i as usize].push(s); + } + } + let mut in_sets: Vec> = vec![Vec::new(); nd]; + for (t, x) in sets.iter().enumerate() { + for &i in x { + in_sets[i as usize].push(t); + } + } + let mut flop_size: Vec = flop_sets.iter().map(|f| set_size(f, dims)).collect(); + let mut set_sz: Vec = sets.iter().map(|x| set_size(x, dims)).collect(); + let mut sliced: Vec = Vec::new(); + let mut slices = 1.0; + let mut is_sliced = vec![false; nd]; + loop { + let biggest = set_sz.iter().copied().fold(0.0, f64::max); + if biggest <= max_size { + return sliced; + } + let per: f64 = flop_size.iter().sum(); + // Candidates: indices of every tensor still over the cap, never open ones. + let mut cands: Vec = Vec::new(); + for (t, x) in sets.iter().enumerate() { + if set_sz[t] > max_size { + cands.extend(x.iter().copied().filter(|i| !is_sliced[*i as usize] && !open.contains(i) && dims[*i as usize] > 1)); + } + } + cands.sort_unstable(); + cands.dedup(); + let mut best: Option<(f64, u32)> = None; + for &i in &cands { + let d = dims[i as usize] as f64; + let saved: f64 = in_flops[i as usize].iter().map(|&s| flop_size[s]).sum::() * (1.0 - 1.0 / d); + let total = (per - saved) * slices * d; + if best.is_none_or(|(b, _)| total < b) { + best = Some((total, i)); + } + } + let Some((_, i)) = best else { return sliced }; + let d = dims[i as usize] as f64; + for &s in &in_flops[i as usize] { + flop_size[s] /= d; + } + for &t in &in_sets[i as usize] { + set_sz[t] /= d; + } + is_sliced[i as usize] = true; + slices *= d; + sliced.push(i); + } +} + +/// Settings for [`find_path`]. +#[derive(Clone, Copy, Debug)] +pub struct PathOptions { + /// Candidate trees to try. + pub trials: usize, + pub seed: u64, + /// Largest intermediate allowed (elements); slice to fit. + pub max_size: Option, + /// Frontier size for subtree reconfiguration (0 = none). + pub reconfigure: usize, + /// Which tree families to sample. + pub families: Families, +} + +/// The tree families the hyper-optimiser samples. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Families { + Both, + Greedy, + Partition, +} + +impl Default for PathOptions { + fn default() -> Self { + PathOptions { trials: 64, seed: 0, max_size: None, reconfigure: 8, families: Families::Both } + } +} + +/// A contraction plan: the tree, the sliced indices, and what it costs. +#[derive(Clone, Debug, PartialEq)] +pub struct Plan { + pub path: Path, + pub sliced: Vec, + /// Cost of one slice. + pub per_slice: Cost, + /// Number of slices. + pub slices: f64, +} + +impl Plan { + /// Total operations over every slice. + pub fn total_flops(&self) -> f64 { + self.per_slice.flops * self.slices + } +} + +/// Seeded hyper-optimisation: `trials` trees from both families with +/// randomised parameters, each sliced to the memory cap and scored by total +/// operations. The best is then reconfigured and re-sliced. +pub fn find_path(shape: &[Vec], dims: &[usize], open: &[u32], opts: &PathOptions) -> Plan { + let mut best = search(shape, dims, open, opts); + // With slices, search again on the network with some of the first + // round's slices fixed: a tree found for the sliced network routes around + // the sliced bonds instead of being patched after the fact. Try fixing all + // of them and the more valuable half (the slicer picks in order of value), + // letting the second round choose the rest; keep the cheaper. + if opts.max_size.is_some() && best.slices > 1.0 { + let first = best.sliced.clone(); + let mut fixings = vec![first.clone()]; + if first.len() > 1 { + fixings.push(first[..first.len().div_ceil(2)].to_vec()); + } + for (k, pre) in fixings.into_iter().enumerate() { + let fixed = sliced_dims(dims, &pre); + let again = PathOptions { seed: opts.seed ^ (0x5ee_d0f5_11ce + k as u64), trials: opts.trials.div_ceil(2), ..*opts }; + let mut second = search(shape, &fixed, open, &again); + let pre_slices: f64 = pre.iter().map(|&i| dims[i as usize] as f64).product(); + let mut sliced = pre; + sliced.extend(second.sliced.iter().copied()); + second.sliced = sliced; + second.slices *= pre_slices; + if second.total_flops() < best.total_flops() { + best = second; + } + } + } + best +} + +/// One round of the hyper-optimiser, on dimensions that may already have +/// indices fixed (dimension one). +fn search(shape: &[Vec], dims: &[usize], open: &[u32], opts: &PathOptions) -> Plan { + let mut rng = Rng::new(opts.seed); + let score = |path: &Path| -> Plan { + let sliced = match opts.max_size { + Some(m) => slice(shape, dims, open, path, m), + None => Vec::new(), + }; + let d = sliced_dims(dims, &sliced); + let per_slice = path_cost(shape, &d, open, path); + let slices = sliced.iter().map(|&i| dims[i as usize] as f64).product(); + Plan { path: path.clone(), sliced, per_slice, slices } + }; + const LN_1000: f64 = 6.907_755_278_982_137; + let trials = opts.trials.max(1); + // Best (total flops, parameters) per family: greedy (costmod, temperature), + // partition (imbalance, cutoff, costmod, temperature). + let mut best_greedy: Option<(f64, f64, f64)> = None; + let mut best_part: Option<(f64, PartitionParams)> = None; + // The best few plans, kept for reconfiguration. + let mut top: Vec = Vec::new(); + let keep = 4; + for t in 0..trials { + let seed = rng.next(); + let greedy = match opts.families { + Families::Greedy => true, + Families::Partition => false, + // Alternate while warming up, then favour the family doing better. + Families::Both => { + if t < 8 { + t % 2 == 0 + } else { + let greedy_ahead = match (best_greedy, best_part) { + (Some(g), Some(p)) => g.0 <= p.0, + (Some(_), None) => true, + _ => false, + }; + (rng.unit() < 0.75) == greedy_ahead + } + } + }; + // In the second half, half the trials explore around the best so far. + let local = t >= trials / 2 && rng.unit() < 0.5; + let path; + if greedy { + let (mut costmod, mut temperature) = + (0.1 + 3.9 * rng.unit(), if t == 0 { 0.0 } else { crate::repro::exp(-LN_1000 * rng.unit()) }); + if let (true, Some((_, c, tp))) = (local, best_greedy) { + costmod = (c + 0.6 * (rng.unit() - 0.5)).clamp(0.1, 4.0); + temperature = (tp.max(1e-3) * crate::repro::exp(2.0 * rng.unit() - 1.0)).min(1.0); + } + path = greedy_path(shape, dims, open, costmod, temperature, seed); + let plan = score(&path); + if best_greedy.is_none_or(|b| plan.total_flops() < b.0) { + best_greedy = Some((plan.total_flops(), costmod, temperature)); + } + push_top(&mut top, plan, keep); + } else { + let mut p = PartitionParams { + imbalance: 0.02 + 0.6 * rng.unit(), + cutoff: 4 + rng.below(12), + passes: 2 + rng.below(4), + costmod: 0.1 + 3.9 * rng.unit(), + temperature: crate::repro::exp(-LN_1000 * rng.unit()), + }; + if let (true, Some((_, b))) = (local, best_part) { + p.imbalance = (b.imbalance + 0.2 * (rng.unit() - 0.5)).clamp(0.01, 0.9); + p.cutoff = (b.cutoff as isize + rng.below(5) as isize - 2).clamp(3, 20) as usize; + p.costmod = (b.costmod + 0.6 * (rng.unit() - 0.5)).clamp(0.1, 4.0); + p.temperature = b.temperature; + } + path = partition_path(shape, dims, open, p, seed); + let plan = score(&path); + if best_part.is_none_or(|b| plan.total_flops() < b.0) { + best_part = Some((plan.total_flops(), p)); + } + push_top(&mut top, plan, keep); + } + } + // Reconfigure each of the best few to convergence; keep the best result. + // With slicing, each round reconfigures the per-slice tree (the sliced + // indices fixed) and then slices afresh, so the tree adapts to its slices. + let mut best = top[0].clone(); + if opts.reconfigure >= 3 { + for cand in top { + let mut cur = cand; + for _ in 0..8 { + let d = sliced_dims(dims, &cur.sliced); + let plan = score(&reconfigure(shape, &d, open, &cur.path, opts.reconfigure)); + if plan.total_flops() < cur.total_flops() * (1.0 - 1e-9) { + cur = plan; + } else { + break; + } + } + if cur.total_flops() < best.total_flops() { + best = cur; + } + } + } + best +} + +/// Keep the `keep` cheapest plans, cheapest first (earlier wins ties). +fn push_top(top: &mut Vec, plan: Plan, keep: usize) { + let at = top.iter().position(|p| plan.total_flops() < p.total_flops()).unwrap_or(top.len()); + if at < keep { + top.insert(at, plan); + top.truncate(keep); + } +} + +/// Contract a network along a path; returns the final tensor, with its +/// indices in the network's open order. +pub fn contract_path(net: &Network, path: &Path) -> Result { + let n = net.tensors.len(); + if path.steps.len() + 1 != n.max(1) { + return Err(TnError::PlanMismatch); + } + let mut live: Vec> = net.tensors.iter().cloned().map(Some).collect(); + for &(a, b) in &path.steps { + let ta = live.get_mut(a).and_then(Option::take).ok_or(TnError::PlanMismatch)?; + let tb = live.get_mut(b).and_then(Option::take).ok_or(TnError::PlanMismatch)?; + live.push(Some(contract(&ta, &tb, &net.dims, &net.open))); + } + let last = live.into_iter().rev().flatten().next().ok_or(TnError::PlanMismatch)?; + let order = net.open.clone(); + let data = permute(&last, &order, &net.dims); + Ok(Tensor { inds: order, data }) +} + +/// Fix the sliced indices of every tensor to one assignment. +fn fix(net: &Network, sliced: &[u32], values: &[usize]) -> Network { + let mut tensors = Vec::with_capacity(net.tensors.len()); + for t in &net.tensors { + let mut cur = t.clone(); + for (si, &i) in sliced.iter().enumerate() { + if let Some(pos) = cur.inds.iter().position(|&j| j == i) { + let r = cur.inds.len(); + let dim = net.dims[i as usize]; + let inner: usize = cur.inds[pos + 1..].iter().map(|&j| net.dims[j as usize]).product(); + let outer: usize = cur.inds[..pos].iter().map(|&j| net.dims[j as usize]).product(); + let mut data = Vec::with_capacity(outer * inner); + for o in 0..outer { + let base = (o * dim + values[si]) * inner; + data.extend_from_slice(&cur.data[base..base + inner]); + } + let mut inds = cur.inds.clone(); + inds.remove(pos); + let _ = r; + cur = Tensor { inds, data }; + } + } + tensors.push(cur); + } + Network { tensors, dims: net.dims.clone(), open: net.open.clone() } +} + +/// Memory the slice threads may hold together, in bytes. +#[cfg(not(target_arch = "wasm32"))] +const SLICE_MEMORY: f64 = 8.0 * 1024.0 * 1024.0 * 1024.0; + +/// Contract a network by a plan: every slice along the plan's path, the +/// slices summed in order. Slices run on several threads where available, +/// as many as fit 8 GiB at about six intermediates of the plan's largest size +/// each. The sum is in slice order whatever the threading, so +/// the bits do not depend on it. +pub fn contract_plan(net: &Network, plan: &Plan) -> Result { + let dims: Vec = plan.sliced.iter().map(|&i| net.dims[i as usize]).collect(); + let count: usize = dims.iter().product(); + let assignment = |mut s: usize| -> Vec { + let mut v = vec![0; dims.len()]; + for k in (0..dims.len()).rev() { + v[k] = s % dims[k]; + s /= dims[k]; + } + v + }; + let run = |s: usize| contract_path(&fix(net, &plan.sliced, &assignment(s)), &plan.path); + #[cfg(not(target_arch = "wasm32"))] + let parts: Vec> = { + let per_thread = 6.0 * plan.per_slice.max_size * core::mem::size_of::() as f64; + let fit = (SLICE_MEMORY / per_thread).floor().max(1.0) as usize; + let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(count).min(fit).max(1); + let mut parts: Vec>> = (0..count).map(|_| None).collect(); + std::thread::scope(|sc| { + let chunks: Vec<_> = parts.chunks_mut(count.div_ceil(threads)).enumerate().collect(); + for (c, chunk) in chunks { + let run = &run; + let base = c * count.div_ceil(threads); + sc.spawn(move || { + for (k, slot) in chunk.iter_mut().enumerate() { + *slot = Some(run(base + k)); + } + }); + } + }); + parts.into_iter().map(|p| p.expect("every slice ran")).collect() + }; + #[cfg(target_arch = "wasm32")] + let parts: Vec> = (0..count).map(run).collect(); + let mut total: Option = None; + for p in parts { + let t = p?; + total = Some(match total { + None => t, + Some(mut acc) => { + for (x, y) in acc.data.iter_mut().zip(&t.data) { + x.re += y.re; + x.im += y.im; + } + acc + } + }); + } + total.ok_or(TnError::PlanMismatch) +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + + /// A random unitary on `k` qubits from the QR of a random matrix + /// (Gram–Schmidt on columns): enough to exercise the contractions. + #[allow(clippy::needless_range_loop)] + pub(crate) fn random_unitary(k: usize, rng: &mut Rng) -> Vec { + let d = 1 << k; + let mut cols: Vec> = (0..d) + .map(|_| (0..d).map(|_| C::new(rng.unit() - 0.5, rng.unit() - 0.5)).collect()) + .collect(); + for j in 0..d { + for i in 0..j { + let mut dot = C::ZERO; + for r in 0..d { + dot = dot.add(cols[i][r].conj().mul(cols[j][r])); + } + for r in 0..d { + let t = cols[i][r].mul(dot); + cols[j][r] = cols[j][r].sub(t); + } + } + let norm = cols[j].iter().map(|c| c.norm2()).sum::().sqrt(); + for r in 0..d { + cols[j][r] = cols[j][r].scale(1.0 / norm); + } + } + let mut m = vec![C::ZERO; d * d]; + for r in 0..d { + for c in 0..d { + m[r * d + c] = cols[c][r]; + } + } + m + } + + /// Dense reference: apply gates to |0…0⟩, qubit 0 most significant. + pub(crate) fn dense(n: u32, gates: &[Gate]) -> Vec { + let dim = 1usize << n; + let mut psi = vec![C::ZERO; dim]; + psi[0] = C::ONE; + for g in gates { + let k = g.qubits.len(); + let mut out = vec![C::ZERO; dim]; + for (x, amp) in psi.iter().enumerate() { + if amp.re == 0.0 && amp.im == 0.0 { + continue; + } + let col = g.qubits.iter().fold(0usize, |c, &q| (c << 1) | ((x >> (n - 1 - q)) & 1)); + for row in 0..1usize << k { + let mut y = x; + for (i, &q) in g.qubits.iter().enumerate() { + let bit = (row >> (k - 1 - i)) & 1; + let shift = n - 1 - q; + y = (y & !(1 << shift)) | (bit << shift); + } + out[y] = out[y].add(g.matrix[row * (1 << k) + col].mul(*amp)); + } + } + psi = out; + } + psi + } + + pub(crate) fn random_circuit(n: u32, depth: usize, seed: u64) -> Vec { + let mut rng = Rng::new(seed); + let mut gates = Vec::new(); + for layer in 0..depth { + for q in 0..n { + gates.push(Gate { qubits: vec![q], matrix: random_unitary(1, &mut rng) }); + } + let mut q = (layer % 2) as u32; + while q + 1 < n { + gates.push(Gate { qubits: vec![q, q + 1], matrix: random_unitary(2, &mut rng) }); + q += 2; + } + // An occasional long-range gate, so the network is not just a ladder. + let a = rng.below(n as usize) as u32; + let b = rng.below(n as usize) as u32; + if a != b { + gates.push(Gate { qubits: vec![b, a], matrix: random_unitary(2, &mut rng) }); + } + } + gates + } + + fn close(a: C, b: C) -> bool { + (a.re - b.re).abs() < 1e-12 && (a.im - b.im).abs() < 1e-12 + } + + #[test] + fn contraction_matches_the_state_vector() { + for seed in 0..6 { + let n = 7; + let gates = random_circuit(n, 6, seed); + let psi = dense(n, &gates); + for x in [0usize, 5, 77, 127] { + let bits: Vec = (0..n).map(|q| ((x >> (n - 1 - q)) & 1) as u8).collect(); + let mut net = Network::amplitude(n, &gates, &bits, &[]).unwrap(); + net.simplify(); + let path = greedy_path(&net.shape(), &net.dims, &net.open, 1.0, 0.0, 0); + let t = contract_path(&net, &path).unwrap(); + assert_eq!(t.data.len(), 1); + assert!(close(t.data[0], psi[x]), "seed {seed} x {x}: {:?} vs {:?}", t.data[0], psi[x]); + } + } + } + + #[test] + fn open_outputs_give_a_batch_of_amplitudes() { + let n = 6; + let gates = random_circuit(n, 5, 9); + let psi = dense(n, &gates); + // Fix qubits 0, 2, 4 to 1, 0, 1; leave 5 and 1 open (5 most significant). + let bits = [1u8, 0, 0, 0, 1, 0]; + let mut net = Network::amplitude(n, &gates, &bits, &[5, 1]).unwrap(); + net.simplify(); + let path = greedy_path(&net.shape(), &net.dims, &net.open, 1.0, 0.3, 4); + let t = contract_path(&net, &path).unwrap(); + assert_eq!(t.data.len(), 4); + for (j, got) in t.data.iter().enumerate() { + let (b5, b1) = ((j >> 1) & 1, j & 1); + let x = (1 << 5) | (b1 << 4) | (1 << 1) | b5; + assert!(close(*got, psi[x]), "j {j}"); + } + } + + #[test] + fn every_path_gives_the_same_amplitude_and_the_same_bits_twice() { + let n = 8; + let gates = random_circuit(n, 6, 3); + let bits = [0u8, 1, 1, 0, 1, 0, 0, 1]; + let mut net = Network::amplitude(n, &gates, &bits, &[]).unwrap(); + net.simplify(); + let reference = contract_path(&net, &greedy_path(&net.shape(), &net.dims, &net.open, 1.0, 0.0, 0)).unwrap().data[0]; + for seed in 0..8 { + let p = greedy_path(&net.shape(), &net.dims, &net.open, 0.5 + 0.2 * seed as f64, 0.5, seed); + let a = contract_path(&net, &p).unwrap().data[0]; + assert!(close(a, reference)); + let b = contract_path(&net, &p).unwrap().data[0]; + assert_eq!((a.re.to_bits(), a.im.to_bits()), (b.re.to_bits(), b.im.to_bits())); + } + } + + /// A brickwork of random two-qubit gates on a `w × h` grid, alternating + /// horizontal and vertical couplings: a network whose cost depends + /// strongly on the order. + pub(crate) fn grid_circuit(w: u32, h: u32, depth: usize, seed: u64) -> Vec { + let mut rng = Rng::new(seed); + let q = |x: u32, y: u32| y * w + x; + let mut gates = Vec::new(); + for layer in 0..depth { + for i in 0..w * h { + gates.push(Gate { qubits: vec![i], matrix: random_unitary(1, &mut rng) }); + } + let (horizontal, offset) = ((layer / 2) % 2 == 0, (layer % 2) as u32); + if horizontal { + for y in 0..h { + let mut x = offset; + while x + 1 < w { + gates.push(Gate { qubits: vec![q(x, y), q(x + 1, y)], matrix: random_unitary(2, &mut rng) }); + x += 2; + } + } + } else { + for x in 0..w { + let mut y = offset; + while y + 1 < h { + gates.push(Gate { qubits: vec![q(x, y), q(x, y + 1)], matrix: random_unitary(2, &mut rng) }); + y += 2; + } + } + } + } + gates + } + + #[test] + fn every_tree_family_gives_the_amplitude() { + let (w, h) = (3, 3); + let n = w * h; + let gates = grid_circuit(w, h, 8, 21); + let psi = dense(n, &gates); + let bits: Vec = (0..n).map(|q| (q % 3 == 0) as u8).collect(); + let x = bits.iter().fold(0usize, |a, &b| (a << 1) | b as usize); + let mut net = Network::amplitude(n, &gates, &bits, &[]).unwrap(); + net.simplify(); + let shape = net.shape(); + for seed in 0..6 { + let p = PartitionParams { imbalance: 0.3, cutoff: 3 + seed as usize, passes: 3, costmod: 1.0, temperature: 0.1 }; + let path = partition_path(&shape, &net.dims, &net.open, p, seed); + assert!(close(contract_path(&net, &path).unwrap().data[0], psi[x]), "partition seed {seed}"); + let r = reconfigure(&shape, &net.dims, &net.open, &path, 8); + let (c0, c1) = (path_cost(&shape, &net.dims, &net.open, &path), path_cost(&shape, &net.dims, &net.open, &r)); + assert!(c1.flops <= c0.flops, "reconfiguration made it costlier"); + assert!(close(contract_path(&net, &r).unwrap().data[0], psi[x]), "reconfigured seed {seed}"); + } + } + + #[test] + fn hyper_optimisation_beats_plain_greedy() { + let gates = grid_circuit(4, 4, 10, 2); + let n = 16; + let mut net = Network::amplitude(n, &gates, &vec![0; n as usize], &[]).unwrap(); + net.simplify(); + let shape = net.shape(); + let plain = path_cost(&shape, &net.dims, &net.open, &greedy_path(&shape, &net.dims, &net.open, 1.0, 0.0, 0)); + let plan = find_path(&shape, &net.dims, &net.open, &PathOptions { trials: 24, seed: 1, max_size: None, reconfigure: 6, families: Families::Both }); + assert!(plan.total_flops() < plain.flops, "{} vs {}", plan.total_flops(), plain.flops); + // The same seed gives the same plan. + let again = find_path(&shape, &net.dims, &net.open, &PathOptions { trials: 24, seed: 1, max_size: None, reconfigure: 6, families: Families::Both }); + assert_eq!(plan, again); + } + + #[test] + fn slices_fit_the_cap_and_sum_to_the_amplitude() { + let (w, h) = (4, 3); + let n = w * h; + let gates = grid_circuit(w, h, 10, 5); + let psi = dense(n, &gates); + let bits: Vec = (0..n).map(|q| (q % 2) as u8).collect(); + let x = bits.iter().fold(0usize, |a, &b| (a << 1) | b as usize); + let mut net = Network::amplitude(n, &gates, &bits, &[]).unwrap(); + net.simplify(); + let shape = net.shape(); + let free = find_path(&shape, &net.dims, &net.open, &PathOptions { trials: 12, seed: 3, max_size: None, reconfigure: 6, families: Families::Both }); + let cap = free.per_slice.max_size / 16.0; + let plan = find_path(&shape, &net.dims, &net.open, &PathOptions { trials: 12, seed: 3, max_size: Some(cap), reconfigure: 6, families: Families::Both }); + assert!(plan.per_slice.max_size <= cap); + assert!(plan.slices >= 2.0); + let t = contract_plan(&net, &plan).unwrap(); + assert!(close(t.data[0], psi[x]), "{:?} vs {:?}", t.data[0], psi[x]); + // Bit-identical however the slices are scheduled: run twice. + let u = contract_plan(&net, &plan).unwrap(); + assert_eq!((t.data[0].re.to_bits(), t.data[0].im.to_bits()), (u.data[0].re.to_bits(), u.data[0].im.to_bits())); + } + + /// On a network of at most eight tensors, reconfiguration with a frontier + /// of eight sees the whole tree, so it must reach the optimum, found here + /// independently by trying every split of every subset. + #[test] + fn reconfiguration_reaches_the_optimum_on_small_networks() { + fn optimum(shape: &[Vec], dims: &[usize], set: u32, memo: &mut BTreeMap)>) -> (f64, Vec) { + if set.count_ones() == 1 { + return (0.0, shape[set.trailing_zeros() as usize].clone()); + } + if let Some(v) = memo.get(&set) { + return v.clone(); + } + let mut best: Option<(f64, Vec)> = None; + let mut a = (set - 1) & set; + while a > 0 { + let b = set ^ a; + if a < b { + let (ca, sa) = optimum(shape, dims, a, memo); + let (cb, sb) = optimum(shape, dims, b, memo); + let c = ca + cb + set_size(&union(&sa, &sb), dims); + if best.as_ref().is_none_or(|x| c < x.0) { + best = Some((c, sym_diff(&sa, &sb, &[]))); + } + } + a = (a - 1) & set; + } + let v = best.unwrap(); + memo.insert(set, v.clone()); + v + } + let mut rng = Rng::new(77); + for trial in 0..40 { + let n = 5 + rng.below(4); // 5..8 tensors + // A random connected graph: a spanning chain plus extra bonds. + let mut shape: Vec> = vec![Vec::new(); n]; + let mut dims = Vec::new(); + let bond = |a: usize, b: usize, shape: &mut Vec>, dims: &mut Vec, rng: &mut Rng| { + let i = dims.len() as u32; + dims.push(2 + rng.below(3)); + shape[a].push(i); + shape[b].push(i); + }; + for t in 1..n { + let u = rng.below(t); + bond(u, t, &mut shape, &mut dims, &mut rng); + } + for _ in 0..n { + let (a, b) = (rng.below(n), rng.below(n)); + if a != b { + bond(a, b, &mut shape, &mut dims, &mut rng); + } + } + let path = greedy_path(&shape, &dims, &[], 1.0, 1.0, trial); + let r = reconfigure(&shape, &dims, &[], &path, 8); + let got = path_cost(&shape, &dims, &[], &r).flops; + let (want, _) = optimum(&shape, &dims, (1u32 << n) - 1, &mut BTreeMap::new()); + assert!((got - want).abs() <= 1e-9 * want, "trial {trial}: reconfigured {got}, optimum {want}"); + } + } + + /// The slicer prices each candidate incrementally; a full recount of the + /// whole path for every candidate must choose the same indices. + #[test] + fn slicing_prices_match_a_full_recount() { + for (seed, cap_div) in [(1u64, 8.0), (2, 64.0), (3, 512.0)] { + let gates = grid_circuit(4, 4, 9, seed); + let mut net = Network::amplitude(16, &gates, &[0; 16], &[]).unwrap(); + net.simplify(); + let (shape, dims, open) = (net.shape(), net.dims.clone(), net.open.clone()); + let path = greedy_path(&shape, &dims, &open, 2.0, 0.2, seed); + let cap = path_cost(&shape, &dims, &open, &path).max_size / cap_div; + // The reference: recount everything for every candidate. + let mut want: Vec = Vec::new(); + loop { + let d = sliced_dims(&dims, &want); + if path_cost(&shape, &d, &open, &path).max_size <= cap { + break; + } + let mut sets: Vec> = shape.clone(); + for &(a, b) in &path.steps { + sets.push(sym_diff(&sets[a], &sets[b], &open)); + } + let mut cands: Vec = + sets.iter().filter(|x| set_size(x, &d) > cap).flatten().copied().filter(|i| d[*i as usize] > 1).collect(); + cands.sort_unstable(); + cands.dedup(); + let mut best: Option<(f64, u32)> = None; + for &i in &cands { + let mut t = want.clone(); + t.push(i); + let total = path_cost(&shape, &sliced_dims(&dims, &t), &open, &path).flops * (1u64 << t.len()) as f64; + if best.is_none_or(|(b, _)| total < b) { + best = Some((total, i)); + } + } + want.push(best.unwrap().1); + } + assert_eq!(slice(&shape, &dims, &open, &path, cap), want, "seed {seed}"); + } + } + + #[test] + fn bad_input_is_refused() { + let g = vec![Gate { qubits: vec![0, 0], matrix: vec![C::ONE; 16] }]; + assert_eq!(Network::amplitude(2, &g, &[0, 0], &[]).unwrap_err(), TnError::BadGate(0)); + assert_eq!(Network::amplitude(2, &[], &[0], &[]).unwrap_err(), TnError::BadOutput); + assert_eq!(Network::amplitude(2, &[], &[0, 2], &[]).unwrap_err(), TnError::BadOutput); + } +} diff --git a/wai/wai-quantum/src/repro.rs b/wai/wai-quantum/src/repro.rs index d95b88f82..0b922ac40 100644 --- a/wai/wai-quantum/src/repro.rs +++ b/wai/wai-quantum/src/repro.rs @@ -30,28 +30,28 @@ pub use trig::sin_cos; // Natural logarithm // --------------------------------------------------------------------------- -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LN2_HI: f64 = 6.931_471_803_691_238_164_90e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LN2_LO: f64 = 1.908_214_929_270_587_700_02e-10; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG1: f64 = 6.666_666_666_666_735_130e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG2: f64 = 3.999_999_999_940_941_908e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG3: f64 = 2.857_142_874_366_239_149e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG4: f64 = 2.222_219_843_214_978_396e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG5: f64 = 1.818_357_216_161_805_012e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG6: f64 = 1.531_383_769_920_937_332e-1; -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] const LG7: f64 = 1.479_819_860_511_658_591e-1; /// The natural logarithm, within an ulp of the true value: `−∞` at zero, NaN /// below it. -#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] +#[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] pub fn ln(x: f64) -> f64 { let mut bits = x.to_bits(); let mut hx = (bits >> 32) as u32; @@ -95,7 +95,7 @@ pub fn ln(x: f64) -> f64 { // Exponential // --------------------------------------------------------------------------- -#[cfg(feature = "quantum_frame")] +#[cfg(any(feature = "quantum_frame", feature = "quantum_tn"))] const EXP_P: [f64; 5] = [ 1.666_666_666_666_660_190_37e-1, -2.777_777_777_701_559_338_42e-3, @@ -107,7 +107,7 @@ const EXP_P: [f64; 5] = [ /// `eˣ`, bit-identical on every IEEE-754 platform: fdlibm's algorithm /// (reduction by `ln 2` in two parts, a rational approximation on /// `[−ln2/2, ln2/2]`), using only `+ − × ÷`. Within an ulp of the true value. -#[cfg(feature = "quantum_frame")] +#[cfg(any(feature = "quantum_frame", feature = "quantum_tn"))] pub fn exp(x: f64) -> f64 { const LN2_HI_E: f64 = 6.931_471_803_691_238_164_90e-1; const LN2_LO_E: f64 = 1.908_214_929_270_587_700_02e-10; @@ -308,7 +308,7 @@ mod tests { #[allow(unused_imports)] use super::*; - #[cfg(feature = "quantum_frame")] + #[cfg(any(feature = "quantum_frame", feature = "quantum_tn"))] #[test] fn exp_is_within_an_ulp_of_the_platform() { let mut s = 0x9e37_79b9_7f4a_7c15u64; @@ -351,7 +351,7 @@ mod tests { }) } - #[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame"))] + #[cfg(any(feature = "quantum_vml", feature = "quantum_kernel", feature = "quantum_resource", feature = "quantum_frame", feature = "quantum_tn"))] #[test] fn ln_is_within_an_ulp_of_the_platform() { for x in inputs().filter(|&x| x > 0.0) { diff --git a/wai/wai-quantum/tests/bit_identity.rs b/wai/wai-quantum/tests/bit_identity.rs index e28da2d84..05103a5e1 100644 --- a/wai/wai-quantum/tests/bit_identity.rs +++ b/wai/wai-quantum/tests/bit_identity.rs @@ -16,7 +16,8 @@ feature = "quantum_spd", feature = "quantum_bptn", feature = "quantum_resource", - feature = "quantum_frame" + feature = "quantum_frame", + feature = "quantum_tn" ))] use wai_quantum::*; @@ -172,3 +173,50 @@ fn frame_sampling_and_matching() { assert_eq!(digest(&v), 0xcbb3_99f9_6d5c_e8f0); } +/// A contraction plan found by the seeded hyper-optimiser, and the amplitude +/// its sliced, threaded contraction gives: the search's random stream and +/// logarithms, and the fixed-order complex arithmetic. +#[test] +fn tensor_network_plan_and_amplitude() { + use quantum_tn::{complex, contract_plan, find_path, Families, Gate, Network, PathOptions}; + // A 3 × 4 grid, ten layers of fixed two-qubit gates and phased rotations. + let (w, h) = (4u32, 3u32); + let mut gates = Vec::new(); + let mut k = 0.0f64; + for layer in 0..10u32 { + for q in 0..w * h { + k += 0.37; + let (c, s) = (1.0 / (1.0 + k * k).sqrt(), k / (1.0 + k * k).sqrt()); + gates.push(Gate { qubits: vec![q], matrix: vec![complex(c, 0.0), complex(0.0, -s), complex(0.0, -s), complex(c, 0.0)] }); + } + let horizontal = layer % 2 == 0; + for y in 0..h { + for x in 0..w { + let (a, b) = if horizontal { (y * w + x, y * w + x + 1) } else { (y * w + x, (y + 1) * w + x) }; + let ok = if horizontal { x + 1 < w && (x + layer / 2) % 2 == 0 } else { y + 1 < h && (y + layer / 2) % 2 == 0 }; + if ok { + // An iSWAP-like gate with a controlled phase. + let z = complex(0.0, 0.0); + let one = complex(1.0, 0.0); + let mi = complex(0.0, -1.0); + let ph = complex(0.5, -0.866_025_403_784_438_6); + gates.push(Gate { qubits: vec![a, b], matrix: vec![one, z, z, z, z, z, mi, z, z, mi, z, z, z, z, z, ph] }); + } + } + } + } + let n = w * h; + let bits: Vec = (0..n).map(|q| (q % 3 == 1) as u8).collect(); + let mut net = Network::amplitude(n, &gates, &bits, &[]).unwrap(); + net.simplify(); + let shape = net.shape(); + let free = find_path(&shape, &net.dims, &net.open, &PathOptions { trials: 12, seed: 5, max_size: None, reconfigure: 6, families: Families::Both }); + let cap = free.per_slice.max_size / 8.0; + let plan = find_path(&shape, &net.dims, &net.open, &PathOptions { trials: 12, seed: 5, max_size: Some(cap), reconfigure: 6, families: Families::Both }); + let amp = contract_plan(&net, &plan).unwrap().data[0]; + let steps: Vec = plan.path.steps.iter().flat_map(|&(a, b)| [a as f64, b as f64]).collect(); + let mut v = vec![amp.re, amp.im, plan.per_slice.flops, plan.per_slice.max_size, plan.slices]; + v.extend(steps); + assert_eq!(digest(&v), 0x8954_0c3a_d1fc_a04d); +} + diff --git a/wai/wai-rs/Cargo.lock b/wai/wai-rs/Cargo.lock index 44da8e1a9..61e27e5b6 100644 --- a/wai/wai-rs/Cargo.lock +++ b/wai/wai-rs/Cargo.lock @@ -3900,7 +3900,7 @@ dependencies = [ [[package]] name = "wai-quantum" -version = "0.3.25" +version = "0.3.26" dependencies = [ "blake3", "ed25519-dalek", diff --git a/wai/wai-rs/Cargo.toml b/wai/wai-rs/Cargo.toml index 4bf6b742f..18a044e3e 100644 --- a/wai/wai-rs/Cargo.toml +++ b/wai/wai-rs/Cargo.toml @@ -351,6 +351,8 @@ quantum_resource_receipt = ["quantum_resource", "wai-quantum/quantum_resource_re # Noisy Clifford circuits at QEC scale: frame sampling, detector error models, # exact matching, measured code models. quantum_frame = ["dep:wai-quantum", "wai-quantum/quantum_frame"] +# Tensor-network contraction: hyper-optimised trees, slicing, circuit amplitudes. +quantum_tn = ["dep:wai-quantum", "wai-quantum/quantum_tn"] quantum_qir = ["quantum_toolchain", "wai-quantum/quantum_qir"] # Variational quantum machine learning (extensions/quantum-sim § VML, # `wai.quantum.vml`): the interference-based ML the thermodynamic-computing pivot @@ -480,6 +482,7 @@ required-features = [ "quantum_bptn", "quantum_resource_receipt", "quantum_frame", + "quantum_tn", ] [[bin]] diff --git a/wai/wai-rs/src/bin/wai_quantum.rs b/wai/wai-rs/src/bin/wai_quantum.rs index bf2935145..167c62e79 100644 --- a/wai/wai-rs/src/bin/wai_quantum.rs +++ b/wai/wai-rs/src/bin/wai_quantum.rs @@ -161,6 +161,8 @@ fn usage() -> ! { wai-quantum qec --file CIRCUIT [--shots N] (the ecosystem's circuit text)\n \ wai-quantum qec [--circuit | --dem] [-d D | --file CIRCUIT]\n \ wai-quantum qec --fit [--shots N] (measure a and p* of the code model)\n\n\ + TENSOR NETWORKS:\n \ + wai-quantum tn --file CIRCUIT.jsonl [-n N] [--bits B] [--open q,q] [--trials T] [--seed S] [--cap-log2 L] [--plan-only] [--json]\n\n\ Channels: identity, bitflip, dephasing, depolarizing, damping.\n\ Run `wai-quantum list` for the algorithm catalog." ); @@ -759,6 +761,99 @@ fn main() { } } + // A circuit's amplitudes by tensor-network contraction: a seeded + // hyper-optimised contraction tree, sliced to a memory cap, contracted + // in a fixed order. Circuits are JSON lines, one gate per line: + // {"q":[0,1],"m":[re,im, …]} (rows are outputs, first qubit most + // significant). + "tn" => { + use wai::quantum_tn::{complex, contract_plan, find_path, Families, Gate, Network, PathOptions}; + let path = opt(args, "--file", "").unwrap_or_else(|| die("--file CIRCUIT.jsonl is required")); + let text = std::fs::read_to_string(&path).unwrap_or_else(|e| die(&format!("{path}: {e}"))); + let mut gates = Vec::new(); + let mut n: u32 = 0; + for (i, line) in text.lines().enumerate().filter(|(_, l)| !l.trim().is_empty()) { + let v: serde_json::Value = serde_json::from_str(line).unwrap_or_else(|_| die(&format!("line {}: not JSON", i + 1))); + let qubits: Vec = v["q"] + .as_array() + .unwrap_or_else(|| die(&format!("line {}: missing \"q\"", i + 1))) + .iter() + .map(|x| x.as_u64().unwrap_or_else(|| die("qubits are integers")) as u32) + .collect(); + let m: Vec = v["m"] + .as_array() + .unwrap_or_else(|| die(&format!("line {}: missing \"m\"", i + 1))) + .iter() + .map(|x| x.as_f64().unwrap_or_else(|| die("matrix entries are numbers"))) + .collect(); + n = n.max(qubits.iter().max().map_or(0, |q| q + 1)); + gates.push(Gate { qubits, matrix: m.chunks(2).map(|c| complex(c[0], *c.get(1).unwrap_or(&0.0))).collect() }); + } + let n: u32 = num(args, "--qubits", "-n", n); + let bits: Vec = match opt(args, "--bits", "") { + Some(b) => b.chars().map(|c| if c == '1' { 1 } else if c == '0' { 0 } else { die("--bits is a 0/1 string") }).collect(), + None => vec![0; n as usize], + }; + let open: Vec = opt(args, "--open", "") + .map(|l| l.split(',').map(|q| q.trim().parse().unwrap_or_else(|_| die("bad --open qubit"))).collect()) + .unwrap_or_default(); + if open.len() > 20 { + die("at most 20 open qubits (2^20 amplitudes)"); + } + let mut net = Network::amplitude(n, &gates, &bits, &open).unwrap_or_else(|e| die(&format!("{e:?}"))); + let raw = net.tensors.len(); + net.simplify(); + let shape = net.shape(); + let cap: Option = opt(args, "--cap-log2", "").map(|v| 2f64.powf(v.parse().unwrap_or_else(|_| die("bad --cap-log2")))); + let options = PathOptions { + trials: num(args, "--trials", "", 64), + seed: num(args, "--seed", "", 0), + max_size: cap, + reconfigure: 8, + families: Families::Both, + }; + let plan = find_path(&shape, &net.dims, &net.open, &options); + let flops = plan.total_flops(); + let plan_line = if flops > 0.0 { + format!("10^{:.3} operations, largest intermediate 2^{:.1}, {} slices", flops.log10(), plan.per_slice.max_size.log2(), plan.slices) + } else { + "no contractions needed".to_string() + }; + if flag(args, "--plan-only") || flops > 1e13 { + if !flag(args, "--plan-only") { + eprintln!("wai-quantum: the plan costs 10^{:.1} operations; printing the plan only (pass a smaller circuit, or --plan-only)", flops.log10()); + } + if flag(args, "--json") { + println!( + "{{\"flops\":{flops},\"log2_max_size\":{},\"slices\":{},\"tensors\":{}}}", + plan.per_slice.max_size.log2(), + plan.slices, + shape.len() + ); + } else { + println!("network {raw} tensors, {} after simplification", shape.len()); + println!("plan {plan_line}"); + } + return; + } + let t = contract_plan(&net, &plan).unwrap_or_else(|e| die(&format!("{e:?}"))); + if flag(args, "--json") { + let amps: Vec = t.data.iter().map(|c| format!("[{},{}]", c.re, c.im)).collect(); + println!( + "{{\"amplitudes\":[{}],\"flops\":{flops},\"log2_max_size\":{},\"slices\":{}}}", + amps.join(","), + plan.per_slice.max_size.log2(), + plan.slices + ); + } else { + println!("network {raw} tensors, {} after simplification", shape.len()); + println!("plan {plan_line}"); + for (i, c) in t.data.iter().enumerate() { + println!("amp[{i}] {:+.12e} {:+.12e}i", c.re, c.im); + } + } + } + "entropy" => { let c = circuit_from(args, algo); let sv = c.simulate().unwrap_or_else(|e| die(&format!("simulate: {e:?}"))); diff --git a/wai/wai-rs/src/lib.rs b/wai/wai-rs/src/lib.rs index 2ee536b01..797a3faa4 100644 --- a/wai/wai-rs/src/lib.rs +++ b/wai/wai-rs/src/lib.rs @@ -624,6 +624,9 @@ pub use wai_quantum::quantum_frame; #[cfg(feature = "quantum_frame")] pub use wai_quantum::quantum_match; + +#[cfg(feature = "quantum_tn")] +pub use wai_quantum::quantum_tn; #[cfg(feature = "quantum_qir")] pub use wai_quantum::quantum_qir; diff --git a/wai/wai-web/Cargo.lock b/wai/wai-web/Cargo.lock index 4453ab014..70acda408 100644 --- a/wai/wai-web/Cargo.lock +++ b/wai/wai-web/Cargo.lock @@ -734,7 +734,7 @@ dependencies = [ [[package]] name = "wai-quantum" -version = "0.3.25" +version = "0.3.26" dependencies = [ "blake3", "ed25519-dalek",