[Mlir-commits] [mlir] [mlir][dataflow] disallow outside use of propagateIfChanged for DataFlowSolver (PR #120885)

[Hongren Zheng]

https://github.com/ZenithalHourlyRate updated https://github.com/llvm/llvm-project/pull/120885

>From 2dc277b04fd8922caf183207bfb5d9ed1128f55e Mon Sep 17 00:00:00 2001
From: Zenithal <i at zenithal.me>
Date: Sun, 22 Dec 2024 09:42:48 +0000
Subject: [PATCH] [mlir][dataflow] Only allow propagateIfChanged when
 DataFlowSolver is running

---
 mlir/include/mlir/Analysis/DataFlowFramework.h |  5 +++++
 mlir/lib/Analysis/DataFlowFramework.cpp        | 14 ++++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index dfd358e7017a4e..ee776e4b4cbef0 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -403,6 +403,8 @@ class DataFlowSolver {
 
   /// Propagate an update to an analysis state if it changed by pushing
   /// dependent work items to the back of the queue.
+  /// This should only be used when DataFlowSolver is running.
+  /// Otherwise, the solver won't process the work items.
   void propagateIfChanged(AnalysisState *state, ChangeResult changed);
 
   /// Get the configuration of the solver.
@@ -412,6 +414,9 @@ class DataFlowSolver {
   /// Configuration of the dataflow solver.
   DataFlowConfig config;
 
+  // The solver is working on the worklist.
+  bool isRunning = false;
+
   /// The solver's work queue. Work items can be inserted to the front of the
   /// queue to be processed greedily, speeding up computations that otherwise
   /// quickly degenerate to quadratic due to propagation of state updates.
diff --git a/mlir/lib/Analysis/DataFlowFramework.cpp b/mlir/lib/Analysis/DataFlowFramework.cpp
index 7e83668c067652..62287c2703c867 100644
--- a/mlir/lib/Analysis/DataFlowFramework.cpp
+++ b/mlir/lib/Analysis/DataFlowFramework.cpp
@@ -104,12 +104,16 @@ Location LatticeAnchor::getLoc() const {
 //===----------------------------------------------------------------------===//
 
 LogicalResult DataFlowSolver::initializeAndRun(Operation *top) {
+  // Enable enqueue to the worklist.
+  isRunning = true;
   // Initialize the analyses.
   for (DataFlowAnalysis &analysis : llvm::make_pointee_range(childAnalyses)) {
     DATAFLOW_DEBUG(llvm::dbgs()
                    << "Priming analysis: " << analysis.debugName << "\n");
-    if (failed(analysis.initialize(top)))
+    if (failed(analysis.initialize(top))) {
+      isRunning = false;
       return failure();
+    }
   }
 
   // Run the analysis until fixpoint.
@@ -121,19 +125,25 @@ LogicalResult DataFlowSolver::initializeAndRun(Operation *top) {
 
       DATAFLOW_DEBUG(llvm::dbgs() << "Invoking '" << analysis->debugName
                                   << "' on: " << point << "\n");
-      if (failed(analysis->visit(point)))
+      if (failed(analysis->visit(point))) {
+        isRunning = false;
         return failure();
+      }
     }
 
     // Iterate until all states are in some initialized state and the worklist
     // is exhausted.
   } while (!worklist.empty());
 
+  // Prevent further updates to the worklist
+  isRunning = false;
   return success();
 }
 
 void DataFlowSolver::propagateIfChanged(AnalysisState *state,
                                         ChangeResult changed) {
+  assert(isRunning &&
+         "DataFlowSolver is not running, should not use propagateIfChanged");
   if (changed == ChangeResult::Change) {
     DATAFLOW_DEBUG(llvm::dbgs() << "Propagating update to " << state->debugName
                                 << " of " << state->anchor << "\n"