warriors/trisha/scripts/generate_quantum_baseline.py

#!/usr/bin/env python3
"""Independent finite-field matrix formulas, emitted directly as stack algebra.

This generator never reads compiler output. Coordinates follow source order,
first deepest; each pure routine consumes its arguments and returns all results.
"""
from pathlib import Path
from quantum_ram import apply_single_gate

ROOT = Path(__file__).resolve().parents[1]


def var(index):
    return ("var", index)


def add(a, b):
    return ("add", a, b)


def neg(a):
    return ("neg", a)


def mul(a, b):
    return ("mul", a, b)


def complex_product(a, b):
    return [add(mul(a[0], b[0]), neg(mul(a[1], b[1]))),
            add(mul(a[0], b[1]), mul(a[1], b[0]))]


def routine(name, width, results):
    code = [f"__{name}:"]
    depth = width

    def expression(expr):
        nonlocal depth
        if isinstance(expr, int):
            code.append(f"    push {expr}")
            depth += 1
        elif expr[0] == "var":
            distance = depth - 1 - expr[1]
            assert 0 <= distance <= 15, (name, distance)
            code.append(f"    dup {distance}")
            depth += 1
        elif expr[0] == "neg":
            expression(expr[1])
            code.extend(["    push -1", "    mul"])
        else:
            expression(expr[1])
            expression(expr[2])
            code.append(f"    {expr[0]}")
            depth -= 1

    for result in results:
        expression(result)
    for _ in range(width):
        code.extend([f"    pick {len(results)}", "    pop 1"])
    return "\n".join(code + ["    return", ""])


def permutation(name, indices, negative=()):
    width = len(indices)
    order = list(range(width))
    code = [f"__{name}:"]
    for position, wanted in enumerate(indices):
        old = order.index(wanted)
        if position == old:
            continue
        if old == width - 1:
            code.append(f"    swap {width - 1 - position}")
        else:
            code.extend(f"    swap {width - 1 - i}" for i in [position, old, position])
        order[position], order[old] = order[old], order[position]
    for position in negative:
        distance = width - 1 - position
        if distance:
            code.append(f"    swap {distance}")
        code.extend(["    push -1", "    mul"])
        if distance:
            code.append(f"    swap {distance}")
    return "\n".join(code + ["    return", ""])


def pure_routines():
    a, b, c, d = map(var, range(4))
    code = [routine("complex_zero", 0, [0, 0]),
            routine("complex_one", 0, [1, 0]),
            routine("complex_add", 4, [add(a, c), add(b, d)]),
            routine("complex_sub", 4, [add(a, neg(c)), add(b, neg(d))]),
            routine("complex_mul", 4, complex_product([a, b], [c, d])),
            routine("complex_scale", 3, [mul(a, b), mul(a, c)]),
            permutation("complex_conj", [0, 1], [1]),
            routine("complex_norm_sq", 2, [add(mul(a, a), mul(b, b))]),
            routine("init_zero", 0, [1, 0, 0, 0]),
            routine("init_one", 0, [0, 0, 1, 0]),
            permutation("paulix", [2, 3, 0, 1]),
            permutation("pauliz", [0, 1, 2, 3], [2, 3]),
            permutation("pauliy", [3, 2, 1, 0], [1, 2]),
            routine("hadamard", 4, [add(a, c), add(b, d), add(a, neg(c)), add(b, neg(d))]),
            permutation("sgate", [0, 1, 3, 2], [2]),
            routine("tgate", 4, [a, b, add(c, neg(d)), add(c, d)])]
    tensor = []
    for ai, bi in [(0, 4), (0, 6), (2, 4), (2, 6)]:
        tensor.extend(complex_product([var(ai), var(ai+1)], [var(bi), var(bi+1)]))
    code.extend([routine("two_qubit_product", 8, tensor),
                 permutation("cnot", [0, 1, 2, 3, 6, 7, 4, 5]),
                 permutation("cz", list(range(8)), [6, 7]),
                 permutation("swap", [0, 1, 4, 5, 2, 3, 6, 7])])
    # The source API's documented signed-field comparison surrogate, not Born sampling.
    diff = add(add(mul(a, a), mul(b, b)), neg(add(mul(c, c), mul(d, d))))
    comparison = routine("measure_deterministic", 4, [diff])
    comparison = comparison.replace("    return", "    split\n    pop 1\n    push 2147483647\n    swap 1\n    lt\n    return")
    code.append(comparison)
    return "\n".join(code)


if __name__ == "__main__":
    target = ROOT / "baselines/triton/std/quantum/gates.tasm"
    target.write_text("// Independent finite-field gate algebra. Generated by scripts/generate_quantum_baseline.py.\n"
                      "// Arguments/results use source coordinate order: first deepest.\n\n"
                      + pure_routines() + "\n" + apply_single_gate())