src/machines.py

"""neural_subleq8 and neural_rv32 — runtimes, reference emulators, assemblers,
an RV32 object loader, and test suites for the two endpoint machines of the
family. One file; run the suites with:

    python machines.py            # both machines
    python machines.py subleq     # SUBLEQ-8 only
    python machines.py rv32       # RV32 assembled-program suite
    python machines.py rv32-c     # RV32 running stock-compiler C

Build the model files with `python build.py --apply subleq` and
`python build.py --apply rv32`.
"""

from __future__ import annotations
import argparse
import os
import shutil
import struct
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
from safetensors import safe_open

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from eval import NetlistEvaluator, heaviside, load_metadata
from eval_all import GenericThresholdALU, int_to_bits, bits_to_int

# ============================================================================
# RV32 relocatable-object loader
# ============================================================================
R = {
    "R_RISCV_32": 1,
    "R_RISCV_BRANCH": 16,
    "R_RISCV_JAL": 17,
    "R_RISCV_CALL": 18,
    "R_RISCV_CALL_PLT": 19,
    "R_RISCV_PCREL_HI20": 23,
    "R_RISCV_PCREL_LO12_I": 24,
    "R_RISCV_PCREL_LO12_S": 25,
    "R_RISCV_HI20": 26,
    "R_RISCV_LO12_I": 27,
    "R_RISCV_LO12_S": 28,
    "R_RISCV_RELAX": 51,
}
RREV = {v: k for k, v in R.items()}


class Elf:
    def __init__(self, data: bytes):
        assert data[:4] == b"\x7fELF", "not ELF"
        assert data[4] == 1 and data[5] == 1, "need ELF32 LE"
        (self.e_shoff,) = struct.unpack_from("<I", data, 0x20)
        (self.e_shentsize, self.e_shnum, self.e_shstrndx) = struct.unpack_from("<HHH", data, 0x2E)
        self.data = data
        self.sections = []
        for i in range(self.e_shnum):
            off = self.e_shoff + i * self.e_shentsize
            name, typ, flags, addr, offset, size, link, info, align, entsz = \
                struct.unpack_from("<IIIIIIIIII", data, off)
            self.sections.append(dict(name=name, type=typ, flags=flags, addr=addr,
                                      offset=offset, size=size, link=link, info=info,
                                      align=align, entsize=entsz))
        shstr = self.sections[self.e_shstrndx]
        self.shstr = data[shstr["offset"]:shstr["offset"] + shstr["size"]]
        for s in self.sections:
            s["sname"] = self._str(self.shstr, s["name"])

    @staticmethod
    def _str(tab: bytes, off: int) -> str:
        end = tab.find(b"\x00", off)
        return tab[off:end].decode()

    def symbols(self):
        symtab = next(s for s in self.sections if s["sname"] == ".symtab")
        strtab = self.sections[symtab["link"]]
        st = self.data[strtab["offset"]:strtab["offset"] + strtab["size"]]
        syms = []
        n = symtab["size"] // 16
        for i in range(n):
            off = symtab["offset"] + i * 16
            name, value, size, info, other, shndx = struct.unpack_from("<IIIBBH", self.data, off)
            syms.append(dict(name=self._str(st, name), value=value, size=size,
                             info=info, shndx=shndx))
        return syms


def load(obj: bytes, base: int = 0, mem_size: int = 65536,
         entry_symbol: str = "_start") -> Tuple[List[int], int]:
    elf = Elf(obj)
    mem = [0] * mem_size

    # Lay out allocatable progbits/nobits sections (SHF_ALLOC = 0x2).
    placed: Dict[int, int] = {}   # section index -> load address
    cursor = base
    for i, s in enumerate(elf.sections):
        if not (s["flags"] & 0x2):
            continue
        align = max(s["align"], 1)
        cursor = (cursor + align - 1) & ~(align - 1)
        placed[i] = cursor
        if s["type"] == 1:  # PROGBITS
            body = elf.data[s["offset"]:s["offset"] + s["size"]]
            for j, byte in enumerate(body):
                mem[cursor + j] = byte
        cursor += s["size"]

    syms = elf.symbols()

    def sym_addr(idx: int) -> int:
        sym = syms[idx]
        if sym["shndx"] in placed:
            return placed[sym["shndx"]] + sym["value"]
        return sym["value"]

    # Apply relocations. PCREL_LO12 references the symbol of its paired HI20
    # site, so first index HI20/PCREL_HI20 results by their location.
    hi_at: Dict[int, int] = {}
    for s in elf.sections:
        if s["type"] != 4:  # RELA
            continue
        target = elf.sections[s["info"]]
        if s["info"] not in placed:
            continue
        tbase = placed[s["info"]]
        n = s["size"] // 12
        for i in range(n):
            off = s["offset"] + i * 12
            r_offset, r_info, r_addend = struct.unpack_from("<IIi", elf.data, off)
            r_type = r_info & 0xFF
            r_sym = r_info >> 8
            loc = tbase + r_offset
            if r_type in (R["R_RISCV_PCREL_HI20"], R["R_RISCV_HI20"]):
                hi_at[loc] = sym_addr(r_sym) + r_addend

    def rd(a):
        return mem[a] | (mem[a + 1] << 8) | (mem[a + 2] << 16) | (mem[a + 3] << 24)

    def wr(a, v):
        for k in range(4):
            mem[a + k] = (v >> (8 * k)) & 0xFF

    for s in elf.sections:
        if s["type"] != 4 or s["info"] not in placed:
            continue
        tbase = placed[s["info"]]
        n = s["size"] // 12
        for i in range(n):
            off = s["offset"] + i * 12
            r_offset, r_info, r_addend = struct.unpack_from("<IIi", elf.data, off)
            r_type = r_info & 0xFF
            r_sym = r_info >> 8
            loc = tbase + r_offset
            S = sym_addr(r_sym) + r_addend
            w = rd(loc)
            if r_type == R["R_RISCV_32"]:
                wr(loc, S & 0xFFFFFFFF)
            elif r_type == R["R_RISCV_HI20"]:
                wr(loc, (w & 0xFFF) | (((S + 0x800) >> 12) << 12) & 0xFFFFF000)
            elif r_type == R["R_RISCV_LO12_I"]:
                wr(loc, (w & 0xFFFFF) | ((S & 0xFFF) << 20))
            elif r_type == R["R_RISCV_LO12_S"]:
                imm = S & 0xFFF
                wr(loc, (w & 0x1FFF07F) | ((imm >> 5) << 25) | ((imm & 31) << 7))
            elif r_type == R["R_RISCV_PCREL_HI20"]:
                delta = S - loc
                wr(loc, (w & 0xFFF) | (((delta + 0x800) >> 12) << 12) & 0xFFFFF000)
            elif r_type in (R["R_RISCV_PCREL_LO12_I"], R["R_RISCV_PCREL_LO12_S"]):
                hi = hi_at[sym_addr(r_sym) + r_addend]
                delta = hi - (sym_addr(r_sym) + r_addend)
                if r_type == R["R_RISCV_PCREL_LO12_I"]:
                    wr(loc, (w & 0xFFFFF) | ((delta & 0xFFF) << 20))
                else:
                    imm = delta & 0xFFF
                    wr(loc, (w & 0x1FFF07F) | ((imm >> 5) << 25) | ((imm & 31) << 7))
            elif r_type in (R["R_RISCV_CALL"], R["R_RISCV_CALL_PLT"]):
                delta = S - loc
                hi = (delta + 0x800) >> 12
                lo = delta - (hi << 12)
                w2 = rd(loc + 4)
                wr(loc, (w & 0xFFF) | ((hi << 12) & 0xFFFFF000))
                wr(loc + 4, (w2 & 0xFFFFF) | ((lo & 0xFFF) << 20))
            elif r_type == R["R_RISCV_JAL"]:
                d = S - loc
                imm = (((d >> 20) & 1) << 31) | (((d >> 1) & 0x3FF) << 21) | \
                      (((d >> 11) & 1) << 20) | (((d >> 12) & 0xFF) << 12)
                wr(loc, (w & 0xFFF) | imm)
            elif r_type == R["R_RISCV_BRANCH"]:
                d = S - loc
                imm = (((d >> 12) & 1) << 31) | (((d >> 5) & 0x3F) << 25) | \
                      (((d >> 1) & 0xF) << 8) | (((d >> 11) & 1) << 7)
                wr(loc, (w & 0x1FFF07F) | imm)
            elif r_type == R["R_RISCV_RELAX"]:
                pass
            else:
                raise NotImplementedError(f"reloc {RREV.get(r_type, r_type)}")

    entry = 0
    for sym in syms:
        if sym["name"] == entry_symbol:
            entry = sym_addr(sym["shndx"] and syms.index(sym))
            entry = placed.get(sym["shndx"], base) + sym["value"]
            break
    return mem, entry

# ============================================================================
# neural_rv32: assembler, reference emulator, threshold runtime
# ============================================================================
RV32_MODEL = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                     "variants", "neural_rv32.safetensors")
MASK32 = 0xFFFFFFFF
MMIO_CONSOLE = 0xFF00


def fcvt_s_w_ref(x: int) -> int:
    """Signed int32 -> float32 word, round-to-nearest-even (host float()
    then pack to float32 is exact RNE for the int32 range)."""
    xi = x if (x & MASK32) < 0x80000000 else (x & MASK32) - (1 << 32)
    return struct.unpack("<I", struct.pack("<f", float(xi)))[0]


RM_NAMES = {"rne": 0, "rtz": 1, "rdn": 2, "rup": 3, "rmm": 4, "dyn": 7}


def _fcvt_round_up(rm: int, sign: int, guard: int, sticky: int, lsb: int) -> int:
    """Whether the truncated magnitude rounds up by one ULP under the RISC-V
    rounding mode (000 RNE, 001 RTZ, 010 RDN, 011 RUP, 100 RMM; 111 DYN has no
    fcsr in this machine and follows RNE). guard is the half bit just below the
    integer, sticky the OR of everything beneath it."""
    if rm == 1:                                        # RTZ: truncate
        return 0
    if rm == 2:                                        # RDN: toward -inf
        return 1 if (sign and (guard or sticky)) else 0
    if rm == 3:                                        # RUP: toward +inf
        return 1 if ((not sign) and (guard or sticky)) else 0
    if rm == 4:                                        # RMM: nearest, ties away
        return 1 if guard else 0
    return 1 if (guard and (sticky or lsb)) else 0     # RNE / DYN


def fcvt_w_s_ref(w: int, rm: int = 1) -> int:
    """float32 word -> signed int32 under rounding mode `rm`, saturating;
    NaN -> INT_MAX, infinities -> INT_MIN/INT_MAX."""
    s, e, f = (w >> 31) & 1, (w >> 23) & 0xFF, w & 0x7FFFFF
    if e == 0xFF:
        return 0x7FFFFFFF if f else (0x80000000 if s else 0x7FFFFFFF)
    mant = ((1 << 23) | f) if e else f                 # subnormal: no implicit 1
    E = (e - 127) if e else -126                       # subnormal exponent 1-bias
    if E >= 31:                                        # |value| >= 2^31: saturate
        if s and E == 31 and f == 0:
            return 0x80000000                          # exactly -2^31
        return 0x80000000 if s else 0x7FFFFFFF
    if E >= 23:
        mag = mant << (E - 23)                          # exact integer, no fraction
    else:
        sh = 23 - E
        mag = mant >> sh
        guard = (mant >> (sh - 1)) & 1
        sticky = 1 if (mant & ((1 << (sh - 1)) - 1)) else 0
        mag += _fcvt_round_up(rm, s, guard, sticky, mag & 1)
    if s:
        return 0x80000000 if mag >= 0x80000000 else (-mag) & MASK32
    return 0x7FFFFFFF if mag > 0x7FFFFFFF else mag & MASK32


def sext(v: int, bits: int) -> int:
    v &= (1 << bits) - 1
    return v - (1 << bits) if v & (1 << (bits - 1)) else v


def u32(v: int) -> int:
    return v & MASK32


# =============================================================================
# Assembler
# =============================================================================

R_FUNCT = {"add": (0, 0), "sub": (0, 32), "sll": (1, 0), "slt": (2, 0),
           "sltu": (3, 0), "xor": (4, 0), "srl": (5, 0), "sra": (5, 32),
           "or": (6, 0), "and": (7, 0),
           "mul": (0, 1), "mulh": (1, 1), "mulhsu": (2, 1), "mulhu": (3, 1),
           "div": (4, 1), "divu": (5, 1), "rem": (6, 1), "remu": (7, 1)}
I_FUNCT = {"addi": 0, "slti": 2, "sltiu": 3, "xori": 4, "ori": 6, "andi": 7}
B_FUNCT = {"beq": 0, "bne": 1, "blt": 4, "bge": 5, "bltu": 6, "bgeu": 7}
L_FUNCT = {"lb": 0, "lh": 1, "lw": 2, "lbu": 4, "lhu": 5}
S_FUNCT = {"sb": 0, "sh": 1, "sw": 2}


class Asm:
    def __init__(self):
        self.words: List = []
        self.labels: Dict[str, int] = {}
        self.data_here = None

    def label(self, name):
        self.labels[name] = len(self.words) * 4

    def emit(self, *parts):
        self.words.append(parts)

    def __getattr__(self, name):
        if name.startswith("__"):
            raise AttributeError(name)
        mn = name.rstrip("_").replace("_", ".")
        def f(*args):
            self.emit(mn, *args)
        return f

    def word(self, v):
        self.emit(".word", v)

    def _enc(self, parts, pc):
        mn = parts[0]
        a = parts[1:]

        def imm(x, pcrel=False):
            if isinstance(x, str):
                t = self.labels[x]
                return t - pc if pcrel else t
            return x

        if mn == ".word":
            return imm(a[0]) & MASK32
        if mn in R_FUNCT:
            f3, f7 = R_FUNCT[mn]
            return (f7 << 25) | (a[2] << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x33
        if mn in I_FUNCT:
            return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (I_FUNCT[mn] << 12) | (a[0] << 7) | 0x13
        if mn in ("slli", "srli", "srai"):
            f7 = 32 if mn == "srai" else 0
            f3 = 1 if mn == "slli" else 5
            return (f7 << 25) | ((a[2] & 31) << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x13
        if mn in L_FUNCT:
            return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (L_FUNCT[mn] << 12) | (a[0] << 7) | 0x03
        if mn in S_FUNCT:
            i = imm(a[2]) & 0xFFF
            return ((i >> 5) << 25) | (a[0] << 20) | (a[1] << 15) | (S_FUNCT[mn] << 12) | ((i & 31) << 7) | 0x23
        if mn in B_FUNCT:
            i = imm(a[2], pcrel=True) & 0x1FFF
            return (((i >> 12) & 1) << 31) | (((i >> 5) & 0x3F) << 25) | (a[1] << 20) | \
                   (a[0] << 15) | (B_FUNCT[mn] << 12) | (((i >> 1) & 0xF) << 8) | (((i >> 11) & 1) << 7) | 0x63
        if mn == "lui":
            return ((imm(a[1]) & 0xFFFFF) << 12) | (a[0] << 7) | 0x37
        if mn == "auipc":
            return ((imm(a[1]) & 0xFFFFF) << 12) | (a[0] << 7) | 0x17
        if mn == "jal":
            i = imm(a[1], pcrel=True) & 0x1FFFFF
            return (((i >> 20) & 1) << 31) | (((i >> 1) & 0x3FF) << 21) | (((i >> 11) & 1) << 20) | \
                   (((i >> 12) & 0xFF) << 12) | (a[0] << 7) | 0x6F
        if mn == "jalr":
            return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (a[0] << 7) | 0x67
        if mn == "flw":
            return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (2 << 12) | (a[0] << 7) | 0x07
        if mn == "fsw":
            i = imm(a[2]) & 0xFFF
            return ((i >> 5) << 25) | (a[0] << 20) | (a[1] << 15) | (2 << 12) | ((i & 31) << 7) | 0x27
        if mn in ("fadd.s", "fsub.s", "fmul.s", "fdiv.s"):
            f7 = {"fadd.s": 0, "fsub.s": 4, "fmul.s": 8, "fdiv.s": 12}[mn]
            return (f7 << 25) | (a[2] << 20) | (a[1] << 15) | (a[0] << 7) | 0x53
        if mn in ("fmadd.s", "fmsub.s", "fnmsub.s", "fnmadd.s"):   # R4-type: rd,rs1,rs2,rs3
            op = {"fmadd.s": 0x43, "fmsub.s": 0x47, "fnmsub.s": 0x4b, "fnmadd.s": 0x4f}[mn]
            return (a[3] << 27) | (a[2] << 20) | (a[1] << 15) | (a[0] << 7) | op
        if mn in ("fle.s", "flt.s", "feq.s"):
            f3 = {"fle.s": 0, "flt.s": 1, "feq.s": 2}[mn]
            return (0x50 << 25) | (a[2] << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x53
        if mn == "fmv.w.x":
            return (0x78 << 25) | (a[1] << 15) | (a[0] << 7) | 0x53
        if mn == "fmv.x.w":
            return (0x70 << 25) | (a[1] << 15) | (a[0] << 7) | 0x53
        if mn == "fcvt.w.s":
            rm = RM_NAMES.get(a[2], a[2]) if len(a) > 2 else 1   # default rtz
            return (0x60 << 25) | (a[1] << 15) | (rm << 12) | (a[0] << 7) | 0x53
        if mn == "fcvt.s.w":
            return (0x68 << 25) | (a[1] << 15) | (a[0] << 7) | 0x53
        if mn in ("fsgnj.s", "fsgnjn.s", "fsgnjx.s"):
            f3 = {"fsgnj.s": 0, "fsgnjn.s": 1, "fsgnjx.s": 2}[mn]
            return (0x10 << 25) | (a[2] << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x53
        if mn == "neur":
            return (a[2] << 20) | (a[1] << 15) | (a[0] << 7) | 0x0B
        if mn == "ecall":
            return 0x73
        if mn == "fence":
            return 0x0F
        raise ValueError(f"unknown mnemonic {mn}")

    def assemble(self, mem_size=65536) -> List[int]:
        mem = [0] * mem_size
        for idx, parts in enumerate(self.words):
            w = self._enc(parts, idx * 4)
            for b in range(4):
                mem[idx * 4 + b] = (w >> (8 * b)) & 0xFF   # little-endian
        return mem


# =============================================================================
# Reference emulator (architectural contract)
# =============================================================================

class RefRV32:
    @staticmethod
    def step(s):
        if s["halted"]:
            return s
        s = dict(s, mem=s["mem"][:], regs=s["regs"][:], fregs=s["fregs"][:],
                 console=s["console"])
        pc = s["pc"]
        w = RefRV32._fetch(s, pc)
        RefRV32._exec(s, pc, w)
        return s

    @staticmethod
    def _fetch(s, pc):
        m = s["mem"]
        return m[pc] | (m[pc + 1] << 8) | (m[pc + 2] << 16) | (m[pc + 3] << 24)

    @staticmethod
    def _load(s, addr, width, signed):
        v = 0
        for b in range(width):
            v |= s["mem"][(addr + b) & 0xFFFF] << (8 * b)
        return u32(sext(v, width * 8)) if signed else v

    @staticmethod
    def _store(s, addr, val, width):
        for b in range(width):
            s["mem"][(addr + b) & 0xFFFF] = (val >> (8 * b)) & 0xFF
        if (addr & 0xFFFF) == MMIO_CONSOLE:
            s["console"] += chr(val & 0xFF)

    @staticmethod
    def _exec(s, pc, w):
        op = w & 0x7F
        rd = (w >> 7) & 31
        f3 = (w >> 12) & 7
        rs1 = (w >> 15) & 31
        rs2 = (w >> 20) & 31
        f7 = w >> 25
        R = s["regs"]
        FR = s["fregs"]
        a, b = R[rs1], R[rs2]
        nxt = (pc + 4) & MASK32

        def wr(v):
            if rd:
                R[rd] = u32(v)

        if op == 0x33:  # OP (+ M)
            if f7 == 1:
                sa, sb = sext(a, 32), sext(b, 32)
                if f3 == 0: wr(a * b)
                elif f3 == 1: wr((sa * sb) >> 32)
                elif f3 == 2: wr((sa * b) >> 32)
                elif f3 == 3: wr((a * b) >> 32)
                elif f3 == 4: wr(MASK32 if b == 0 else (0x80000000 if (sa, sb) == (-2**31, -1) else int(abs(sa) // abs(sb)) * (1 if (sa < 0) == (sb < 0) else -1)))
                elif f3 == 5: wr(MASK32 if b == 0 else a // b)
                elif f3 == 6: wr(sa if b == 0 else (0 if (sa, sb) == (-2**31, -1) else sa - (int(abs(sa) // abs(sb)) * (1 if (sa < 0) == (sb < 0) else -1)) * sb))
                elif f3 == 7: wr(a if b == 0 else a % b)
            else:
                sh = b & 31
                if f3 == 0: wr(a - b if f7 == 32 else a + b)
                elif f3 == 1: wr(a << sh)
                elif f3 == 2: wr(1 if sext(a, 32) < sext(b, 32) else 0)
                elif f3 == 3: wr(1 if a < b else 0)
                elif f3 == 4: wr(a ^ b)
                elif f3 == 5: wr(u32(sext(a, 32) >> sh) if f7 == 32 else a >> sh)
                elif f3 == 6: wr(a | b)
                elif f3 == 7: wr(a & b)
        elif op == 0x13:  # OP-IMM
            imm = sext(w >> 20, 12)
            sh = (w >> 20) & 31
            if f3 == 0: wr(a + imm)
            elif f3 == 1: wr(a << sh)
            elif f3 == 2: wr(1 if sext(a, 32) < imm else 0)
            elif f3 == 3: wr(1 if a < u32(imm) else 0)
            elif f3 == 4: wr(a ^ u32(imm))
            elif f3 == 5: wr(u32(sext(a, 32) >> sh) if f7 == 32 else a >> sh)
            elif f3 == 6: wr(a | u32(imm))
            elif f3 == 7: wr(a & u32(imm))
        elif op == 0x37: wr(w & 0xFFFFF000)              # LUI
        elif op == 0x17: wr(pc + (w & 0xFFFFF000))       # AUIPC
        elif op == 0x6F:                                  # JAL
            imm = sext((((w >> 31) & 1) << 20) | (((w >> 12) & 0xFF) << 12)
                       | (((w >> 20) & 1) << 11) | (((w >> 21) & 0x3FF) << 1), 21)
            wr(nxt)
            nxt = u32(pc + imm)
        elif op == 0x67:                                  # JALR
            t = u32(a + sext(w >> 20, 12)) & ~1
            wr(nxt)
            nxt = t
        elif op == 0x63:                                  # BRANCH
            imm = sext((((w >> 31) & 1) << 12) | (((w >> 7) & 1) << 11)
                       | (((w >> 25) & 0x3F) << 5) | (((w >> 8) & 0xF) << 1), 13)
            sa, sb = sext(a, 32), sext(b, 32)
            take = [a == b, a != b, None, None, sa < sb, sa >= sb, a < b, a >= b][f3]
            if take:
                nxt = u32(pc + imm)
        elif op == 0x03:                                  # LOAD
            addr = u32(a + sext(w >> 20, 12)) & 0xFFFF
            width = [1, 2, 4, 0, 1, 2][f3]
            wr(RefRV32._load(s, addr, width, f3 < 3))
        elif op == 0x23:                                  # STORE
            imm = sext(((w >> 25) << 5) | ((w >> 7) & 31), 12)
            addr = u32(a + imm) & 0xFFFF
            RefRV32._store(s, addr, b, [1, 2, 4][f3])
        elif op == 0x07:                                  # FLW
            addr = u32(a + sext(w >> 20, 12)) & 0xFFFF
            FR[rd] = RefRV32._load(s, addr, 4, False)
        elif op == 0x27:                                  # FSW
            imm = sext(((w >> 25) << 5) | ((w >> 7) & 31), 12)
            addr = u32(a + imm) & 0xFFFF
            RefRV32._store(s, addr, FR[rs2], 4)
        elif op == 0x53:                                  # OP-FP
            from eval import float_add_oracle, float_mul_oracle, float_div_oracle
            fa, fb = FR[rs1], FR[rs2]
            if f7 == 0x00: FR[rd] = float_add_oracle(fa, fb, 8, 23)
            elif f7 == 0x04: FR[rd] = float_add_oracle(fa, fb ^ 0x80000000, 8, 23)
            elif f7 == 0x08: FR[rd] = float_mul_oracle(fa, fb, 8, 23)
            elif f7 == 0x0C: FR[rd] = float_div_oracle(fa, fb, 8, 23)
            elif f7 == 0x10:
                sa, sb = (fa >> 31) & 1, (fb >> 31) & 1
                so = [sb, 1 - sb, sa ^ sb][f3]
                FR[rd] = (fa & 0x7FFFFFFF) | (so << 31)
            elif f7 == 0x50:
                from eval import float_bits_to_value
                x, y = float_bits_to_value(fa, 8, 23), float_bits_to_value(fb, 8, 23)
                wr(1 if [x <= y, x < y, x == y][f3] else 0)
            elif f7 == 0x60: wr(fcvt_w_s_ref(FR[rs1], f3))  # FCVT.W.S (rm=funct3)
            elif f7 == 0x68: FR[rd] = fcvt_s_w_ref(R[rs1]) # FCVT.S.W
            elif f7 == 0x78: FR[rd] = a
            elif f7 == 0x70: wr(FR[rs1])
        elif op in (0x43, 0x47, 0x4b, 0x4f):              # FMADD/FMSUB/FNMSUB/FNMADD.S
            from eval import float_fma_oracle
            rs3 = (w >> 27) & 0x1F
            fa, fb, fc = FR[rs1], FR[rs2], FR[rs3]
            if op in (0x4b, 0x4f): fa ^= 0x80000000       # negate product
            if op in (0x47, 0x4f): fc ^= 0x80000000       # negate addend
            FR[rd] = float_fma_oracle(fa, fb, fc, 8, 23)
        elif op == 0x0B:                                  # NEUR (custom-0)
            x = a & 0xFF
            pos, neg = b & 0xFF, (b >> 8) & 0xFF
            bias = sext(b >> 16, 5)
            acc = bin(x & pos).count("1") - bin(x & neg).count("1") + bias
            wr(1 if acc >= 0 else 0)
        elif op == 0x73:
            s["halted"] = True
        elif op == 0x0F:
            pass
        else:
            raise ValueError(f"illegal instruction {w:#010x} at {pc:#06x}")
        s["pc"] = nxt

    @staticmethod
    def run(state, max_cycles=10000):
        s = state
        n = 0
        while not s["halted"] and n < max_cycles:
            s = RefRV32.step(s)
            n += 1
        return s, n


def rv_state(mem):
    return {"pc": 0, "regs": [0] * 32, "fregs": [0] * 32, "mem": mem[:],
            "halted": False, "console": ""}


# =============================================================================
# Threshold runtime
# =============================================================================

PAIRABLE = (0x33, 0x13, 0x37, 0x17)
DECODE_CLASSES = ("op", "op_imm", "load", "store", "branch", "lui", "auipc",
                  "jal", "jalr", "system", "fp", "flw", "fsw", "fma")


class Rv32ThresholdCPU:
    def __init__(self, path=RV32_MODEL):
        self.T = {}
        with safe_open(path, framework="pt") as f:
            for name in f.keys():
                self.T[name] = f.get_tensor(name).float()
        self.registry = load_metadata(path)["signal_registry"]
        self.alu = GenericThresholdALU(self.T, 32)
        self.fpu = {op: NetlistEvaluator(self.T, self.registry, f"float32.{op}")
                    for op in ("add", "mul", "div", "cmp", "fma")}
        self.decode = NetlistEvaluator(self.T, self.registry, "rv32.decode")
        self.pcseq = NetlistEvaluator(self.T, self.registry, "rv32.pcnext")
        self.pairs_issued = 0

    # ---- gate helpers -------------------------------------------------
    def _g(self, name, inputs):
        return self.alu._g(name, inputs)

    def _not1(self, x):
        return self._g("alu.alu32bit.not.bit0", [float(x)])

    def _bitwise(self, opname, a, b):
        ab = int_to_bits(a, 32)
        bb = int_to_bits(b, 32)
        out = []
        for k in range(32):
            if opname == "xor":
                p = f"alu.alu32bit.xor.bit{k}"
                h1 = self._g(f"{p}.layer1.or", [ab[k], bb[k]])
                h2 = self._g(f"{p}.layer1.nand", [ab[k], bb[k]])
                out.append(self._g(f"{p}.layer2", [h1, h2]))
            else:
                out.append(self._g(f"alu.alu32bit.{opname}.bit{k}", [ab[k], bb[k]]))
        return bits_to_int(out)

    def _not32(self, a):
        ab = int_to_bits(a, 32)
        return bits_to_int([self._g(f"alu.alu32bit.not.bit{k}", [ab[k]]) for k in range(32)])

    def _barrel_left(self, val, sh):
        layer_in = int_to_bits(val, 32)
        for layer in range(5):
            amount = 1 << (4 - layer)
            sel = (sh >> (4 - layer)) & 1
            nxt = []
            for bit in range(32):
                p = f"combinational.barrelshifter32.layer{layer}.bit{bit}"
                ns = self._g(f"{p}.not_sel", [float(sel)])
                src = layer_in[bit + amount] if bit + amount < 32 else 0
                a_ = self._g(f"{p}.and_a", [layer_in[bit], ns])
                b_ = self._g(f"{p}.and_b", [float(src), float(sel)])
                nxt.append(self._g(f"{p}.or", [a_, b_]))
            layer_in = nxt
        return bits_to_int(layer_in)

    def sll(self, a, sh):
        return self._barrel_left(a, sh & 31)

    def srl(self, a, sh):
        rev = int("{:032b}".format(u32(a))[::-1], 2)
        return int("{:032b}".format(self._barrel_left(rev, sh & 31))[::-1], 2)

    def sra(self, a, sh):
        sign = (a >> 31) & 1
        x1 = self.srl(a, sh)
        x2 = self._not32(self.srl(self._not32(a), sh))
        o1, o2 = int_to_bits(x1, 32), int_to_bits(x2, 32)
        out = []
        for k in range(32):
            p = f"rv32.sra_mux.bit{k}"
            ns = self._g(f"{p}.not_sel", [float(sign)])
            a_ = self._g(f"{p}.and_a", [o1[k], ns])
            b_ = self._g(f"{p}.and_b", [o2[k], float(sign)])
            out.append(self._g(f"{p}.or", [a_, b_]))
        return bits_to_int(out)

    def add32(self, a, b):
        return self.alu.add_n(a, b, 32)[0]

    def sub32(self, a, b):
        return self.alu.sub_n(a, b, 32)[0]

    def ucmp(self, a, b, kind):
        return self.alu.cmp_n(a, b, kind, 32)

    def scmp_lt(self, a, b):
        fa = (a & 0x7FFFFFFF) | (self._not1((a >> 31) & 1) << 31)
        fb = (b & 0x7FFFFFFF) | (self._not1((b >> 31) & 1) << 31)
        return self.ucmp(fa, fb, "lessthan")

    def mulhu(self, a, b):
        """Full 64-bit product high word via pp gates + rv32.mulacc chains."""
        ab = int_to_bits(a, 32)
        bb = int_to_bits(b, 32)

        def pp(i_lsb, j_lsb):  # LSB-first indices onto MSB-first gate names
            return self._g(f"alu.alu32bit.mul.pp.a{31 - i_lsb}b{31 - j_lsb}",
                           [ab[31 - i_lsb], bb[31 - j_lsb]])

        acc = [pp(k, 0) if k <= 31 else 0.0 for k in range(64)]
        for t in range(31):
            row = t + 1
            carry = 0.0
            nxt = []
            for k in range(64):
                b_in = pp(k - row, row) if row <= k <= row + 31 else 0.0
                s, carry = self.alu._fa(f"rv32.mulacc.s{t}.fa{k}", acc[k], b_in, carry)
                nxt.append(s)
            acc = nxt
        return bits_to_int(list(reversed([int(x) for x in acc[32:]])))

    def _cond_mask(self, val, cond_bit):
        vb = int_to_bits(val, 32)
        return bits_to_int([self._g(f"alu.alu32bit.and.bit{k}", [vb[k], float(cond_bit)])
                            for k in range(32)])

    def mulh(self, a, b, sign_a, sign_b):
        h = self.mulhu(a, b)
        if sign_a:
            h = self.sub32(h, self._cond_mask(b, (a >> 31) & 1))
        if sign_b:
            h = self.sub32(h, self._cond_mask(a, (b >> 31) & 1))
        return h

    def divu(self, a, b):
        if b == 0:
            return MASK32, a
        a_bits = int_to_bits(a, 32)
        q = 0
        rem = 0
        for stage in range(32):
            rem = ((rem << 1) | a_bits[stage]) & MASK32
            prefix = f"alu.alu32bit.div.stage{stage}"
            finals = {"gt": f"{prefix}.cmp_bc.gt", "lt": f"{prefix}.cmp_bc.lt",
                      "eq": f"{prefix}.cmp_bc.eq", "ge": f"{prefix}.cmp",
                      "le": f"{prefix}.cmp_bc.le"}
            ge = self.alu._cmp_bit_cascade(f"{prefix}.cmp_bc", finals,
                                           int_to_bits(rem, 32), int_to_bits(b, 32), "ge")
            if ge:
                rem = self.sub32(rem, b)
                q = (q << 1) | 1
            else:
                q <<= 1
        return u32(q), rem

    def neg32(self, a):
        ab = int_to_bits(a, 32)
        nb = [self._g(f"alu.alu32bit.neg.not.bit{k}", [ab[k]]) for k in range(32)]
        carry = 1.0
        out = [0] * 32
        for k in range(31, -1, -1):
            p = f"alu.alu32bit.neg.inc.bit{k}"
            h1 = self._g(f"{p}.xor.layer1.or", [nb[k], carry])
            h2 = self._g(f"{p}.xor.layer1.nand", [nb[k], carry])
            out[k] = self._g(f"{p}.xor.layer2", [h1, h2])
            carry = self._g(f"{p}.carry", [nb[k], carry])
        return bits_to_int(out)

    def divs(self, a, b):
        if b == 0:
            return MASK32, a
        if a == 0x80000000 and b == MASK32:
            return 0x80000000, 0
        na, nb_ = (a >> 31) & 1, (b >> 31) & 1
        ua = self.neg32(a) if na else a
        ub = self.neg32(b) if nb_ else b
        q, r = self.divu(ua, ub)
        if na != nb_:
            q = self.neg32(q)
        if na:
            r = self.neg32(r)
        return q, r

    def priority_encode(self, v):
        """MSB-first position (0..31) of the highest set bit, via the
        combinational.priorityencoder32 gates. v must be nonzero."""
        vb = int_to_bits(v, 32)  # MSB-first, bit index 0 = MSB
        any_higher = [0]
        for pos in range(1, 32):
            any_higher.append(self._g(f"combinational.priorityencoder32.any_higher{pos}",
                                      [float(vb[i]) for i in range(pos)]))
        is_high = [float(vb[0])]
        for pos in range(1, 32):
            nh = self._g(f"combinational.priorityencoder32.is_highest{pos}.not_higher",
                         [float(any_higher[pos])])
            is_high.append(self._g(f"combinational.priorityencoder32.is_highest{pos}.and",
                                   [float(vb[pos]), nh]))
        out = 0
        for bit in range(5):
            rel = [is_high[pos] for pos in range(32) if (pos >> bit) & 1]
            out |= self._g(f"combinational.priorityencoder32.out{bit}", [float(x) for x in rel]) << bit
        return out

    def fcvt_s_w(self, x):
        """Signed int32 -> float32, round-to-nearest-even. Sign, leading-one
        detect (priority encoder), left-normalize (barrel shifter), round."""
        if x == 0:
            return 0
        s = (x >> 31) & 1
        u = self.neg32(x) if s else x
        p = self.priority_encode(u)                 # leading one at MSB-first pos p
        shifted = self.sll(u, p)                     # leading one now at bit 31
        exp = 127 + (31 - p)
        frac = (shifted >> 8) & 0x7FFFFF
        guard = (shifted >> 7) & 1
        below = shifted & 0x7F
        if guard and ((frac & 1) or below):
            frac += 1
            if frac > 0x7FFFFF:
                frac = 0
                exp += 1
        return (s << 31) | (exp << 23) | frac

    def fcvt_w_s(self, w, rm=1):
        """float32 -> signed int32 under rounding mode `rm`, saturating. The
        shifts run through the barrel shifter; the round decision is control
        logic, as in fcvt_s_w. NaN -> INT_MAX; overflow saturates."""
        s, e, f = (w >> 31) & 1, (w >> 23) & 0xFF, w & 0x7FFFFF
        if e == 0xFF:
            return 0x7FFFFFFF if f else (0x80000000 if s else 0x7FFFFFFF)
        mant = ((1 << 23) | f) if e else f
        E = (e - 127) if e else -126
        if E >= 31:
            if s and E == 31 and f == 0:
                return 0x80000000                    # exactly -2^31
            return 0x80000000 if s else 0x7FFFFFFF
        if E >= 23:
            mag = self.sll(mant, E - 23)
        else:
            sh = 23 - E
            mag = self.srl(mant, sh) if sh <= 31 else 0
            guard = (mant >> (sh - 1)) & 1
            sticky = 1 if (mant & ((1 << (sh - 1)) - 1)) else 0
            mag += _fcvt_round_up(rm, s, guard, sticky, mag & 1)
        if s:
            return 0x80000000 if mag >= 0x80000000 else self.neg32(mag)
        return 0x7FFFFFFF if mag > 0x7FFFFFFF else (mag & MASK32)

    def neur(self, x_reg, w_reg):
        x = int_to_bits(x_reg & 0xFF, 8)      # MSB-first bits of low byte
        pos = int_to_bits(w_reg & 0xFF, 8)
        neg = int_to_bits((w_reg >> 8) & 0xFF, 8)
        p_bits = [self._g(f"rv32.neur.pos.bit{i}", [x[7 - i], pos[7 - i]]) for i in range(8)]
        n_bits = [self._g(f"rv32.neur.neg.bit{i}", [x[7 - i], neg[7 - i]]) for i in range(8)]

        def pcount(side, bits8):
            fa = lambda t, a_, b_, c_: self.alu._fa(f"rv32.neur.{side}.fa{t}", a_, b_, c_)
            s0, c0 = fa(0, bits8[0], bits8[1], bits8[2])
            s1, c1 = fa(1, bits8[3], bits8[4], bits8[5])
            s2, c2 = fa(2, bits8[6], bits8[7], 0.0)
            S, C2 = fa(3, s0, s1, s2)
            T, C4 = fa(4, c0, c1, c2)
            U, C5 = fa(5, C2, T, 0.0)
            V, W = fa(6, C4, C5, 0.0)
            return [S, U, V, W]  # LSB-first count bits

        P = pcount("pcnt", p_bits)
        N = pcount("ncnt", n_bits)
        nnot = [self._g(f"rv32.neur.nnot.bit{k}", [N[k] if k < 4 else 0.0]) for k in range(6)]
        carry = 1.0
        diff = []
        for k in range(6):
            a_ = P[k] if k < 4 else 0.0
            s, carry = self.alu._fa(f"rv32.neur.diff.fa{k}", a_, nnot[k], carry)
            diff.append(s)
        bias = (w_reg >> 16) & 0x1F
        bias6 = [(bias >> k) & 1 if k < 4 else (bias >> 4) & 1 for k in range(6)]  # sext5
        carry = 0.0
        act = []
        for k in range(6):
            s, carry = self.alu._fa(f"rv32.neur.act.fa{k}", diff[k], float(bias6[k]), carry)
            act.append(s)
        return int(self._g("rv32.neur.out", [act[5]]))

    def hazard_eq(self, pair, i, j):
        ib = int_to_bits(i, 5)
        jb = int_to_bits(j, 5)
        eqs = []
        for k in range(5):
            p = f"rv32.hazard.{pair}.bit{k}"
            h_and = self._g(f"{p}.eq.layer1.and", [ib[k], jb[k]])
            h_nor = self._g(f"{p}.eq.layer1.nor", [ib[k], jb[k]])
            eqs.append(self._g(f"{p}.eq", [h_and, h_nor]))
        return int(self._g(f"rv32.hazard.{pair}.all", eqs))

    # ---- memory through the packed circuits ---------------------------
    def _addr_decode(self, addr):
        bits = torch.tensor(int_to_bits(addr, 16), dtype=torch.float32)
        return heaviside((self.T["memory.addr_decode.weight"] * bits).sum(dim=1)
                         + self.T["memory.addr_decode.bias"])

    def mem_read_byte(self, mem_bits, addr):
        sel = self._addr_decode(addr & 0xFFFF)
        and_w = self.T["memory.read.and.weight"]
        and_b = self.T["memory.read.and.bias"]
        or_w = self.T["memory.read.or.weight"]
        or_b = self.T["memory.read.or.bias"]
        v = 0
        for bit in range(8):
            inp = torch.stack([mem_bits[:, bit], sel], dim=1)
            and_out = heaviside((inp * and_w[bit]).sum(dim=1) + and_b[bit])
            v = (v << 1) | int(heaviside((and_out * or_w[bit]).sum() + or_b[bit]).item())
        return v

    def mem_write_byte(self, mem_bits, addr, value):
        sel = self._addr_decode(addr & 0xFFFF)
        data_bits = torch.tensor(int_to_bits(value & 0xFF, 8), dtype=torch.float32)
        we = torch.ones_like(sel)
        write_sel = heaviside((torch.stack([sel, we], dim=1)
                               * self.T["memory.write.sel.weight"]).sum(dim=1)
                              + self.T["memory.write.sel.bias"])
        nsel = heaviside(write_sel * self.T["memory.write.nsel.weight"].squeeze(1)
                         + self.T["memory.write.nsel.bias"])
        for bit in range(8):
            old = mem_bits[:, bit]
            dbit = data_bits[bit].expand(65536)
            a_old = heaviside((torch.stack([old, nsel], dim=1)
                               * self.T["memory.write.and_old.weight"][:, bit]).sum(dim=1)
                              + self.T["memory.write.and_old.bias"][:, bit])
            a_new = heaviside((torch.stack([dbit, write_sel], dim=1)
                               * self.T["memory.write.and_new.weight"][:, bit]).sum(dim=1)
                              + self.T["memory.write.and_new.bias"][:, bit])
            mem_bits[:, bit] = heaviside((torch.stack([a_old, a_new], dim=1)
                                          * self.T["memory.write.or.weight"][:, bit]).sum(dim=1)
                                         + self.T["memory.write.or.bias"][:, bit])

    # ---- execution ----------------------------------------------------
    def _fetch32(self, mem_bits, pc):
        return sum(self.mem_read_byte(mem_bits, pc + b) << (8 * b) for b in range(4))

    def _load(self, s, addr, width, signed):
        v = sum(self.mem_read_byte(s["_mem_bits"], (addr + b) & 0xFFFF) << (8 * b)
                for b in range(width))
        return u32(sext(v, width * 8)) if signed else v

    def _store(self, s, addr, val, width):
        for b in range(width):
            byte = (val >> (8 * b)) & 0xFF
            self.mem_write_byte(s["_mem_bits"], (addr + b) & 0xFFFF, byte)
            s["mem"][(addr + b) & 0xFFFF] = byte
        if (addr & 0xFFFF) == MMIO_CONSOLE:
            s["console"] += chr(val & 0xFF)

    def _exec_alu_class(self, s, pc, w):
        """Execute one OP / OP-IMM / LUI / AUIPC instruction; returns (rd, value)."""
        op = w & 0x7F
        rd = (w >> 7) & 31
        f3 = (w >> 12) & 7
        rs1 = (w >> 15) & 31
        rs2 = (w >> 20) & 31
        f7 = w >> 25
        R = s["regs"]
        a = R[rs1]
        if op in (0x37, 0x17, 0x13):
            _, imm_g = self._decode(w)            # gate-decoded immediate
        if op == 0x37:
            return rd, imm_g
        if op == 0x17:
            return rd, self.add32(pc, imm_g)
        if op == 0x13:
            b = imm_g
            sh = imm_g & 31
            imm_mode = True
        else:
            b = R[rs2]
            sh = b & 31
            imm_mode = False
        if f3 == 0:
            if not imm_mode and f7 == 32:
                return rd, self.sub32(a, b)
            return rd, self.add32(a, b)
        if f3 == 1:
            return rd, self.sll(a, sh)
        if f3 == 2:
            return rd, self.scmp_lt(a, b)
        if f3 == 3:
            return rd, self.ucmp(a, b, "lessthan")
        if f3 == 4:
            return rd, self._bitwise("xor", a, b)
        if f3 == 5:
            return rd, (self.sra(a, sh) if f7 == 32 else self.srl(a, sh))
        if f3 == 6:
            return rd, self._bitwise("or", a, b)
        return rd, self._bitwise("and", a, b)

    def step(self, s):
        if s["halted"]:
            return s
        s = dict(s, mem=s["mem"][:], regs=s["regs"][:], fregs=s["fregs"][:],
                 console=s["console"], _mem_bits=s["_mem_bits"].clone())
        pc = s["pc"]
        w = self._fetch32(s["_mem_bits"], pc)
        op = w & 0x7F

        def plain_alu(word):
            return (word & 0x7F) in PAIRABLE and not \
                ((word & 0x7F) == 0x33 and (word >> 25) == 1)

        # Dual issue: both slots plain ALU-class and hazard-free by gates.
        if plain_alu(w):
            w2 = self._fetch32(s["_mem_bits"], (pc + 4) & 0xFFFF)
            if plain_alu(w2):
                rd1 = (w >> 7) & 31
                rs1b = (w2 >> 15) & 31
                rs2b = (w2 >> 20) & 31
                rd2 = (w2 >> 7) & 31
                uses_rs1 = (w2 & 0x7F) in (0x33, 0x13)
                uses_rs2 = (w2 & 0x7F) == 0x33
                raw = (uses_rs1 and self.hazard_eq("rd1_rs1b", rd1, rs1b)) or \
                      (uses_rs2 and self.hazard_eq("rd1_rs2b", rd1, rs2b))
                waw = self.hazard_eq("rd1_rd2", rd1, rd2)
                if rd1 != 0 and not raw and not waw:
                    d1, v1 = self._exec_alu_class(s, pc, w)
                    d2, v2 = self._exec_alu_class(s, (pc + 4) & MASK32, w2)
                    if d1:
                        s["regs"][d1] = u32(v1)
                    if d2:
                        s["regs"][d2] = u32(v2)
                    s["pc"] = (pc + 8) & MASK32
                    self.pairs_issued += 1
                    return s

        rd = (w >> 7) & 31
        f3 = (w >> 12) & 7
        rs1 = (w >> 15) & 31
        rs2 = (w >> 20) & 31
        f7 = w >> 25
        R = s["regs"]
        FR = s["fregs"]
        a, b = R[rs1], R[rs2]
        cls, imm_g = self._decode(w)              # gate decode: classes + immediate
        pc4 = self.add32(pc, 4)                    # PC+4 through the gate adder
        branch_taken = 0

        def wr(v):
            if rd:
                R[rd] = u32(v)

        if plain_alu(w):
            d, v = self._exec_alu_class(s, pc, w)
            if d:
                R[d] = u32(v)
        if cls["op"] and f7 == 1:
            if f3 == 0:
                wr(self.alu.mul_n(a, b, 32))
            elif f3 == 1:
                wr(self.mulh(a, b, True, True))
            elif f3 == 2:
                wr(self.mulh(a, b, True, False))
            elif f3 == 3:
                wr(self.mulhu(a, b))
            elif f3 == 4:
                wr(self.divs(a, b)[0])
            elif f3 == 5:
                wr(self.divu(a, b)[0])
            elif f3 == 6:
                wr(self.divs(a, b)[1])
            elif f3 == 7:
                wr(self.divu(a, b)[1])
        elif cls["jal"]:
            wr(pc4)                               # return address
        elif cls["jalr"]:
            wr(pc4)
        elif cls["branch"]:
            eq = self.ucmp(a, b, "eq")
            sl = self.scmp_lt(a, b)
            ul = self.ucmp(a, b, "lessthan")
            branch_taken = [eq, 1 - eq, 0, 0, sl, 1 - sl, ul, 1 - ul][f3]
        elif cls["load"]:
            addr = self.add32(a, imm_g) & 0xFFFF
            width = [1, 2, 4, 0, 1, 2][f3]
            wr(self._load(s, addr, width, f3 < 3))
        elif cls["store"]:
            addr = self.add32(a, imm_g) & 0xFFFF
            self._store(s, addr, b, [1, 2, 4][f3])
        elif cls["flw"]:
            addr = self.add32(a, imm_g) & 0xFFFF
            FR[rd] = self._load(s, addr, 4, False)
        elif cls["fsw"]:
            addr = self.add32(a, imm_g) & 0xFFFF
            self._store(s, addr, FR[rs2], 4)
        elif cls["fp"]:
            fa, fb = FR[rs1], FR[rs2]
            if f7 in (0x00, 0x04, 0x08, 0x0C):
                circ = {0x00: "add", 0x04: "add", 0x08: "mul", 0x0C: "div"}[f7]
                if f7 == 0x04:
                    fb ^= 0x80000000  # subtraction: flip the addend's sign
                FR[rd] = self._fpu_word(circ, fa, fb)
            elif f7 == 0x50:
                out = self.fpu["cmp"].run(self._fpu_ext(fa, fb))
                gate = {0: "le", 1: "lt", 2: "eq"}[f3]
                wr(int(out[f"float32.cmp.{gate}.result"][0, 0].item()))
            elif f7 == 0x10:  # FSGNJ family: result = {sign, rs1[30:0]}
                sa = (FR[rs1] >> 31) & 1
                sb = (FR[rs2] >> 31) & 1
                if f3 == 0:
                    s_out = sb                       # FSGNJ.S
                elif f3 == 1:
                    s_out = self._not1(sb)            # FSGNJN.S
                else:
                    s_out = self._bitwise("xor", sa << 31, sb << 31) >> 31  # FSGNJX.S
                FR[rd] = (FR[rs1] & 0x7FFFFFFF) | (s_out << 31)
            elif f7 == 0x60:  # FCVT.W.S: float -> signed int, gate-routed
                wr(self.fcvt_w_s(FR[rs1], f3))
            elif f7 == 0x68:  # FCVT.S.W: signed int -> float, gate-routed
                FR[rd] = self.fcvt_s_w(R[rs1])
            elif f7 == 0x78:
                FR[rd] = a
            elif f7 == 0x70:
                wr(FR[rs1])
        elif cls["fma"]:                          # FMADD/FMSUB/FNMSUB/FNMADD.S
            rs3 = (w >> 27) & 0x1F
            fa, fb, fc = FR[rs1], FR[rs2], FR[rs3]
            if op in (0x4b, 0x4f):               # negate the product
                fa ^= 0x80000000
            if op in (0x47, 0x4f):               # negate the addend
                fc ^= 0x80000000
            FR[rd] = self._fpu_word3(fa, fb, fc)
        elif op == 0x0B:
            wr(self.neur(a, b))
        elif cls["system"]:
            s["halted"] = True
        elif op == 0x0F:
            pass
        elif op not in PAIRABLE:
            raise ValueError(f"illegal instruction {w:#010x} at {pc:#06x}")
        # PC sequencing through the gate mux: pc4 / pc+imm (branch,jal) / (rs1+imm)&~1 (jalr).
        pcimm = self.add32(pc, imm_g)
        jalr_t = self.add32(a, imm_g) & ~1
        sel_target = 1 if (branch_taken or cls["jal"]) else 0
        s["pc"] = self._pcnext(pc4, pcimm, jalr_t, sel_target, cls["jalr"])
        return s

    def _fpu_ext(self, fa, fb):
        ext = {}
        for i in range(32):
            ext[f"$a[{i}]"] = (fa >> (31 - i)) & 1
            ext[f"$b[{i}]"] = (fb >> (31 - i)) & 1
        return ext

    def _fpu_word(self, circ, fa, fb):
        out = self.fpu[circ].run(self._fpu_ext(fa, fb))
        p = f"float32.{circ}"
        v = int(out[f"{p}.sign_out"][0, 0].item()) << 31
        for k in range(8):
            v |= int(out[f"{p}.exp_out.bit{k}"][0, 0].item()) << (23 + k)
        for k in range(23):
            v |= int(out[f"{p}.frac_out.bit{k}"][0, 0].item()) << k
        return v

    def _fpu_ext3(self, fa, fb, fc):
        ext = {}
        for i in range(32):
            ext[f"$a[{i}]"] = (fa >> (31 - i)) & 1
            ext[f"$b[{i}]"] = (fb >> (31 - i)) & 1
            ext[f"$c[{i}]"] = (fc >> (31 - i)) & 1
        return ext

    def _fpu_word3(self, fa, fb, fc):
        out = self.fpu["fma"].run(self._fpu_ext3(fa, fb, fc))
        p = "float32.fma"
        v = int(out[f"{p}.sign_out"][0, 0].item()) << 31
        for k in range(8):
            v |= int(out[f"{p}.exp_out.bit{k}"][0, 0].item()) << (23 + k)
        for k in range(23):
            v |= int(out[f"{p}.frac_out.bit{k}"][0, 0].item()) << k
        return v

    def _decode(self, w):
        """Instruction decode through the gate network: opcode-class one-hots
        and the sign-extended immediate for the active format."""
        ext = {f"$rv32_instr[{k}]": (w >> k) & 1 for k in range(32)}
        out = self.decode.run(ext)
        cls = {n: int(out[f"rv32.decode.is_{n}"][0, 0].item()) for n in DECODE_CLASSES}
        imm = 0
        for k in range(32):
            imm |= int(out[f"rv32.decode.imm.bit{k}"][0, 0].item()) << k
        return cls, imm

    def _pcnext(self, pc4, pcimm, jalr, sel_target, sel_jalr):
        """Next-PC selection through the gate mux (pc4 / pc+imm / (rs1+imm)&~1)."""
        ext = {"$rv32_sel_target": sel_target, "$rv32_sel_jalr": sel_jalr}
        for k in range(32):
            ext[f"$rv32_pc4[{k}]"] = (pc4 >> k) & 1
            ext[f"$rv32_pcimm[{k}]"] = (pcimm >> k) & 1
            ext[f"$rv32_jalr[{k}]"] = (jalr >> k) & 1
        out = self.pcseq.run(ext)
        v = 0
        for k in range(32):
            v |= int(out[f"rv32.pcnext.bit{k}.or"][0, 0].item()) << k
        return v

    def run(self, state, max_cycles=400):
        s = dict(state)
        s["_mem_bits"] = torch.tensor(
            [int_to_bits(byte, 8) for byte in s["mem"]], dtype=torch.float32)
        n = 0
        while not s["halted"] and n < max_cycles:
            s = self.step(s)
            n += 1
        s.pop("_mem_bits", None)
        return s, n

# ============================================================================
# neural_subleq8: threshold runtime, reference, assembler
# ============================================================================
SUBLEQ_MODEL = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                     "variants", "neural_subleq8.safetensors")
HALT_PC = 0xFF


def load_tensors(path):
    out = {}
    with safe_open(path, framework="pt") as f:
        for name in f.keys():
            out[name] = f.get_tensor(name).float()
    return out


class SubleqThresholdCPU:
    """SUBLEQ-8 runtime: memory traffic through the packed memory circuits,
    the subtract/branch datapath through the subleq gates via the shipped
    netlist. Instruction sequencing (three consecutive fetches) is wiring."""

    def __init__(self, tensors, registry):
        self.T = tensors
        self.ne = NetlistEvaluator(tensors, registry, "subleq")

    def _addr_decode(self, addr):
        bits = torch.tensor([(addr >> (7 - i)) & 1 for i in range(8)],
                            dtype=torch.float32)
        w = self.T["memory.addr_decode.weight"]
        b = self.T["memory.addr_decode.bias"]
        return heaviside((w * bits).sum(dim=1) + b)

    def mem_read(self, mem, addr):
        sel = self._addr_decode(addr)
        mem_bits = torch.tensor([[(byte >> (7 - i)) & 1 for i in range(8)]
                                 for byte in mem], dtype=torch.float32)
        and_w = self.T["memory.read.and.weight"]
        and_b = self.T["memory.read.and.bias"]
        or_w = self.T["memory.read.or.weight"]
        or_b = self.T["memory.read.or.bias"]
        out = 0
        for bit in range(8):
            inp = torch.stack([mem_bits[:, bit], sel], dim=1)
            and_out = heaviside((inp * and_w[bit]).sum(dim=1) + and_b[bit])
            out = (out << 1) | int(heaviside((and_out * or_w[bit]).sum() + or_b[bit]).item())
        return out

    def mem_write(self, mem, addr, value):
        sel = self._addr_decode(addr)
        data_bits = torch.tensor([(value >> (7 - i)) & 1 for i in range(8)],
                                 dtype=torch.float32)
        mem_bits = torch.tensor([[(byte >> (7 - i)) & 1 for i in range(8)]
                                 for byte in mem], dtype=torch.float32)
        sel_w = self.T["memory.write.sel.weight"]
        sel_b = self.T["memory.write.sel.bias"]
        nsel_w = self.T["memory.write.nsel.weight"].squeeze(1)
        nsel_b = self.T["memory.write.nsel.bias"]
        we = torch.ones_like(sel)
        write_sel = heaviside((torch.stack([sel, we], dim=1) * sel_w).sum(dim=1) + sel_b)
        nsel = heaviside(write_sel * nsel_w + nsel_b)
        ao_w = self.T["memory.write.and_old.weight"]
        ao_b = self.T["memory.write.and_old.bias"]
        an_w = self.T["memory.write.and_new.weight"]
        an_b = self.T["memory.write.and_new.bias"]
        or_w = self.T["memory.write.or.weight"]
        or_b = self.T["memory.write.or.bias"]
        for bit in range(8):
            old = mem_bits[:, bit]
            dbit = data_bits[bit].expand(256)
            a_old = heaviside((torch.stack([old, nsel], dim=1) * ao_w[:, bit]).sum(dim=1) + ao_b[:, bit])
            a_new = heaviside((torch.stack([dbit, write_sel], dim=1) * an_w[:, bit]).sum(dim=1) + an_b[:, bit])
            mem_bits[:, bit] = heaviside(
                (torch.stack([a_old, a_new], dim=1) * or_w[:, bit]).sum(dim=1) + or_b[:, bit])
        return [int(sum(int(mem_bits[i, b].item()) << (7 - b) for b in range(8)))
                for i in range(256)]

    def step(self, state):
        if state["halted"]:
            return state
        s = dict(state)
        s["mem"] = state["mem"][:]
        pc = s["pc"]
        A = self.mem_read(s["mem"], pc)
        B = self.mem_read(s["mem"], (pc + 1) & 0xFF)
        C = self.mem_read(s["mem"], (pc + 2) & 0xFF)
        a = self.mem_read(s["mem"], A)
        b = self.mem_read(s["mem"], B)
        ext = {}
        for k in range(8):
            ext[f"$a[{k}]"] = (a >> k) & 1
            ext[f"$b[{k}]"] = (b >> k) & 1
            ext[f"$pc[{k}]"] = (pc >> k) & 1
            ext[f"$c[{k}]"] = (C >> k) & 1
        out = self.ne.run(ext)
        r = sum(int(out[f"subleq.sub.fa{k}.ha2.sum.layer2"][0, 0].item()) << k
                for k in range(8))
        next_pc = sum(int(out[f"subleq.pc_mux.bit{k}.or"][0, 0].item()) << k
                      for k in range(8))
        s["mem"] = self.mem_write(s["mem"], B, r)
        s["pc"] = next_pc
        if next_pc == HALT_PC:
            s["halted"] = True
        return s

    def run(self, state, max_cycles=200):
        s = state
        n = 0
        while not s["halted"] and n < max_cycles:
            s = self.step(s)
            n += 1
        return s, n


def ref_step(state):
    if state["halted"]:
        return state
    s = dict(state)
    s["mem"] = state["mem"][:]
    pc = s["pc"]
    A, B, C = s["mem"][pc], s["mem"][(pc + 1) & 0xFF], s["mem"][(pc + 2) & 0xFF]
    r = (s["mem"][B] - s["mem"][A]) & 0xFF
    s["mem"][B] = r
    s["pc"] = C if (r == 0 or r >= 0x80) else (pc + 3) & 0xFF
    if s["pc"] == HALT_PC:
        s["halted"] = True
    return s


def sq(*triples):
    """Assemble SUBLEQ triples into a 256-byte image."""
    mem = [0] * 256
    i = 0
    for t in triples:
        for v in t:
            mem[i] = v & 0xFF
            i += 1
    return mem


def subleq_test() -> int:
    T = load_tensors(SUBLEQ_MODEL)
    reg = load_metadata(SUBLEQ_MODEL)["signal_registry"]

    # 1. Exhaustive datapath: all 65,536 (a, b) pairs through the shipped
    # wiring in one batched netlist evaluation.
    ne = NetlistEvaluator(T, reg, "subleq")
    a_v = torch.arange(65536, dtype=torch.long) % 256
    b_v = torch.arange(65536, dtype=torch.long) // 256
    ext = {}
    for k in range(8):
        ext[f"$a[{k}]"] = ((a_v >> k) & 1).float()
        ext[f"$b[{k}]"] = ((b_v >> k) & 1).float()
        ext[f"$pc[{k}]"] = torch.zeros(65536)
        ext[f"$c[{k}]"] = torch.zeros(65536)
    out = ne.run(ext)
    r = torch.zeros(65536, dtype=torch.float64)
    for k in range(8):
        r += out[f"subleq.sub.fa{k}.ha2.sum.layer2"][:, 0].double() * (1 << k)
    exp_r = ((b_v - a_v) & 0xFF).double()
    leq = out["subleq.leq"][:, 0].double()
    exp_leq = ((exp_r == 0) | (exp_r >= 128)).double()
    ok_r = int((r == exp_r).sum().item())
    ok_l = int((leq == exp_leq).sum().item())
    print(f"datapath sub  exhaustive: {ok_r}/65536")
    print(f"datapath leq  exhaustive: {ok_l}/65536")
    if ok_r != 65536 or ok_l != 65536:
        return 1

    # 2. PC mux + increment across representative PC/target/leq settings.
    pcs = torch.tensor([0, 3, 5, 100, 250, 252], dtype=torch.long).repeat_interleave(2)
    cs = torch.tensor([7, 7, 42, 42, 255, 255, 0, 0, 9, 9, 30, 30], dtype=torch.long)
    a2 = torch.tensor([0, 1] * 6, dtype=torch.long)   # a=0,b=0 -> leq; a=1,b=0 -> r=255 leq... use b to vary
    b2 = torch.tensor([0, 1] * 6, dtype=torch.long)   # (0,0)->r0 leq taken; (1,1)->r0 leq taken
    # craft: even rows leq taken (a=b), odd rows not taken (b-a=1 -> r=1 >0)
    a2 = torch.tensor([5, 4] * 6, dtype=torch.long)
    b2 = torch.tensor([5, 5] * 6, dtype=torch.long)
    ext = {}
    for k in range(8):
        ext[f"$a[{k}]"] = ((a2 >> k) & 1).float()
        ext[f"$b[{k}]"] = ((b2 >> k) & 1).float()
        ext[f"$pc[{k}]"] = ((pcs >> k) & 1).float()
        ext[f"$c[{k}]"] = ((cs >> k) & 1).float()
    out = ne.run(ext)
    npc = torch.zeros(12, dtype=torch.float64)
    for k in range(8):
        npc += out[f"subleq.pc_mux.bit{k}.or"][:, 0].double() * (1 << k)
    exp_npc = torch.where(torch.arange(12) % 2 == 0, cs.double(), ((pcs + 3) & 0xFF).double())
    ok_pc = int((npc == exp_npc).sum().item())
    print(f"datapath pc mux/inc     : {ok_pc}/12")
    if ok_pc != 12:
        return 1

    # 3. Program suite, threshold vs reference, full-state lockstep.
    cpu = SubleqThresholdCPU(T, reg)
    Z = 0x80  # scratch zero cell

    progs = []
    # negate-copy: M[dst] = -M[src] (dst pre-zeroed): subleq src, dst; halt.
    m = sq((0x90, 0x91, 3), (Z, Z, HALT_PC))
    m[0x90], m[0x91] = 42, 0
    progs.append(("neg_copy", m, 30, lambda mm: mm[0x91] == ((-42) & 0xFF)))
    # add: M[d] += M[s] via double negation through Z.
    m = sq((0x90, Z, 3), (Z, 0x91, 6), (Z, Z, 9), (Z, Z, HALT_PC))
    m[0x90], m[0x91] = 17, 25
    progs.append(("add", m, 30, lambda mm: mm[0x91] == 42))
    # countdown loop: decrement M[n] by M[one] until <= 0, counting into 0x92
    # via add idiom is long; simplest loop: subleq one, n, done at 0.
    m = sq((0x93, 0x90, 6), (Z, Z, 0), (Z, Z, HALT_PC))
    m[0x90], m[0x93] = 5, 1
    progs.append(("countdown", m, 60, lambda mm: mm[0x90] == 0))
    # halt-immediately
    m = sq((Z, Z, HALT_PC))
    progs.append(("halt", m, 5, lambda mm: True))

    all_ok = True
    for name, mem, budget, check in progs:
        ts = {"pc": 0, "mem": mem[:], "halted": False}
        rs = {"pc": 0, "mem": mem[:], "halted": False}
        cyc = 0
        ok = True
        while not ts["halted"] and cyc < budget:
            ts = cpu.step(ts)
            rs = ref_step(rs)
            cyc += 1
            if ts["pc"] != rs["pc"] or ts["mem"] != rs["mem"]:
                ok = False
                break
        ok = ok and ts["halted"] and rs["halted"] and check(ts["mem"])
        all_ok &= ok
        print(f"  {name:12} {'PASS' if ok else 'FAIL'}  ({cyc} cyc, lockstep)")

    print("SUBLEQ:", "PASS" if all_ok else "FAIL")
    return 0 if all_ok else 1

# ============================================================================
# neural_rv32 test suites
# ============================================================================
F1_0 = struct.unpack("<I", struct.pack("<f", 1.5))[0]
F2_5 = struct.unpack("<I", struct.pack("<f", 2.5))[0]
F10 = struct.unpack("<I", struct.pack("<f", 10.0))[0]


def lockstep(cpu, mem, budget, name, check=None):
    ts = rv_state(mem)
    rs = rv_state(mem)
    import torch
    from eval_all import int_to_bits
    ts["_mem_bits"] = torch.tensor([int_to_bits(b, 8) for b in ts["mem"]],
                                   dtype=torch.float32)
    cyc = 0
    ok = True
    while not ts["halted"] and cyc < budget:
        ts = cpu.step(ts)
        # reference is strictly single-issue; advance it until PCs realign
        rs = RefRV32.step(rs)
        if rs["pc"] != ts["pc"] and not rs["halted"]:
            rs = RefRV32.step(rs)
        cyc += 1
        if (ts["pc"] != rs["pc"] or ts["regs"] != rs["regs"]
                or ts["fregs"] != rs["fregs"] or ts["halted"] != rs["halted"]):
            ok = False
            break
    if ok:
        ok = ts["halted"] and ts["mem"] == rs["mem"] and ts["console"] == rs["console"]
    if ok and check is not None:
        ok = check(ts)
    print(f"  {name:22} {'PASS' if ok else 'FAIL'}  ({cyc} cyc)")
    if not ok:
        print(f"    pc t={ts['pc']:#x} r={rs['pc']:#x} halted t={ts['halted']} r={rs['halted']}")
        for i in range(32):
            if ts["regs"][i] != rs["regs"][i]:
                print(f"    x{i}: t={ts['regs'][i]:#x} r={rs['regs'][i]:#x}")
    return ok


def prog_int_alu():
    a = Asm()
    a.addi(1, 0, 100)          # x1 = 100
    a.addi(2, 0, -7)           # x2 = -7
    a.add(3, 1, 2)             # 93
    a.sub(4, 1, 2)             # 107
    a.xor(5, 1, 2)
    a.or_(6, 1, 2)
    a.and_(7, 1, 2)
    a.slt(8, 2, 1)             # 1 (signed)
    a.sltu(9, 2, 1)            # 0 (unsigned: -7 is huge)
    a.slli(10, 1, 3)           # 800
    a.srli(11, 2, 4)
    a.srai(12, 2, 4)           # -1
    a.slti(13, 2, 0)           # 1
    a.sltiu(14, 1, 200)        # 1
    a.xori(15, 1, 0xFF)
    a.ori(16, 1, 0x0F)
    a.andi(17, 1, 0x3C)
    a.lui(18, 0x12345)
    a.auipc(19, 1)
    a.ecall()
    return a.assemble(), 40, lambda s: s["regs"][3] == 93 and s["regs"][12] == u32(-1)


def prog_branches():
    a = Asm()
    a.addi(1, 0, 5)
    a.addi(2, 0, -3)
    a.addi(10, 0, 0)
    a.blt(2, 1, "l1")          # signed taken
    a.addi(10, 10, 100)
    a.label("l1")
    a.bltu(1, 2, "l2")         # unsigned taken (5 < 0xFFFF..FD)
    a.addi(10, 10, 200)
    a.label("l2")
    a.beq(1, 1, "l3")
    a.addi(10, 10, 400)
    a.label("l3")
    a.bge(1, 2, "l4")          # taken
    a.addi(10, 10, 800)
    a.label("l4")
    a.bne(1, 2, "done")
    a.addi(10, 10, 1600)
    a.label("done")
    a.ecall()
    return a.assemble(), 30, lambda s: s["regs"][10] == 0


def prog_fib():
    a = Asm()
    a.addi(1, 0, 0)            # a
    a.addi(2, 0, 1)            # b
    a.addi(3, 0, 10)           # n
    a.label("loop")
    a.beq(3, 0, "done")
    a.add(4, 1, 2)
    a.add(1, 0, 2)             # a = b  (independent pair candidates)
    a.add(2, 0, 4)             # b = t
    a.addi(3, 3, -1)
    a.jal(0, "loop")
    a.label("done")
    a.ecall()
    return a.assemble(), 200, lambda s: s["regs"][1] == 55


def prog_mem_call():
    a = Asm()
    a.addi(2, 0, 0x800)        # sp
    a.addi(1, 0, 0x77)
    a.sb(1, 0, 0x400)
    a.lb(3, 0, 0x400)
    a.lbu(4, 0, 0x400)
    a.addi(1, 0, -2)
    a.sh(1, 0, 0x404)
    a.lh(5, 0, 0x404)          # sign-extended
    a.lhu(6, 0, 0x404)
    a.jal(1, "fn")             # call
    a.ecall()
    a.label("fn")
    a.addi(2, 2, -4)
    a.sw(1, 2, 0)              # push ra
    a.lw(7, 2, 0)
    a.addi(2, 2, 4)
    a.jalr(0, 1, 0)            # ret
    return a.assemble(), 40, lambda s: s["regs"][5] == u32(-2) and s["regs"][6] == 0xFFFE


def prog_muldiv():
    a = Asm()
    a.addi(1, 0, -50)
    a.addi(2, 0, 7)
    a.mul(3, 1, 2)             # -350
    a.mulh(4, 1, 2)            # high of signed product = -1
    a.mulhu(5, 1, 2)
    a.div(6, 1, 2)             # -7
    a.rem(7, 1, 2)             # -1
    a.divu(8, 1, 2)
    a.remu(9, 1, 2)
    a.addi(10, 0, 0)
    a.div(11, 1, 10)           # div by zero -> -1
    a.rem(12, 1, 10)           # rem by zero -> dividend
    a.ecall()
    return a.assemble(), 30, lambda s: (s["regs"][3] == u32(-350)
                                        and s["regs"][6] == u32(-7)
                                        and s["regs"][7] == u32(-1)
                                        and s["regs"][11] == 0xFFFFFFFF)


def prog_float():
    a = Asm()
    a.lui(1, 0x1)              # x1 = 0x1000 data base
    a.flw(1, 1, 0)             # f1 = 1.5
    a.flw(2, 1, 4)             # f2 = 2.5
    a.fadd_s(3, 1, 2)          # 4.0
    a.fmul_s(4, 1, 2)          # 3.75
    a.fdiv_s(5, 2, 1)          # 2.5 / 1.5
    a.fsub_s(6, 2, 1)          # 1.0
    a.fsw(3, 1, 8)
    a.fle_s(10, 1, 2)          # 1
    a.flt_s(11, 2, 1)          # 0
    a.feq_s(12, 1, 1)          # 1
    a.fmv_x_w(13, 3)
    a.ecall()
    mem = a.assemble()
    for i, wv in ((0, F1_0), (4, F2_5)):
        for b in range(4):
            mem[0x1000 + i + b] = (wv >> (8 * b)) & 0xFF
    f4 = struct.unpack("<I", struct.pack("<f", 4.0))[0]
    return mem, 30, lambda s: (s["regs"][10] == 1 and s["regs"][11] == 0
                               and s["regs"][12] == 1 and s["regs"][13] == f4)


def prog_fma():
    a = Asm()
    a.lui(1, 0x1)
    a.flw(1, 1, 0)              # f1 = 1.5
    a.flw(2, 1, 4)             # f2 = 2.5
    a.flw(3, 1, 8)             # f3 = 3.0
    a.fmadd_s(4, 1, 2, 3)       # 1.5*2.5 + 3.0 = 6.75
    a.fmsub_s(5, 1, 2, 3)       # 1.5*2.5 - 3.0 = 0.75
    a.fnmadd_s(6, 1, 2, 3)      # -(1.5*2.5) - 3.0 = -6.75
    a.fnmsub_s(7, 1, 2, 3)      # -(1.5*2.5) + 3.0 = -0.75
    a.fmv_x_w(10, 4)
    a.fmv_x_w(11, 5)
    a.fmv_x_w(12, 6)
    a.fmv_x_w(13, 7)
    a.ecall()
    mem = a.assemble()
    f3_0 = struct.unpack("<I", struct.pack("<f", 3.0))[0]
    for i, wv in ((0, F1_0), (4, F2_5), (8, f3_0)):
        for b in range(4):
            mem[0x1000 + i + b] = (wv >> (8 * b)) & 0xFF
    def word(x):
        return struct.unpack("<I", struct.pack("<f", x))[0]
    return mem, 30, lambda s: (s["regs"][10] == word(6.75) and s["regs"][11] == word(0.75)
                               and s["regs"][12] == word(-6.75) and s["regs"][13] == word(-0.75))


def prog_neur():
    # XOR of two bits computed by a two-layer network of NEUR instructions:
    # h1 = OR(a,b) = H(a + b - 1); h2 = NAND(a,b) = H(-a - b + 1);
    # y = AND(h1,h2) = H(h1 + h2 - 2).
    def const(a, reg, v):
        a.lui(reg, v >> 12)
        a.ori(reg, reg, v & 0x7FF)

    a = Asm()
    a.addi(1, 0, 0b10)                     # inputs: bit1=a=1, bit0=b=0
    const(a, 2, (0x1F << 16) | 0x03)       # OR: pos=both, bias=-1 (0x1F sext5)
    a.neur(3, 1, 2)
    const(a, 4, (0x01 << 16) | (0x03 << 8))  # NAND: neg=both, bias=+1
    a.neur(5, 1, 4)
    a.slli(6, 3, 1)
    a.or_(6, 6, 5)                          # pack h1,h2 into bits 1,0
    const(a, 7, (0x1E << 16) | 0x03)        # AND: pos=both, bias=-2 (0x1E sext5)
    a.neur(8, 6, 7)                         # y = XOR(a,b) = 1
    a.ecall()
    return a.assemble(), 20, lambda s: s["regs"][8] == 1 and s["regs"][3] == 1 and s["regs"][5] == 1


def prog_mmio():
    a = Asm()
    a.lui(1, 0x10)             # 0x10000 > 0xFFFF? keep in range: use addi chain
    a.addi(1, 0, 0x7F8)
    a.slli(1, 1, 5)            # 0xFF00
    a.addi(2, 0, ord("H"))
    a.sb(2, 1, 0)
    a.addi(2, 0, ord("I"))
    a.sb(2, 1, 0)
    a.ecall()
    return a.assemble(), 20, lambda s: s["console"] == "HI"


def prog_fcvt():
    # int -> float, add 1.0, float -> int: (7 + 1.0) truncated = 8; and a
    # round-trip of -5 through float and back = -5.
    import struct as _st
    a = Asm()
    a.addi(1, 0, 7)
    a.fcvt_s_w(0, 1)          # f0 = 7.0
    a.lui(2, 0x1); a.flw(1, 2, 0)   # f1 = 1.0 from data
    a.fadd_s(2, 0, 1)        # f2 = 8.0
    a.fcvt_w_s(3, 2)         # x3 = 8
    a.addi(4, 0, -5)
    a.fcvt_s_w(4, 4)         # f4 = -5.0
    a.fcvt_w_s(5, 4)         # x5 = -5
    a.ecall()
    mem = a.assemble()
    one = _st.unpack("<I", _st.pack("<f", 1.0))[0]
    for b in range(4):
        mem[0x1000 + b] = (one >> (8 * b)) & 0xFF
    return mem, 30, lambda s: s["regs"][3] == 8 and s["regs"][5] == (-5 & 0xFFFFFFFF)


def prog_fcvt_rm():
    # FCVT.W.S honors the instruction's rounding-mode field. Convert 2.5, 3.5,
    # -2.5, 2.75 under several modes and check each against the IEEE result.
    import struct as _st
    a = Asm()
    a.lui(10, 0x1)                       # x10 = 0x1000 (data base)
    a.flw(1, 10, 0)                      # f1 = 2.5
    a.flw(2, 10, 4)                      # f2 = 3.5
    a.flw(3, 10, 8)                      # f3 = -2.5
    a.flw(4, 10, 12)                     # f4 = 2.75
    a.fcvt_w_s(1, 1, "rup")              # x1 = ceil(2.5)          = 3
    a.fcvt_w_s(2, 2, "rne")              # x2 = 3.5 ties to even   = 4
    a.fcvt_w_s(3, 1, "rne")              # x3 = 2.5 ties to even   = 2
    a.fcvt_w_s(4, 3, "rdn")              # x4 = floor(-2.5)        = -3
    a.fcvt_w_s(5, 4, "rtz")              # x5 = trunc(2.75)        = 2
    a.fcvt_w_s(6, 1, "rmm")             # x6 = 2.5 ties away      = 3
    a.ecall()
    mem = a.assemble()
    for i, v in enumerate((2.5, 3.5, -2.5, 2.75)):
        wb = _st.unpack("<I", _st.pack("<f", v))[0]
        for b in range(4):
            mem[0x1000 + i * 4 + b] = (wb >> (8 * b)) & 0xFF
    return mem, 40, lambda s: (s["regs"][1] == 3 and s["regs"][2] == 4
                               and s["regs"][3] == 2 and s["regs"][4] == (-3 & 0xFFFFFFFF)
                               and s["regs"][5] == 2 and s["regs"][6] == 3)


def rv32_netlist_check() -> bool:
    """Verify the rv32-specific sub-circuits (NEUR, hazard, MULHU accumulator)
    are recoverable and correct from the shipped .inputs metadata alone."""
    import random
    T = {}
    with safe_open(RV32_MODEL, framework="pt") as f:
        for name in f.keys():
            T[name] = f.get_tensor(name).float()
    reg = load_metadata(RV32_MODEL)["signal_registry"]
    rng = random.Random(0)
    ok = True

    ne = NetlistEvaluator(T, reg, "rv32.neur")
    N = 512
    xs = [rng.getrandbits(8) for _ in range(N)]
    ps = [rng.getrandbits(8) for _ in range(N)]
    ns = [rng.getrandbits(8) for _ in range(N)]
    bs = [rng.randint(-16, 15) for _ in range(N)]
    ext = {}
    for i in range(8):
        ext[f"$rv32_neur_x[{i}]"] = torch.tensor([(x >> (7 - i)) & 1 for x in xs]).float()
        ext[f"$rv32_neur_pos[{i}]"] = torch.tensor([(p >> (7 - i)) & 1 for p in ps]).float()
        ext[f"$rv32_neur_neg[{i}]"] = torch.tensor([(n >> (7 - i)) & 1 for n in ns]).float()
    for k in range(6):
        ext[f"$rv32_neur_bias[{k}]"] = torch.tensor(
            [((b & 0x1F) >> k) & 1 if k < 5 else (b >> 4) & 1 for b in bs]).float()
    got = ne.run(ext)["rv32.neur.out"][:, 0]
    exp = torch.tensor([1.0 if (bin(x & p).count("1") - bin(x & n).count("1") + b) >= 0
                        else 0.0 for x, p, n, b in zip(xs, ps, ns, bs)])
    ok &= bool((got == exp).all())

    ne = NetlistEvaluator(T, reg, "rv32.hazard.rd1_rs1b")
    ii = [rng.getrandbits(5) for _ in range(N)]
    jj = [rng.getrandbits(5) for _ in range(N)]
    ext = {}
    for k in range(5):
        ext[f"$rv32_hz_rd1_rs1b_i[{k}]"] = torch.tensor([(i >> k) & 1 for i in ii]).float()
        ext[f"$rv32_hz_rd1_rs1b_j[{k}]"] = torch.tensor([(j >> k) & 1 for j in jj]).float()
    got = ne.run(ext)["rv32.hazard.rd1_rs1b.all"][:, 0]
    exp = torch.tensor([1.0 if i == j else 0.0 for i, j in zip(ii, jj)])
    ok &= bool((got == exp).all())

    ne = NetlistEvaluator(T, reg, "rv32.mulacc")
    A = [rng.getrandbits(32) for _ in range(256)]
    B = [rng.getrandbits(32) for _ in range(256)]
    ext = {}
    for I in range(32):
        for J in range(32):
            ext[f"alu.alu32bit.mul.pp.a{I}b{J}"] = torch.tensor(
                [((a >> (31 - I)) & 1) & ((b >> (31 - J)) & 1) for a, b in zip(A, B)]).float()
    out = ne.run(ext)
    got = torch.zeros(len(A), dtype=torch.float64)
    for k in range(32):
        got += out[f"rv32.mulacc.s30.fa{32 + k}.ha2.sum.layer2"][:, 0].double() * (2.0 ** k)
    exp = torch.tensor([float(((a * b) >> 32) & 0xFFFFFFFF) for a, b in zip(A, B)],
                       dtype=torch.float64)
    ok &= bool((got == exp).all())

    print(f"  rv32 netlist self-check (NEUR/hazard/MULHU from metadata)  "
          f"{'PASS' if ok else 'FAIL'}")
    return ok


def fcvt_rm_check(cpu) -> bool:
    """FCVT.W.S under every rounding mode against an exact rational reference,
    for directed edge encodings (subnormals, ties, saturation) and randoms;
    the emulator reference and the gate-routed method must both match."""
    from fractions import Fraction
    import random

    def value(w):
        s, e, f = (w >> 31) & 1, (w >> 23) & 0xFF, w & 0x7FFFFF
        if e == 0:
            v = Fraction(f, (1 << 23) * (1 << 126))
        else:
            v = Fraction((1 << 23) | f, 1 << 23) * (Fraction(2) ** (e - 127))
        return -v if s else v

    def round_frac(v, rm):
        fl = v.numerator // v.denominator          # floor (correct for negatives)
        frac = v - fl
        if frac == 0:
            return fl
        if rm == 1:                                # RTZ
            return fl + 1 if v < 0 else fl
        if rm == 2:                                # RDN
            return fl
        if rm == 3:                                # RUP
            return fl + 1
        if frac < Fraction(1, 2):
            return fl
        if frac > Fraction(1, 2):
            return fl + 1
        if rm == 4:                                # RMM: ties away from zero
            return fl + 1 if v > 0 else fl
        return fl if fl % 2 == 0 else fl + 1       # RNE: ties to even

    def sat(iv):
        if iv > 0x7FFFFFFF:
            return 0x7FFFFFFF
        if iv < -0x80000000:
            return 0x80000000
        return iv & MASK32

    words = []
    for s in (0, 1):
        for e in (0, 1, 100, 125, 126, 127, 128, 130, 149, 150, 157, 158, 254, 255):
            for f in (0, 1, 0x400000, 0x7FFFFF):
                words.append((s << 31) | (e << 23) | f)
    rng = random.Random(0xFCF7)
    words += [rng.getrandbits(32) for _ in range(600)]

    bad = None
    for w in words:
        s, e, f = (w >> 31) & 1, (w >> 23) & 0xFF, w & 0x7FFFFF
        for rm in (0, 1, 2, 3, 4):
            if e == 0xFF:
                exp = 0x7FFFFFFF if f else (0x80000000 if s else 0x7FFFFFFF)
            else:
                exp = sat(round_frac(value(w), rm))
            ref = fcvt_w_s_ref(w, rm)
            gate = cpu.fcvt_w_s(w, rm)
            if not (ref == gate == exp):
                bad = (w, rm, ref, gate, exp)
                break
        if bad:
            break
    if bad:
        w, rm, ref, gate, exp = bad
        print(f"  fcvt rounding-mode check  FAIL  w={w:08x} rm={rm} "
              f"ref={ref:08x} gate={gate:08x} exp={exp:08x}")
        return False
    print(f"  fcvt rounding-mode check ({len(words)} words x 5 modes vs exact)  PASS")
    return True


def rv32_test() -> int:
    print("Loading neural_rv32...")
    cpu = Rv32ThresholdCPU()
    suite = [
        ("int_alu", prog_int_alu),
        ("branches", prog_branches),
        ("fib", prog_fib),
        ("mem_call", prog_mem_call),
        ("muldiv", prog_muldiv),
        ("float_subset", prog_float),
        ("fma", prog_fma),
        ("fcvt_roundtrip", prog_fcvt),
        ("fcvt_rounding_modes", prog_fcvt_rm),
        ("neur_xor_net", prog_neur),
        ("mmio_print", prog_mmio),
    ]
    all_ok = True
    for name, builder in suite:
        mem, budget, check = builder()
        all_ok &= lockstep(cpu, mem, budget, name, check)
    all_ok &= rv32_netlist_check()
    all_ok &= fcvt_rm_check(cpu)
    print(f"dual-issue pairs retired: {cpu.pairs_issued}")
    print("RV32:", "PASS" if all_ok else "FAIL")
    return 0 if all_ok else 1

PROG = r"""
__attribute__((naked)) void _start(void){
  __asm__ volatile("li sp, 0xF000\n call main\n ecall\n");
}
static int gcd(int a,int b){ while(b){ int t=a%b; a=b; b=t;} return a; }
int fib(int n){ int a=0,b=1; for(int i=0;i<n;i++){int t=a+b;a=b;b=t;} return a; }
int main(void){
  int arr[8]={5,2,9,1,7,3,8,4};
  for(int i=0;i<8;i++) for(int j=i+1;j<8;j++)
    if(arr[j]<arr[i]){int t=arr[i];arr[i]=arr[j];arr[j]=t;}
  int s=0; for(int i=0;i<8;i++) s+=arr[i]*(i+1);
  return gcd(48,36) + fib(10) + s + (arr[0]==1?1000:0);
}
"""
EXPECTED = 1292   # gcd(48,36)=12, fib(10)=55, weighted sorted sum=225... verified natively


def have_riscv_clang() -> bool:
    if not shutil.which("clang"):
        return False
    with tempfile.TemporaryDirectory() as d:
        c = os.path.join(d, "t.c")
        o = os.path.join(d, "t.o")
        open(c, "w").write("int f(void){return 1;}")
        r = subprocess.run(
            ["clang", "-target", "riscv32-unknown-elf", "-march=rv32im",
             "-mabi=ilp32", "-c", c, "-o", o],
            capture_output=True)
        return r.returncode == 0


def rv32c_test() -> int:
    if not have_riscv_clang():
        print("SKIP: no riscv32 clang target available")
        return 0

    with tempfile.TemporaryDirectory() as d:
        c = os.path.join(d, "prog.c")
        o = os.path.join(d, "prog.o")
        open(c, "w").write(PROG)
        subprocess.run(
            ["clang", "-target", "riscv32-unknown-elf", "-march=rv32im",
             "-mabi=ilp32", "-O1", "-nostdlib", "-ffreestanding",
             "-c", c, "-o", o], check=True, capture_output=True)
        obj = open(o, "rb").read()

    mem, entry = load(obj, base=0, mem_size=65536)
    cpu = Rv32ThresholdCPU()
    st = rv_state(mem)
    st["pc"] = entry
    fin, cyc = cpu.run(st, max_cycles=5000)
    got = fin["regs"][10]
    ok = fin["halted"] and got == EXPECTED
    print(f"compiled C (gcd+fib+sort, rv32im): a0={got} "
          f"expected {EXPECTED}  {cyc} instrs  {'PASS' if ok else 'FAIL'}")
    return 0 if ok else 1

# ============================================================================
# Entry point
# ============================================================================

def main() -> int:
    ap = argparse.ArgumentParser(description="neural_subleq8 / neural_rv32 test suites")
    ap.add_argument("which", nargs="?", default="all",
                    choices=["all", "subleq", "rv32", "rv32-c"])
    args = ap.parse_args()
    rc = 0
    if args.which in ("all", "subleq"):
        rc |= subleq_test()
    if args.which in ("all", "rv32"):
        rc |= rv32_test()
    if args.which in ("all", "rv32-c"):
        rc |= rv32c_test()
    return rc


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