From 5bedd675d741b177606dae9f13cc7c1531b6b954 Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles@arm.com>
Date: Fri, 23 Dec 2022 21:03:14 +0000
Subject: [PATCH] [mlir] Allow overriding
 AbstractDenseDataFlowAnalysis::visitOperation

AbstractDenseDataFlowAnalysis::visitOperation controls how the dataflow
analysis proceeds around control flow. In particular, conservative
assumptions are made about call operations which can prevent some
analysis from succeeding.

The motivating case for this change is https://reviews.llvm.org/D140415,
for which it is correct and necessary for the lattice to be preserved
after call operations.

Some renaming was necessary to avoid confusion with
DenseDataFlowAnalysis::visitOperation.
AbstractDenseDataFlowAnalysis::visitRegionBranchOperation and
DenseDataFlowAnalysis::visitOperationImpl are also made protected
to allow implementation of AbstractDenseDataFlowAnalysis::visitOperation,
although I did not need these to be virtual.

Differential Revision: https://reviews.llvm.org/D140879
---
 mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h | 13 ++++++-------
 mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp        |  6 +++---
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
index 32733a1..9b71ae1 100644
--- a/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h
@@ -93,16 +93,11 @@ protected:
     propagateIfChanged(lhs, lhs->join(rhs));
   }
 
-private:
   /// Visit an operation. If this is a call operation or region control-flow
   /// operation, then the state after the execution of the operation is set by
   /// control-flow or the callgraph. Otherwise, this function invokes the
   /// operation transfer function.
-  void visitOperation(Operation *op);
-
-  /// Visit a block. The state at the start of the block is propagated from
-  /// control-flow predecessors or callsites
-  void visitBlock(Block *block);
+  virtual void processOperation(Operation *op);
 
   /// Visit a program point within a region branch operation with predecessors
   /// in it. This can either be an entry block of one of the regions of the
@@ -110,6 +105,11 @@ private:
   void visitRegionBranchOperation(ProgramPoint point,
                                   RegionBranchOpInterface branch,
                                   AbstractDenseLattice *after);
+
+private:
+  /// Visit a block. The state at the start of the block is propagated from
+  /// control-flow predecessors or callsites
+  void visitBlock(Block *block);
 };
 
 //===----------------------------------------------------------------------===//
@@ -148,7 +148,6 @@ protected:
     setToEntryState(static_cast<LatticeT *>(lattice));
   }
 
-private:
   /// Type-erased wrappers that convert the abstract dense lattice to a derived
   /// lattice and invoke the virtual hooks operating on the derived lattice.
   void visitOperationImpl(Operation *op, const AbstractDenseLattice &before,
diff --git a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
index c7bc02c..6450891 100644
--- a/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp
@@ -20,7 +20,7 @@ using namespace mlir::dataflow;
 
 LogicalResult AbstractDenseDataFlowAnalysis::initialize(Operation *top) {
   // Visit every operation and block.
-  visitOperation(top);
+  processOperation(top);
   for (Region &region : top->getRegions()) {
     for (Block &block : region) {
       visitBlock(&block);
@@ -34,7 +34,7 @@ LogicalResult AbstractDenseDataFlowAnalysis::initialize(Operation *top) {
 
 LogicalResult AbstractDenseDataFlowAnalysis::visit(ProgramPoint point) {
   if (auto *op = point.dyn_cast<Operation *>())
-    visitOperation(op);
+    processOperation(op);
   else if (auto *block = point.dyn_cast<Block *>())
     visitBlock(block);
   else
@@ -42,7 +42,7 @@ LogicalResult AbstractDenseDataFlowAnalysis::visit(ProgramPoint point) {
   return success();
 }
 
-void AbstractDenseDataFlowAnalysis::visitOperation(Operation *op) {
+void AbstractDenseDataFlowAnalysis::processOperation(Operation *op) {
   // If the containing block is not executable, bail out.
   if (!getOrCreateFor<Executable>(op, op->getBlock())->isLive())
     return;
-- 
2.7.4