[MLIR] DataFlowAnalysis: Use a queue to maintain the worklist

[vaivaswatha]

Summary

Since the analysis is described to be suitable for a forward data-flow analysis, maintaining the worklist as a queue is closer to the RPO ordering of block visits, thus reaching the fixpoint earlier.

For example,

// mlir-opt -mlir-disable-threading -allow-unregistered-dialect -pass-pipeline="builtin.func(sccp)" sccp-iterations.mlir
func @simple_control_flow(%arg0 : i32, %arg1 : i1) -> i32 {
  cond_br %arg1, ^bb1, ^bb3

^bb1:
  %1 = arith.constant 1 : i32
  br ^bb2(%1 : i32)

^bb3:
  %3 = arith.constant 2 : i32
  br ^bb2(%3 : i32)

^bb2(%arg : i32):
  %2 = arith.constant 3 : i32
  %4 = arith.addi %arg, %2 : i32
  return %4 : i32
}

Running SCCP on this code results in the following visits when using a stack (currently) to maintain the worklist (I modified visitOperation in SCCP.cpp to print the operation being visited):

(SCCP): Visiting %c2_i32 = arith.constant 2 : i32
(SCCP): Visiting %c3_i32 = arith.constant 3 : i32                                                                                                                                                                  
(SCCP): Visiting %1 = arith.addi %0, %c3_i32 : i32
(SCCP): Visiting %c1_i32 = arith.constant 1 : i32
(SCCP): Visiting %1 = arith.addi %0, %c3_i32 : i32
module  {                                         
  func @simple_control_flow(%arg0: i32, %arg1: i1) -> i32 {
    %c3_i32 = arith.constant 3 : i32                                                                     
    %c2_i32 = arith.constant 2 : i32
    %c1_i32 = arith.constant 1 : i32
    cond_br %arg1, ^bb1, ^bb2       
  ^bb1:  // pred: ^bb0       
    br ^bb3(%c1_i32 : i32)
  ^bb2:  // pred: ^bb0    
    br ^bb3(%c2_i32 : i32)
  ^bb3(%0: i32):  // 2 preds: ^bb1, ^bb2
    %1 = arith.addi %0, %c3_i32 : i32   
    return %1 : i32                  
  }                
}

But with the proposed change to use a queue instead, the visits are:

(SCCP): Visiting %c1_i32 = arith.constant 1 : i32
(SCCP): Visiting %c2_i32 = arith.constant 2 : i32
(SCCP): Visiting %c3_i32 = arith.constant 3 : i32
(SCCP): Visiting %1 = arith.addi %0, %c3_i32 : i32
module  {
  func @simple_control_flow(%arg0: i32, %arg1: i1) -> i32 {
    %c3_i32 = arith.constant 3 : i32
    %c2_i32 = arith.constant 2 : i32
    %c1_i32 = arith.constant 1 : i32
    cond_br %arg1, ^bb1, ^bb2
  ^bb1:  // pred: ^bb0
    br ^bb3(%c1_i32 : i32)
  ^bb2:  // pred: ^bb0
    br ^bb3(%c2_i32 : i32)
  ^bb3(%0: i32):  // 2 preds: ^bb1, ^bb2
    %1 = arith.addi %0, %c3_i32 : i32
    return %1 : i32
  }
}

On the longer run, we could explore more complex strategies (not to mention the simple improvement of checking if a node is already in the queue before inserting it) as is described in "Iterative Data-flow Analysis - Revisited" (Keith D. Cooper, Timothy J. Harvey, and Ken Kennedy).

I'm not sure what's the right way to add a test for this. Perhaps add a statistic to count when an Operation is visited, print that and CHECK for that in the unit test?

[rriddle]

Hmmm, I'm slightly fine with submitting this without an explicit test given that it is changing a data structure (and not a crazy feature/behavior change) and seems relatively straightforward. Do you mind adding in the logging that you added (potentially as a separate commit)? It would save time for future debugging. Ideally we could CHECK the debug output, but can't think of any examples of that offhand.

mlir/lib/Analysis/DataFlowAnalysis.cpp

9

System includes should go after the others:

https://llvm.org/docs/CodingStandards.html#include-style

[vaivaswatha]

Incorporating review comments.

[vaivaswatha]

@rriddle I've added the debug statement I had locally. If this looks fine to you, I'll go ahead and push the change.

mlir/lib/Analysis/DataFlowAnalysis.cpp

 //===- DataFlowAnalysis.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/DataFlowAnalysis.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
+#include <queue>
+
 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
 /// This class contains various state used when computing the lattice elements
 /// of a callable operation.
 class CallableLatticeState {
 public:
private:
   template <typename ValuesT>
   void markAllPessimisticFixpoint(ValuesT values) {
     for (auto value : values)
       markPessimisticFixpoint(value);
   }
   template <typename ValuesT>
   void markAllPessimisticFixpoint(Operation *op, ValuesT values) {
     markAllPessimisticFixpoint(values);
-    opWorklist.push_back(op);
+    opWorklist.push(op);
   }
   template <typename ValuesT>
   void markAllPessimisticFixpointAndVisitUsers(ValuesT values) {
     for (auto value : values) {
       AbstractLatticeElement &lattice = analysis.getLatticeElement(value);
       if (lattice.markPessimisticFixpoint() == ChangeResult::Change)
         visitUsers(value);
     }
private:
 
   /// The set of blocks that are known to execute, or are intrinsically live.
   SmallPtrSet<Block *, 16> executableBlocks;
 
   /// The set of control flow edges that are known to execute.
   DenseSet<std::pair<Block *, Block *>> executableEdges;
 
   /// A worklist containing blocks that need to be processed.
-  SmallVector<Block *, 64> blockWorklist;
+  std::queue<Block *> blockWorklist;
 
   /// A worklist of operations that need to be processed.
-  SmallVector<Operation *, 64> opWorklist;
+  std::queue<Operation *> opWorklist;
 
   /// The callable operations that have their argument/result state tracked.
   DenseMap<Operation *, CallableLatticeState> callableLatticeState;
 
   /// A map between a call operation and the resolved symbol callable. This
   /// avoids re-resolving symbol references during propagation. Value based
   /// callables are trivial to resolve, so they can be done in-place.
   DenseMap<Operation *, Operation *> callToSymbolCallable;
for (Region &region : op->getRegions()) {
     markEntryBlockExecutable(&region, /*markPessimisticFixpoint=*/true);
   }
   initializeSymbolCallables(op);
 }
 
 void ForwardDataFlowSolver::solve() {
   while (!blockWorklist.empty() || !opWorklist.empty()) {
     // Process any operations in the op worklist.
-    while (!opWorklist.empty())
-      visitUsers(*opWorklist.pop_back_val());
+    while (!opWorklist.empty()) {
+      Operation *nextOp = opWorklist.front();
+      opWorklist.pop();
+      visitUsers(*nextOp);
+    }
 
     // Process any blocks in the block worklist.
-    while (!blockWorklist.empty())
-      visitBlock(blockWorklist.pop_back_val());
+    while (!blockWorklist.empty()) {
+      Block *nextBlock = blockWorklist.front();
+      blockWorklist.pop();
+      visitBlock(nextBlock);
+    }
   }
 }
 
 void ForwardDataFlowSolver::initializeSymbolCallables(Operation *op) {
   // Initialize the set of symbol callables that can have their state tracked.
   // This tracks which symbol callable operations we can propagate within and
   // out of.
   auto walkFn = [&](Operation *symTable, bool allUsesVisible) {
void ForwardDataFlowSolver::visitOperation(Operation *op) {
   // If all of the results of this operation are already resolved, bail out
   // early.
   auto isAtFixpointFn = [&](Value value) { return isAtFixpoint(value); };
   if (llvm::all_of(op->getResults(), isAtFixpointFn))
     return;
 
   // Visit the current operation.
   if (analysis.visitOperation(op, operandLattices) == ChangeResult::Change)
-    opWorklist.push_back(op);
+    opWorklist.push(op);
 
   // `visitOperation` is required to define all of the result lattices.
   assert(llvm::none_of(
              op->getResults(),
              [&](Value value) {
                return analysis.getLatticeElement(value).isUninitialized();
              }) &&
          "expected `visitOperation` to define all result lattices");
if (!region) {
         continue;
 
       // Mark the results outside of the input range as having reached the
       // pessimistic fixpoint.
       // TODO: This isn't exactly ideal. There may be situations in which a
       // region operation can provide information for certain results that
       // aren't part of the control flow.
       if (succArgs.size() != results.size()) {
-        opWorklist.push_back(parentOp);
+        opWorklist.push(parentOp);
         if (succArgs.empty()) {
           markAllPessimisticFixpoint(results);
           continue;
         }
 
         unsigned firstResIdx = succArgs[0].cast<OpResult>().getResultNumber();
         markAllPessimisticFixpoint(results.take_front(firstResIdx));
         markAllPessimisticFixpoint(
if (!region->empty()) {
     return markBlockExecutable(&region->front());
   }
   return ChangeResult::NoChange;
 }
 
 ChangeResult ForwardDataFlowSolver::markBlockExecutable(Block *block) {
   bool marked = executableBlocks.insert(block).second;
   if (marked)
-    blockWorklist.push_back(block);
+    blockWorklist.push(block);
   return marked ? ChangeResult::Change : ChangeResult::NoChange;
 }
 
 bool ForwardDataFlowSolver::isBlockExecutable(Block *block) const {
   return executableBlocks.count(block);
 }
 
 void ForwardDataFlowSolver::markEdgeExecutable(Block *from, Block *to) {
bool ForwardDataFlowSolver::isAtFixpoint(Value value) const {
   if (auto *lattice = analysis.lookupLatticeElement(value))
     return lattice->isAtFixpoint();
   return false;
 }
 
 void ForwardDataFlowSolver::join(Operation *owner, AbstractLatticeElement &to,
                                  const AbstractLatticeElement &from) {
   if (to.join(from) == ChangeResult::Change)
-    opWorklist.push_back(owner);
+    opWorklist.push(owner);
 }
 
 //===----------------------------------------------------------------------===//
 // AbstractLatticeElement
 //===----------------------------------------------------------------------===//
 
 AbstractLatticeElement::~AbstractLatticeElement() = default;
 

mlir/lib/Transforms/SCCP.cpp

 #include "mlir/Analysis/DataFlowAnalysis.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/Passes.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "sccp"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // SCCP Analysis
 //===----------------------------------------------------------------------===//
 
 namespace {

 
 struct SCCPAnalysis : public ForwardDataFlowAnalysis<SCCPLatticeValue> {
   using ForwardDataFlowAnalysis<SCCPLatticeValue>::ForwardDataFlowAnalysis;
   ~SCCPAnalysis() override = default;
 
   ChangeResult
   visitOperation(Operation *op,
                  ArrayRef<LatticeElement<SCCPLatticeValue> *> operands) final {
+
+    LLVM_DEBUG(llvm::dbgs() << "SCCP: Visiting operation: " << *op << "\n");
+
     // Don't try to simulate the results of a region operation as we can't
     // guarantee that folding will be out-of-place. We don't allow in-place
     // folds as the desire here is for simulated execution, and not general
     // folding.
     if (op->getNumRegions())
       return markAllPessimisticFixpoint(op->getResults());
 
     SmallVector<Attribute> constantOperands(