[mlir][vector] Support multiple result types in vector.mask

[springerm]

Summary

The verifier already had support for multiple result types, but the op definition assumed a single, optional result.

[dcaballe]

Thanks for extending this! I would've loved to have time to review this before landing. Limiting vector.mask to single result ops was intentional. To support maskable ops with multiple results we also have to allow multiple passthru values. We have to verify that results and passthru values match both in number and type. Could you also add a test with multiple results? You can use an operation from the test dialect, I guess. That's mostly why I didn't add it in first place: I couldn't find a actual use case.

[springerm]

In D141683#4052448, @dcaballe wrote:

Thanks for extending this! I would've loved to have time to review this before landing. Limiting vector.mask to single result ops was intentional. To support maskable ops with multiple results we also have to allow multiple passthru values. We have to verify that results and passthru values match both in number and type. Could you also add a test with multiple results? You can use an operation from the test dialect, I guess. That's mostly why I didn't add it in first place: I couldn't find a actual use case.

I don't have a use case for multiple vector results. But this op should return all results of the masked operation. In particular, for bufferization purposes, we need any tensor values that are returned by the masked op.

Imagine a hypothetical vector op that reads a vector, writes a new vector to the location, then returns the original vector.

%r = vector.mask { "vector.read_update_and_write"(%v) : (tensor<?xf32>, vector<5xf32>) -> (vector<5xf32>, tensor<?xf32>) }

Now we have two return values. But only one vector result. I didn't think about the mask, passthru values when making this change. How about we limit the number of vector results to 1 for now? https://reviews.llvm.org/D141786

[dcaballe]

Yep, I couldn't find a use case or way to test it either so that's way I didn't add it. Also didn't want to deal with the multi-passthrough case. The mask should still be one for multi return cases, I hope.
Yep, it sounds good. We could support multiple results in the API and restrict to one in the verifier. In that way the API is more ergonomic. Thanks for doing this!

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

def Vector_MaskOp : Vector_Op<"mask", [
   DeclareOpInterfaceMethods<MaskingOpInterface>,
   RecursiveMemoryEffects, NoRegionArguments
 ]> {
   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 implicitly `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`. The
+    terminator yields all results of the maskable operation to the result of
+    this operation.
 
     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.
 
let description = [{
 
     ```
       %0 = vector.mask %mask, %passthru { arith.divsi %a, %b : vector<8xi32> } : vector<8xi1> -> vector<8xi32>
     ```
 
     ```
       vector.mask %mask { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, memref<?xf32> } : vector<16xi1>
     ```
+
+    ```
+      vector.mask %mask { vector.transfer_write %val, %t0[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
+    ```
   }];
 
   // TODO: Support multiple results and passthru values.
   let arguments = (ins VectorOf<[I1]>:$mask,
                    Optional<AnyType>:$passthru);
-  let results = (outs Optional<AnyType>:$results);
+  let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$maskRegion);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "Value":$mask,
                    CArg<"function_ref<void(OpBuilder &, Location)>",
                         "buildTerminatedBody">:$maskRegion)>,
-    OpBuilder<(ins "Type":$resultType, "Value":$mask,
+    OpBuilder<(ins "TypeRange":$resultTypes, "Value":$mask,
                    CArg<"function_ref<void(OpBuilder &, Location)>",
                         "buildTerminatedBody">:$maskRegion)>,
-    OpBuilder<(ins "Type":$resultType, "Value":$mask,
+    OpBuilder<(ins "TypeRange":$resultTypes, "Value":$mask,
                    "Value":$passthru,
                    CArg<"function_ref<void(OpBuilder &, Location)>",
                         "buildTerminatedBody">:$maskRegion)>
   ];
 
   let extraClassDeclaration = [{
     static void ensureTerminator(Region &region, Builder &builder, Location loc);
   }];

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

void MaskOp::build(
   result.addOperands(mask);
   OpBuilder::InsertionGuard guard(builder);
   Region *maskRegion = result.addRegion();
   builder.createBlock(maskRegion);
   maskRegionBuilder(builder, result.location);
 }
 
 void MaskOp::build(
-    OpBuilder &builder, OperationState &result, Type resultType, Value mask,
-    function_ref<void(OpBuilder &, Location)> maskRegionBuilder) {
-  build(builder, result, resultType, mask, /*passthru=*/Value(),
+    OpBuilder &builder, OperationState &result, TypeRange resultTypes,
+    Value mask, function_ref<void(OpBuilder &, Location)> maskRegionBuilder) {
+  build(builder, result, resultTypes, mask, /*passthru=*/Value(),
         maskRegionBuilder);
 }
 
 void MaskOp::build(
-    OpBuilder &builder, OperationState &result, Type resultType, Value mask,
-    Value passthru,
+    OpBuilder &builder, OperationState &result, TypeRange resultTypes,
+    Value mask, Value passthru,
     function_ref<void(OpBuilder &, Location)> maskRegionBuilder) {
   build(builder, result, mask, maskRegionBuilder);
   if (passthru)
     result.addOperands(passthru);
-  result.addTypes(resultType);
+  result.addTypes(resultTypes);
 }
 
 ParseResult MaskOp::parse(OpAsmParser &parser, OperationState &result) {
   // Create the op region.
   result.regions.reserve(1);
   Region &maskRegion = *result.addRegion();
 
   auto &builder = parser.getBuilder();