wai-quantum 0.6.0

A deterministic quantum stack in pure Rust: byte-exact circuit simulation (statevector / stabilizer / tensor-network MPS / sparse-Pauli backends), sparse Pauli dynamics at utility scale (arbitrary angles, 1024 qubits), belief-propagation tensor networks on the hardware graph, error mitigation, qLDPC decoding, noise learning, circuit-equivalence proofs, a phasor interference-ML layer, information-theoretic limits, noisy channels and state tomography, and signed energy-accounted receipts. No QPU, no cloud, no system libraries in the default features — identical results native, in the browser, and as a WASI component at the edge; the opt-in WebGPU member reaches the platform's GPU interface.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

"""The tile array's program for wai-quantum's npu-cores member.

One compute core runs the kernel (npu_cores.cc) over a dispatch of LINES
lines of LINE cases each: the array's data movers bring each line into the
core's memory and take its outputs back to the host's buffer. The member runs
this script with the open toolchain's Python to build the program, then loads
it itself; nothing here runs on the device.

    python npu_cores_design.py --kernel npu_cores.cc --rnd 536870912 --sh 30 --last-add 0 \\
        --line 256 --lines 4 --device npu2 --xclbin final.xclbin --insts insts.bin

With --versions it prints the toolchain's versions, one per line, and exits.
"""

import argparse
import importlib.metadata as md
import os
import site
import sys

HEADER = 16  # words before a line's cases
CASE_IN = 4  # words per case in
CASE_OUT = 8  # words per case out: four 64-bit values


def versions():
    for name in ("mlir_aie", "llvm-aie"):
        try:
            print(f"{name} {md.version(name)}")
        except md.PackageNotFoundError:
            print(f"{name} absent")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--versions", action="store_true")
    ap.add_argument("--kernel")
    ap.add_argument("--rnd", type=int)
    ap.add_argument("--sh", type=int)
    ap.add_argument("--last-add", type=int, choices=[0, 1])
    ap.add_argument("--line", type=int)
    ap.add_argument("--lines", type=int)
    ap.add_argument("--device", default="npu2")
    ap.add_argument("--xclbin")
    ap.add_argument("--insts")
    a = ap.parse_args()
    if a.versions:
        versions()
        return
    # The core compiler ships in its own package beside this interpreter's.
    if "PEANO_INSTALL_DIR" not in os.environ:
        for d in site.getsitepackages():
            if os.path.isdir(os.path.join(d, "llvm-aie", "bin")):
                os.environ["PEANO_INSTALL_DIR"] = os.path.join(d, "llvm-aie")
                break

    import numpy as np
    import aie.iron as iron
    from aie.iron import CompileTime, In, ObjectFifo, Out, Program, Runtime, Worker
    from aie.iron.device import from_name
    from aie.iron.kernel import ExternalFunction

    source = open(a.kernel).read()
    line, lines = a.line, a.lines
    flags = [f"-DRND={a.rnd}", f"-DSH={a.sh}", f"-DLAST_ADD={a.last_add}"]

    @iron.jit
    def design(inp: In, out: Out, *, n: CompileTime[int]):
        lin = np.ndarray[(HEADER + CASE_IN * line,), np.dtype[np.int32]]
        lout = np.ndarray[(CASE_OUT * line,), np.dtype[np.int32]]
        vin = np.ndarray[(n * (HEADER + CASE_IN * line),), np.dtype[np.int32]]
        vout = np.ndarray[(n * CASE_OUT * line,), np.dtype[np.int32]]
        of_in = ObjectFifo(lin, name="in")
        of_out = ObjectFifo(lout, name="out")
        kernel = ExternalFunction("wq_line", source_string=source, arg_types=[lin, lout, np.int32], compile_flags=flags)

        def core(of_in, of_out, kernel):
            o = of_out.acquire(1)
            i = of_in.acquire(1)
            kernel(i, o, line)
            of_in.release(1)
            of_out.release(1)

        worker = Worker(core, [of_in.cons(), of_out.prod(), kernel])

        def sequence(x, y, in_h, out_h):
            in_h.fill(x)
            out_h.drain(y, wait=True)

        rt = Runtime(sequence, [vin, vout, of_in.prod(), of_out.cons()])
        return Program(iron.get_current_device(), rt, workers=[worker]).resolve_program()

    iron.set_current_device(from_name(a.device, n_cols=1))
    design.specialize(n=lines).compile(xclbin_path=a.xclbin, inst_path=a.insts)


if __name__ == "__main__":
    sys.exit(main())