EnzymeAD / EnzymeAD/Enzyme-JAX
Drop the block-merge workarounds once llvm/llvm-project#215036 lands
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
Two passes pin the greedy driver's region simplification to `GreedySimplifyRegionLevel::Normal` so that aggressive block merging cannot thread non-LLVM-typed values (index, memref) through an `llvm.invoke`/`llvm.cond_br` terminator, which produces IR that fails the terminator's own operand-type constraints:
- `src/enzyme_ad/jax/Passes/LowerAlignedAffineAccesses.cpp`
- `src/enzyme_ad/jax/Passes/ParallelSerialization.cpp`
(introduced in #2813; the failure it dodges is a generic upstream bug — `mergeIdenticalBlocks` only checks that predecessors implement `BranchOpInterface`, never whether they can forward the new operand types.)
llvm/llvm-project#215036 fixes this at the source: `BranchOpInterface` gains `mayForwardTypeToSuccessor`, block merging refuses clusters a predecessor cannot forward, and the LLVM dialect terminators answer with `isCompatibleType`. Once Enzyme-JAX's LLVM is at or past that commit, both `setRegionSimplificationLevel(Normal)` calls (and their TODO comments) should be dropped, restoring aggressive region simplification in those passes. `lower_aligned_eh.mlir` and `parallel_serialization_eh.mlir` stay as the regression tests either way.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
https://claude.ai/code/session_016zErYp7upmqr4NHfhod9UD
Contributor guide
No contributing guide indexed for this repository
Research direction
Check the LLVM version used by Enzyme-JAX against llvm/llvm-project#215036. In src/enzyme_ad/jax/Passes/LowerAlignedAffineAccesses.cpp and src/enzyme_ad/jax/Passes/ParallelSerialization.cpp, inspect the setRegionSimplificationLevel(Normal) calls and their TODO comments. Remove them once the dependency is new enough, then run lower_aligned_eh.mlir and parallel_serialization_eh.mlir as regression tests.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Refactor
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100