[Mlir-commits] [mlir] fe72b96 - [mlir][dataflow] Deduplicate data-flow solver worklist items (#221983)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 14 05:30:53 PDT 2026
Author: Yoni Katzuv
Date: 2026-09-14T15:30:48+03:00
New Revision: fe72b9678b037eaa71627a503dcc17288f9c3811
URL: https://github.com/llvm/llvm-project/commit/fe72b9678b037eaa71627a503dcc17288f9c3811
DIFF: https://github.com/llvm/llvm-project/commit/fe72b9678b037eaa71627a503dcc17288f9c3811.diff
LOG: [mlir][dataflow] Deduplicate data-flow solver worklist items (#221983)
The solver enqueues a work item every time a state it depends on
changes, so the same item can appear in the worklist many times.
A duplicate work item is re-processed even though its inputs have not
changed since the previous processing. For example, a dense analysis can
visit a join point with N predecessors up to N times in a row.
To deduplicate the worklist, pending items are tracked using a
`DenseSet`, and an item that is already pending is not pushed again.
Items are removed from the set when dequeued, before being visited, so
an analysis can still re-enqueue the point it is currently visiting.
Since analyses already have to be monotone to reach a fixed point, this
change does not affect analysis results, only the number of visits until
convergence. The test `FooState` used XOR for the join operation, which
is not monotone; the join operation was switched to OR and dependent
tests were updated.
A test that counts how many times a join block with several predecessors
is visited was added to check the deduplication.
Added:
Modified:
mlir/include/mlir/Analysis/DataFlowFramework.h
mlir/lib/Analysis/DataFlowFramework.cpp
mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
mlir/test/Analysis/test-foo-analysis.mlir
mlir/test/lib/Analysis/TestDataFlowFramework.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index 30fed25befa5d..69485fe4c53f8 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -18,6 +18,7 @@
#include "mlir/IR/Operation.h"
#include "mlir/Support/StorageUniquer.h"
+#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/EquivalenceClasses.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/STLFunctionalExtras.h"
@@ -408,9 +409,23 @@ class DataFlowSolver {
/// Each item is processed by invoking the child analysis at the program
/// point.
using WorkItem = std::pair<ProgramPoint *, DataFlowAnalysis *>;
- /// Push a work item onto the worklist.
- void enqueue(WorkItem item) { worklist.push(std::move(item)); }
+ /// Push a work item onto the worklist, if it is not already there.
+ void enqueue(WorkItem item) {
+ if (pending.insert(item).second)
+ worklist.push(item);
+ }
+private:
+ /// Pop a work item off the worklist, which must be nonempty.
+ WorkItem dequeue() {
+ assert(!worklist.empty());
+ auto item = worklist.front();
+ worklist.pop();
+ pending.erase(item);
+ return item;
+ }
+
+public:
/// Get the state associated with the given lattice anchor. If it does not
/// exist, create an uninitialized state.
template <typename StateT, typename AnchorT>
@@ -450,6 +465,9 @@ class DataFlowSolver {
/// queue to be processed greedily, speeding up computations that otherwise
/// quickly degenerate to quadratic due to propagation of state updates.
std::queue<WorkItem> worklist;
+ /// Keep track of pending work items, to prevent redundant processing of items
+ /// that were enqueued multiple times.
+ DenseSet<WorkItem> pending;
/// Type-erased instances of the children analyses.
SmallVector<std::unique_ptr<DataFlowAnalysis>> childAnalyses;
diff --git a/mlir/lib/Analysis/DataFlowFramework.cpp b/mlir/lib/Analysis/DataFlowFramework.cpp
index e51ae7a1d7ca7..a2aa380749068 100644
--- a/mlir/lib/Analysis/DataFlowFramework.cpp
+++ b/mlir/lib/Analysis/DataFlowFramework.cpp
@@ -146,8 +146,7 @@ LogicalResult DataFlowSolver::initializeAndRun(
// Iterate until all states are in some initialized state and the worklist
// is exhausted.
while (!worklist.empty()) {
- auto [point, analysis] = worklist.front();
- worklist.pop();
+ auto [point, analysis] = dequeue();
DATAFLOW_DEBUG(LDBG() << "Invoking '" << analysis->debugName
<< "' on: " << *point);
diff --git a/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir b/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
index 740b648078f75..bb27ac285dfb9 100644
--- a/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
+++ b/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
@@ -15,38 +15,38 @@ func.func @linear() {
// converges.
//
// Under the current `FooAnalysis` implementation:
-// - entry op after-state is 0 xor 7 = 7
-// - bb0 terminator after-state is 7 xor 1 = 6
+// - entry op after-state is 0 or 1 = 1
+// - bb0 terminator after-state is 1 or 2 = 3
// - when the join block is first visited, only bb0 has contributed, so the
-// join op transiently sees 6 xor 2 = 4
+// join op transiently sees 3 or 12 = 15
// - once the other predecessor arrives, revisiting the join updates the
-// final staged `foo_state` to 7 for the first op in the join block and it
-// stays 7 for the following op
+// final staged `foo_state` to 31 for the first op in the join block and it
+// stays 31 for the following op
//
// But if a non-staged `BarAnalysis` observed bb2 after only bb0 had reached
-// it, bb2's first tagged op would transiently see 6 xor 2 = 4 and latch
+// it, bb2's first tagged op would transiently see 3 or 12 = 15 and latch
// `bar_state = false`, poisoning later points. The staged run below must use
// only the converged `FooState`, so `bar_state` stays true.
//
// CHECK-LABEL: func.func @requires_staged_bar()
func.func @requires_staged_bar() {
- // CHECK: "test.branch"()[^bb{{[0-9]+}}, ^bb{{[0-9]+}}] {bar_state = true, foo = 7 : ui64, foo_state = 7 : i64, tag = "annotate"} : () -> ()
- "test.branch"() [^bb0, ^bb2] {tag = "annotate", foo = 7 : ui64} : () -> ()
+ // CHECK: "test.branch"()[^bb{{[0-9]+}}, ^bb{{[0-9]+}}] {bar_state = true, foo = 1 : ui64, foo_state = 1 : i64, tag = "annotate"} : () -> ()
+ "test.branch"() [^bb0, ^bb2] {tag = "annotate", foo = 1 : ui64} : () -> ()
^bb0:
- // CHECK: "test.branch"()[^bb{{[0-9]+}}] {bar_state = true, foo = 1 : ui64, foo_state = 6 : i64, tag = "annotate"} : () -> ()
- "test.branch"() [^bb1] {tag = "annotate", foo = 1 : ui64} : () -> ()
+ // CHECK: "test.branch"()[^bb{{[0-9]+}}] {bar_state = true, foo = 2 : ui64, foo_state = 3 : i64, tag = "annotate"} : () -> ()
+ "test.branch"() [^bb1] {tag = "annotate", foo = 2 : ui64} : () -> ()
^bb1:
- // CHECK: "test.foo"() {bar_state = true, foo = 2 : ui64, foo_state = 7 : i64, tag = "annotate"} : () -> ()
- "test.foo"() {tag = "annotate", foo = 2 : ui64} : () -> ()
- // CHECK: "test.foo"() {bar_state = true, foo_state = 7 : i64, tag = "annotate"} : () -> ()
+ // CHECK: "test.foo"() {bar_state = true, foo = 12 : ui64, foo_state = 31 : i64, tag = "annotate"} : () -> ()
+ "test.foo"() {tag = "annotate", foo = 12 : ui64} : () -> ()
+ // CHECK: "test.foo"() {bar_state = true, foo_state = 31 : i64, tag = "annotate"} : () -> ()
"test.foo"() {tag = "annotate"} : () -> ()
return
^bb2:
- // CHECK: "test.branch"()[^bb{{[0-9]+}}] {bar_state = true, foo = 2 : ui64, foo_state = 5 : i64, tag = "annotate"} : () -> ()
- "test.branch"() [^bb1] {tag = "annotate", foo = 2 : ui64} : () -> ()
+ // CHECK: "test.branch"()[^bb{{[0-9]+}}] {bar_state = true, foo = 16 : ui64, foo_state = 17 : i64, tag = "annotate"} : () -> ()
+ "test.branch"() [^bb1] {tag = "annotate", foo = 16 : ui64} : () -> ()
}
// -----
diff --git a/mlir/test/Analysis/test-foo-analysis.mlir b/mlir/test/Analysis/test-foo-analysis.mlir
index edd153d433878..bb53f9fc9097e 100644
--- a/mlir/test/Analysis/test-foo-analysis.mlir
+++ b/mlir/test/Analysis/test-foo-analysis.mlir
@@ -26,8 +26,8 @@ func.func @test_two_join() -> () {
"test.foo"() {tag = "a"} : () -> ()
// CHECK: b -> 1
"test.foo"() {tag = "b", foo = 1 : ui64} : () -> ()
- // CHECK: c -> 0
- "test.foo"() {tag = "c", foo = 1 : ui64} : () -> ()
+ // CHECK: c -> 3
+ "test.foo"() {tag = "c", foo = 2 : ui64} : () -> ()
return
}
@@ -47,7 +47,7 @@ func.func @test_fork() -> () {
"test.branch"() [^bb2] {tag = "b", foo = 4 : ui64} : () -> ()
^bb2:
- // CHECK: end -> 6
+ // CHECK: end -> 7
"test.foo"() {tag = "end"} : () -> ()
return
@@ -61,7 +61,7 @@ func.func @test_simple_loop() -> () {
"test.branch"() [^bb0] {tag = "init", foo = 1 : ui64} : () -> ()
^bb0:
- // CHECK: a -> 1
+ // CHECK: a -> 3
"test.foo"() {tag = "a", foo = 3 : ui64} : () -> ()
"test.branch"() [^bb0, ^bb1] : () -> ()
@@ -79,17 +79,51 @@ func.func @test_double_loop() -> () {
"test.branch"() [^bb0] {tag = "init", foo = 2 : ui64} : () -> ()
^bb0:
- // CHECK: a -> 1
+ // CHECK: a -> 7
"test.foo"() {tag = "a", foo = 3 : ui64} : () -> ()
"test.branch"() [^bb0, ^bb1] : () -> ()
^bb1:
- // CHECK: b -> 4
+ // CHECK: b -> 7
"test.foo"() {tag = "b", foo = 5 : ui64} : () -> ()
"test.branch"() [^bb0, ^bb2] : () -> ()
^bb2:
- // CHECK: end -> 4
+ // CHECK: end -> 7
"test.foo"() {tag = "end"} : () -> ()
return
}
+
+// -----
+
+// The join block is enqueued once per predecessor while `initialize` walks the
+// rest of the blocks. The solver collapses those into a single work item, so
+// the join block is visited exactly once (in addition to the initialization
+// visit). Without worklist deduplication it would be visited once per
+// predecessor.
+
+// CHECK-LABEL: function: @wide_fan_in
+func.func @wide_fan_in() {
+ "test.branch"() [^bb1, ^bb2, ^bb3, ^bb4, ^bb5] : () -> ()
+
+^bb0:
+ // CHECK: join block visits -> 2
+ // CHECK: join -> 31
+ "test.foo"() {tag = "join"} : () -> ()
+ return
+
+^bb1:
+ "test.branch"() [^bb0] {foo = 1 : ui64} : () -> ()
+
+^bb2:
+ "test.branch"() [^bb0] {foo = 2 : ui64} : () -> ()
+
+^bb3:
+ "test.branch"() [^bb0] {foo = 4 : ui64} : () -> ()
+
+^bb4:
+ "test.branch"() [^bb0] {foo = 8 : ui64} : () -> ()
+
+^bb5:
+ "test.branch"() [^bb0] {foo = 16 : ui64} : () -> ()
+}
diff --git a/mlir/test/lib/Analysis/TestDataFlowFramework.cpp b/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
index 98fbcd8c9b12c..afe614b317c4d 100644
--- a/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
+++ b/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
@@ -10,6 +10,7 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
#include <optional>
using namespace mlir;
@@ -20,7 +21,7 @@ constexpr char kFooAttrName[] = "foo";
constexpr char kFooStateAttrName[] = "foo_state";
constexpr char kBarStateAttrName[] = "bar_state";
-/// This analysis state represents an integer that is XOR'd with other states.
+/// This analysis state represents an integer that is OR'd with other states.
class FooState : public AnalysisState {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FooState)
@@ -39,7 +40,7 @@ class FooState : public AnalysisState {
}
/// Join the state with another. If either is uninitialized, take the
- /// initialized value. Otherwise, XOR the integer values.
+ /// initialized value. Otherwise, OR the integer values.
ChangeResult join(const FooState &rhs) {
if (rhs.isUninitialized())
return ChangeResult::NoChange;
@@ -51,7 +52,7 @@ class FooState : public AnalysisState {
return ChangeResult::Change;
}
uint64_t before = *state;
- state = before ^ value;
+ state = before | value;
return before == *state ? ChangeResult::NoChange : ChangeResult::Change;
}
@@ -72,7 +73,7 @@ class FooState : public AnalysisState {
};
/// This analysis computes `FooState` across operations and control-flow edges.
-/// If an op specifies a `foo` integer attribute, the contained value is XOR'd
+/// If an op specifies a `foo` integer attribute, the contained value is OR'd
/// with the value before the operation.
class FooAnalysis : public DataFlowAnalysis {
public:
@@ -87,15 +88,19 @@ class FooAnalysis : public DataFlowAnalysis {
LogicalResult initialize(Operation *top) override;
LogicalResult visit(ProgramPoint *point) override;
+ unsigned lookupVisits(const Block *block) const;
+
private:
void visitBlock(Block *block);
void visitOperation(Operation *op);
+
+ DenseMap<const Block *, unsigned> blockVisits; // including initial visit
};
/// This analysis state stores whether all previously observed `FooState`
/// values at tagged program points along the CFG leading to the current point
-/// have been non-multiples of 4. Once the state becomes false at some point,
-/// all later points reachable from it also remain false.
+/// have not been positive multiples of 5. Once the state becomes false at some
+/// point, all later points reachable from it also remain false.
class BarState : public AnalysisState {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BarState)
@@ -138,10 +143,10 @@ class BarState : public AnalysisState {
/// This analysis is intended to be loaded after `FooAnalysis` has converged.
/// It records whether every observed `FooState` on or before a given tagged
-/// program point has been non-divisible by 4. Because the state only ever
-/// transitions from true to false, observing a transient divisible-by-4
-/// `FooState` before `FooAnalysis` converges can permanently poison the
-/// result.
+/// program point has not been a positive multiple of 5. Because the state only
+/// ever transitions from true to false, observing a transient positive
+/// divisible-by-5 `FooState` before `FooAnalysis` converges can permanently
+/// poison the result.
class BarAnalysis : public DataFlowAnalysis {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BarAnalysis)
@@ -210,7 +215,12 @@ LogicalResult FooAnalysis::visit(ProgramPoint *point) {
return success();
}
+unsigned FooAnalysis::lookupVisits(const Block *const block) const {
+ return blockVisits.lookup(block);
+}
+
void FooAnalysis::visitBlock(Block *block) {
+ ++blockVisits[block];
if (block->isEntryBlock()) {
// This is the initial state. Let the framework default-initialize it.
return;
@@ -305,7 +315,8 @@ void BarAnalysis::visitOperation(Operation *op) {
const FooState *fooState = getOrCreateFor<FooState>(point, point);
if (fooState->isUninitialized())
return;
- result |= state->join((fooState->getValue() & 0x3) != 0);
+ const auto value = fooState->getValue();
+ result |= state->join(value > 0 && value % 5 != 0);
}
propagateIfChanged(state, result);
}
@@ -313,7 +324,7 @@ void BarAnalysis::visitOperation(Operation *op) {
void TestFooAnalysisPass::runOnOperation() {
func::FuncOp func = getOperation();
DataFlowSolver solver;
- solver.load<FooAnalysis>();
+ auto &analysis = *solver.load<FooAnalysis>();
if (failed(solver.initializeAndRun(func)))
return signalPassFailure();
@@ -324,6 +335,9 @@ void TestFooAnalysisPass::runOnOperation() {
auto tag = op->getDiscardableAttrOfType<StringAttr>(kTagAttrName);
if (!tag)
return;
+ if (auto *const block = op->getBlock(); &block->front() == op)
+ os << tag.getValue() << " block visits -> "
+ << analysis.lookupVisits(block) << "\n";
const FooState *state =
solver.lookupState<FooState>(solver.getProgramPointAfter(op));
assert(state && !state->isUninitialized());
More information about the Mlir-commits
mailing list