[mlir] [dataflow] unify semantics of program point (llvm#110344)

The concept of a 'program point' in the original data flow framework is
ambiguous. It can refer to either an operation or a block itself. This
representation has different interpretations in forward and backward
data-flow analysis. In forward data-flow analysis, the program point of
an operation represents the state after the operation, while in backward
data flow analysis, it represents the state before the operation. When
using forward or backward data-flow analysis, it is crucial to carefully
handle this distinction to ensure correctness.

This patch refactors the definition of program point, unifying the
interpretation of program points in both forward and backward data-flow
analysis.

How to integrate this patch?

For dense forward data-flow analysis and other analysis (except dense
backward data-flow analysis), the program point corresponding to the
original operation can be obtained by `getProgramPointAfter(op)`, and
the program point corresponding to the original block can be obtained by
`getProgramPointBefore(block)`.

For dense backward data-flow analysis, the program point corresponding
to the original operation can be obtained by
`getProgramPointBefore(op)`, and the program point corresponding to the
original block can be obtained by `getProgramPointAfter(block)`.

NOTE: If you need to get the lattice of other data-flow analyses in
dense backward data-flow analysis, you should still use the dense
forward data-flow approach. For example, to get the Executable state of
a block in dense backward data-flow analysis and add the dependency of
the current operation, you should write:

``getOrCreateFor<Executable>(getProgramPointBefore(op),
getProgramPointBefore(block))``

In case above, we use getProgramPointBefore(op) because the analysis we
rely on is dense backward data-flow, and we use
getProgramPointBefore(block) because the lattice we query is the result
of a non-dense backward data flow computation.

related dsscussion:
https://discourse.llvm.org/t/rfc-unify-the-semantics-of-program-points/80671/8
corresponding PSA:
https://discourse.llvm.org/t/psa-program-point-semantics-change/81479

Author: cxy-1993

flang/lib/Optimizer/Transforms/StackArrays.cpp

@@ -375,7 +375,7 @@ mlir::LogicalResult AllocationAnalysis::visitOperation(
     }
   } else if (mlir::isa<fir::ResultOp>(op)) {
     mlir::Operation *parent = op->getParentOp();
-    LatticePoint *parentLattice = getLattice(parent);
+    LatticePoint *parentLattice = getLattice(getProgramPointAfter(parent));
     assert(parentLattice);
     mlir::ChangeResult parentChanged = parentLattice->join(*after);
     propagateIfChanged(parentLattice, parentChanged);
@@ -396,28 +396,29 @@ void AllocationAnalysis::setToEntryState(LatticePoint *lattice) {
 /// Mostly a copy of AbstractDenseLattice::processOperation - the difference
 /// being that call operations are passed through to the transfer function
 mlir::LogicalResult AllocationAnalysis::processOperation(mlir::Operation *op) {
+  mlir::ProgramPoint *point = getProgramPointAfter(op);
   // If the containing block is not executable, bail out.
-  if (!getOrCreateFor<mlir::dataflow::Executable>(op, op->getBlock())->isLive())
+  if (op->getBlock() != nullptr &&
+      !getOrCreateFor<mlir::dataflow::Executable>(
+           point, getProgramPointBefore(op->getBlock()))
+           ->isLive())
     return mlir::success();
 
   // Get the dense lattice to update
-  mlir::dataflow::AbstractDenseLattice *after = getLattice(op);
+  mlir::dataflow::AbstractDenseLattice *after = getLattice(point);
 
   // If this op implements region control-flow, then control-flow dictates its
   // transfer function.
   if (auto branch = mlir::dyn_cast<mlir::RegionBranchOpInterface>(op)) {
-    visitRegionBranchOperation(op, branch, after);
+    visitRegionBranchOperation(point, branch, after);
     return mlir::success();
   }
 
   // pass call operations through to the transfer function
 
   // Get the dense state before the execution of the op.
-  const mlir::dataflow::AbstractDenseLattice *before;
-  if (mlir::Operation *prev = op->getPrevNode())
-    before = getLatticeFor(op, prev);
-  else
-    before = getLatticeFor(op, op->getBlock());
+  const mlir::dataflow::AbstractDenseLattice *before =
+      getLatticeFor(point, getProgramPointBefore(op));
 
   /// Invoke the operation transfer function
   return visitOperationImpl(op, *before, after);
@@ -452,9 +453,10 @@ StackArraysAnalysisWrapper::analyseFunction(mlir::Operation *func) {
     return mlir::failure();
   }
 
-  LatticePoint point{func};
+  LatticePoint point{solver.getProgramPointAfter(func)};
   auto joinOperationLattice = [&](mlir::Operation *op) {
-    const LatticePoint *lattice = solver.lookupState<LatticePoint>(op);
+    const LatticePoint *lattice =
+        solver.lookupState<LatticePoint>(solver.getProgramPointAfter(op));
     // there will be no lattice for an unreachable block
     if (lattice)
       (void)point.join(*lattice);

mlir/include/mlir/Analysis/DataFlow/DeadCodeAnalysis.h

@@ -182,7 +182,7 @@ class DeadCodeAnalysis : public DataFlowAnalysis {
 
   /// Visit an operation with control-flow semantics and deduce which of its
   /// successors are live.
-  LogicalResult visit(ProgramPoint point) override;
+  LogicalResult visit(ProgramPoint *point) override;
 
 private:
   /// Find and mark symbol callables with potentially unknown callsites as

mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.h

@@ -36,8 +36,7 @@ enum class CallControlFlowAction { EnterCallee, ExitCallee, ExternalCallee };
 //===----------------------------------------------------------------------===//
 
 /// This class represents a dense lattice. A dense lattice is attached to
-/// operations to represent the program state after their execution or to blocks
-/// to represent the program state at the beginning of the block. A dense
+/// program point to represent the program state at the program point.
 /// lattice is propagated through the IR by dense data-flow analysis.
 class AbstractDenseLattice : public AnalysisState {
 public:
@@ -59,15 +58,13 @@ class AbstractDenseLattice : public AnalysisState {
 //===----------------------------------------------------------------------===//
 
 /// Base class for dense forward data-flow analyses. Dense data-flow analysis
-/// attaches a lattice between the execution of operations and implements a
-/// transfer function from the lattice before each operation to the lattice
-/// after. The lattice contains information about the state of the program at
-/// that point.
+/// attaches a lattice to program points and implements a transfer function from
+/// the lattice before each operation to the lattice after. The lattice contains
+/// information about the state of the program at that program point.
 ///
-/// In this implementation, a lattice attached to an operation represents the
-/// state of the program after its execution, and a lattice attached to block
-/// represents the state of the program right before it starts executing its
-/// body.
+/// Visit a program point in forward dense data-flow analysis will invoke the
+/// transfer function of the operation preceding the program point iterator.
+/// Visit a program point at the begining of block will visit the block itself.
 class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
 public:
   using DataFlowAnalysis::DataFlowAnalysis;
@@ -76,13 +73,14 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// may modify the program state; that is, every operation and block.
   LogicalResult initialize(Operation *top) override;
 
-  /// Visit a program point that modifies the state of the program. If this is a
-  /// block, then the state is propagated from control-flow predecessors or
-  /// callsites. 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.
-  LogicalResult visit(ProgramPoint point) override;
+  /// Visit a program point that modifies the state of the program. If the
+  /// program point is at the beginning of a block, then the state is propagated
+  /// from control-flow predecessors or callsites.  If the operation before
+  /// program point iterator 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 before the program point iterator.
+  LogicalResult visit(ProgramPoint *point) override;
 
 protected:
   /// Propagate the dense lattice before the execution of an operation to the
@@ -91,15 +89,14 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
                                            const AbstractDenseLattice &before,
                                            AbstractDenseLattice *after) = 0;
 
-  /// Get the dense lattice after the execution of the given lattice anchor.
+  /// Get the dense lattice on the given lattice anchor.
   virtual AbstractDenseLattice *getLattice(LatticeAnchor anchor) = 0;
 
-  /// Get the dense lattice after the execution of the given program point and
-  /// add it as a dependency to a lattice anchor. That is, every time the
-  /// lattice after anchor is updated, the dependent program point must be
-  /// visited, and the newly triggered visit might update the lattice after
-  /// dependent.
-  const AbstractDenseLattice *getLatticeFor(ProgramPoint dependent,
+  /// Get the dense lattice on the given lattice anchor and add dependent as its
+  /// dependency. That is, every time the lattice after anchor is updated, the
+  /// dependent program point must be visited, and the newly triggered visit
+  /// might update the lattice on dependent.
+  const AbstractDenseLattice *getLatticeFor(ProgramPoint *dependent,
                                             LatticeAnchor anchor);
 
   /// Set the dense lattice at control flow entry point and propagate an update
@@ -153,7 +150,7 @@ class AbstractDenseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// 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
   /// parent operation itself.
-  void visitRegionBranchOperation(ProgramPoint point,
+  void visitRegionBranchOperation(ProgramPoint *point,
                                   RegionBranchOpInterface branch,
                                   AbstractDenseLattice *after);
 
@@ -294,14 +291,12 @@ class DenseForwardDataFlowAnalysis
 //===----------------------------------------------------------------------===//
 
 /// Base class for dense backward dataflow analyses. Such analyses attach a
-/// lattice between the execution of operations and implement a transfer
-/// function from the lattice after the operation ot the lattice before it, thus
-/// propagating backward.
+/// lattice to program point and implement a transfer function from the lattice
+/// after the operation to the lattice before it, thus propagating backward.
 ///
-/// In this implementation, a lattice attached to an operation represents the
-/// state of the program before its execution, and a lattice attached to a block
-/// represents the state of the program before the end of the block, i.e., after
-/// its terminator.
+/// Visit a program point in dense backward data-flow analysis will invoke the
+/// transfer function of the operation following the program point iterator.
+/// Visit a program point at the end of block will visit the block itself.
 class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
 public:
   /// Construct the analysis in the given solver. Takes a symbol table
@@ -321,9 +316,9 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// operations, the state is propagated using the transfer function
   /// (visitOperation).
   ///
-  /// Note: the transfer function is currently *not* invoked for operations with
-  /// region or call interface, but *is* invoked for block terminators.
-  LogicalResult visit(ProgramPoint point) override;
+  /// Note: the transfer function is currently *not* invoked before operations
+  /// with region or call interface, but *is* invoked before block terminators.
+  LogicalResult visit(ProgramPoint *point) override;
 
 protected:
   /// Propagate the dense lattice after the execution of an operation to the
@@ -337,10 +332,11 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// block.
   virtual AbstractDenseLattice *getLattice(LatticeAnchor anchor) = 0;
 
-  /// Get the dense lattice before the execution of the program point in
-  /// `anchor` and declare that the `dependent` program point must be updated
-  /// every time `point` is.
-  const AbstractDenseLattice *getLatticeFor(ProgramPoint dependent,
+  /// Get the dense lattice on the given lattice anchor and add dependent as its
+  /// dependency. That is, every time the lattice after anchor is updated, the
+  /// dependent program point must be visited, and the newly triggered visit
+  /// might update the lattice before dependent.
+  const AbstractDenseLattice *getLatticeFor(ProgramPoint *dependent,
                                             LatticeAnchor anchor);
 
   /// Set the dense lattice before at the control flow exit point and propagate
@@ -400,7 +396,7 @@ class AbstractDenseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// (from which the state is propagated) in or after it. `regionNo` indicates
   /// the region that contains the successor, `nullopt` indicating the successor
   /// of the branch operation itself.
-  void visitRegionBranchOperation(ProgramPoint point,
+  void visitRegionBranchOperation(ProgramPoint *point,
                                   RegionBranchOpInterface branch,
                                   RegionBranchPoint branchPoint,
                                   AbstractDenseLattice *before);

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

@@ -179,18 +179,22 @@ class Lattice : public AbstractSparseLattice {
 /// operands to the lattices of the results. This analysis will propagate
 /// lattices across control-flow edges and the callgraph using liveness
 /// information.
+///
+/// Visit a program point in sparse forward data-flow analysis will invoke the
+/// transfer function of the operation preceding the program point iterator.
+/// Visit a program point at the begining of block will visit the block itself.
 class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
 public:
   /// Initialize the analysis by visiting every owner of an SSA value: all
   /// operations and blocks.
   LogicalResult initialize(Operation *top) override;
 
-  /// Visit a program point. If this is a block and all control-flow
-  /// predecessors or callsites are known, then the arguments lattices are
-  /// propagated from them. If this is a call operation or an operation with
-  /// region control-flow, then its result lattices are set accordingly.
-  /// Otherwise, the operation transfer function is invoked.
-  LogicalResult visit(ProgramPoint point) override;
+  /// Visit a program point. If this is at beginning of block and all
+  /// control-flow predecessors or callsites are known, then the arguments
+  /// lattices are propagated from them. If this is after call operation or an
+  /// operation with region control-flow, then its result lattices are set
+  /// accordingly.  Otherwise, the operation transfer function is invoked.
+  LogicalResult visit(ProgramPoint *point) override;
 
 protected:
   explicit AbstractSparseForwardDataFlowAnalysis(DataFlowSolver &solver);
@@ -221,7 +225,7 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
 
   /// Get a read-only lattice element for a value and add it as a dependency to
   /// a program point.
-  const AbstractSparseLattice *getLatticeElementFor(ProgramPoint point,
+  const AbstractSparseLattice *getLatticeElementFor(ProgramPoint *point,
                                                     Value value);
 
   /// Set the given lattice element(s) at control flow entry point(s).
@@ -251,7 +255,8 @@ class AbstractSparseForwardDataFlowAnalysis : public DataFlowAnalysis {
   /// operation `branch`, which can either be the entry block of one of the
   /// regions or the parent operation itself, and set either the argument or
   /// parent result lattices.
-  void visitRegionSuccessors(ProgramPoint point, RegionBranchOpInterface branch,
+  void visitRegionSuccessors(ProgramPoint *point,
+                             RegionBranchOpInterface branch,
                              RegionBranchPoint successor,
                              ArrayRef<AbstractSparseLattice *> lattices);
 };
@@ -312,7 +317,7 @@ class SparseForwardDataFlowAnalysis
 
   /// Get the lattice element for a value and create a dependency on the
   /// provided program point.
-  const StateT *getLatticeElementFor(ProgramPoint point, Value value) {
+  const StateT *getLatticeElementFor(ProgramPoint *point, Value value) {
     return static_cast<const StateT *>(
         AbstractSparseForwardDataFlowAnalysis::getLatticeElementFor(point,
                                                                     value));
@@ -377,10 +382,10 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// under it.
   LogicalResult initialize(Operation *top) override;
 
-  /// Visit a program point. If this is a call operation or an operation with
+  /// Visit a program point. If it is after call operation or an operation with
   /// block or region control-flow, then operand lattices are set accordingly.
   /// Otherwise, invokes the operation transfer function (`visitOperationImpl`).
-  LogicalResult visit(ProgramPoint point) override;
+  LogicalResult visit(ProgramPoint *point) override;
 
 protected:
   explicit AbstractSparseBackwardDataFlowAnalysis(
@@ -445,14 +450,14 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
   /// Get the lattice element for a value, and also set up
   /// dependencies so that the analysis on the given ProgramPoint is re-invoked
   /// if the value changes.
-  const AbstractSparseLattice *getLatticeElementFor(ProgramPoint point,
+  const AbstractSparseLattice *getLatticeElementFor(ProgramPoint *point,
                                                     Value value);
 
   /// Get the lattice elements for a range of values, and also set up
   /// dependencies so that the analysis on the given ProgramPoint is re-invoked
   /// if any of the values change.
   SmallVector<const AbstractSparseLattice *>
-  getLatticeElementsFor(ProgramPoint point, ValueRange values);
+  getLatticeElementsFor(ProgramPoint *point, ValueRange values);
 
   SymbolTableCollection &symbolTable;
 };
@@ -465,6 +470,10 @@ class AbstractSparseBackwardDataFlowAnalysis : public DataFlowAnalysis {
 /// backwards across the IR by implementing transfer functions for operations.
 ///
 /// `StateT` is expected to be a subclass of `AbstractSparseLattice`.
+///
+/// Visit a program point in sparse backward data-flow analysis will invoke the
+/// transfer function of the operation preceding the program point iterator.
+/// Visit a program point at the begining of block will visit the block itself.
 template <typename StateT>
 class SparseBackwardDataFlowAnalysis
     : public AbstractSparseBackwardDataFlowAnalysis {