Add run_mlir.jl: pre_xla MLIR → standalone XLA execute script#2581
Add run_mlir.jl: pre_xla MLIR → standalone XLA execute script#2581gbaraldi wants to merge 3 commits intoEnzymeAD:mainfrom
Conversation
… script A code-generation tool for testing the XLA compilation/execution pipeline on pre_xla_compile MLIR dumps without needing the full model setup. Phase 1 (generator — no Reactant dependency): julia run_mlir.jl first.mlir loop.mlir output.jl - Parses func.func @main signatures (types, shapes, sdy shardings) - Extracts mesh spec, num_partitions/replicas from module attributes - Detects grid constants vs aliased state variables - Emits a self-contained Julia script with hardcoded shapes/shardings Phase 2 (generated script): julia --project=Reactant.jl output.jl [--cpu] - Creates mock ConcreteRArrays with correct NamedSharding - Loads MLIR modules and XLA-compiles with shardy partitioner - Executes first_time_step, marshals outputs → loop inputs, executes loop - Syncs all results for accurate timing - --cpu flag: sets CUDA_VISIBLE_DEVICES="" and XLA_FLAGS="--xla_force_host_platform_device_count=N" to run on virtual CPU devices Includes example MLIR files (1-device, stablehlo.add) for testing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Creates inputs with known values (grid=ones, dt=3, state=zeros), executes both MLIR modules, and asserts the outputs match expected values (first: [1,1,1,1], loop: [2,2,2,2]). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
@avik-pal Maybe this is useful for reproducing gb-25 bugs. Also might be useful for executing as well since it can obliviate a lot of the loop (well the logic) |
Codecov Report✅ All modified and coverable lines are covered by tests.
Additional details and impacted files@@ Coverage Diff @@
## main #2581 +/- ##
===========================================
- Coverage 68.16% 32.30% -35.87%
===========================================
Files 109 174 +65
Lines 11779 28567 +16788
===========================================
+ Hits 8029 9228 +1199
- Misses 3750 19339 +15589 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
scripts/run_mlir.jl
Outdated
| @@ -0,0 +1,594 @@ | |||
| #!/usr/bin/env julia | |||
| """ | |||
| run_mlir.jl — Generator: parse pre_xla MLIR files → emit standalone execute script. | |||
There was a problem hiding this comment.
why do we need this to not have reactant dependency? Seems like parsing would be much simpler with that one dep
Use sdyTensorShardingAttr*, sdyMeshAttr*, IR.FunctionType, IR.julia_type, and IR.size to introspect MLIR modules programmatically instead of regex. Zero regex patterns remain in the generator. The only string-level check is a dict_get helper for dictionary attribute lookup (TODO: upstream haskey/get for dict Attributes into Reactant's MLIR IR bindings). Also uses randn instead of zeros for mock array data to exercise non-trivial computation paths, and syncs results for accurate timing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
| run_mlir_v2.jl — Generator (Reactant-dependent): MLIR IR introspection → emit execute script. | ||
|
|
||
| Usage: | ||
| julia --project=Reactant.jl run_mlir_v2.jl [first.mlir] [loop.mlir] [output.jl] |
| # Main: analyze MLIR via IR APIs, emit script | ||
| # ────────────────────────────────────────────────────────────── | ||
|
|
||
| function main() |
There was a problem hiding this comment.
can this just be added as a reactant-internal utility function in the serialization subpackage that takes a list of mlir files [in this case it would be first and loop], and an output path?
There was a problem hiding this comment.
even nicer would be a cli app on 1.12+ that calls the serialization function
There was a problem hiding this comment.
The question is how do you choose how to pipe outputs into inputs etc.
3 participants
Summary
scripts/run_mlir.jl, a two-phase code-generation tool for testing the XLA compilation/execution pipeline on*_pre_xla_compile.mlirdumps — without needing the full model (e.g. Oceananigans/GB-25) setup--cpuflag to run on virtual CPU devices viaXLA_FLAGS="--xla_force_host_platform_device_count=N"How it works
Phase 1 — Generator (no Reactant dependency):
func.func @main(...)signatures (tensor types, shapes,sdy.shardingannotations)mhlo.num_partitions/num_replicasfrom module attributestf.aliasing_output)Phase 2 — Generated script:
ConcreteRArray/ConcreteRNumberwith correctNamedShardingsynced_buffer) for accurate timingTest plan
["x"=2,"y"=2])stablehlo.add) end-to-end on GPU--cpu--cpu(4 virtual devices)🤖 Generated with Claude Code