[Mlir-commits] [mlir] [mlir][Vector] Improve `vector.mask` verifier (PR #139823)

[Diego Caballero]

================
@@ -6543,29 +6543,31 @@ void mlir::vector::MaskOp::print(OpAsmPrinter &p) {
 }
 
 void MaskOp::ensureTerminator(Region &region, Builder &builder, Location loc) {
-  OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
-      MaskOp>::ensureTerminator(region, builder, loc);
-  // Keep the default yield terminator if the number of masked operations is not
-  // the expected. This case will trigger a verification failure.
-  Block &block = region.front();
-  if (block.getOperations().size() != 2)
+  // Create default terminator if there are no ops to mask.
+  if (region.empty() || region.front().empty()) {
+    OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
+        MaskOp>::ensureTerminator(region, builder, loc);
     return;
+  }
 
-  // Replace default yield terminator with a new one that returns the results
-  // from the masked operation.
-  OpBuilder opBuilder(builder.getContext());
-  Operation *maskedOp = &block.front();
-  Operation *oldYieldOp = &block.back();
-  assert(isa<vector::YieldOp>(oldYieldOp) && "Expected vector::YieldOp");
+  // If  region has an explicit terminator, we don't modify it.
+  Block &block = region.front();
+  if (isa<vector::YieldOp>(block.back()))
+    return;
 
-  // Empty vector.mask op.
-  if (maskedOp == oldYieldOp)
+  // Create default terminator if the number of masked operations is not
+  // one. This case will trigger a verification failure.
+  if (block.getOperations().size() != 1) {
+    OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
+        MaskOp>::ensureTerminator(region, builder, loc);
     return;
+  }
 
-  opBuilder.setInsertionPoint(oldYieldOp);
+  // Create a terminator that yields the results from the masked operation.
+  OpBuilder opBuilder(builder.getContext());
+  Operation *maskedOp = &block.front();
+  opBuilder.setInsertionPointToEnd(&block);
----------------
dcaballe wrote:

I highlighted the three use cases

https://github.com/llvm/llvm-project/pull/139823