[eudsl-llvmpy] Python-driven register allocators and instruction schedulers (MIR codegen)
- Dominant language
- Python
- Stars
- 79
- Forks
- 14
- Avg merge
- 11h 43m
- Merged PRs (30d)
- 72
Description
## 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&)`.
- **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(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(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 (#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).
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.