[mlir][IR] Fix Block::without_terminator for blocks without terminator

[matthias-springer]

Blocks without a terminator are not handled correctly by Block::without_terminator: the last operation is excluded, even though it is not a terminator. With this commit, only terminators are excluded. If the last operation is unregistered, it is included for safety.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

Blocks without a terminator are not handled correctly by Block::without_terminator: the last operation is excluded, even though it is not a terminator. With this commit, only terminators are excluded. If the last operation is unregistered, it is included for safety.

---

Full diff: https://github.com/llvm/llvm-project/pull/154498.diff

2 Files Affected:

• (modified) mlir/include/mlir/IR/Block.h (+11-1)
• (modified) mlir/lib/IR/Block.cpp (+9)

diff --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index e486bb627474d..07281b0ad9ba0 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -205,10 +205,15 @@ class alignas(8) Block : public IRObjectWithUseList<BlockOperand>,
   }
 
   /// Return an iterator range over the operation within this block excluding
-  /// the terminator operation at the end.
+  /// the terminator operation at the end. If the block has no terminator,
+  /// return an iterator range over the entire block. If it is unknown if the
+  /// block has a terminator (i.e., last block operation is unregistered), also
+  /// return an iterator range over the entire block.
   iterator_range<iterator> without_terminator() {
     if (begin() == end())
       return {begin(), end()};
+    if (hasNoTerminator())
+      return {begin(), end()};
     auto endIt = --end();
     return {begin(), endIt};
   }
@@ -224,6 +229,11 @@ class alignas(8) Block : public IRObjectWithUseList<BlockOperand>,
   /// Check whether this block might have a terminator.
   bool mightHaveTerminator();
 
+  /// Return "true" if this block has no terminator. Return "false" if the last
+  /// operation is unregistered. In that case, the presence of a terminator
+  /// cannot be determined.
+  bool hasNoTerminator();
+
   //===--------------------------------------------------------------------===//
   // Predecessors and successors.
   //===--------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index 57825d9b42178..383847f5d9a66 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -251,6 +251,15 @@ bool Block::mightHaveTerminator() {
   return !empty() && back().mightHaveTrait<OpTrait::IsTerminator>();
 }
 
+/// Check whether this block has no terminator.
+bool Block::hasNoTerminator() {
+  if (empty())
+    return true;
+  if (!back().isRegistered())
+    return false;
+  return back().hasTrait<OpTrait::IsTerminator>();
+}
+
 // Indexed successor access.
 unsigned Block::getNumSuccessors() {
   return empty() ? 0 : back().getNumSuccessors();

[joker-eph]

When does this happen? How could this be exercised upstream?

[matthias-springer]

This issue came up when I was debugging code in a downstream project. The code there iterated over all ops, except for the terminator (for dubious reasons...). The block belonged to an op with NoTerminator. Block::without_terminator did not include the last operation (when it should have).

I found a way to test this via sortTopologically: based on this API change, the implementation there can be simplified.