[mlir][vector][NFC] Remove terminator from vector.mask

[springerm]

Summary

Use GraphRegionNoTerminator.traits instead.

[springerm]

There is one problem with this change: The wrapped op could DCE away.

[dcaballe]

Thanks, @springerm! I haven't used this trait before. What happens when the maskable op has results? Are they identified as results of the vector.mask op?
There are also cases where we may accidentally and temporarily have multiple operations within the same vector.mask, until the vector.mask is "legalized" (see reviews.llvm.org/D141559). Would those cases still work with this approach?

[springerm]

This works with results but the op could DCE away because there is no user of the results. So i think we actually need the terminator.

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

 //===----------------------------------------------------------------------===//
 
 #ifndef VECTOR_OPS
 #define VECTOR_OPS
 
 include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.td"
 include "mlir/Dialect/Vector/Interfaces/MaskingOpInterface.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/IR/RegionKindInterface.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/VectorInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 def Vector_Dialect : Dialect {
def Vector_CreateMaskOp :
   }];
 
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
   let assemblyFormat = "$operands attr-dict `:` type(results)";
 }
 
 def Vector_MaskOp : Vector_Op<"mask", [
-  SingleBlockImplicitTerminator<"vector::YieldOp">,
   DeclareOpInterfaceMethods<MaskingOpInterface>,
   RecursiveMemoryEffects, NoRegionArguments
-]> {
+] # GraphRegionNoTerminator.traits> {
   let summary = "Predicates a maskable vector operation";
   let description = [{
     The `vector.mask` is a `MaskingOpInterface` operation that predicates the
     execution of another operation. It takes an `i1` vector mask and an
     optional passthru vector as arguments.
-    A `vector.yield`-terminated region encloses the operation to be masked.
-    Values used within the region are captured from above. Only one *maskable*
-    operation can be masked with a `vector.mask` operation at a time. An
-    operation is *maskable* if it implements the `MaskableOpInterface`.
+    A region encloses exactly one operation to be masked. Values used within
+    the region are captured from above. Only one *maskable* operation can be
+    masked with a `vector.mask` operation at a time. An operation is *maskable*
+    if it implements the `MaskableOpInterface`.
 
     The vector mask argument holds a bit for each vector lane and determines
     which vector lanes should execute the maskable operation and which ones
     should not. The `vector.mask` operation returns the value produced by the
     masked execution of the nested operation, if any. The masked-off lanes in
     the result vector are taken from the corresponding lanes of the pass-thru
     argument, if provided, or left unmodified, otherwise.
 

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

void MaskOp::build(
   result.addTypes(resultType);
 }
 
 ParseResult MaskOp::parse(OpAsmParser &parser, OperationState &result) {
   // Create the op region.
   result.regions.reserve(1);
   Region &maskRegion = *result.addRegion();
 
-  auto &builder = parser.getBuilder();
-
   // Parse all the operands.
   OpAsmParser::UnresolvedOperand mask;
   if (parser.parseOperand(mask))
     return failure();
 
   // Optional passthru operand.
   OpAsmParser::UnresolvedOperand passthru;
   ParseResult parsePassthru = parser.parseOptionalComma();
   if (parsePassthru.succeeded() && parser.parseOperand(passthru))
     return failure();
 
   // Parse op region.
   if (parser.parseRegion(maskRegion, /*arguments=*/{}, /*argTypes=*/{}))
     return failure();
-
-  MaskOp::ensureTerminator(maskRegion, builder, result.location);
+  // Make sure that there is at least one block so that we generate a better
+  // error message if the region is empty.
+  if (maskRegion.empty())
+    maskRegion.emplaceBlock();
 
   // Parse the optional attribute list.
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   // Parse all the types.
   Type maskType;
   if (parser.parseColonType(maskType))
ParseResult MaskOp::parse(OpAsmParser &parser, OperationState &result) {
 
   return success();
 }
 
 void mlir::vector::MaskOp::print(OpAsmPrinter &p) {
   p << " " << getMask();
   if (getPassthru())
     p << ", " << getPassthru();
-
+  p << " ";
   // Print single masked operation and skip terminator.
   p << " { ";
   Block *singleBlock = &getMaskRegion().getBlocks().front();
-  if (singleBlock && singleBlock->getOperations().size() > 1)
+  if (singleBlock && singleBlock->getOperations().size() > 0)
     p.printCustomOrGenericOp(&singleBlock->front());
   p << " }";
 
   p.printOptionalAttrDict(getOperation()->getAttrs());
 
   p << " : " << getMask().getType();
   if (getNumResults() > 0)
     p << " -> " << getResultTypes();
 }
 
-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.
-  if (region.front().getOperations().size() != 2)
-    return;
-
-  // Replace default yield terminator with a new one that returns the results
-  // from the masked operation.
-  OpBuilder opBuilder(builder.getContext());
-  Operation *maskedOp = &region.front().front();
-  Operation *oldYieldOp = &region.front().back();
-  assert(isa<vector::YieldOp>(oldYieldOp) && "Expected vector::YieldOp");
-
-  opBuilder.setInsertionPoint(oldYieldOp);
-  opBuilder.create<vector::YieldOp>(maskedOp->getLoc(), maskedOp->getResults());
-  oldYieldOp->dropAllReferences();
-  oldYieldOp->erase();
-}
-
 LogicalResult MaskOp::verify() {
   // Structural checks.
   Block &block = getMaskRegion().getBlocks().front();
-  if (block.getOperations().size() < 2)
+  if (block.getOperations().size() == 0)
     return emitOpError("expects an operation to mask");
-  if (block.getOperations().size() > 2)
+  if (block.getOperations().size() > 1)
     return emitOpError("expects only one operation to mask");
 
   auto maskableOp = dyn_cast<MaskableOpInterface>(block.front());
   if (!maskableOp)
     return emitOpError("expects a maskable operation");
 
   // Result checks.
   if (maskableOp->getNumResults() != getNumResults())