[mlir][IR] Add ForwardDominanceIterator for IR walkers

[springerm]

Summary

This iterator is similar to ForwardIterator but enumerates blocks according to their successor relationship. As a first use case, this new iterator is utilized in the dialect conversion framework.

Depends On: D144887

[springerm]

add comment

[springerm]

This does not yet handle the case correctly when a block is taken from remainingBlocks: a block without predecessors must be selected. In case there is no such block, there is loop and any block can be taken. It would be more efficient to pre-compute a topologic sorting instead if following successors during the traversal. This would also allow for removing blocks during the traversal.

[zero9178]

mlir/include/mlir/IR/Iterators.h

23

Please don't use using namespace in headers

62

I believe the not equals is automatically inherited from llvm::iterator_facade_base and therefor not necessary

[springerm]

ForwardDominanceIterator does not enumerate unreachable blocks. They do not post-dominate the entry block.

[springerm]

address comments

[rriddle]

Why can't we just use the ReversePostOrderTraversal infra in LLVM?

[springerm]

address comments: use llvm::depth_first

[springerm]

In D144888#4167587, @rriddle wrote:

Why can't we just use the ReversePostOrderTraversal infra in LLVM?

Oh right, for the forward traversal we could use llvm::depth_first.

[springerm]

mlir/lib/Transforms/Utils/DialectConversion.cpp

73–75

Note: This error is no longer produced. Unreachable blocks DCE away so it should no matter if there are unreachable blocks.

[rriddle]

mlir/include/mlir/IR/Iterators.h

67–68

Use llvm::make_pointee_range instead.

mlir/lib/IR/RegionKindInterface.cpp

21–26 ↗

(On Diff #502949)

mlir/include/mlir/IR/Iterators.h

+//===- Iterators.h - IR iterators for IR visitors ---------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// The iterators defined in this file can be used together with IR visitors.
+// Note: These iterators cannot be defined in Visitors.h because that would
+// introduce a cyclic header dependency due to Operation.h.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_ITERATORS_H
+#define MLIR_IR_ITERATORS_H
+
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SetVector.h"
+
+using namespace mlir;
+
+namespace mlir {
+/// This iterator enumerates elements in "reverse" order. It is a wrapper around
+/// llvm::reverse.
+struct ReverseIterator {
+  // llvm::reverse uses RangeT::rbegin and RangeT::rend.
+  template <typename RangeT>
+  static constexpr auto makeIterable(RangeT &&range) {
+    return llvm::reverse(
+        ForwardIterator::makeIterable(std::forward<RangeT>(range)));
+  }
+};
+
+/// An iterator that enumerates blocks according to their successor
+/// relationship.
+class ForwardDominanceBlockIterator final
+    : public llvm::iterator_facade_base<ForwardDominanceBlockIterator,
+                                        std::forward_iterator_tag, Block> {
+public:
+  /// Create an "end" iterator.
+  ForwardDominanceBlockIterator();
+
+  explicit ForwardDominanceBlockIterator(
+      iterator_range<Region::iterator> blocks);
+
+  explicit ForwardDominanceBlockIterator(Region *region)
+      : ForwardDominanceBlockIterator(region->getBlocks()) {}
+
+  using llvm::iterator_facade_base<ForwardDominanceBlockIterator,
+                                   std::forward_iterator_tag,
+                                   Block>::operator++;
+
+  ForwardDominanceBlockIterator &operator++();
+  Block *operator->() const { return worklist.back(); }
+  Block &operator*() const { return *worklist.back(); }
+
+  /// Compare this iterator with another iterator.
+  bool operator==(const ForwardDominanceBlockIterator &rhs) const;
+  bool operator!=(const ForwardDominanceBlockIterator &rhs) const {
+    return !(*this == rhs);
+  }
+
+private:
+  /// A worklist that contains the blocks that will be enumerated next upon
+  /// iterator increment. The last block on the list is the current position
+  /// of the iterator. An iterator with an empty worklist is an "end" iterator.
+  SmallVector<Block *> worklist;
+
+  /// Set of blocks that have not been visited yet.
+  SetVector<Block *> remainingBlocks;
+};
+
+/// This iterator enumerates elements according to their dominance relationship.
+/// Operations and regions are enumerated in "forward" order. Blocks are
+/// enumerated according to their successor relationship.
+struct ForwardDominanceIterator {
+  static Block &makeIterable(Block &range) {
+    return ForwardIterator::makeIterable(range);
+  }
+
+  static auto makeIterable(Region &region) {
+    return llvm::make_range(ForwardDominanceBlockIterator(&region), {});
+  }
+
+  static auto makeIterable(iterator_range<Region::iterator> range) {
+    return llvm::make_range(ForwardDominanceBlockIterator(range), {});
+  }
+
+  static MutableArrayRef<Region> makeIterable(Operation &range) {
+    return ForwardIterator::makeIterable(range);
+  }
+};
+} // namespace mlir
+
+#endif // MLIR_IR_ITERATORS_H

mlir/include/mlir/IR/Visitors.h

struct ForwardIterator {
   static MutableArrayRef<Region> makeIterable(Operation &range);
 
   /// Regions and block are already iterable.
   template <typename T> static constexpr T &makeIterable(T &range) {
     return range;
   }
 };
 
-/// This iterator enumerates elements in "reverse" order. It is a wrapper around
-/// llvm::reverse.
-struct ReverseIterator {
-  template <typename RangeT>
-  static constexpr auto makeIterable(RangeT &&range) {
-    // llvm::reverse uses RangeT::rbegin and RangeT::rend.
-    return llvm::reverse(
-        ForwardIterator::makeIterable(std::forward<RangeT>(range)));
-  }
-};
-
 /// A utility class to encode the current walk stage for "generic" walkers.
 /// When walking an operation, we can either choose a Pre/Post order walker
 /// which invokes the callback on an operation before/after all its attached
 /// regions have been visited, or choose a "generic" walker where the callback
 /// is invoked on the operation N+1 times where N is the number of regions
 /// attached to that operation. The `WalkStage` class below encodes the current
 /// stage of the walk, i.e., which regions have already been visited, and the
 /// callback accepts an additional argument for the current stage. Such

mlir/lib/IR/CMakeLists.txt

add_mlir_library(MLIRIR
   Diagnostics.cpp
   Dialect.cpp
   DialectResourceBlobManager.cpp
   Dominance.cpp
   ExtensibleDialect.cpp
   FunctionImplementation.cpp
   FunctionInterfaces.cpp
   IntegerSet.cpp
+  Iterators.cpp
   Location.cpp
   MLIRContext.cpp
   Operation.cpp
   OperationSupport.cpp
   PatternMatch.cpp
   Region.cpp
   RegionKindInterface.cpp
   SymbolTable.cpp

mlir/lib/IR/Iterators.cpp

+//===- Iterators.cpp - IR iterators for IR visitors -----------------------===//
+//
+// 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/IR/Iterators.h"
+
+using namespace mlir;
+
+ForwardDominanceBlockIterator::ForwardDominanceBlockIterator() = default;
+
+ForwardDominanceBlockIterator::ForwardDominanceBlockIterator(
+    iterator_range<Region::iterator> blocks) {
+  if (blocks.empty())
+    return;
+
+  // Always start iterating with the first block.
+  worklist.push_back(&*blocks.begin());
+  for (Block &b : llvm::drop_begin(blocks))
+    remainingBlocks.insert(&b);
+}
+
+ForwardDominanceBlockIterator &ForwardDominanceBlockIterator::operator++() {
+  assert(!worklist.empty() && "cannot increment iterator past end");
+  Block *block = worklist.pop_back_val();
+
+  // Add successors to worklist.
+  for (Block *succ : block->getSuccessors()) {
+    auto it = llvm::find(remainingBlocks, succ);
+    if (it != remainingBlocks.end()) {
+      remainingBlocks.erase(it);
+      worklist.push_back(succ);
+    }
+  }
+
+  // Take the next block from the set of unvisited blocks if we run out of
+  // blocks.
+  if (worklist.empty() && !remainingBlocks.empty()) {
+    worklist.push_back(remainingBlocks.pop_back_val());
+  }
+
+  return *this;
+}
+
+bool ForwardDominanceBlockIterator::operator==(
+    const ForwardDominanceBlockIterator &rhs) const {
+  return llvm::equal(worklist, rhs.worklist) &&
+         llvm::equal(remainingBlocks, rhs.remainingBlocks);
+}

mlir/lib/Transforms/Utils/DialectConversion.cpp

 //===- DialectConversion.cpp - MLIR dialect conversion generic pass -------===//
 //
 // 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/Transforms/DialectConversion.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Iterators.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::detail;
 
 #define DEBUG_TYPE "dialect-conversion"
 
-/// Recursively collect all of the operations to convert from within 'region'.
-/// If 'target' is nonnull, operations that are recursively legal have their
-/// regions pre-filtered to avoid considering them for legalization.
-static LogicalResult
-computeConversionSet(iterator_range<Region::iterator> region,
-                     Location regionLoc,
-                     SmallVectorImpl<Operation *> &toConvert,
-                     ConversionTarget *target = nullptr) {
-  if (region.empty())
-    return success();
-
-  // Traverse starting from the entry block.
-  SmallVector<Block *, 16> worklist(1, &*region.begin());
-  DenseSet<Block *> visitedBlocks;
-  visitedBlocks.insert(worklist.front());
-  while (!worklist.empty()) {
-    Block *block = worklist.pop_back_val();
-
-    // Compute the conversion set of each of the nested operations.
-    for (Operation &op : *block) {
-      toConvert.emplace_back(&op);
-
-      // Don't check this operation's children for conversion if the operation
-      // is recursively legal.
-      auto legalityInfo =
-          target ? target->isLegal(&op)
-                 : std::optional<ConversionTarget::LegalOpDetails>();
-      if (legalityInfo && legalityInfo->isRecursivelyLegal)
-        continue;
-      for (auto &region : op.getRegions()) {
-        if (failed(computeConversionSet(region.getBlocks(), region.getLoc(),
-                                        toConvert, target)))
-          return failure();
-      }
-    }
-
-    // Recurse to children that haven't been visited.
-    for (Block *succ : block->getSuccessors())
-      if (visitedBlocks.insert(succ).second)
-        worklist.push_back(succ);
-  }
-
-  // Check that all blocks in the region were visited.
-  if (llvm::any_of(llvm::drop_begin(region, 1),
-                   [&](Block &block) { return !visitedBlocks.count(&block); }))
-    return emitError(regionLoc, "unreachable blocks were not converted");
-  return success();
-}
-
 /// A utility function to log a successful result for the given reason.
 template <typename... Args>
 static void logSuccess(llvm::ScopedPrinter &os, StringRef fmt, Args &&...args) {
   LLVM_DEBUG({
     os.unindent();
     os.startLine() << "} -> SUCCESS";
     if (!fmt.empty())
       os.getOStream() << " : "
struct ConversionPatternRewriterImpl {
 
   /// Notifies that `block` is being merged with `srcBlock`.
   void notifyBlocksBeingMerged(Block *block, Block *srcBlock);
 
   /// Notifies that the blocks of a region are about to be moved.
   void notifyRegionIsBeingInlinedBefore(Region &region, Region &parent,
                                         Region::iterator before);
 
-  /// Notifies that the blocks of a region were cloned into another.
-  void notifyRegionWasClonedBefore(iterator_range<Region::iterator> &blocks,
-                                   Location origRegionLoc);
-
   /// Notifies that a pattern match failed for the given reason.
   LogicalResult
   notifyMatchFailure(Location loc,
                      function_ref<void(Diagnostic &)> reasonCallback);
 
   //===--------------------------------------------------------------------===//
   // State
   //===--------------------------------------------------------------------===//
void ConversionPatternRewriterImpl::notifyRegionIsBeingInlinedBefore(
   for (auto &earlierBlock : llvm::drop_begin(llvm::reverse(region), 1)) {
     blockActions.push_back(
         BlockAction::getMove(laterBlock, {&region, &earlierBlock}));
     laterBlock = &earlierBlock;
   }
   blockActions.push_back(BlockAction::getMove(laterBlock, {&region, nullptr}));
 }
 
-void ConversionPatternRewriterImpl::notifyRegionWasClonedBefore(
-    iterator_range<Region::iterator> &blocks, Location origRegionLoc) {
-  for (Block &block : blocks)
-    blockActions.push_back(BlockAction::getCreate(&block));
-
-  // Compute the conversion set for the inlined region.
-  auto result = computeConversionSet(blocks, origRegionLoc, createdOps);
-
-  // This original region has already had its conversion set computed, so there
-  // shouldn't be any new failures.
-  (void)result;
-  assert(succeeded(result) && "expected region to have no unreachable blocks");
-}
-
 LogicalResult ConversionPatternRewriterImpl::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
     if (notifyCallback)
       notifyCallback(diag);

 }
 
 void ConversionPatternRewriter::cloneRegionBefore(Region &region,
                                                   Region &parent,
                                                   Region::iterator before,
                                                   IRMapping &mapping) {
   if (region.empty())
     return;
+
   PatternRewriter::cloneRegionBefore(region, parent, before, mapping);
 
-  // Collect the range of the cloned blocks.
-  auto clonedBeginIt = mapping.lookup(&region.front())->getIterator();
-  auto clonedBlocks = llvm::make_range(clonedBeginIt, before);
-  impl->notifyRegionWasClonedBefore(clonedBlocks, region.getLoc());
+  for (Block &b : ForwardDominanceIterator::makeIterable(region)) {
+    Block *cloned = mapping.lookup(&b);
+    impl->notifyCreatedBlock(cloned);
+    cloned->walk<WalkOrder::PreOrder, ForwardDominanceIterator>(
+        [&](Operation *op) { notifyOperationInserted(op); });
+  }
 }
 
 void ConversionPatternRewriter::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     impl->logger.startLine()
         << "** Insert  : '" << op->getName() << "'(" << op << ")\n";
   });
   impl->createdOps.push_back(op);
LogicalResult OperationConverter::convertOperations(
     function_ref<void(Diagnostic &)> notifyCallback) {
   if (ops.empty())
     return success();
   ConversionTarget &target = opLegalizer.getTarget();
 
   // Compute the set of operations and blocks to convert.
   SmallVector<Operation *> toConvert;
   for (auto *op : ops) {
-    toConvert.emplace_back(op);
-    for (auto &region : op->getRegions())
-      if (failed(computeConversionSet(region.getBlocks(), region.getLoc(),
-                                      toConvert, &target)))
-        return failure();
+    op->walk<WalkOrder::PreOrder, ForwardDominanceIterator>([&](Operation *op) {
+      toConvert.push_back(op);
+      // Don't check this operation's children for conversion if the operation
+      // is recursively legal.
+      auto legalityInfo = target.isLegal(op);
+      if (legalityInfo && legalityInfo->isRecursivelyLegal)
+        return WalkResult::skip();
+      return WalkResult::advance();
+    });
   }
 
   // Convert each operation and discard rewrites on failure.
   ConversionPatternRewriter rewriter(ops.front()->getContext());
   ConversionPatternRewriterImpl &rewriterImpl = rewriter.getImpl();
   rewriterImpl.notifyCallback = notifyCallback;
 
   for (auto *op : toConvert)

mlir/test/IR/visitors.mlir

 // CHECK:       Erasing op 'func.return'
 
 // CHECK-LABEL: Block post-order erasures (no skip)
 // CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
 // CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
 // CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
 // CHECK:       Erasing block ^bb0 from region 0 from operation 'func.func'
 // CHECK:       Erasing block ^bb0 from region 0 from operation 'builtin.module'
+
+// -----
+
+func.func @unordered_cfg_with_loop() {
+  "regionOp0"() ({
+    ^bb0:
+      %c = "op0"() : () -> (i1)
+      cf.cond_br %c, ^bb2, ^bb3
+    ^bb1:
+      "op1"(%val) : (i32) -> ()
+      cf.br ^bb5
+    ^bb2:
+      %val = "op2"() : () -> (i32)
+      cf.br ^bb1
+    ^bb3:
+      "op3"() : () -> ()
+      cf.br ^bb2
+    ^bb4:
+      "op4"() : () -> ()
+      cf.br ^bb2
+    ^bb5:
+      "op5"() : () -> ()
+      cf.br ^bb7
+    ^bb6:
+      "op6"() : () -> ()
+      cf.br ^bb6
+    ^bb7:
+      "op7"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+//      4
+//      |
+//      v
+// 0 -> 2 --> 1 --> 5 --> 7
+// |    ^
+// |    |      6 --
+// |   /       ^   \
+// |  /         \  /
+// v /           --
+//  3
+
+// CHECK-LABEL: Op forward dominance post-order visits
+// CHECK:       Visiting op 'op0'
+// CHECK:       Visiting op 'cf.cond_br'
+// CHECK:       Visiting op 'op3'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op2'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op1'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op5'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op7'
+// CHECK:       Visiting op 'op6'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'op4'
+// CHECK:       Visiting op 'cf.br'
+// CHECK:       Visiting op 'regionOp0'
+// CHECK:       Visiting op 'func.return'
+// CHECK:       Visiting op 'func.func'
+// CHECK:       Visiting op 'builtin.module'
+
+// CHECK-LABEL: Block forward dominance post-order visits
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb3 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb2 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb1 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb5 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb7 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb6 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb4 from region 0 from operation 'regionOp0'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
+// CHECK:       Visiting block ^bb0 from region 0 from operation 'builtin.module'
+
+// CHECK-LABEL: Region forward dominance post-order visits
+// CHECK:       Visiting region 0 from operation 'regionOp0'
+// CHECK:       Visiting region 0 from operation 'func.func'
+// CHECK:       Visiting region 0 from operation 'builtin.module'
+
+// CHECK-LABEL: Block pre-order erasures (skip)
+// CHECK:       Erasing block ^bb0 from region 0 from operation 'regionOp0'
+// CHECK:       Cannot erase block ^bb0 from region 0 from operation 'regionOp0', still has uses
+// CHECK:       Cannot erase block ^bb1 from region 0 from operation 'regionOp0', still has uses
+// CHECK:       Erasing block ^bb2 from region 0 from operation 'regionOp0'
+// CHECK:       Erasing block ^bb2 from region 0 from operation 'regionOp0'
+// CHECK:       Cannot erase block ^bb2 from region 0 from operation 'regionOp0', still has uses
+// CHECK:       Cannot erase block ^bb3 from region 0 from operation 'regionOp0', still has uses
+// CHECK:       Cannot erase block ^bb4 from region 0 from operation 'regionOp0', still has uses

mlir/test/lib/IR/TestVisitors.cpp

 //===- TestIRVisitors.cpp - Pass to test the IR visitors ------------------===//
 //
 // 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/IR/BuiltinOps.h"
+#include "mlir/IR/Iterators.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 static void printRegion(Region *region) {
   llvm::outs() << "region " << region->getRegionNumber() << " from operation '"
                << region->getParentOp()->getName() << "'";
 }
llvm::outs() << "Op reverse post-order visits"
                << "\n";
   op->walk<WalkOrder::PostOrder, ReverseIterator>(opPure);
   llvm::outs() << "Block reverse post-order visits"
                << "\n";
   op->walk<WalkOrder::PostOrder, ReverseIterator>(blockPure);
   llvm::outs() << "Region reverse post-order visits"
                << "\n";
   op->walk<WalkOrder::PostOrder, ReverseIterator>(regionPure);
+
+  llvm::outs() << "Op forward dominance post-order visits"
+               << "\n";
+  op->walk<WalkOrder::PostOrder, ForwardDominanceIterator>(opPure);
+  llvm::outs() << "Block forward dominance post-order visits"
+               << "\n";
+  op->walk<WalkOrder::PostOrder, ForwardDominanceIterator>(blockPure);
+  llvm::outs() << "Region forward dominance post-order visits"
+               << "\n";
+  op->walk<WalkOrder::PostOrder, ForwardDominanceIterator>(regionPure);
 }
 
 /// Tests erasure callbacks that skip the walk.
 static void testSkipErasureCallbacks(Operation *op) {
   auto skipOpErasure = [](Operation *op) {
     // Do not erase module and module children operations. Otherwise, there
     // wouldn't be too much to test in pre-order.
     if (isa<ModuleOp>(op) || isa<ModuleOp>(op->getParentOp()))
       return WalkResult::advance();
 
     llvm::outs() << "Erasing ";
     printOperation(op);
     llvm::outs() << "\n";
     op->dropAllUses();
     op->erase();
     return WalkResult::skip();
   };
   auto skipBlockErasure = [](Block *block) {
     // Do not erase module and module children blocks. Otherwise there wouldn't
     // be too much to test in pre-order.
     Operation *parentOp = block->getParentOp();
     if (isa<ModuleOp>(parentOp) || isa<ModuleOp>(parentOp->getParentOp()))
       return WalkResult::advance();
 
+    if (block->use_empty()) {
       llvm::outs() << "Erasing ";
       printBlock(block);
       llvm::outs() << "\n";
       block->erase();
       return WalkResult::skip();
+    } else {
+      llvm::outs() << "Cannot erase ";
+      printBlock(block);
+      llvm::outs() << ", still has uses\n";
+      return WalkResult::advance();
+    }
   };
 
   llvm::outs() << "Op pre-order erasures (skip)"
                << "\n";
   Operation *cloned = op->clone();
   cloned->walk<WalkOrder::PreOrder>(skipOpErasure);
   cloned->erase();
 
static void testNoSkipErasureCallbacks(Operation *op) {
   auto noSkipOpErasure = [](Operation *op) {
     llvm::outs() << "Erasing ";
     printOperation(op);
     llvm::outs() << "\n";
     op->dropAllUses();
     op->erase();
   };
   auto noSkipBlockErasure = [](Block *block) {
+    if (block->use_empty()) {
       llvm::outs() << "Erasing ";
       printBlock(block);
       llvm::outs() << "\n";
       block->erase();
+    } else {
+      llvm::outs() << "Cannot erase ";
+      printBlock(block);
+      llvm::outs() << ", still has uses\n";
+    }
   };
 
   llvm::outs() << "Op post-order erasures (no skip)"
                << "\n";
   Operation *cloned = op->clone();
   cloned->walk<WalkOrder::PostOrder>(noSkipOpErasure);
 
   llvm::outs() << "Block post-order erasures (no skip)"