Add execution-based numerical tests for MLIR lit tests - #3011
Draft
Pangoraw wants to merge 5 commits into
Draft
Conversation
This reverts commit e0f08d5.
…table memory The existing scf_for*mutable_memory* tests only FileCheck the generated IR shape; they can't catch numerical bugs like the checkpointing forward-loop induction variable double-scaling fixed in the previous commit. Add mlir-runner-executed variants (checkpointing off, uniform, and binomial) that JIT-run the differentiated function and check the printed primal/gradient values match across all three checkpointing strategies. This requires two small additions: - %mlir-opt/%mlir-runner/%mlir_runner_utils/%mlir_c_runner_utils lit substitutions, gated behind a `mlir-runner` feature so the new tests are skipped (not failed) wherever these tools aren't built. - Lowering enzyme.load/enzyme.store (dynamic-size-annotated memref load/store) in --convert-enzyme-to-memref, which previously only handled init/push/pop/get/set; the uniform checkpointing path emits load/store for its per-iteration caches, so this is needed to fully lower to executable IR.
Pangoraw
marked this pull request as draft
July 27, 2026 17:00
Pangoraw
changed the title
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
1 participant
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.
Note
Experiment since this has been useful to find bugs in the checkpointing implementations.
Summary
Stacked on #3009. The tests added there (
scf_for*mutable_memory*) are FileCheck-only — they assert on the shape of the generated IR, not on actual numerical correctness. That's what let the forward-loop induction-variable double-scaling bug (fixed inalso fix inner iv rematerialization) go unnoticed: the IR looked plausible but computed wrong indices for any checkpoint period that didn't evenly divide the loop trip count.This PR adds real execution tests: they run the differentiated function through
mlir-runnerand check the printed primal/gradient values, for all three checkpointing strategies (off, uniform, binomial). All three now agree exactly, which they didn't before the fix.test/MLIR/Integration/ReverseMode/{scf_for_mutable_memory,scf_for_checkpointing_mutable_memory,scf_for_checkpointing_binomial_mutable_memory}_exec.mlir: new execution tests.test/lit.site.cfg.py.in: adds%mlir-opt/%mlir-runner/%mlir_runner_utils/%mlir_c_runner_utilssubstitutions, gated behind an available-only-if-presentmlir-runnerfeature so these tests cleanly skip (not fail) wherever those tools aren't built (e.g. the plain CMake build against system LLVM, which doesn't build MLIR tools at all).Enzyme/MLIR/Passes/EnzymeToMemRef.cpp: adds lowering patterns forenzyme.load/enzyme.store(previously onlyinit/push/pop/get/setwere handled) — needed because the uniform checkpointing path emits these for its per-iteration caches, and they have to be gone before the IR can reachmlir-opt's LLVM lowering pipeline.Test plan
bazel test //test/MLIR/Integration/ReverseMode:all— 3/3 new tests passbazel test //test/MLIR/...— 147/147 tests pass (no regressions from theEnzymeToMemRef.cppadditions)UNSUPPORTED(not built/tested) since the CMake config here doesn't build MLIR tools — confirm this is acceptable, or wire up MLIR tool builds in CMake CI as a follow-up🤖 Generated with Claude Code