[mlir][DataFlow] Add visitBlockTransfer hook to dense analyses

[fabianmcg]

Add a customizable visitBlockTransfer method to dense forward and backward dataflow analyses, allowing subclasses to customize lattice propagation behavior along control flow edges between blocks. Default implementation preserves existing join/meet semantics.

This change mirrors the exiting structure of both dense dataflow classes, where RegionBranchOpInterface and callables are allowed to be customized by subclasses.

The use case motivating this change is dense liveness analysis. Currently, without the customization hook the block transfer function produces incorrect results. The issue is the current logic doesn't remove the successor block arguments from the live set, as it only meets the successor state with the predecessor state (ie. set union).
With this change is now possible to compute the correct result by specifying the correct logic in visitBlockTransfer.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Fabian Mora (fabianmcg)

Changes

Add a customizable visitBlockTransfer method to dense forward and backward dataflow analyses, allowing subclasses to customize lattice propagation behavior along control flow edges between blocks. Default implementation preserves existing join/meet semantics.

This change mirrors the exiting structure of both dense dataflow classes, where RegionBranchOpInterface and callables are allowed to be customized by subclasses.

---

Full diff: https://github.com/llvm/llvm-project/pull/166263.diff

2 Files Affected:

• (modified) mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h (+70)
• (modified) mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp (+6-3)

diff --git a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
index 3c87c453a4cf0..d8c2490810ecf 100644
--- a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
@@ -127,6 +127,19 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// them into the same equivalent class.
   virtual void buildOperationEquivalentLatticeAnchor(Operation *op) {}
 
+  /// Visit a block and propagate the dense lattice forward along the control
+  /// flow edge from predecessor to block. `point` corresponds to the program
+  /// point before `block`. The default implementation merges in the state from
+  /// the predecessor's terminator.
+  virtual void visitBlockTransfer(Block *block, ProgramPoint *point,
+                                  Block *predecessor,
+                                  const AbstractDenseLattice &before,
+                                  AbstractDenseLattice *after) {
+    // Merge in the state from the predecessor's terminator.
+    join(after, *getLatticeFor(
+                    point, getProgramPointAfter(predecessor->getTerminator())));
+  }
+
   /// Propagate the dense lattice forward along the control flow edge from
   /// `regionFrom` to `regionTo` regions of the `branch` operation. `nullopt`
   /// values correspond to control flow branches originating at or targeting the
@@ -259,6 +272,22 @@ class DenseForwardDataFlowAnalysis
         branch, regionFrom, regionTo, before, after);
   }
 
+  /// Hook for customizing the behavior of lattice propagation along the control
+  /// flow edges between blocks. The control flows from `predecessor` to
+  /// `block`. The lattice is propagated forward along this edge. The lattices
+  /// are as follows:
+  ///   - `before` is the lattice at the end of the predecessor block;
+  ///   - `after` is the lattice at the beginning of the block.
+  /// By default, the `after` state is simply joined with the `before` state.
+  /// Concrete analyses can override this behavior or delegate to the parent
+  /// call for the default behavior.
+  virtual void visitBlockTransfer(Block *block, ProgramPoint *point,
+                                  Block *predecessor, const LatticeT &before,
+                                  LatticeT *after) {
+    AbstractDenseForwardDataFlowAnalysis::visitBlockTransfer(
+        block, point, predecessor, before, after);
+  }
+
 protected:
   /// Get the dense lattice on this lattice anchor.
   LatticeT *getLattice(LatticeAnchor anchor) override {
@@ -306,6 +335,13 @@ class DenseForwardDataFlowAnalysis
                                          static_cast<const LatticeT &>(before),
                                          static_cast<LatticeT *>(after));
   }
+  void visitBlockTransfer(Block *block, ProgramPoint *point, Block *predecessor,
+                          const AbstractDenseLattice &before,
+                          AbstractDenseLattice *after) final {
+    visitBlockTransfer(block, point, predecessor,
+                       static_cast<const LatticeT &>(before),
+                       static_cast<LatticeT *>(after));
+  }
 };
 
 //===----------------------------------------------------------------------===//
@@ -388,6 +424,17 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// them into the same equivalent class.
   virtual void buildOperationEquivalentLatticeAnchor(Operation *op) {}
 
+  /// Visit a block and propagate the dense lattice backward along the control
+  /// flow edge from successor to block. `point` corresponds to the program
+  /// point after `block`. The default implementation merges in the state from
+  /// the successor's first operation or the block itself when empty.
+  virtual void visitBlockTransfer(Block *block, ProgramPoint *point,
+                                  Block *successor,
+                                  const AbstractDenseLattice &after,
+                                  AbstractDenseLattice *before) {
+    meet(before, *getLatticeFor(point, getProgramPointBefore(successor)));
+  }
+
   /// Propagate the dense lattice backwards along the control flow edge from
   /// `regionFrom` to `regionTo` regions of the `branch` operation. `nullopt`
   /// values correspond to control flow branches originating at or targeting the
@@ -531,6 +578,22 @@ class DenseBackwardDataFlowAnalysis
         branch, regionFrom, regionTo, after, before);
   }
 
+  /// Hook for customizing the behavior of lattice propagation along the control
+  /// flow edges between blocks. The control flows from `successor` to
+  /// `block`. The lattice is propagated back along this edge. The lattices
+  /// are as follows:
+  ///   - `after` is the lattice at the beginning of the successor block;
+  ///   - `before` is the lattice at the end of the block.
+  /// By default, the `before` state is simply met with the `after` state.
+  /// Concrete analyses can override this behavior or delegate to the parent
+  /// call for the default behavior.
+  virtual void visitBlockTransfer(Block *block, ProgramPoint *point,
+                                  Block *successor, const LatticeT &after,
+                                  LatticeT *before) {
+    AbstractDenseBackwardDataFlowAnalysis::visitBlockTransfer(
+        block, point, successor, after, before);
+  }
+
 protected:
   /// Get the dense lattice at the given lattice anchor.
   LatticeT *getLattice(LatticeAnchor anchor) override {
@@ -577,6 +640,13 @@ class DenseBackwardDataFlowAnalysis
                                          static_cast<const LatticeT &>(after),
                                          static_cast<LatticeT *>(before));
   }
+  void visitBlockTransfer(Block *block, ProgramPoint *point, Block *successor,
+                          const AbstractDenseLattice &after,
+                          AbstractDenseLattice *before) final {
+    visitBlockTransfer(block, point, successor,
+                       static_cast<const LatticeT &>(after),
+                       static_cast<LatticeT *>(before));
+  }
 };
 
 } // end namespace dataflow
diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index 0682e5f26785a..22bc0b32a9bd1 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -266,9 +266,10 @@ void AbstractDenseForwardDataFlowAnalysis::visitBlock(Block *block) {
     }
 
     LDBG() << "    Joining state from predecessor " << predecessor;
+    const AbstractDenseLattice &before = *getLatticeFor(
+        point, getProgramPointAfter(predecessor->getTerminator()));
     // Merge in the state from the predecessor's terminator.
-    join(after, *getLatticeFor(
-                    point, getProgramPointAfter(predecessor->getTerminator())));
+    visitBlockTransfer(block, point, predecessor, before, after);
   }
 }
 
@@ -614,7 +615,9 @@ void AbstractDenseBackwardDataFlowAnalysis::visitBlock(Block *block) {
     LDBG() << "    Meeting state from successor " << successor;
     // Merge in the state from the successor: either the first operation, or the
     // block itself when empty.
-    meet(before, *getLatticeFor(point, getProgramPointBefore(successor)));
+    visitBlockTransfer(block, point, successor,
+                       *getLatticeFor(point, getProgramPointBefore(successor)),
+                       before);
   }
 }
 

[Copilot]

Pull Request Overview

This PR introduces new customization hooks for dense data flow analysis by adding visitBlockTransfer methods that allow subclasses to customize lattice propagation behavior along control flow edges between blocks. The changes refactor the existing inline lattice joining/meeting logic into overridable virtual methods.

• Added visitBlockTransfer virtual methods to both forward and backward dense analysis classes
• Refactored call sites in DenseAnalysis.cpp to invoke the new transfer functions
• Provided default implementations that preserve existing behavior

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h	Added new visitBlockTransfer virtual method declarations and implementations for both forward and backward analysis classes, along with documentation
mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp	Refactored visitBlock methods to call the new visitBlockTransfer hooks instead of directly joining/meeting lattices

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ftynse]

Looks good in principle though I have two remarks:

1. This is not an NFC as it specifically adds a new functionality.
2. It would be better to have an example for when such a customization is necessary.

[fabianmcg]

It would be better to have an example for when such a customization is necessary.

I do have one example, using dense backward dataflow for liveness analysis.
Without the customization point, the analysis is incorrect because the block transfer function is not removing the successor block arguments from the set.

The reason I didn't include that analysis in this PR is because there's already 2 versions of liveness in upstream. However, none of them suited my use case.