Refactor JIT codegen into submodules (gradual)#2
Merged
Conversation
This is a partial refactor: moved JIT quantum/profiling state and type helpers into dedicated submodules (codegen/jit_types.rs, codegen/quantum.rs, codegen/registry.rs).
There was a problem hiding this comment.
Pull request overview
This PR refactors the Python JIT codegen implementation into smaller, focused submodules and adjusts JIT logging to use tracing instead of direct stderr printing.
Changes:
- Split the monolithic
src/py/jit/codegen/mod.rsinto dedicated modules (jit_types,compiler,gen_expr,lowering,jit_module,registry,quantum). - Switch
jit_logoutput fromeprintln!totracing::info!. - Remove a couple of config re-exports from
src/py/jit/mod.rsthat are no longer needed.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/py/jit/mod.rs | Drops re-exports for two quantum config helpers (now referenced directly where needed). |
| src/py/jit/config.rs | Routes JIT logging through tracing and keeps the test hook path. |
| src/py/jit/codegen/mod.rs | Becomes a façade that wires together the new codegen submodules and re-exports their APIs. |
| src/py/jit/codegen/jit_types.rs | New: centralizes shared JIT type definitions (entry types, lowered kernels, strategies, profile structs). |
| src/py/jit/codegen/compiler.rs | New: holds top-level compilation functions extracted from the old monolith. |
| src/py/jit/codegen/gen_expr.rs | New: extracted expression → Cranelift IR lowering. |
| src/py/jit/codegen/lowering.rs | New: extracted lowered-kernel detection helpers. |
| src/py/jit/codegen/jit_module.rs | New: extracted TLS JIT module setup and execution profiling state. |
| src/py/jit/codegen/registry.rs | New: extracted registry and iris_jit_invoke_* helper exports. |
| src/py/jit/codegen/quantum.rs | New: extracted quantum profiling / variant selection state machine and helpers. |
You can also share your feedback on Copilot code review. Take the survey.
|
|
||
| use crate::py::jit::heuristics; | ||
| use crate::py::jit::parser::Expr; | ||
|
|
| } | ||
|
|
||
| eprintln!("{}", msg()); | ||
| // Use `tracing` when available; falls back to stderr if no subscriber is set. |
| @@ -0,0 +1,149 @@ | |||
| // src/py/jit/codegen/types.rs | |||
Comment on lines
+47
to
+50
| fn invoke_named_entry(func_ptr: i64, args: &[f64]) -> f64 { | ||
| let f: extern "C" fn(*const f64) -> f64 = unsafe { std::mem::transmute(func_ptr as usize) }; | ||
| f(args.as_ptr()) | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a modular refactor and significant enhancements to the JIT (Just-In-Time) code generation pipeline for the Iris project. The changes separate key components into their own files, add new type definitions, and implement logic for lowering parsed expressions into optimized kernel representations. The main goals are to improve maintainability, enable better kernel detection for vectorization, and provide a foundation for profiling and variant strategies.
The most important changes are:
JIT Kernel Lowering and Detection:
src/py/jit/codegen/lowering.rswith logic to analyze parsed expressions (Expr) and detect when they can be lowered into unary or binary kernels, or further into a generic lowered expression tree. This enables the JIT to recognize and optimize simple operations for vectorization and efficient execution.arg_index_of_var,normalize_intrinsic_name, anddetect_lowered_exprto support the lowering process.Type Definitions and Kernel Metadata:
src/py/jit/codegen/jit_types.rs(notably astypes.rsin the file header) containing shared type definitions for the JIT pipeline, including enums for reduction modes, return types, kernel types, variant strategies, and profiling structures. This centralizes and clarifies the metadata associated with JIT-compiled functions.JIT Compilation and Variant Strategy:
src/py/jit/codegen/compiler.rsto use the new lowering logic, and to support compiling multiple speculative variants of the same expression (e.g., optimized, scalar fallback, and experimental fast trig). This supports dynamic selection of the best-performing variant at runtime.JIT Module Initialization and Profiling:
src/py/jit/codegen/jit_module.rsto encapsulate thread-local JIT module setup, backend selection (including SIMD capabilities), and function pointer profiling. This improves initialization robustness and enables per-thread JIT state management.JitExecProfileenum and thread-local storage for profiling execution signatures of JIT-compiled functions.These changes lay the groundwork for more advanced JIT optimizations, better profiling, and easier extension of the codegen pipeline.