[MLIR][NFC] Minor cleanup for BufferDeallocation pass.

[jurahul]

Summary

• Change walkReturnOperations() to be a non-template and look at block terminator for ReturnLike trait.
• Clarify description of validateSupportedControlFlow
• Eliminate unused argument in Backedges::recurse.
• Eliminate repeated calls to getFunction()
• Fix wording for non-SCF loop failure

[pifon2a]

Thank you!

mlir/lib/Transforms/BufferDeallocation.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferUtils.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 
 /// Walks over all immediate return-like terminators in the given region.
-template <typename FuncT>
-static void walkReturnOperations(Region *region, const FuncT &func) {
-  for (Block &block : *region)
-    for (Operation &operation : block) {
-      // Skip non-return-like terminators.
-      if (operation.hasTrait<OpTrait::ReturnLike>())
-        func(&operation);
-    }
+static void walkReturnOperations(Region *region,
+                                 std::function<void(Operation *)> func) {
+  for (Block &block : *region) {
+    Operation *terminator = block.getTerminator();
+    if (terminator->hasTrait<OpTrait::ReturnLike>())
+      func(terminator);
 }
 
-/// Checks if all operations in a given region have at least one attached region
-/// that implements the RegionBranchOpInterface. This is not required in edge
+/// Checks if all operations in a given region that have at least one attached
+/// region implement the RegionBranchOpInterface. This is not required in edge
 /// cases, where we have a single attached region and the parent operation has
 /// no results.
 static bool validateSupportedControlFlow(Region &region) {
   bool success = true;
   region.walk([&success](Operation *operation) {
     auto regions = operation->getRegions();
     // Walk over all operations in a region and check if the operation has at
     // least one region and implements the RegionBranchOpInterface. If there

 /// explicit control flow.
 class Backedges {
 public:
   using BlockSetT = SmallPtrSet<Block *, 16>;
   using BackedgeSetT = llvm::DenseSet<std::pair<Block *, Block *>>;
 
 public:
   /// Constructs a new backedges analysis using the op provided.
-  Backedges(Operation *op) { recurse(op, op->getBlock()); }
+  Backedges(Operation *op) { recurse(op); }
 
   /// Returns the number of backedges formed by explicit control flow.
   size_t size() const { return edgeSet.size(); }
 
   /// Returns the start iterator to loop over all backedges.
   BackedgeSetT::const_iterator begin() const { return edgeSet.begin(); }
 
   /// Returns the end iterator to loop over all backedges.
bool enter(Block &current, Block *predecessor) {
     return inserted;
   }
 
   /// Leaves the current block.
   void exit(Block &current) { visited.erase(&current); }
 
   /// Recurses into the given operation while taking all attached regions into
   /// account.
-  void recurse(Operation *op, Block *predecessor) {
+  void recurse(Operation *op) {
     Block *current = op->getBlock();
     // If the current op implements the `BranchOpInterface`, there can be
     // cycles in the scope of all successor blocks.
     if (isa<BranchOpInterface>(op)) {
       for (Block *succ : current->getSuccessors())
         recurse(*succ, current);
     }
     // Recurse into all distinct regions and check for explicit control-flow
private:
   void recurse(Block &block, Block *predecessor) {
     // Try to enter the current block. If this is not possible, we are
     // currently processing this block and can safely return here.
     if (!enter(block, predecessor))
       return;
 
     // Recurse into all operations and successor blocks.
     for (Operation &op : block.getOperations())
-      recurse(&op, predecessor);
+      recurse(&op);
 
     // Leave the current block.
     exit(block);
   }
 
   /// Stores all blocks that are currently visited and on the processing stack.
   BlockSetT visited;
 
private:
   void placeDeallocs() const {
     // Move or insert deallocs using the previously computed information.
     // These deallocations will be linked to their associated allocation nodes
     // since they don't have any aliases that can (potentially) increase their
     // liveness.
     for (const BufferPlacementAllocs::AllocEntry &entry : allocs) {
       Value alloc = std::get<0>(entry);
       auto aliasesSet = aliases.resolve(alloc);
-      assert(aliasesSet.size() > 0 && "must contain at least one alias");
+      assert(!aliasesSet.empty() && "must contain at least one alias");
 
       // Determine the actual block to place the dealloc and get liveness
       // information.
       Block *placementBlock =
           findCommonDominator(alloc, aliasesSet, postDominators);
       const LivenessBlockInfo *livenessInfo =
           liveness.getLiveness(placementBlock);
 

 
 /// The actual buffer deallocation pass that inserts and moves dealloc nodes
 /// into the right positions. Furthermore, it inserts additional clones if
 /// necessary. It uses the algorithm described at the top of the file.
 struct BufferDeallocationPass : BufferDeallocationBase<BufferDeallocationPass> {
 
   void runOnFunction() override {
     // Ensure that there are supported loops only.
-    Backedges backedges(getFunction());
+    FuncOp func = getFunction();
+    Backedges backedges(func);
     if (backedges.size()) {
-      getFunction().emitError(
-          "Structured control-flow loops are supported only.");
+      func.emitError("Only structured control-flow loops are supported.");
       return signalPassFailure();
     }
 
     // Check that the control flow structures are supported.
-    if (!validateSupportedControlFlow(getFunction().getRegion())) {
+    if (!validateSupportedControlFlow(func.getRegion())) {
       return signalPassFailure();
     }
 
     // Place all required temporary clone and dealloc nodes.
-    BufferDeallocation deallocation(getFunction());
+    BufferDeallocation deallocation(func);
     deallocation.deallocate();
   }
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createBufferDeallocationPass() {
   return std::make_unique<BufferDeallocationPass>();
 }

mlir/test/Transforms/buffer-deallocation.mlir

 // CHECK-NEXT: memref.dealloc %[[VAL_7]]
 
 // -----
 
 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
 // The BufferDeallocation transformation should fail on this explicit
 // control-flow loop since they are not supported.
 
-// expected-error@+1 {{Structured control-flow loops are supported only}}
+// expected-error@+1 {{Only structured control-flow loops are supported}}
 func @loop_dynalloc(
   %arg0 : i32,
   %arg1 : i32,
   %arg2: memref<?xf32>,
   %arg3: memref<?xf32>) {
   %const0 = constant 0 : i32
   br ^loopHeader(%const0, %arg2 : i32, memref<?xf32>)
 

 }
 
 // -----
 
 // Test Case: explicit control-flow loop with a dynamically allocated buffer.
 // The BufferDeallocation transformation should fail on this explicit
 // control-flow loop since they are not supported.
 
-// expected-error@+1 {{Structured control-flow loops are supported only}}
+// expected-error@+1 {{Only structured control-flow loops are supported}}
 func @do_loop_alloc(
   %arg0 : i32,
   %arg1 : i32,
   %arg2: memref<2xf32>,
   %arg3: memref<2xf32>) {
   %const0 = constant 0 : i32
   br ^loopBody(%const0, %arg2 : i32, memref<2xf32>)