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

[Hongren Zheng]

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

Detailed writeup is in https://github.com/google/heir/issues/1153. See also https://github.com/llvm/llvm-project/pull/120881. In short, `propagateIfChanged` is used outside of the `DataFlowAnalysis` scope, because it is public, but it does not propagate as expected as the `DataFlowSolver` has stopped running.

To solve such misuse, `propagateIfChanged` should be made protected/private.

For downstream users affected by this, to correctly propagate the change, the Analysis should be re-run (check #120881) instead of just a `propagateIfChanged`

The change to `IntegerRangeAnalysis` is just a expansion of the `solver->propagateIfChanged`. The `Lattice` has already been updated by the `join`. Propagation is done by `onUpdate`.

Cc @Mogball for review

>From 66f5e03d285fc71317af7f73f374f528f12e86ba 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] disallow outside use of propagateIfChanged
 for DataFlowSolver

---
 mlir/include/mlir/Analysis/DataFlowFramework.h  | 10 +++++++---
 .../Analysis/DataFlow/IntegerRangeAnalysis.cpp  | 17 ++++++++++++-----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index 969664dc7a4fe3..1dcb32f762c003 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -394,13 +394,17 @@ class DataFlowSolver {
   template <typename StateT, typename AnchorT>
   StateT *getOrCreateState(AnchorT anchor);
 
+  /// Get the configuration of the solver.
+  const DataFlowConfig &getConfig() const { return config; }
+
+protected:
   /// 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 by DataFlowAnalysis instances.
+  /// When used outside of DataFlowAnalysis, the solver won't process
+  /// the work items
   void propagateIfChanged(AnalysisState *state, ChangeResult changed);
 
-  /// Get the configuration of the solver.
-  const DataFlowConfig &getConfig() const { return config; }
-
 private:
   /// Configuration of the dataflow solver.
   DataFlowConfig config;
diff --git a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
index 9e9411e5ede12c..60ae7d00c0bbdb 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp
@@ -45,9 +45,13 @@ void IntegerValueRangeLattice::onUpdate(DataFlowSolver *solver) const {
   std::optional<APInt> constant = getValue().getValue().getConstantValue();
   auto value = cast<Value>(anchor);
   auto *cv = solver->getOrCreateState<Lattice<ConstantValue>>(value);
-  if (!constant)
-    return solver->propagateIfChanged(
-        cv, cv->join(ConstantValue::getUnknownConstant()));
+  if (!constant) {
+    auto changed = cv->join(ConstantValue::getUnknownConstant());
+    if (changed == ChangeResult::Change) {
+      cv->onUpdate(solver);
+    }
+    return;
+  }
 
   Dialect *dialect;
   if (auto *parent = value.getDefiningOp())
@@ -56,8 +60,11 @@ void IntegerValueRangeLattice::onUpdate(DataFlowSolver *solver) const {
     dialect = value.getParentBlock()->getParentOp()->getDialect();
 
   Type type = getElementTypeOrSelf(value);
-  solver->propagateIfChanged(
-      cv, cv->join(ConstantValue(IntegerAttr::get(type, *constant), dialect)));
+  auto changed =
+      cv->join(ConstantValue(IntegerAttr::get(type, *constant), dialect));
+  if (changed == ChangeResult::Change) {
+    cv->onUpdate(solver);
+  }
 }
 
 LogicalResult IntegerRangeAnalysis::visitOperation(