[Mlir-commits] [mlir] [mlir] `int-range-optmizations`: Fix referencing of deleted ops (PR #91807)

[Felix Schneider]

https://github.com/ubfx created https://github.com/llvm/llvm-project/pull/91807

The pass runs a `DataFlowSolver` and collects state information on the input IR. Then, the rewrite driver and folding is applied. During pattern application and folding it can happen that an Op from the input IR is deleted and a new Op is created at the same address. When the newly created Ops is looked up in the `DataFlowSolver` state memory, the state of the original Op is returned.

This patch adds a method to `DataFlowSolver` which removes all state related to a `ProgramPoint`. It also adds a listener to the Pass which clears the state information of deleted Ops from the `DataFlowSolver`.

This doesn't seem like the prettiest solution, maybe someone can advise on how to fix this in a nicer way? This pattern of using a `DataFlowSolver` throughout a pass with multiple pattern applications/foldings happens in other Passes, too. So it seems worth to find a proper solution for this. The bugs and crashes resulting from this are elusive and annoying to debug because they are dependent on how things are allocated.

Fix https://github.com/llvm/llvm-project/issues/81228

>From fb6ed097e646ad9710619ea5568e2f90c1a7e49a Mon Sep 17 00:00:00 2001
From: Felix Schneider <fx.schn at gmail.com>
Date: Fri, 10 May 2024 22:41:16 +0200
Subject: [PATCH] [mlir] `int-range-optmizations`: Fix referencing of deleted
 ops

The pass runs a `DataFlowSolver` and collects state information on the
input IR. Then, the rewrite driver and folding is applied. During pattern
application and folding it can happen that an Op from the input IR is
deleted and a new Op is created at the same address. When the newly created
Ops is looked up in the `DataFlowSolver` state memory, the state of the
original Op is returned.

This patch adds a method to `DataFlowSolver` which removes all state
related to a `ProgramPoint`. It also adds a listener to the Pass which
clears the state information of deleted Ops from the `DataFlowSolver`.

Fix https://github.com/llvm/llvm-project/issues/81228
---
 .../include/mlir/Analysis/DataFlowFramework.h | 11 ++++++++
 .../Transforms/IntRangeOptimizations.cpp      | 25 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index c76cfac07fc77..2580ec28b5190 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -242,6 +242,17 @@ class DataFlowSolver {
     return static_cast<const StateT *>(it->second.get());
   }
 
+  /// Erase any analysis state associated with the given program point.
+  template <typename PointT>
+  void eraseState(PointT point) {
+    ProgramPoint pp(point);
+
+    for (auto it = analysisStates.begin(); it != analysisStates.end(); ++it) {
+      if (it->first.first == pp)
+        analysisStates.erase(it);
+    }
+  }
+
   /// Get a uniqued program point instance. If one is not present, it is
   /// created with the provided arguments.
   template <typename PointT, typename... Args>
diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 92cad7cd1ef26..2473169962b95 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -102,6 +102,24 @@ static FailureOr<bool> handleUge(ConstantIntRanges lhs, ConstantIntRanges rhs) {
 }
 
 namespace {
+/// This class listens on IR transformations performed during a pass relying on
+/// information from a `DataflowSolver`. It erases state associated with the
+/// erased operation and its results from the `DataFlowSolver` so that Patterns
+/// do not accidentally query old state information for newly created Ops.
+class DataFlowListener : public RewriterBase::Listener {
+public:
+  DataFlowListener(DataFlowSolver &s) : s(s) {}
+
+protected:
+  void notifyOperationErased(Operation *op) override {
+    s.eraseState(op);
+    for (Value res : op->getResults())
+      s.eraseState(res);
+  }
+
+  DataFlowSolver &s;
+};
+
 struct ConvertCmpOp : public OpRewritePattern<arith::CmpIOp> {
 
   ConvertCmpOp(MLIRContext *context, DataFlowSolver &s)
@@ -167,10 +185,15 @@ struct IntRangeOptimizationsPass
     if (failed(solver.initializeAndRun(op)))
       return signalPassFailure();
 
+    DataFlowListener listener(solver);
+
     RewritePatternSet patterns(ctx);
     populateIntRangeOptimizationsPatterns(patterns, solver);
 
-    if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
+    GreedyRewriteConfig config;
+    config.listener = &listener;
+
+    if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns), config)))
       signalPassFailure();
   }
 };