[Mlir-commits] [mlir] [mlir][dataflow] Add merge-site widening hook for infinite-height lattices (PR #193338)
Ivan Butygin
llvmlistbot at llvm.org
Fri May 8 07:59:34 PDT 2026
https://github.com/Hardcode84 updated https://github.com/llvm/llvm-project/pull/193338
>From 57f343b1ddf382b50919760cf6f7ef045518284d Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Mon, 20 Apr 2026 00:00:10 +0200
Subject: [PATCH 1/8] [mlir][dataflow] Add merge-site widening hook for
infinite-height lattices
IntegerRangeAnalysis can fail to converge on scf.while loops with dynamic
bounds when a loop-carried range is propagated through nested region ops
(e.g. scf.if). The range ratchets [0,0]->[0,1]->[0,2]->... one step per
worklist visit, requiring O(2^31) iterations for i32.
Rather than bolting per-analysis loop-visit caps onto IntegerRangeAnalysis
itself (which duplicates base-class control-flow plumbing and solves only
the one analysis's variant of the problem), this patch adds a generic
widening hook to the framework:
- DataFlowSolver::enableWidening<StateT>(budget, fn) registers a per-state
widen callback keyed by TypeID. Called by analyses at construction time.
- AbstractSparseForwardDataFlowAnalysis::join is the single choke point
for all merge-site joins (block args, region successors, callable args).
It now consults the solver to bump a per-state change counter and, once
the budget is exhausted, invokes the registered widen fn to force the
state to a sound over-approximation.
- tryWidenAtMerge returns only the widening-induced delta; the caller
unions it with its own ChangeResult. This makes it structurally
impossible for the hook to downgrade a preceding Change, which would
silently drop a dataflow update.
- hasWideningEnabled is derived from the configs map (no redundant bool).
Zero runtime cost for analyses that don't opt in: one well-predicted
branch on !wideningConfigs.empty() in the join hot path, and a DenseMap
that stays empty. Finite-height analyses (constant-prop, liveness,
dead-code, SCCP) are unaffected.
IntegerRangeAnalysis opts in with budget 16, widening to
IntegerValueRange::getMaxRange. Full MLIR regression suite passes.
Dense and sparse-backward variants are intentionally not wired up; no
in-tree analysis of either shape has an infinite-height lattice yet.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../Analysis/DataFlow/IntegerRangeAnalysis.h | 6 +-
.../mlir/Analysis/DataFlow/SparseAnalysis.h | 19 +++++-
.../include/mlir/Analysis/DataFlowFramework.h | 59 +++++++++++++++++++
.../DataFlow/IntegerRangeAnalysis.cpp | 15 +++++
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 15 ++++-
mlir/lib/Analysis/DataFlowFramework.cpp | 26 ++++++++
6 files changed, 134 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
index 5b6ae9bf84265..ba45fa083b227 100644
--- a/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h
@@ -40,7 +40,11 @@ class IntegerValueRangeLattice : public Lattice<IntegerValueRange> {
class IntegerRangeAnalysis
: public SparseForwardDataFlowAnalysis<IntegerValueRangeLattice> {
public:
- using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
+ /// Construct the analysis and opt into merge-site widening. Without this,
+ /// `scf.while` loops with dynamic bounds and nested region ops can keep the
+ /// solver ratcheting a loop-carried range by +1 per visit for up to 2^31
+ /// iterations on i32.
+ explicit IntegerRangeAnalysis(DataFlowSolver &solver);
/// At an entry point, we cannot reason about integer value ranges.
void setToEntryState(IntegerValueRangeLattice *lattice) override {
diff --git a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
index df50d8d193aeb..fb9c73ef9789a 100644
--- a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
@@ -194,7 +194,17 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
LogicalResult visit(ProgramPoint *point) override;
protected:
- explicit AbstractSparseForwardDataFlowAnalysis(DataFlowSolver &solver);
+ /// Construct the analysis. `stateTypeID` identifies the concrete lattice
+ /// type used by this analysis; it routes the merge-site widening lookup in
+ /// the solver. The CRTP wrapper `SparseForwardDataFlowAnalysis<StateT>`
+ /// fills it in from the template parameter.
+ AbstractSparseForwardDataFlowAnalysis(DataFlowSolver &solver,
+ TypeID stateTypeID);
+
+ /// Legacy constructor. Using this disables widening for this analysis
+ /// because the state TypeID is unknown. Prefer the two-argument ctor.
+ explicit AbstractSparseForwardDataFlowAnalysis(DataFlowSolver &solver)
+ : AbstractSparseForwardDataFlowAnalysis(solver, TypeID()) {}
/// The operation transfer function. Given the operand lattices, this
/// function is expected to set the result lattices.
@@ -286,6 +296,11 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
visitRegionSuccessors(ProgramPoint *point, RegionBranchOpInterface branch,
RegionSuccessor successor,
ArrayRef<AbstractSparseLattice *> lattices);
+
+ /// TypeID of the concrete lattice type used by this analysis. Threaded
+ /// through from the CRTP derived class so that merge-site widening lookups
+ /// know which `WideningConfig` to consult. Empty TypeID disables widening.
+ TypeID stateTypeID;
};
//===----------------------------------------------------------------------===//
@@ -305,7 +320,7 @@ class SparseForwardDataFlowAnalysis
public:
explicit SparseForwardDataFlowAnalysis(DataFlowSolver &solver)
- : AbstractSparseForwardDataFlowAnalysis(solver) {}
+ : AbstractSparseForwardDataFlowAnalysis(solver, TypeID::get<StateT>()) {}
/// Visit an operation with the lattices of its operands. This function is
/// expected to set the lattices of the operation's results.
diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index 25506645f2f26..2913fde31f108 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -439,10 +439,65 @@ class DataFlowSolver {
/// Get the configuration of the solver.
const DataFlowConfig &getConfig() const { return config; }
+ //===--------------------------------------------------------------------===//
+ // Widening (opt-in convergence cap for infinite-height lattices)
+ //===--------------------------------------------------------------------===//
+
+ /// A widening function forces a state to a sound over-approximation (lattice
+ /// top) once its merge-site visit budget is exhausted. Invoked by
+ /// `tryWidenAtMerge`. Must return `Change` iff the state was modified.
+ using WidenFn = std::function<ChangeResult(AnalysisState *)>;
+
+ /// Enable widening for states of type `StateT`. After `budget` merge-site
+ /// joins that each produced a change, subsequent joins at the same state
+ /// trigger `widen`. Must be called before `initializeAndRun`. Analyses with
+ /// finite-height lattices should not call this; they converge naturally.
+ template <typename StateT>
+ void enableWidening(unsigned budget, WidenFn widen) {
+ wideningConfigs.try_emplace(TypeID::get<StateT>(),
+ WideningConfig{budget, std::move(widen)});
+ }
+
+ /// Fast-path predicate: true if any analysis loaded in this solver has
+ /// opted into widening. Callers on the hot path should check this before
+ /// doing any further widening work. Derived from the configs map rather
+ /// than a redundant bool so there is only one source of truth.
+ bool hasWideningEnabled() const { return !wideningConfigs.empty(); }
+
+ /// If widening is configured for `stateTypeID` and this state's merge-site
+ /// change budget is exhausted, invoke the registered widen function and
+ /// return the change it produced. Otherwise, record `changeFromJoin`
+ /// against the budget (incrementing only on `Change`) and return
+ /// `NoChange`. The caller is responsible for unioning this result into its
+ /// own `ChangeResult`:
+ ///
+ /// ChangeResult c = lhs->join(rhs);
+ /// c |= solver.tryWidenAtMerge(lhs, typeID, c);
+ ///
+ /// Returning only the *additional* change from widening (rather than the
+ /// combined result) makes it structurally impossible for this function to
+ /// downgrade a preceding `Change` into `NoChange`, which would silently
+ /// lose a dataflow update.
+ ChangeResult tryWidenAtMerge(AnalysisState *state, TypeID stateTypeID,
+ ChangeResult changeFromJoin);
+
private:
/// Configuration of the dataflow solver.
DataFlowConfig config;
+ /// Per-state-type widening policy. Emptiness of this map is the canonical
+ /// answer to "is widening enabled for this solver"; see
+ /// `hasWideningEnabled`.
+ struct WideningConfig {
+ unsigned budget;
+ WidenFn widen;
+ };
+ DenseMap<TypeID, WideningConfig> wideningConfigs;
+
+ /// Per-state merge-change counter. Populated only for states whose type has
+ /// widening enabled; empty for any solver that never calls `enableWidening`.
+ DenseMap<AnalysisState *, uint32_t> mergeChangeCounts;
+
/// The solver is working on the worklist.
bool isRunning = false;
@@ -716,6 +771,10 @@ class DataFlowAnalysis {
/// Return the configuration of the solver used for this analysis.
const DataFlowConfig &getSolverConfig() const { return solver.getConfig(); }
+ /// Return the solver that owns this analysis. Needed by sparse framework
+ /// code in subclasses (e.g. the merge-site widening path).
+ DataFlowSolver &getSolver() { return solver; }
+
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
/// When compiling with debugging, keep a name for the analyis.
StringRef debugName;
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index b29fc28131806..6eac43f67c1fd 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -58,6 +58,21 @@ LogicalResult staticallyNonNegative(DataFlowSolver &solver, Operation *op) {
}
} // namespace mlir::dataflow
+/// Budget for merge-site joins per integer-range lattice element before the
+/// solver forces the element to max-range. Sized for realistic DAG nesting
+/// depth; loops with dynamic bounds trip well before 2^31 iterations.
+static constexpr unsigned kIntegerRangeWideningBudget = 16;
+
+IntegerRangeAnalysis::IntegerRangeAnalysis(DataFlowSolver &solver)
+ : SparseForwardDataFlowAnalysis(solver) {
+ solver.enableWidening<IntegerValueRangeLattice>(
+ kIntegerRangeWideningBudget, [](AnalysisState *state) -> ChangeResult {
+ auto *lattice = static_cast<IntegerValueRangeLattice *>(state);
+ return lattice->join(
+ IntegerValueRange::getMaxRange(cast<Value>(lattice->getAnchor())));
+ });
+}
+
LogicalResult IntegerRangeAnalysis::visitOperation(
Operation *op, ArrayRef<const IntegerValueRangeLattice *> operands,
ArrayRef<IntegerValueRangeLattice *> results) {
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 90f2a588d1ca4..a97e91a804764 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -46,8 +46,8 @@ void AbstractSparseLattice::onUpdate(DataFlowSolver *solver) const {
//===----------------------------------------------------------------------===//
AbstractSparseForwardDataFlowAnalysis::AbstractSparseForwardDataFlowAnalysis(
- DataFlowSolver &solver)
- : DataFlowAnalysis(solver) {
+ DataFlowSolver &solver, TypeID stateTypeID)
+ : DataFlowAnalysis(solver), stateTypeID(stateTypeID) {
registerAnchorKind<CFGEdge>();
}
@@ -358,7 +358,16 @@ void AbstractSparseForwardDataFlowAnalysis::setAllToEntryStates(
void AbstractSparseForwardDataFlowAnalysis::join(
AbstractSparseLattice *lhs, const AbstractSparseLattice &rhs) {
- propagateIfChanged(lhs, lhs->join(rhs));
+ ChangeResult changed = lhs->join(rhs);
+ // Single fast-path branch when no analysis in this solver opted into
+ // widening. Well-predicted constant for any solver that never calls
+ // enableWidening, so this path is effectively free for constant-prop,
+ // liveness, dead-code, and SCCP. `tryWidenAtMerge` returns only the
+ // widening-induced delta so it cannot downgrade `changed`.
+ DataFlowSolver &s = getSolver();
+ if (LLVM_UNLIKELY(s.hasWideningEnabled()))
+ changed |= s.tryWidenAtMerge(lhs, stateTypeID, changed);
+ propagateIfChanged(lhs, changed);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Analysis/DataFlowFramework.cpp b/mlir/lib/Analysis/DataFlowFramework.cpp
index e51ae7a1d7ca7..17f817cf56b8a 100644
--- a/mlir/lib/Analysis/DataFlowFramework.cpp
+++ b/mlir/lib/Analysis/DataFlowFramework.cpp
@@ -170,6 +170,32 @@ void DataFlowSolver::propagateIfChanged(AnalysisState *state,
}
}
+ChangeResult DataFlowSolver::tryWidenAtMerge(AnalysisState *state,
+ TypeID stateTypeID,
+ ChangeResult changeFromJoin) {
+ assert(hasWideningEnabled() &&
+ "tryWidenAtMerge called without widening enabled; "
+ "gate callers on hasWideningEnabled()");
+ auto cfgIt = wideningConfigs.find(stateTypeID);
+ if (cfgIt == wideningConfigs.end())
+ return ChangeResult::NoChange;
+
+ const WideningConfig &cfg = cfgIt->second;
+ uint32_t &changes = mergeChangeCounts[state];
+ if (changes >= cfg.budget) {
+ // Budget exhausted; force a sound over-approximation. Only the widen
+ // result is returned — the caller unions it with its own change.
+ ChangeResult widened = cfg.widen(state);
+ DATAFLOW_DEBUG(LDBG() << "Widening state " << state->debugName << " of "
+ << state->anchor << " after " << changes
+ << " merge-site changes\n");
+ return widened;
+ }
+ if (changeFromJoin == ChangeResult::Change)
+ ++changes;
+ return ChangeResult::NoChange;
+}
+
//===----------------------------------------------------------------------===//
// DataFlowAnalysis
//===----------------------------------------------------------------------===//
>From a59120ac3b8bb6ab0b500be84f204a65b7604539 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Mon, 20 Apr 2026 15:34:46 +0200
Subject: [PATCH 2/8] [mlir][dataflow] NFC: unify merge-change counter type
with budget
WideningConfig::budget is `unsigned`; the matching per-state counter was
`uint32_t`. Same width on all supported platforms but inconsistent at
the comparison site. Switch the counter to `unsigned` so the types
match at the declaration.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
mlir/include/mlir/Analysis/DataFlowFramework.h | 2 +-
mlir/lib/Analysis/DataFlowFramework.cpp | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index 2913fde31f108..091a8ee406774 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -496,7 +496,7 @@ class DataFlowSolver {
/// Per-state merge-change counter. Populated only for states whose type has
/// widening enabled; empty for any solver that never calls `enableWidening`.
- DenseMap<AnalysisState *, uint32_t> mergeChangeCounts;
+ DenseMap<AnalysisState *, unsigned> mergeChangeCounts;
/// The solver is working on the worklist.
bool isRunning = false;
diff --git a/mlir/lib/Analysis/DataFlowFramework.cpp b/mlir/lib/Analysis/DataFlowFramework.cpp
index 17f817cf56b8a..ae50c6ab3ec45 100644
--- a/mlir/lib/Analysis/DataFlowFramework.cpp
+++ b/mlir/lib/Analysis/DataFlowFramework.cpp
@@ -181,7 +181,7 @@ ChangeResult DataFlowSolver::tryWidenAtMerge(AnalysisState *state,
return ChangeResult::NoChange;
const WideningConfig &cfg = cfgIt->second;
- uint32_t &changes = mergeChangeCounts[state];
+ unsigned &changes = mergeChangeCounts[state];
if (changes >= cfg.budget) {
// Budget exhausted; force a sound over-approximation. Only the widen
// result is returned — the caller unions it with its own change.
>From b293ce113257f4c1ee1fd4f49e5fa9cfde6b1b57 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Tue, 21 Apr 2026 20:15:36 +0200
Subject: [PATCH 3/8] [mlir][dataflow] Add IntegerRangeAnalysis convergence
test
Test from PR #192235: nested scf.while with dynamic bounds that
previously hung IntegerRangeAnalysis. Converges instantly with the
framework-level merge-site widening added in the preceding commit.
---
.../Arith/int-range-analysis-convergence.mlir | 65 +++++++++++++++++++
1 file changed, 65 insertions(+)
create mode 100644 mlir/test/Dialect/Arith/int-range-analysis-convergence.mlir
diff --git a/mlir/test/Dialect/Arith/int-range-analysis-convergence.mlir b/mlir/test/Dialect/Arith/int-range-analysis-convergence.mlir
new file mode 100644
index 0000000000000..bf4da6bab3a54
--- /dev/null
+++ b/mlir/test/Dialect/Arith/int-range-analysis-convergence.mlir
@@ -0,0 +1,65 @@
+// IntegerRangeAnalysis non-convergence on scf.while with dynamic bounds.
+//
+// The carry range ratchets [0,0]->[0,1]->[0,2]->... without bound.
+// Two nested scf.if layers with differing arith chains (addi, muli)
+// bounded by remui create enough worklist cascade to prevent the
+// solver's back-to-back convergence shortcut from firing.
+//
+// With framework-level merge-site widening, the analysis converges
+// in bounded time.
+//
+// RUN: mlir-opt -int-range-optimizations %s -o /dev/null
+
+func.func @grouped_gemm_while_hang(%n: i32, %flag: i1) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %c3 = arith.constant 3 : i32
+ %c7 = arith.constant 7 : i32
+ %c127 = arith.constant 127 : i32
+ %init = arith.cmpi slt, %c0, %n : i32
+
+ %res:2 = scf.while (%a0 = %c0, %cond = %init) : (i32, i1) -> (i32, i1) {
+ scf.condition(%cond) %a0, %cond : i32, i1
+ } do {
+ ^bb0(%b0: i32, %bc: i1):
+ %t0 = arith.addi %b0, %c1 : i32
+ %ic = arith.cmpi slt, %t0, %n : i32
+
+ %inner:2 = scf.while (%i0 = %t0, %iic = %ic) : (i32, i1) -> (i32, i1) {
+ scf.condition(%iic) %i0, %iic : i32, i1
+ } do {
+ ^bb1(%j0: i32, %jc: i1):
+
+ %L0 = scf.if %flag -> (i32) {
+ %a0_0 = arith.addi %j0, %c1 : i32
+ %a0_1 = arith.muli %a0_0, %c7 : i32
+ %a0_r = arith.remui %a0_1, %c127 : i32
+ scf.yield %a0_r : i32
+ } else {
+ %b0_0 = arith.addi %j0, %c3 : i32
+ %b0_1 = arith.muli %b0_0, %c7 : i32
+ %b0_r = arith.remui %b0_1, %c127 : i32
+ scf.yield %b0_r : i32
+ }
+
+ %L1 = scf.if %flag -> (i32) {
+ %a1_0 = arith.addi %L0, %c1 : i32
+ %a1_1 = arith.muli %a1_0, %c7 : i32
+ %a1_r = arith.remui %a1_1, %c127 : i32
+ scf.yield %a1_r : i32
+ } else {
+ %b1_0 = arith.addi %L0, %c3 : i32
+ %b1_1 = arith.muli %b1_0, %c7 : i32
+ %b1_r = arith.remui %b1_1, %c127 : i32
+ scf.yield %b1_r : i32
+ }
+
+ %nic = arith.cmpi slt, %L1, %n : i32
+ scf.yield %L1, %nic : i32, i1
+ }
+
+ %nc = arith.cmpi slt, %inner#0, %n : i32
+ scf.yield %inner#0, %nc : i32, i1
+ }
+ return %res#0 : i32
+}
>From e9058396761262900627e65f8e5d9e66657b506b Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Tue, 21 Apr 2026 21:58:28 +0200
Subject: [PATCH 4/8] cleanup
---
mlir/include/mlir/Analysis/DataFlowFramework.h | 10 +++-------
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 7 +++----
2 files changed, 6 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index 091a8ee406774..97691dcff9957 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -460,8 +460,7 @@ class DataFlowSolver {
/// Fast-path predicate: true if any analysis loaded in this solver has
/// opted into widening. Callers on the hot path should check this before
- /// doing any further widening work. Derived from the configs map rather
- /// than a redundant bool so there is only one source of truth.
+ /// doing any further widening work.
bool hasWideningEnabled() const { return !wideningConfigs.empty(); }
/// If widening is configured for `stateTypeID` and this state's merge-site
@@ -485,9 +484,7 @@ class DataFlowSolver {
/// Configuration of the dataflow solver.
DataFlowConfig config;
- /// Per-state-type widening policy. Emptiness of this map is the canonical
- /// answer to "is widening enabled for this solver"; see
- /// `hasWideningEnabled`.
+ /// Per-state-type widening policy.
struct WideningConfig {
unsigned budget;
WidenFn widen;
@@ -771,8 +768,7 @@ class DataFlowAnalysis {
/// Return the configuration of the solver used for this analysis.
const DataFlowConfig &getSolverConfig() const { return solver.getConfig(); }
- /// Return the solver that owns this analysis. Needed by sparse framework
- /// code in subclasses (e.g. the merge-site widening path).
+ /// Return the solver that owns this analysis.
DataFlowSolver &getSolver() { return solver; }
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index a97e91a804764..3bf89540c38ef 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -360,12 +360,11 @@ void AbstractSparseForwardDataFlowAnalysis::join(
AbstractSparseLattice *lhs, const AbstractSparseLattice &rhs) {
ChangeResult changed = lhs->join(rhs);
// Single fast-path branch when no analysis in this solver opted into
- // widening. Well-predicted constant for any solver that never calls
- // enableWidening, so this path is effectively free for constant-prop,
- // liveness, dead-code, and SCCP. `tryWidenAtMerge` returns only the
- // widening-induced delta so it cannot downgrade `changed`.
+ // widening.
DataFlowSolver &s = getSolver();
if (LLVM_UNLIKELY(s.hasWideningEnabled()))
+ // `tryWidenAtMerge` returns only the widening-induced delta so it
+ // cannot downgrade `changed`.
changed |= s.tryWidenAtMerge(lhs, stateTypeID, changed);
propagateIfChanged(lhs, changed);
}
>From 2c16fab6c410f7d787c7cfb1287eb08bf953c606 Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Tue, 21 Apr 2026 23:56:02 +0200
Subject: [PATCH 5/8] [mlir][dataflow] Extend merge-site widening hook to all
analysis bases
Wires `DataFlowSolver::tryWidenAtMerge` into the three remaining merge
helpers so the opt-in widening path is uniform across the four
sparse/dense x forward/backward analysis variants:
- `AbstractSparseBackwardDataFlowAnalysis::meet`
- `AbstractDenseForwardDataFlowAnalysis::join` (inline)
- `AbstractDenseBackwardDataFlowAnalysis::meet` (inline)
Each base gains a constructor overload that accepts a `TypeID`; the
CRTP wrappers (`SparseBackwardDataFlowAnalysis<StateT>`,
`DenseForwardDataFlowAnalysis<LatticeT>`, `DenseBackwardDataFlowAnalysis
<LatticeT>`) thread `TypeID::get<...>()` through. Legacy single-arg
constructors are retained and simply store an empty TypeID, which
disables widening for analyses that don't opt in.
No behavior change for analyses that don't call `enableWidening`: the
hot path is the existing single `hasWideningEnabled()` branch on an
empty `DenseMap`.
---
.../mlir/Analysis/DataFlow/DenseAnalysis.h | 62 ++++++++++++++++---
.../mlir/Analysis/DataFlow/SparseAnalysis.h | 21 ++++++-
mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp | 12 +++-
3 files changed, 81 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
index 82f86d06886b6..ffb0136e6ad90 100644
--- a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
@@ -67,7 +67,18 @@ class AbstractDenseLattice : public AnalysisState {
/// Visit a program point at the begining of block will visit the block itself.
class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
public:
- using DataFlowAnalysis::DataFlowAnalysis;
+ /// Construct the analysis. `stateTypeID` identifies the concrete lattice
+ /// type used by this analysis; it routes the merge-site widening lookup in
+ /// the solver. The CRTP wrapper `DenseForwardDataFlowAnalysis<LatticeT>`
+ /// fills it in from the template parameter.
+ AbstractDenseForwardDataFlowAnalysis(DataFlowSolver &solver,
+ TypeID stateTypeID)
+ : DataFlowAnalysis(solver), stateTypeID(stateTypeID) {}
+
+ /// Legacy constructor. Using this disables widening for this analysis
+ /// because the state TypeID is unknown. Prefer the two-argument ctor.
+ explicit AbstractDenseForwardDataFlowAnalysis(DataFlowSolver &solver)
+ : AbstractDenseForwardDataFlowAnalysis(solver, TypeID()) {}
/// Initialize the analysis by visiting every program point whose execution
/// may modify the program state; that is, every operation and block.
@@ -113,7 +124,11 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
/// Join a lattice with another and propagate an update if it changed.
void join(AbstractDenseLattice *lhs, const AbstractDenseLattice &rhs) {
- propagateIfChanged(lhs, lhs->join(rhs));
+ ChangeResult changed = lhs->join(rhs);
+ DataFlowSolver &s = getSolver();
+ if (LLVM_UNLIKELY(s.hasWideningEnabled()))
+ changed |= s.tryWidenAtMerge(lhs, stateTypeID, changed);
+ propagateIfChanged(lhs, changed);
}
/// Visit an operation. If this is a call operation or region control-flow
@@ -189,6 +204,11 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
void visitCallOperation(CallOpInterface call,
const AbstractDenseLattice &before,
AbstractDenseLattice *after);
+
+ /// TypeID of the concrete lattice type used by this analysis. Threaded
+ /// through from the CRTP derived class so that merge-site widening lookups
+ /// know which `WideningConfig` to consult. Empty TypeID disables widening.
+ TypeID stateTypeID;
};
//===----------------------------------------------------------------------===//
@@ -208,8 +228,9 @@ class DenseForwardDataFlowAnalysis
"analysis state class expected to subclass AbstractDenseLattice");
public:
- using AbstractDenseForwardDataFlowAnalysis::
- AbstractDenseForwardDataFlowAnalysis;
+ explicit DenseForwardDataFlowAnalysis(DataFlowSolver &solver)
+ : AbstractDenseForwardDataFlowAnalysis(solver,
+ TypeID::get<LatticeT>()) {}
/// Visit an operation with the dense lattice before its execution. This
/// function is expected to set the dense lattice after its execution and
@@ -358,10 +379,22 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
public:
/// Construct the analysis in the given solver. Takes a symbol table
/// collection that is used to cache symbol resolution in interprocedural part
- /// of the analysis. The symbol table need not be prefilled.
+ /// of the analysis. The symbol table need not be prefilled. `stateTypeID`
+ /// identifies the concrete lattice type used by this analysis; it routes the
+ /// merge-site widening lookup in the solver. The CRTP wrapper
+ /// `DenseBackwardDataFlowAnalysis<LatticeT>` fills it in from the template
+ /// parameter.
+ AbstractDenseBackwardDataFlowAnalysis(DataFlowSolver &solver,
+ SymbolTableCollection &symbolTable,
+ TypeID stateTypeID)
+ : DataFlowAnalysis(solver), symbolTable(symbolTable),
+ stateTypeID(stateTypeID) {}
+
+ /// Legacy constructor. Using this disables widening for this analysis
+ /// because the state TypeID is unknown. Prefer the three-argument ctor.
AbstractDenseBackwardDataFlowAnalysis(DataFlowSolver &solver,
SymbolTableCollection &symbolTable)
- : DataFlowAnalysis(solver), symbolTable(symbolTable) {}
+ : AbstractDenseBackwardDataFlowAnalysis(solver, symbolTable, TypeID()) {}
/// Initialize the analysis by visiting every program point whose execution
/// may modify the program state; that is, every operation and block.
@@ -410,7 +443,11 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
/// Meet a lattice with another lattice and propagate an update if it changed.
void meet(AbstractDenseLattice *lhs, const AbstractDenseLattice &rhs) {
- propagateIfChanged(lhs, lhs->meet(rhs));
+ ChangeResult changed = lhs->meet(rhs);
+ DataFlowSolver &s = getSolver();
+ if (LLVM_UNLIKELY(s.hasWideningEnabled()))
+ changed |= s.tryWidenAtMerge(lhs, stateTypeID, changed);
+ propagateIfChanged(lhs, changed);
}
/// Visit an operation. Dispatches to specialized methods for call or region
@@ -497,6 +534,11 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
/// Symbol table for call-level control flow.
SymbolTableCollection &symbolTable;
+
+ /// TypeID of the concrete lattice type used by this analysis. Threaded
+ /// through from the CRTP derived class so that merge-site widening lookups
+ /// know which `WideningConfig` to consult. Empty TypeID disables widening.
+ TypeID stateTypeID;
};
//===----------------------------------------------------------------------===//
@@ -515,8 +557,10 @@ class DenseBackwardDataFlowAnalysis
"analysis state expected to subclass AbstractDenseLattice");
public:
- using AbstractDenseBackwardDataFlowAnalysis::
- AbstractDenseBackwardDataFlowAnalysis;
+ DenseBackwardDataFlowAnalysis(DataFlowSolver &solver,
+ SymbolTableCollection &symbolTable)
+ : AbstractDenseBackwardDataFlowAnalysis(solver, symbolTable,
+ TypeID::get<LatticeT>()) {}
/// Transfer function. Visits an operation with the dense lattice after its
/// execution. This function is expected to set the dense lattice before its
diff --git a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
index fb9c73ef9789a..4d37279d114ca 100644
--- a/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h
@@ -429,8 +429,19 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
LogicalResult visit(ProgramPoint *point) override;
protected:
+ /// Construct the analysis. `stateTypeID` identifies the concrete lattice
+ /// type used by this analysis; it routes the merge-site widening lookup in
+ /// the solver. The CRTP wrapper `SparseBackwardDataFlowAnalysis<StateT>`
+ /// fills it in from the template parameter.
+ AbstractSparseBackwardDataFlowAnalysis(DataFlowSolver &solver,
+ SymbolTableCollection &symbolTable,
+ TypeID stateTypeID);
+
+ /// Legacy constructor. Using this disables widening for this analysis
+ /// because the state TypeID is unknown. Prefer the three-argument ctor.
explicit AbstractSparseBackwardDataFlowAnalysis(
- DataFlowSolver &solver, SymbolTableCollection &symbolTable);
+ DataFlowSolver &solver, SymbolTableCollection &symbolTable)
+ : AbstractSparseBackwardDataFlowAnalysis(solver, symbolTable, TypeID()) {}
/// The operation transfer function. Given the result lattices, this
/// function is expected to set the operand lattices.
@@ -519,6 +530,11 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
getLatticeElementsFor(ProgramPoint *point, ValueRange values);
SymbolTableCollection &symbolTable;
+
+ /// TypeID of the concrete lattice type used by this analysis. Threaded
+ /// through from the CRTP derived class so that merge-site widening lookups
+ /// know which `WideningConfig` to consult. Empty TypeID disables widening.
+ TypeID stateTypeID;
};
//===----------------------------------------------------------------------===//
@@ -543,7 +559,8 @@ class SparseBackwardDataFlowAnalysis
public:
explicit SparseBackwardDataFlowAnalysis(DataFlowSolver &solver,
SymbolTableCollection &symbolTable)
- : AbstractSparseBackwardDataFlowAnalysis(solver, symbolTable) {}
+ : AbstractSparseBackwardDataFlowAnalysis(solver, symbolTable,
+ TypeID::get<StateT>()) {}
/// Visit an operation with the lattices of its results. This function is
/// expected to set the lattices of the operation's operands.
diff --git a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
index 3bf89540c38ef..c5efea4149875 100644
--- a/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp
@@ -374,8 +374,10 @@ void AbstractSparseForwardDataFlowAnalysis::join(
//===----------------------------------------------------------------------===//
AbstractSparseBackwardDataFlowAnalysis::AbstractSparseBackwardDataFlowAnalysis(
- DataFlowSolver &solver, SymbolTableCollection &symbolTable)
- : DataFlowAnalysis(solver), symbolTable(symbolTable) {
+ DataFlowSolver &solver, SymbolTableCollection &symbolTable,
+ TypeID stateTypeID)
+ : DataFlowAnalysis(solver), symbolTable(symbolTable),
+ stateTypeID(stateTypeID) {
registerAnchorKind<CFGEdge>();
}
@@ -689,5 +691,9 @@ void AbstractSparseBackwardDataFlowAnalysis::setAllToExitStates(
void AbstractSparseBackwardDataFlowAnalysis::meet(
AbstractSparseLattice *lhs, const AbstractSparseLattice &rhs) {
- propagateIfChanged(lhs, lhs->meet(rhs));
+ ChangeResult changed = lhs->meet(rhs);
+ DataFlowSolver &s = getSolver();
+ if (LLVM_UNLIKELY(s.hasWideningEnabled()))
+ changed |= s.tryWidenAtMerge(lhs, stateTypeID, changed);
+ propagateIfChanged(lhs, changed);
}
>From 681ed49398590d6d88028b6477b64193f73022dc Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Tue, 21 Apr 2026 23:56:29 +0200
Subject: [PATCH 6/8] [mlir][dataflow] Add widening hook matrix test
Adds a `TestWideningAnalysis` pass that instantiates an
infinite-height counter lattice (merge = max, +1 per transfer) behind
all four sparse/dense x forward/backward analysis drivers, and a LIT
file that exercises the widening hook for each variant:
- Straight-line code stays below the per-callback budget.
- `scf.while` with dynamic bound trips widening on the loop-carried
lattice element.
- Unstructured `cf.br` cycle also converges via merge-site widening,
demonstrating independence from loop detection or back-edge
identification.
Drives the hook registered via `DataFlowSolver::enableWidening<StateT>`
from a real consumer at each of the four framework choke points so the
uniformity claim is covered by an actual test.
---
.../mlir/Analysis/DataFlow/DenseAnalysis.h | 3 +-
.../test/Analysis/DataFlow/test-widening.mlir | 81 ++++
mlir/test/lib/Analysis/CMakeLists.txt | 1 +
.../DataFlow/TestWideningAnalysis.cpp | 402 ++++++++++++++++++
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
5 files changed, 487 insertions(+), 2 deletions(-)
create mode 100644 mlir/test/Analysis/DataFlow/test-widening.mlir
create mode 100644 mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp
diff --git a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
index ffb0136e6ad90..68266c39ae5c1 100644
--- a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
@@ -229,8 +229,7 @@ class DenseForwardDataFlowAnalysis
public:
explicit DenseForwardDataFlowAnalysis(DataFlowSolver &solver)
- : AbstractDenseForwardDataFlowAnalysis(solver,
- TypeID::get<LatticeT>()) {}
+ : AbstractDenseForwardDataFlowAnalysis(solver, TypeID::get<LatticeT>()) {}
/// Visit an operation with the dense lattice before its execution. This
/// function is expected to set the dense lattice after its execution and
diff --git a/mlir/test/Analysis/DataFlow/test-widening.mlir b/mlir/test/Analysis/DataFlow/test-widening.mlir
new file mode 100644
index 0000000000000..3521beaa0231b
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/test-widening.mlir
@@ -0,0 +1,81 @@
+// Exercises the merge-site widening hook across the four sparse/dense x
+// forward/backward analysis variants. Uses an infinite-height counter
+// lattice with a small widening budget; loops trip widening, straight-line
+// code stays below it.
+//
+// RUN: mlir-opt %s -split-input-file \
+// RUN: --test-widening-analysis=variant=sparse-forward 2>&1 | FileCheck %s --check-prefix=SF
+// RUN: mlir-opt %s -split-input-file \
+// RUN: --test-widening-analysis=variant=sparse-backward 2>&1 | FileCheck %s --check-prefix=SB
+// RUN: mlir-opt %s -split-input-file \
+// RUN: --test-widening-analysis=variant=dense-forward 2>&1 | FileCheck %s --check-prefix=DF
+// RUN: mlir-opt %s -split-input-file \
+// RUN: --test-widening-analysis=variant=dense-backward 2>&1 | FileCheck %s --check-prefix=DB
+
+//===----------------------------------------------------------------------===//
+// Straight-line code: counter stays bounded, no widening.
+//===----------------------------------------------------------------------===//
+
+// SF: tag=straight: results=[count={{[0-9]+}}]
+// SB: tag=straight: operands=[count={{[0-9]+}}, count={{[0-9]+}}]
+// DF: tag=straight: count={{[0-9]+}}
+// DB: tag=straight: count={{[0-9]+}}
+
+func.func @straight(%x: i32) -> i32 {
+ %a = arith.addi %x, %x : i32
+ %b = arith.addi %a, %a : i32
+ %c = arith.addi %b, %b : i32
+ %d = arith.addi %c, %c {tag = "straight"} : i32
+ return %d : i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Structured loop (scf.while with dynamic bound): widening fires on the
+// loop-carried value and propagates through the body.
+//===----------------------------------------------------------------------===//
+
+// SF: tag=scf_body: results=[widened]
+// SB: tag=scf_body: operands=[widened, {{widened|count=[0-9]+}}]
+// DF: tag=scf_body: widened
+// DB: tag=scf_body: widened
+
+func.func @scf_while(%n: i32) -> i32 {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ %init = arith.cmpi slt, %c0, %n : i32
+ %r:2 = scf.while (%acc = %c0, %cond = %init) : (i32, i1) -> (i32, i1) {
+ scf.condition(%cond) %acc, %cond : i32, i1
+ } do {
+ ^bb(%a: i32, %c: i1):
+ %next = arith.addi %a, %c1 {tag = "scf_body"} : i32
+ %nc = arith.cmpi slt, %next, %n : i32
+ scf.yield %next, %nc : i32, i1
+ }
+ return %r#0 : i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Unstructured CFG cycle (cf.br / cf.cond_br): the merge-site hook fires at
+// the block-arg join on the loop head without any loop detection. This is
+// the architectural payoff over back-edge-only widening.
+//===----------------------------------------------------------------------===//
+
+// SF: tag=cfg_body: results=[widened]
+// SB: tag=cfg_body: operands=[widened, {{widened|count=[0-9]+}}]
+// DF: tag=cfg_body: widened
+// DB: tag=cfg_body: widened
+
+func.func @unstructured_cfg(%cond: i1) {
+ %c0 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : i32
+ cf.br ^head(%c0 : i32)
+^head(%a: i32):
+ %next = arith.addi %a, %c1 {tag = "cfg_body"} : i32
+ cf.cond_br %cond, ^head(%next : i32), ^exit
+^exit:
+ return
+}
diff --git a/mlir/test/lib/Analysis/CMakeLists.txt b/mlir/test/lib/Analysis/CMakeLists.txt
index c37671ade37b3..a57ce1f6e230f 100644
--- a/mlir/test/lib/Analysis/CMakeLists.txt
+++ b/mlir/test/lib/Analysis/CMakeLists.txt
@@ -18,6 +18,7 @@ add_mlir_library(MLIRTestAnalysis
DataFlow/TestLivenessAnalysis.cpp
DataFlow/TestSparseBackwardDataFlowAnalysis.cpp
DataFlow/TestStridedMetadataRangeAnalysis.cpp
+ DataFlow/TestWideningAnalysis.cpp
EXCLUDE_FROM_LIBMLIR
diff --git a/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp
new file mode 100644
index 0000000000000..9f87fe77d95d4
--- /dev/null
+++ b/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp
@@ -0,0 +1,402 @@
+//===- TestWideningAnalysis.cpp - Test merge-site widening hook ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Exercises DataFlowSolver::enableWidening across the four sparse/dense x
+// forward/backward analysis variants. Each driver uses the same infinite-
+// height `CounterValue` lattice (max-merge, +1 per transfer), so loops
+// ratchet the counter without bound until widening caps it. The widen
+// callback flips a `widened` bit. Emits the final lattice value for each
+// "tag"-annotated op so FileCheck can distinguish tight-bounded vs widened
+// outcomes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/DataFlow/DenseAnalysis.h"
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/Analysis/DataFlow/Utils.h"
+#include "mlir/Analysis/DataFlowFramework.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+using namespace mlir::dataflow;
+
+namespace {
+
+/// Budget deliberately small so loop tests trip widening quickly. Acyclic
+/// control flow stays well below it.
+static constexpr unsigned kTestBudget = 4;
+
+//===----------------------------------------------------------------------===//
+// Counter value
+//===----------------------------------------------------------------------===//
+
+/// Infinite-height scalar counter. Merge = max on the count; the `widened`
+/// bit saturates (once set, stays set and the count freezes). Without
+/// widening, a loop body's transfer ratchets the count by +1 on every
+/// revisit and never converges; with widening, the bit flips and the value
+/// becomes idempotent.
+struct CounterValue {
+ int count = 0;
+ bool widened = false;
+
+ bool operator==(const CounterValue &rhs) const {
+ return count == rhs.count && widened == rhs.widened;
+ }
+
+ static CounterValue getWidened() {
+ return CounterValue{std::numeric_limits<int>::max(), true};
+ }
+
+ static CounterValue join(const CounterValue &lhs, const CounterValue &rhs) {
+ if (lhs.widened || rhs.widened)
+ return getWidened();
+ return CounterValue{std::max(lhs.count, rhs.count), false};
+ }
+
+ /// For sparse backward we want the same max semantics as forward.
+ static CounterValue meet(const CounterValue &lhs, const CounterValue &rhs) {
+ return join(lhs, rhs);
+ }
+
+ void print(raw_ostream &os) const {
+ if (widened)
+ os << "widened";
+ else
+ os << "count=" << count;
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Sparse lattice
+//===----------------------------------------------------------------------===//
+
+class CounterSparseLattice : public Lattice<CounterValue> {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterSparseLattice)
+ using Lattice::Lattice;
+};
+
+/// Register a widen callback for the sparse lattice: force to `widened`.
+static void enableSparseWidening(DataFlowSolver &solver) {
+ solver.enableWidening<CounterSparseLattice>(
+ kTestBudget, [](AnalysisState *state) -> ChangeResult {
+ auto *lattice = static_cast<CounterSparseLattice *>(state);
+ return lattice->join(CounterValue::getWidened());
+ });
+}
+
+//===----------------------------------------------------------------------===//
+// Dense lattice
+//===----------------------------------------------------------------------===//
+
+class CounterDenseLattice : public AbstractDenseLattice {
+public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterDenseLattice)
+ using AbstractDenseLattice::AbstractDenseLattice;
+
+ ChangeResult join(const AbstractDenseLattice &rhs) override {
+ const CounterValue &rv =
+ static_cast<const CounterDenseLattice &>(rhs).value;
+ return set(CounterValue::join(value, rv));
+ }
+
+ ChangeResult meet(const AbstractDenseLattice &rhs) override {
+ return join(rhs);
+ }
+
+ ChangeResult set(const CounterValue &newValue) {
+ if (value == newValue)
+ return ChangeResult::NoChange;
+ value = newValue;
+ return ChangeResult::Change;
+ }
+
+ const CounterValue &getValue() const { return value; }
+
+ void print(raw_ostream &os) const override { value.print(os); }
+
+private:
+ CounterValue value;
+};
+
+static void enableDenseWidening(DataFlowSolver &solver) {
+ solver.enableWidening<CounterDenseLattice>(
+ kTestBudget, [](AnalysisState *state) -> ChangeResult {
+ auto *lattice = static_cast<CounterDenseLattice *>(state);
+ return lattice->set(CounterValue::getWidened());
+ });
+}
+
+//===----------------------------------------------------------------------===//
+// Sparse forward driver
+//===----------------------------------------------------------------------===//
+
+class CounterSparseForwardAnalysis
+ : public SparseForwardDataFlowAnalysis<CounterSparseLattice> {
+public:
+ explicit CounterSparseForwardAnalysis(DataFlowSolver &solver)
+ : SparseForwardDataFlowAnalysis(solver) {
+ enableSparseWidening(solver);
+ }
+
+ LogicalResult
+ visitOperation(Operation *, ArrayRef<const CounterSparseLattice *> operands,
+ ArrayRef<CounterSparseLattice *> results) override {
+ CounterValue out;
+ for (const CounterSparseLattice *operand : operands)
+ out = CounterValue::join(out, operand->getValue());
+ if (!out.widened)
+ out.count += 1;
+ for (CounterSparseLattice *result : results)
+ propagateIfChanged(result, result->join(out));
+ return success();
+ }
+
+ void setToEntryState(CounterSparseLattice *lattice) override {
+ propagateIfChanged(lattice, lattice->join(CounterValue{}));
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Sparse backward driver
+//===----------------------------------------------------------------------===//
+
+class CounterSparseBackwardAnalysis
+ : public SparseBackwardDataFlowAnalysis<CounterSparseLattice> {
+public:
+ CounterSparseBackwardAnalysis(DataFlowSolver &solver,
+ SymbolTableCollection &symbols)
+ : SparseBackwardDataFlowAnalysis(solver, symbols) {
+ enableSparseWidening(solver);
+ }
+
+ LogicalResult
+ visitOperation(Operation *op, ArrayRef<CounterSparseLattice *> operands,
+ ArrayRef<const CounterSparseLattice *> results) override {
+ CounterValue in;
+ for (const CounterSparseLattice *result : results)
+ in = CounterValue::join(in, result->getValue());
+ if (!in.widened)
+ in.count += 1;
+ // Route through the framework's `meet` helper so the widening hook fires.
+ // `meet` only reads rhs, so a stack-allocated throwaway lattice with an
+ // arbitrary anchor is sufficient.
+ for (CounterSparseLattice *operand : operands) {
+ CounterSparseLattice tmp(operand->getAnchor());
+ (void)tmp.join(in);
+ meet(operand, tmp);
+ }
+ return success();
+ }
+
+ void visitBranchOperand(OpOperand &) override {}
+ void visitCallOperand(OpOperand &) override {}
+ void visitNonControlFlowArguments(RegionSuccessor &,
+ ArrayRef<BlockArgument>) override {}
+
+ void setToExitState(CounterSparseLattice *lattice) override {
+ propagateIfChanged(lattice, lattice->join(CounterValue{}));
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Dense forward driver
+//===----------------------------------------------------------------------===//
+
+class CounterDenseForwardAnalysis
+ : public DenseForwardDataFlowAnalysis<CounterDenseLattice> {
+public:
+ explicit CounterDenseForwardAnalysis(DataFlowSolver &solver)
+ : DenseForwardDataFlowAnalysis(solver) {
+ enableDenseWidening(solver);
+ }
+
+ LogicalResult visitOperation(Operation *, const CounterDenseLattice &before,
+ CounterDenseLattice *after) override {
+ CounterValue v = before.getValue();
+ if (!v.widened)
+ v.count += 1;
+ // Route through the framework's `join` helper so the widening hook can
+ // fire on the transfer update. `join` at merge sites is the only hook
+ // the framework exposes for widening; for a max-monotone lattice,
+ // joining-in the computed value is equivalent to assigning it.
+ CounterDenseLattice tmp(after->getAnchor());
+ (void)tmp.set(v);
+ join(after, tmp);
+ return success();
+ }
+
+ void setToEntryState(CounterDenseLattice *lattice) override {
+ propagateIfChanged(lattice, lattice->set(CounterValue{}));
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Dense backward driver
+//===----------------------------------------------------------------------===//
+
+class CounterDenseBackwardAnalysis
+ : public DenseBackwardDataFlowAnalysis<CounterDenseLattice> {
+public:
+ CounterDenseBackwardAnalysis(DataFlowSolver &solver,
+ SymbolTableCollection &symbols)
+ : DenseBackwardDataFlowAnalysis(solver, symbols) {
+ enableDenseWidening(solver);
+ }
+
+ LogicalResult visitOperation(Operation *, const CounterDenseLattice &after,
+ CounterDenseLattice *before) override {
+ CounterValue v = after.getValue();
+ if (!v.widened)
+ v.count += 1;
+ CounterDenseLattice tmp(before->getAnchor());
+ (void)tmp.set(v);
+ meet(before, tmp);
+ return success();
+ }
+
+ void setToExitState(CounterDenseLattice *lattice) override {
+ propagateIfChanged(lattice, lattice->set(CounterValue{}));
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Test pass
+//===----------------------------------------------------------------------===//
+
+enum class WideningVariant {
+ SparseForward,
+ SparseBackward,
+ DenseForward,
+ DenseBackward,
+};
+
+struct TestWideningAnalysisPass
+ : public PassWrapper<TestWideningAnalysisPass, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestWideningAnalysisPass)
+
+ TestWideningAnalysisPass() = default;
+ TestWideningAnalysisPass(const TestWideningAnalysisPass &other)
+ : PassWrapper(other) {
+ variant = other.variant;
+ }
+
+ StringRef getArgument() const override { return "test-widening-analysis"; }
+ StringRef getDescription() const override {
+ return "Exercise DataFlowSolver::enableWidening across the four analysis "
+ "variants; reflect the resulting counter value.";
+ }
+
+ Option<WideningVariant> variant{
+ *this, "variant", llvm::cl::desc("Which analysis variant to run"),
+ llvm::cl::init(WideningVariant::SparseForward),
+ llvm::cl::values(clEnumValN(WideningVariant::SparseForward,
+ "sparse-forward", "Sparse forward"),
+ clEnumValN(WideningVariant::SparseBackward,
+ "sparse-backward", "Sparse backward"),
+ clEnumValN(WideningVariant::DenseForward,
+ "dense-forward", "Dense forward"),
+ clEnumValN(WideningVariant::DenseBackward,
+ "dense-backward", "Dense backward"))};
+
+ void runOnOperation() override {
+ Operation *top = getOperation();
+ DataFlowSolver solver;
+ loadBaselineAnalyses(solver);
+ SymbolTableCollection symbols;
+
+ switch (variant) {
+ case WideningVariant::SparseForward:
+ solver.load<CounterSparseForwardAnalysis>();
+ break;
+ case WideningVariant::SparseBackward:
+ solver.load<CounterSparseBackwardAnalysis>(symbols);
+ break;
+ case WideningVariant::DenseForward:
+ solver.load<CounterDenseForwardAnalysis>();
+ break;
+ case WideningVariant::DenseBackward:
+ solver.load<CounterDenseBackwardAnalysis>(symbols);
+ break;
+ }
+
+ if (failed(solver.initializeAndRun(top)))
+ return signalPassFailure();
+
+ raw_ostream &os = llvm::outs();
+ top->walk([&](Operation *op) {
+ auto tag = op->getAttrOfType<StringAttr>("tag");
+ if (!tag)
+ return;
+ os << "tag=" << tag.getValue() << ": ";
+ switch (variant) {
+ case WideningVariant::SparseForward:
+ printSparseResults(os, solver, op);
+ break;
+ case WideningVariant::SparseBackward:
+ printSparseOperands(os, solver, op);
+ break;
+ case WideningVariant::DenseForward:
+ printDenseAt(os, solver, solver.getProgramPointAfter(op));
+ break;
+ case WideningVariant::DenseBackward:
+ printDenseAt(os, solver, solver.getProgramPointBefore(op));
+ break;
+ }
+ os << "\n";
+ });
+ }
+
+private:
+ static void printSparseResults(raw_ostream &os, DataFlowSolver &solver,
+ Operation *op) {
+ os << "results=[";
+ llvm::interleaveComma(op->getResults(), os, [&](Value v) {
+ if (const auto *lat = solver.lookupState<CounterSparseLattice>(v))
+ lat->getValue().print(os);
+ else
+ os << "<unset>";
+ });
+ os << "]";
+ }
+
+ static void printSparseOperands(raw_ostream &os, DataFlowSolver &solver,
+ Operation *op) {
+ os << "operands=[";
+ llvm::interleaveComma(op->getOperands(), os, [&](Value v) {
+ if (const auto *lat = solver.lookupState<CounterSparseLattice>(v))
+ lat->getValue().print(os);
+ else
+ os << "<unset>";
+ });
+ os << "]";
+ }
+
+ static void printDenseAt(raw_ostream &os, DataFlowSolver &solver,
+ ProgramPoint *point) {
+ if (const auto *lat = solver.lookupState<CounterDenseLattice>(point))
+ lat->getValue().print(os);
+ else
+ os << "<unset>";
+ }
+};
+
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestWideningAnalysisPass() {
+ PassRegistration<TestWideningAnalysisPass>();
+}
+} // namespace test
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index c4754b3a08551..db81efc519b3e 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -161,6 +161,7 @@ void registerTestPassStateExtensionCommunication();
void registerTestVectorLowerings();
void registerTestVectorReductionToSPIRVDotProd();
void registerTestVulkanRunnerPipeline();
+void registerTestWideningAnalysisPass();
void registerTestWrittenToPass();
void registerTestXeGPULowerings();
#if MLIR_ENABLE_PDL_IN_PATTERNMATCH
@@ -310,6 +311,7 @@ static void registerTestPasses() {
mlir::test::registerTestVectorLowerings();
mlir::test::registerTestVectorReductionToSPIRVDotProd();
mlir::test::registerTestVulkanRunnerPipeline();
+ mlir::test::registerTestWideningAnalysisPass();
mlir::test::registerTestWrittenToPass();
mlir::test::registerTestXeGPULowerings();
#if MLIR_ENABLE_PDL_IN_PATTERNMATCH
>From 57f7f58b748e841c1731194e04add261a36afd6d Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Fri, 8 May 2026 15:17:17 +0200
Subject: [PATCH 7/8] missing typeid macros
---
mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp | 8 ++++++++
1 file changed, 8 insertions(+)
diff --git a/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp
index 9f87fe77d95d4..17fdd12c0514c 100644
--- a/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp
+++ b/mlir/test/lib/Analysis/DataFlow/TestWideningAnalysis.cpp
@@ -144,6 +144,8 @@ static void enableDenseWidening(DataFlowSolver &solver) {
class CounterSparseForwardAnalysis
: public SparseForwardDataFlowAnalysis<CounterSparseLattice> {
public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterSparseForwardAnalysis)
+
explicit CounterSparseForwardAnalysis(DataFlowSolver &solver)
: SparseForwardDataFlowAnalysis(solver) {
enableSparseWidening(solver);
@@ -174,6 +176,8 @@ class CounterSparseForwardAnalysis
class CounterSparseBackwardAnalysis
: public SparseBackwardDataFlowAnalysis<CounterSparseLattice> {
public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterSparseBackwardAnalysis)
+
CounterSparseBackwardAnalysis(DataFlowSolver &solver,
SymbolTableCollection &symbols)
: SparseBackwardDataFlowAnalysis(solver, symbols) {
@@ -216,6 +220,8 @@ class CounterSparseBackwardAnalysis
class CounterDenseForwardAnalysis
: public DenseForwardDataFlowAnalysis<CounterDenseLattice> {
public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterDenseForwardAnalysis)
+
explicit CounterDenseForwardAnalysis(DataFlowSolver &solver)
: DenseForwardDataFlowAnalysis(solver) {
enableDenseWidening(solver);
@@ -248,6 +254,8 @@ class CounterDenseForwardAnalysis
class CounterDenseBackwardAnalysis
: public DenseBackwardDataFlowAnalysis<CounterDenseLattice> {
public:
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CounterDenseBackwardAnalysis)
+
CounterDenseBackwardAnalysis(DataFlowSolver &solver,
SymbolTableCollection &symbols)
: DenseBackwardDataFlowAnalysis(solver, symbols) {
>From bb3ec7327a2f7019546f6b4a30bd29241c9c6f9d Mon Sep 17 00:00:00 2001
From: Ivan Butygin <ivan.butygin at gmail.com>
Date: Fri, 8 May 2026 15:42:17 +0200
Subject: [PATCH 8/8] increase kIntegerRangeWideningBudget
---
mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp | 9 ++++++---
1 file changed, 6 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index 6eac43f67c1fd..b6ddffc62e528 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -59,9 +59,12 @@ LogicalResult staticallyNonNegative(DataFlowSolver &solver, Operation *op) {
} // namespace mlir::dataflow
/// Budget for merge-site joins per integer-range lattice element before the
-/// solver forces the element to max-range. Sized for realistic DAG nesting
-/// depth; loops with dynamic bounds trip well before 2^31 iterations.
-static constexpr unsigned kIntegerRangeWideningBudget = 16;
+/// solver forces the element to max-range. Sized as 2 * max supported integer
+/// bitwidth (i64): the longest legitimate strictly-monotone join chain a
+/// ConstantIntRanges lattice can produce is one bit per step on either the lo
+/// or hi end, so any chain longer than 2 * bitwidth is sub-bit-per-iteration
+/// growth (i.e. the +1 ratchet pathology that this hook exists to cut off).
+static constexpr unsigned kIntegerRangeWideningBudget = 128;
IntegerRangeAnalysis::IntegerRangeAnalysis(DataFlowSolver &solver)
: SparseForwardDataFlowAnalysis(solver) {
More information about the Mlir-commits
mailing list