EnzymeAD / EnzymeAD/Enzyme-JAX
dFEM TUs take ~10min to compile: 58% is 18 single-threaded module-scope canonicalize runs over a 62MB module
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
## Where the ~10 minutes go (measured on the worst TU)
`tests/unit/dfem/test_functional_gradient.cpp`, merged-main toolchain, plain `-O3`:
| stage | time |
|---|---|
| plain clang -O3 (no plugin), whole TU | 38.6s |
| **with plugin, whole TU** | **656s (10:56)** |
| — MLIR pipeline (offline replay, `--mlir-timing`) | **443s** |
| — remainder (import + translate-back + link + clang backend on merged module) | ~175s |
The TU imports as a **62MB module, 3,029 functions, 188 `__enzyme_*` call sites** (grows to 112MB after inlining).
**Pass ranking** (of 443s): `CanonicalizerPass` **257s (58%)** across 18 module-scope runs (~8.5s each); `LLVMToAffineAccess` 61s (14%); `AffineCFG` 17s; the Differentiate pass itself only 8.7s. Everything else ≤2%.
**Inside canonicalize the profile is flat** — no hot pattern: storage-uniquer/hashing ~5%, malloc/free ~4.5%, region-simplification liveness (`propagateLiveness`/`eraseUnreachableBlocks`/`SuccessorRange`) ~6%, the rest spread across the greedy driver. The cost is simply *module size × run count*, and — decisive — **User ≈ Wall (462s vs 443s): the entire pipeline is effectively single-threaded**, because every canonicalize (and the other heavy passes) runs at module scope.
## Fix directions, ranked by measured/estimated payoff
1. **Run the heavy passes nested per-function so the pass manager parallelizes them.** `builtin.module(llvm.func(canonicalize), func.func(canonicalize))` measured 5.78s wall / 8.19s user vs 8.22s/8.52s module-scope on the same dump (both including ~3.5s single-threaded parse; the canonicalize portion roughly halves, in-process would gain more). Speedup is currently capped by a few giant dFEM kernel functions — but combined with (2)/(3) and 192 cores this is the largest lever. Needs an audit of which of the 18 runs actually require module scope (canonicalize itself doesn't; interplay is with `symbol-dce`, which can stay module-level).
2. **Halve the run count.** Several of the 18 are guard runs after cheap local passes (`polygeist-mem2reg`, `convert-llvm-to-cf`, `enzyme-lift-cf-to-scf`, …) that could share one canonicalize. Consolidating to ~9 saves ~100–120s on this TU.
3. **Cheapen intermediate runs**: `canonicalize{region-simplify=disabled max-iterations=1 top-down=true}` measured 8.5→6.3s per run (−25%); keep one full-strength run before AD and one at the end.
4. **Slice the conversion module.** The round trip re-optimizes the whole TU, but the link-back already merges into the original module — functions the raising never changes don't need to ride through 18 canonicalizes *or* translate-back. Payoff is modest on dFEM test TUs (the template closure is most of the module) but large for library TUs and for the ~175s non-pipeline remainder.
5. `LLVMToAffineAccess` (61s) is likewise single-threaded and per-function in nature — same nesting treatment applies.
For scale: (2)+(3) alone are ~150s off a 656s compile; (1) on top plausibly brings the pipeline under 2 minutes; (4) attacks the remaining 175s.
Repro artifacts: imported module + pipeline dumped via `DEBUG_REACTANT`; offline replay `enzymexlamlir-opt --pass-pipeline='builtin.module()' --mlir-timing`.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with tests/unit/dfem/test_functional_gradient.cpp and the DEBUG_REACTANT replay artifacts. Run enzymexlamlir-opt with the dumped pipeline and --mlir-timing, then compare pass scope, run count, and intermediate canonicalize settings against the reported timings. Done means the dFEM TU compiles substantially faster while preserving the existing pipeline behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100