[Mlir-commits] [mlir] [mlir][CSE] Add cleanup PostDominanceInfo logic to cse (PR #192628)

[lonely eagle]

https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/192628

>From 3cb609ab579e4d5555f913e4221549d1f85df3b7 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Wed, 29 Apr 2026 06:10:26 +0000
Subject: [PATCH] rebase main.

---
 mlir/include/mlir/Transforms/CSE.h |  7 +++++--
 mlir/lib/Transforms/CSE.cpp        |  9 ++++++++-
 mlir/lib/Transforms/Utils/CSE.cpp  | 21 ++++++++++++++-------
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Transforms/CSE.h b/mlir/include/mlir/Transforms/CSE.h
index 4a87d585e0eb9..9190911a9a743 100644
--- a/mlir/include/mlir/Transforms/CSE.h
+++ b/mlir/include/mlir/Transforms/CSE.h
@@ -18,6 +18,7 @@
 namespace mlir {
 
 class DominanceInfo;
+class PostDominanceInfo;
 class Operation;
 class Region;
 class RewriterBase;
@@ -32,7 +33,8 @@ void eliminateCommonSubExpressions(RewriterBase &rewriter,
                                    DominanceInfo &domInfo, Operation *op,
                                    bool *changed = nullptr,
                                    int64_t *numCSE = nullptr,
-                                   int64_t *numDCE = nullptr);
+                                   int64_t *numDCE = nullptr,
+                                   PostDominanceInfo *postDomInfo = nullptr);
 
 /// Eliminate common subexpressions within the given region.
 ///
@@ -41,7 +43,8 @@ void eliminateCommonSubExpressions(RewriterBase &rewriter,
 /// DCE counts are needed.
 void eliminateCommonSubExpressions(RewriterBase &rewriter,
                                    DominanceInfo &domInfo, Region &region,
-                                   bool *changed = nullptr);
+                                   bool *changed = nullptr,
+                                   PostDominanceInfo *postDomInfo = nullptr);
 
 } // namespace mlir
 
diff --git a/mlir/lib/Transforms/CSE.cpp b/mlir/lib/Transforms/CSE.cpp
index f7afa03e2f02b..508a7eca1792f 100644
--- a/mlir/lib/Transforms/CSE.cpp
+++ b/mlir/lib/Transforms/CSE.cpp
@@ -35,13 +35,20 @@ struct CSE : public impl::CSEPassBase<CSE> {
 void CSE::runOnOperation() {
   IRRewriter rewriter(&getContext());
   auto &domInfo = getAnalysis<DominanceInfo>();
+  // The CSE implementation does not rely on PostDominanceInfo. However,
+  // since we mark it as preserved at the end, if a cached PostDominanceInfo
+  // exists, we need to update it within CSE.
+  PostDominanceInfo *postDomInfo = nullptr;
+  if (auto dominate = getCachedAnalysis<PostDominanceInfo>())
+    postDomInfo = &dominate.value().get();
+
   bool changed = false;
   // `numCSE` / `numDCE` are `llvm::Statistic` objects, not raw `int64_t`, so
   // the public API's out-parameters cannot point at them directly.
   int64_t cseCount = 0;
   int64_t dceCount = 0;
   eliminateCommonSubExpressions(rewriter, domInfo, getOperation(), &changed,
-                                &cseCount, &dceCount);
+                                &cseCount, &dceCount, postDomInfo);
 
   numCSE = cseCount;
   numDCE = dceCount;
diff --git a/mlir/lib/Transforms/Utils/CSE.cpp b/mlir/lib/Transforms/Utils/CSE.cpp
index 90444e6201891..45f0564070f52 100644
--- a/mlir/lib/Transforms/Utils/CSE.cpp
+++ b/mlir/lib/Transforms/Utils/CSE.cpp
@@ -52,8 +52,9 @@ namespace {
 /// Simple common sub-expression elimination.
 class CSEDriver {
 public:
-  CSEDriver(RewriterBase &rewriter, DominanceInfo *domInfo)
-      : rewriter(rewriter), domInfo(domInfo) {}
+  CSEDriver(RewriterBase &rewriter, DominanceInfo *domInfo,
+            PostDominanceInfo *postDomInfo)
+      : rewriter(rewriter), domInfo(domInfo), postDomInfo(postDomInfo) {}
 
   /// Simplify all operations within the given op.
   void simplify(Operation *op, bool *changed = nullptr);
@@ -118,6 +119,7 @@ class CSEDriver {
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
   DominanceInfo *domInfo = nullptr;
+  PostDominanceInfo *postDomInfo = nullptr;
   MemEffectsCache memEffectsCache;
 
   // Various statistics.
@@ -389,8 +391,11 @@ void CSEDriver::eraseDeadOps(bool *changed) {
   // Erase any operations that were marked as dead during simplification, and
   // remove their associated dominator trees.
   for (auto *op : opsToErase) {
-    for (Region &region : op->getRegions())
+    for (Region &region : op->getRegions()) {
       domInfo->invalidate(&region);
+      if (postDomInfo != nullptr)
+        postDomInfo->invalidate(&region);
+    }
     rewriter.eraseOp(op);
   }
   if (changed)
@@ -418,8 +423,9 @@ void CSEDriver::simplify(Region &region, bool *changed) {
 void mlir::eliminateCommonSubExpressions(RewriterBase &rewriter,
                                          DominanceInfo &domInfo, Operation *op,
                                          bool *changed, int64_t *numCSE,
-                                         int64_t *numDCE) {
-  CSEDriver driver(rewriter, &domInfo);
+                                         int64_t *numDCE,
+                                         PostDominanceInfo *postDomInfo) {
+  CSEDriver driver(rewriter, &domInfo, postDomInfo);
   driver.simplify(op, changed);
   if (numCSE)
     *numCSE = driver.getNumCSE();
@@ -429,7 +435,8 @@ void mlir::eliminateCommonSubExpressions(RewriterBase &rewriter,
 
 void mlir::eliminateCommonSubExpressions(RewriterBase &rewriter,
                                          DominanceInfo &domInfo, Region &region,
-                                         bool *changed) {
-  CSEDriver driver(rewriter, &domInfo);
+                                         bool *changed,
+                                         PostDominanceInfo *postDomInfo) {
+  CSEDriver driver(rewriter, &domInfo, postDomInfo);
   driver.simplify(region, changed);
 }