[mlir] Add hasTerminator helper function to Block

[springerm]

Summary

Such functionality is used in various places.

[mehdi_amini]

mlir/lib/IR/Block.cpp

264

Should we name this mightHaveTerminator()?

[springerm]

mlir/lib/IR/Block.cpp

264

What's your reasoning? A block either has a terminator or not.

[springerm]

update

[springerm]

mlir/lib/IR/Block.cpp

264

I see, to account for unregistered ops. All code locations that I found are actually looking for terminators, not may-be-terminators.

[mehdi_amini]

mlir/lib/IR/Block.cpp

264

Yes, it’s tricky because there is no definitive answer for unregistered ops…

For example right now !hasTerminator() does not mean no terminator , so I’d name it hasRegisteredTerminator to be exact

[springerm]

update

mlir/include/mlir/IR/Block.h

iterator_range<iterator> without_terminator() {
     return {begin(), endIt};
   }
 
   //===--------------------------------------------------------------------===//
   // Terminator management
   //===--------------------------------------------------------------------===//
 
   /// Get the terminator operation of this block. This function asserts that
-  /// the block has a valid terminator operation.
+  /// the block has a valid terminator operation: an op with the IsTerminator
+  /// trait or an unregistered op.
   Operation *getTerminator();
 
+  /// Return true if this block has a terminator op with the IsTerminator trait.
+  bool hasRegisteredTerminator();
+
   //===--------------------------------------------------------------------===//
   // Predecessors and successors.
   //===--------------------------------------------------------------------===//
 
   // Predecessor iteration.
   using pred_iterator = PredecessorIterator;
   pred_iterator pred_begin() {
     return pred_iterator((BlockOperand *)getFirstUse());

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, OpBuilder &b,
 
     // Create new TiledLoopOp.
     auto newTiledLoopOp = rewriter.create<TiledLoopOp>(
         tiledLoopOp.getLoc(), tiledLoopOp.lowerBound(),
         tiledLoopOp.upperBound(), tiledLoopOp.step(), newInputs, newOutputs,
         tiledLoopOp.iterator_types(), tiledLoopOp.distribution_types());
 
     // Remove terminator.
-    if (!newTiledLoopOp.getBody()->empty())
+    if (newTiledLoopOp.getBody()->hasRegisteredTerminator())
       rewriter.eraseOp(tiledLoopOp.getBody()->getTerminator());
 
     // Compute new loop body arguments.
     SmallVector<Value> newBlockArgs, newRegionInOutArgs, oldRegionInOutArgs;
     ValueRange newInductionVars = newTiledLoopOp.getInductionVars();
     newBlockArgs.append(newInductionVars.begin(), newInductionVars.end());
 
     ValueRange newRegionInArgs = newTiledLoopOp.getRegionInputArgs();

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, OpBuilder &b,
     }
 
     // Create new op.
     auto newIfOp =
         rewriter.create<scf::IfOp>(ifOp.getLoc(), newTypes, ifOp.getCondition(),
                                    /*withElseRegion=*/true);
 
     // Remove terminators.
-    if (!newIfOp.thenBlock()->empty()) {
+    if (newIfOp.thenBlock()->hasRegisteredTerminator()) {
       rewriter.eraseOp(newIfOp.thenBlock()->getTerminator());
       rewriter.eraseOp(newIfOp.elseBlock()->getTerminator());
     }
 
     // Move over then/else blocks.
     rewriter.mergeBlocks(ifOp.thenBlock(), newIfOp.thenBlock());
     rewriter.mergeBlocks(ifOp.elseBlock(), newIfOp.elseBlock());
 
LogicalResult bufferize(Operation *op, OpBuilder & /*b*/,
     rewriter.setInsertionPointToStart(loopBody);
     SmallVector<Value> iterArgs =
         convert(newForOp.getRegionIterArgs(), [&](Value val, int64_t index) {
           return rewriter.create<bufferization::ToTensorOp>(val.getLoc(), val);
         });
     iterArgs.insert(iterArgs.begin(), newForOp.getInductionVar());
 
     // Erase terminator if present.
-    if (iterArgs.size() == 1)
+    if (loopBody->hasRegisteredTerminator())
       rewriter.eraseOp(loopBody->getTerminator());
 
     // Move loop body to new loop.
     rewriter.mergeBlocks(oldLoopBody, loopBody, iterArgs);
 
     // Update scf.yield of new loop.
     auto yieldOp = cast<scf::YieldOp>(loopBody->getTerminator());
     rewriter.setInsertionPoint(yieldOp);

mlir/lib/IR/AsmPrinter.cpp

if (printBlockArgs) {
         if (printerFlags.shouldPrintDebugInfo())
           // TODO: Allow deferring argument locations.
           initializer.visit(arg.getLoc(), /*canBeDeferred=*/false);
       }
     }
 
     // Consider the operations within this block, ignoring the terminator if
     // requested.
-    bool hasTerminator =
-        !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
     auto range = llvm::make_range(
         block->begin(),
         std::prev(block->end(),
-                  (!hasTerminator || printBlockTerminator) ? 0 : 1));
+                  (!block->hasRegisteredTerminator() || printBlockTerminator)
+                      ? 0
+                      : 1));
     for (Operation &op : range)
       print(&op);
   }
 
   /// Print the given region.
   void printRegion(Region &region, bool printEntryBlockArgs,
                    bool printBlockTerminators,
                    bool printEmptyBlock = false) override {
if (!block->getParent()) {
       interleaveComma(predIDs, [&](std::pair<unsigned, Block *> pred) {
         printBlockName(pred.second);
       });
     }
     os << newLine;
   }
 
   currentIndent += indentWidth;
-  bool hasTerminator =
-      !block->empty() && block->back().hasTrait<OpTrait::IsTerminator>();
   auto range = llvm::make_range(
       block->begin(),
-      std::prev(block->end(),
-                (!hasTerminator || printBlockTerminator) ? 0 : 1));
+      std::prev(
+          block->end(),
+          (!block->hasRegisteredTerminator() || printBlockTerminator) ? 0 : 1));
   for (auto &op : range) {
     print(&op);
     os << newLine;
   }
   currentIndent -= indentWidth;
 }
 
 void OperationPrinter::printValueID(Value value, bool printResultNo,

mlir/lib/IR/Block.cpp

void Block::eraseArguments(function_ref<bool(BlockArgument)> shouldEraseFn) {
   arguments.erase(firstDead, arguments.end());
 }
 
 //===----------------------------------------------------------------------===//
 // Terminator management
 //===----------------------------------------------------------------------===//
 
 /// Get the terminator operation of this block. This function asserts that
-/// the block has a valid terminator operation.
+/// the block has a valid terminator operation: an op with the IsTerminator
+/// trait or an unregistered op.
 Operation *Block::getTerminator() {
   assert(!empty() && back().mightHaveTrait<OpTrait::IsTerminator>());
   return &back();
 }
 
+/// Return true if this block has a terminator op with the IsTerminator trait.
+bool Block::hasRegisteredTerminator() {
+  return !empty() && back().hasTrait<OpTrait::IsTerminator>();
+}
+
 // Indexed successor access.
 unsigned Block::getNumSuccessors() {
   return empty() ? 0 : back().getNumSuccessors();
 }
 
 Block *Block::getSuccessor(unsigned i) {
   assert(i < getNumSuccessors());
   return getTerminator()->getSuccessor(i);

mlir/lib/IR/Operation.cpp

 void impl::ensureRegionTerminator(
     Region &region, OpBuilder &builder, Location loc,
     function_ref<Operation *(OpBuilder &, Location)> buildTerminatorOp) {
   OpBuilder::InsertionGuard guard(builder);
   if (region.empty())
     builder.createBlock(&region);
 
   Block &block = region.back();
-  if (!block.empty() && block.back().hasTrait<OpTrait::IsTerminator>())
+  if (block.hasRegisteredTerminator())
     return;
 
   builder.setInsertionPointToEnd(&block);
   builder.insert(buildTerminatorOp(builder, loc));
 }
 
 /// Create a simple OpBuilder and forward to the OpBuilder version of this
 /// function.
 void impl::ensureRegionTerminator(
     Region &region, Builder &builder, Location loc,
     function_ref<Operation *(OpBuilder &, Location)> buildTerminatorOp) {
   OpBuilder opBuilder(builder.getContext());
   ensureRegionTerminator(region, opBuilder, loc, buildTerminatorOp);
 }