[Mlir-commits] [mlir] [mlir] [dataflow] further optimize dataflow compile time (PR #149804)

[donald chen]

https://github.com/cxy-1993 created https://github.com/llvm/llvm-project/pull/149804

Optimize dataflow compilation time by skipping initialization of irrelevant operations

>From dee817a671e2c1601e3deafdc1a9457aa0093835 Mon Sep 17 00:00:00 2001
From: donald chen <chenxunyu1993 at gmail.com>
Date: Mon, 21 Jul 2025 12:26:24 +0000
Subject: [PATCH] [mlir] [dataflow] further optimize dataflow compile time

Optimize dataflow compilation time by skipping initialization
of irrelevant operations
---
 mlir/include/mlir/Analysis/DataFlowFramework.h           | 9 +++++++++
 mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp             | 8 ++++++++
 .../DataFlow/TestDenseBackwardDataFlowAnalysis.cpp       | 1 +
 .../DataFlow/TestDenseForwardDataFlowAnalysis.cpp        | 1 +
 4 files changed, 19 insertions(+)

diff --git a/mlir/include/mlir/Analysis/DataFlowFramework.h b/mlir/include/mlir/Analysis/DataFlowFramework.h
index 49862927caff2..67d593a7bfad4 100644
--- a/mlir/include/mlir/Analysis/DataFlowFramework.h
+++ b/mlir/include/mlir/Analysis/DataFlowFramework.h
@@ -658,6 +658,12 @@ class DataFlowAnalysis {
     return solver.getOrCreateState<StateT>(anchor);
   }
 
+  /// Add irrelevant program point.
+  template <typename PointT>
+  void addIrrelevantPoint(PointT point) {
+    irrelevantPoints.insert(ProgramPoint(point));
+  }
+
   /// Get a read-only analysis state for the given point and create a dependency
   /// on `dependent`. If the return state is updated elsewhere, this analysis is
   /// re-invoked on the dependent.
@@ -695,6 +701,9 @@ class DataFlowAnalysis {
   StringRef debugName;
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
+  /// Program points shouldn't analyzed by this analysis.
+  DenseSet<ProgramPoint> irrelevantPoints;
+
 private:
   /// The parent data-flow solver.
   DataFlowSolver &solver;
diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index d05374f667a51..0d0d841b0bff8 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -104,6 +104,10 @@ void AbstractDenseForwardDataFlowAnalysis::visitCallOperation(
 
 LogicalResult
 AbstractDenseForwardDataFlowAnalysis::processOperation(Operation *op) {
+  // Skip irrelavant program points.
+  if (irrelevantPoints.contains(ProgramPoint(op)))
+    return;
+
   ProgramPoint *point = getProgramPointAfter(op);
   // If the containing block is not executable, bail out.
   if (op->getBlock() != nullptr &&
@@ -333,6 +337,10 @@ void AbstractDenseBackwardDataFlowAnalysis::visitCallOperation(
 
 LogicalResult
 AbstractDenseBackwardDataFlowAnalysis::processOperation(Operation *op) {
+  // Skip irrelavant program points.
+  if (irrelevantPoints.contains(ProgramPoint(op)))
+    return;
+
   ProgramPoint *point = getProgramPointBefore(op);
   // If the containing block is not executable, bail out.
   if (op->getBlock() != nullptr &&
diff --git a/mlir/test/lib/Analysis/DataFlow/TestDenseBackwardDataFlowAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestDenseBackwardDataFlowAnalysis.cpp
index d57b41c41de64..f73430fb78c58 100644
--- a/mlir/test/lib/Analysis/DataFlow/TestDenseBackwardDataFlowAnalysis.cpp
+++ b/mlir/test/lib/Analysis/DataFlow/TestDenseBackwardDataFlowAnalysis.cpp
@@ -147,6 +147,7 @@ LogicalResult NextAccessAnalysis::visitOperation(Operation *op,
 
 void NextAccessAnalysis::buildOperationEquivalentLatticeAnchor(Operation *op) {
   if (isMemoryEffectFree(op)) {
+    addIrrelevantPoint(op);
     unionLatticeAnchors<NextAccess>(getProgramPointBefore(op),
                                     getProgramPointAfter(op));
   }
diff --git a/mlir/test/lib/Analysis/DataFlow/TestDenseForwardDataFlowAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestDenseForwardDataFlowAnalysis.cpp
index a88ed7f8dea8b..ea5614c24a6bf 100644
--- a/mlir/test/lib/Analysis/DataFlow/TestDenseForwardDataFlowAnalysis.cpp
+++ b/mlir/test/lib/Analysis/DataFlow/TestDenseForwardDataFlowAnalysis.cpp
@@ -154,6 +154,7 @@ LogicalResult LastModifiedAnalysis::visitOperation(
 void LastModifiedAnalysis::buildOperationEquivalentLatticeAnchor(
     Operation *op) {
   if (isMemoryEffectFree(op)) {
+    addIrrelevantPoint(op);
     unionLatticeAnchors<LastModification>(getProgramPointBefore(op),
                                           getProgramPointAfter(op));
   }