EnzymeAD / EnzymeAD/Enzyme-JAX

Drop the block-merge workarounds once llvm/llvm-project#215036 lands

Open Beginner friendly
#2,815 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.