[eudsl-llvmpy] Python-driven register allocators and instruction schedulers (MIR codegen)

[makslevental]

Context

eudsl-llvmpy already drives the codegen pass pipeline from Python: llvm.mir.run_codegen_to_mir builds a TargetPassConfig and runs ISel (src/MIR/Machine.cpp:1060), and emit_object runs the back half of codegen — register allocation + emission — over already-selected MIR (src/MIR/Machine.cpp:950, back-half at 979-1042, using a set+restore of the process-global -start-after cl::opt at 982-1002). Register allocation and instruction scheduling are the codegen stages inside that pipeline, but neither can currently be driven by user logic.

This tracks adding Python-driven register allocators and instruction schedulers — a custom allocator/scheduler whose core decisions are made by a Python callable.

Feasibility (already investigated): both are registerable and both can be Python-driven. It is C++-only (no llvm-c C API for either), which is fine — the extension links LLVM's C++ codegen libs and can subclass the interfaces + construct the registry objects directly. Registration points:

• Reg alloc: RegisterRegAlloc::setDefault(&createFn) (in-process, no CLI) (RegAllocRegistry.h:54); selected by TargetPassConfig::createRegAllocPass (TargetPassConfig.cpp:1415). The pass is a MachineFunctionPass filling VirtRegMap; idiomatic base RegAllocBase (RegAllocBase.h) with decision point selectOrSplit(const LiveInterval&, SmallVectorImpl<Register>&).
• Scheduler: MachineSchedRegistry + -misched= (MachineScheduler.h:173, selected in MachineScheduler.cpp:525), or the TargetMachine::createMachineScheduler hook (TargetMachine.h:167). Custom strategy implements MachineSchedStrategy (MachineScheduler.h:246): initialize, pickNode(bool&IsTopNode), schedNode, releaseTopNode, releaseBottomNode; wrap via createSchedLive<Strategy>(C).

Binding shape follows MLIR's callback trampoline (mlir/lib/Bindings/Python/Rewrite.cpp:87-118) but simpler — hold the nb::callable as a member of the C++ subclass and override the virtual directly.

Companion findings doc: REGALLOCA_ISCHED_PASSES.md. Sibling issue for IR passes: #599.

Conventions

• New TU (e.g. src/MIR/PythonCodegen.cpp) wired into nanobind_add_module (CMakeLists.txt:105-129) + a populate_* into the llvm.mir submodule (src/eudslllvm_ext.cpp:26,64-65), mirroring src/MIR/Machine.cpp.
• Bind the decision body as nb::callable (the legitimate exception to CLAUDE.md's "prefer specific types"); everything else uses specific nanobind types (mirror LLVM names).
• Tests use parse_assembly / MIR builders, assert_no_leaks(), one behavior per test, TDD (see tests/mir/test_emit_jit.py, test_mir_passes.py).
• Reuse the existing back-half emission path (Machine.cpp:979-1042) and diagnostic capture (ScopedDiagnosticCapture, Machine.cpp:1093-1101) rather than rebuilding the pipeline.
• Free-threading caveat: the module is built FREE_THREADED (CMakeLists.txt:106). The trampoline's refcount/call must be correct under a free-threaded interpreter — do not assume a process-wide GIL. Acquire the GIL in the decision callback with nb::gil_scoped_acquire; do not release it around pm->run().
• cl::opt safety: selecting a scheduler via -misched= follows the existing set+restore-under-GIL pattern used for -start-after (Machine.cpp:982-1002), since these options are process-global and unlocked.

Recommended order

Do the scheduler pickNode slice first — it is safer than reg alloc (a bad selectOrSplit must still produce a valid assignment-or-spill or codegen aborts). Prove each vertical slice against the existing AArch64 back-half pipeline before generalizing.

Commits (atomic, test-driven)

Scheduler track

• 1 — Foundation: custom-scheduler registration + selection, native trivial policy. Add class PyMachineSchedStrategy : public MachineSchedStrategy implementing the 5 virtuals with a trivial correct C++ policy (e.g. pick the first ready node), registered via a MachineSchedRegistry factory using createSchedLive<PyMachineSchedStrategy>(C). Add a binding (flag on run_codegen_to_mir or a new entry point) that selects it by name through the -misched= cl::opt (set+restore per Machine.cpp:982-1002). No Python callback yet. Test: codegen with the custom scheduler selected produces valid MIR/object; runs end to end.

• 2 — Route pickNode into a Python callable. Store an nb::callable; pickNode acquires the GIL, hands the ready set to the callable, and uses its choice; the release/sched hooks stay native. inc_ref/dec_ref for lifetime. Test: a Python scheduler mimicking the Commit-1 policy yields the same result; callable invoked; assert_no_leaks().

• 3 — Marshal richer schedule state + exception propagation. Expose SUnit/ready-node views the callable needs; a Python exception in pickNode surfaces as a codegen failure via ScopedDiagnosticCapture (Machine.cpp:1093-1101), not a crash. Test: pytest.raises on a raising callback; module/context leak-clean after; valid schedule when the callback picks legally.

Register-allocator track

• 4 — Foundation: custom reg alloc registration, native trivial policy. Add class PyRegAlloc : public MachineFunctionPass, public RegAllocBase with a trivial-but-correct selectOrSplit (+ spiller/enqueueImpl/dequeue), registered via RegisterRegAlloc::setDefault(&createFn). Reuse the back-half emission path (Machine.cpp:979-1042). No Python callback yet. Test: valid object emitted with the custom allocator selected.

• 5 — Route selectOrSplit into a Python callable. Store an nb::callable; translate its answer into either Matrix->assign(...) or a spill/split list. GIL acquire + refcount. Test: a Python allocator produces a valid assignment for a small function; callable invoked; leaks clean.

• 6 — Marshal LiveInterval + candidate physregs + exception propagation. Give the callback the interval and legal candidate registers (reuse/extend MIR downcast/caster infra); surface exceptions via ScopedDiagnosticCapture. Test: exception surfaces as error; correct assignment when the callback returns a legal physreg.

Docs

• 7 — Docs. Update README.md + add PYTHON_CODEGEN_PASSES.md: API for both, the decision-callback contracts, cl::opt/GIL/free-threading caveats, and that this is the codegen sibling of the IR-pass work (Python-driven LLVM IR passes (new PassManager) #599).

Files

• Create: src/MIR/PythonCodegen.cpp (trampolines + populate_*) wired into CMakeLists.txt:105-129; forward decl + submodule call in src/eudslllvm_ext.cpp.
• Create: tests/mir/test_python_scheduler.py, tests/mir/test_python_regalloc.py.
• Create: PYTHON_CODEGEN_PASSES.md; edit README.md.
• Edit: src/MIR/Machine.cpp only if a shared back-half/selection helper is factored out.
• Reference (no edit): src/IR/Casters.cpp/Kinds.h (downcast), and LLVM headers cited above under third_party/llvm-project/.

Risks / open questions

• Free-threaded correctness of the callable refcount + reentrancy across the pipeline (Commits 2/5).
• Reg alloc correctness contract: a Python selectOrSplit must always yield a legal assignment or a valid spill, or codegen aborts under NDEBUG — the native trivial policy (Commit 4) de-risks this before Python enters.
• PostRA scheduler has no -misched= override (only the TargetMachine hook) — decide whether to cover PostRA at all in this issue or defer.
• Performance: a per-node/per-interval Python callback is slow; this is about expressiveness/experimentation, not production codegen — document that.

Verification

• Build compiles the new TU; extension imports.
• pytest tests/mir/ -q green incl. the new scheduler/regalloc tests; every test ends with assert_no_leaks().
• Manual: a Python scheduler and a Python allocator each drive the AArch64 back-half pipeline to a valid object; verified against the default pipeline's output shape.
• Each commit builds and passes tests independently (atomic); native-policy foundation commits (1, 4) land before their Python-callback commits (2, 5).

[makslevental]

Learnings from the IR-pass sprint (#599) that should shape this MIR work

The IR-pass stack landed (#602/#603/#604/#605 for Python passes + registration, #608 for object-level mutation, #609/#610 for the example passes). Several things we hit there transfer directly to the scheduler/regalloc plan here — some sharpen the existing caveats, a couple change the risk calculus.

1. The -fno-exceptions unwinding hazard is real, and worse for codegen than for IR

LLVM is built LLVM_ENABLE_EH=OFF (-fno-exceptions). A Python exception thrown across LLVM's pass-manager frames skips their destructors (dangling PrettyStackTraceHead, leaked PreservedAnalyses) and can std::terminate on targets without async unwind tables. The ScopedDiagnosticCapture note in the plan is necessary but not sufficient — the mechanism that actually worked for IR was:

• a thread_local std::exception_ptr; the trampoline does try { ...callback... } catch (...) { stash; return <benign>; } so it always returns normally into LLVM's frames;
• a wrapper around pm->run() that re-raises the stashed exception after the pipeline returns, so the throw only ever crosses our own -fexceptions frames.

Do this for pickNode (Commit 3) and selectOrSplit (Commit 6). Caveat that's sharper here: for IR we could gate later optional passes off after a stash (registerShouldRunOptionalPassCallback), but codegen passes are almost all isRequired() — they can't be skipped. So after a stashed error, the rest of codegen may still run on half-produced MIR before the re-raise. Decide deliberately: either re-raise at the earliest safe boundary, or accept that the benign return value from a failed pickNode/selectOrSplit must itself be legal (a valid node / a valid assignment-or-spill) so the remainder of the required pipeline doesn't abort before you get to re-raise. This ties directly to the plan's own "must always yield a legal assignment or codegen aborts under NDEBUG" risk.

Honesty note we learned the hard way: don't add a defensive catch for "a C++ exception escaped pm->run()" — we did, then realized it's dead code (LLVM can't throw; our trampolines stash rather than throw), and removed it. Verify a guard is reachable before writing it; otherwise it's just LCOV_EXCL noise fighting the 100% gate.

2. GIL-guard placement has a specific correct answer

Put nb::gil_scoped_acquire outside the try, because the catch stashes via std::current_exception(), which for an nb::python_error touches Python refcounts and must run under the held GIL. PyGILState_Ensure doesn't throw, so nothing is lost. Add a comment saying so, or someone will "fix" it into a refcount bug under the free-threaded build.

3. Decision-value errors should use nb::raise_from, not bare nb::python_error

When a callback's return can't be interpreted (e.g. the truthiness of an IR-pass result), wrap it: nb::raise_from(err, PyExc_ValueError, "could not ... ") gives a descriptive message and chains the original as __cause__. For MIR, the analogue is a callback that returns a node/register that isn't in the legal candidate set — surface that with a clear message rather than letting a downstream assert fire.

4. cl::opt set+restore must be exception-safe

The plan already notes the -misched= / -start-after set+restore-under-GIL pattern. Combine it with (1): if the callback raises and you stash, the cl::opt restore still has to happen. Use RAII for the restore so it runs on every path, including the stash path.

5. If you bind MI-level mutation, reuse the erase-poison fix

#608 hit the dangling-handle problem: after erase_from_parent() the C++ object is freed but nanobind keeps the wrapper's pointer, so a later access is a use-after-free. Fix: after erasing, nb::inst_set_state(self, /*ready=*/false, /*destruct=*/false) poisons the wrapper so any later use raises TypeError instead of dereferencing freed memory. MachineInstr/MachineBasicBlock erasure has the identical hazard — apply the same pattern. (It poisons only the called handle; aliasing/transitive invalidation is tracked in #615, which a MIR live-value registry should share.)

6. Static registries are the #614 free-threaded-race pattern

RegisterRegAlloc::setDefault(&createFn) and MachineSchedRegistry are process-global mutable state, same shape as the caster/pass registries flagged in #614 (unsynchronized under FREE_THREADED). Whatever synchronization decision #614 reaches should cover these too — note it here rather than adding a third unsynchronized registry.

7. Test honesty: structural checks are not enough — execute

The example-pass PRs taught this concretely: #609 (SLP) initially went green on a semantically wrong vectorization because string-shape + verify() are satisfied by any well-formed vector op regardless of lane/operand wiring. The fix was to JIT-execute and assert computed results. This project already emits objects + JIT-executes (test_emit_jit.py) — lean on it: for the native trivial-policy commits (1, 4), make the test fail if the policy were a no-op, and for the Python-callback commits, assert the executed output, not just that valid MIR/object was produced. Also, when a guard genuinely can't fire in the test's regime (we had a purity guard over integer adds that are always pure), document that honestly instead of contriving a test that fakes coverage.

8. 100% per-PR coverage gate, and the local-tooling gotcha

Coverage is enforced per PR over src/IR + src/MIR, so every new binding line (including the native trivial policies in Commits 1/4) needs a test in the same PR or a justified LCOV_EXCL. Running it locally requires llvm-profdata/llvm-cov from the same toolchain that built the .so (the bundled mlir_wheel tools are a newer LLVM and reject the system profraw) — see the ## C++ coverage note now in projects/eudsl-llvmpy/CLAUDE.md (macOS xcrun -f, Linux versioned llvm-*-N).

9. Stacked-PR mechanics

gh stack worked well for the 10-PR IR stack, with one recurring friction: an edit low in the stack that touches the same populate_* region as a later PR conflicts on rebase every time. Keep each MIR commit's edits localized (separate TU PythonCodegen.cpp as planned helps) to minimize this.

Bonus

The MCP server proposed in #616 (in-package llvm.mcp) is the agent-facing surface for all of this; a Python scheduler/allocator that can be authored, run, and JIT-checked from a session is squarely in scope for it once this lands.