[mlir][Vector] Add folding for vector.mask with all-true masks

[dcaballe]

Summary

This patch removes vector.mask operations with all-true masks (i.e.,
all lanes enabled).

[nicolasvasilache]

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

5690

Should this just be a folding pattern or do you see particular reasons to not always apply this folding?

[hanchung]

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

5690

+1, it looks more like a folding pattern to me

[dcaballe]

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

5690

I tried that but the folding patters are pretty limited. I can't use a rewriter or move instructions around. I tried returning the maskable operations to replace the vector.mask one but it crashes as the maskable operation is nested within the operation to be replaced. Let me give it another try. Let me know if you have any suggestions in this regard.

[dcaballe]

Move pattern to fold.

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

5690

I figure it out :)

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

```
       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.
+  // TODO: Support multiple passthru values.
   let arguments = (ins VectorOf<[I1]>:$mask,
                    Optional<AnyType>:$passthru);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$maskRegion);
 
   let skipDefaultBuilders = 1;
   let builders = [
     OpBuilder<(ins "Value":$mask, "Operation *":$maskableOp,
                    CArg<"function_ref<void(OpBuilder &, Operation *)>">:$maskRegion)>,
     OpBuilder<(ins "TypeRange":$resultTypes, "Value":$mask, "Operation *":$maskableOp,
                    CArg<"function_ref<void(OpBuilder &, Operation *)>">:$maskRegion)>,
     OpBuilder<(ins "TypeRange":$resultTypes, "Value":$mask, "Value":$passthru,
                    "Operation *":$maskableOp,
                    CArg<"function_ref<void(OpBuilder &, Operation *)>">:$maskRegion)>
   ];
 
   let extraClassDeclaration = [{
     Block *getMaskBlock() { return &getMaskRegion().front(); }
-    static void ensureTerminator(Region &region, Builder &builder, Location loc);
+
+    /// Returns true if mask op is not masking any operation.
+    bool isEmpty() {
+      Block *block = getMaskBlock();
+      if (block->getOperations().size() > 1)
+        return false;
+      return true;
+    }
+
+    static void ensureTerminator(Region &region, Builder &builder,
+                                 Location loc);
   }];
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def Vector_TransposeOp :
   Vector_Op<"transpose", [Pure,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>,
     PredOpTrait<"operand and result have same element type",

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

if (passthru) {
 
     if (passthru.getType() != maskableOp->getResultTypes()[0])
       return emitOpError("expects passthru type to match result type");
   }
 
   return success();
 }
 
+/// Folds vector.mask ops with an all-true mask.
+LogicalResult MaskOp::fold(FoldAdaptor adaptor,
+                           SmallVectorImpl<OpFoldResult> &results) {
+  MaskFormat maskFormat = getMaskFormat(getMask());
+  if (isEmpty())
+    return failure();
+
+  if (maskFormat != MaskFormat::AllTrue)
+    return failure();
+
+  // Move maskable operation outside of the `vector.mask` region.
+  Operation *maskableOp = getMaskableOp();
+  maskableOp->dropAllUses();
+  maskableOp->moveBefore(getOperation());
+
+  results.push_back(maskableOp->getResult(0));
+  return success();
+}
+
 // Elides empty vector.mask operations with or without return values. Propagates
 // the yielded values by the vector.yield terminator, if any, or erases the op,
 // otherwise.
 class ElideEmptyMaskOp : public OpRewritePattern<MaskOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(MaskOp maskOp,
                                 PatternRewriter &rewriter) const override {
     auto maskingOp = cast<MaskingOpInterface>(maskOp.getOperation());
     if (maskingOp.getMaskableOp())
       return failure();
 
-    Block *block = maskOp.getMaskBlock();
-    if (block->getOperations().size() > 1)
+    if (!maskOp.isEmpty())
       return failure();
 
+    Block *block = maskOp.getMaskBlock();
     auto terminator = cast<vector::YieldOp>(block->front());
     if (terminator.getNumOperands() == 0)
       rewriter.eraseOp(maskOp);
     else
       rewriter.replaceOp(maskOp, terminator.getOperands());
 
     return success();
   }
 };
 
 void MaskOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
   results.add<ElideEmptyMaskOp>(context);
 }
 
 // MaskingOpInterface definitions.
 
 /// Returns the operation masked by this 'vector.mask'.
 Operation *MaskOp::getMaskableOp() {
   Block *block = getMaskBlock();

mlir/test/Dialect/Vector/canonicalize.mlir

 //  CHECK-SAME:     %[[IN:.*]]: vector<8xf32>
 func.func @empty_vector_mask_with_return(%a : vector<8xf32>, %mask : vector<8xi1>) -> vector<8xf32> {
 //   CHECK-NOT:   vector.mask
 //       CHECK:   return %[[IN]] : vector<8xf32>
   %0 = vector.mask %mask { vector.yield %a : vector<8xf32> } : vector<8xi1> -> vector<8xf32>
   return %0 : vector<8xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @all_true_vector_mask
+//  CHECK-SAME:     %[[IN:.*]]: vector<3x4xf32>
+func.func @all_true_vector_mask(%a : vector<3x4xf32>) -> vector<3x4xf32> {
+//   CHECK-NOT:   vector.mask
+//       CHECK:   %[[ADD:.*]] = arith.addf %[[IN]], %[[IN]] : vector<3x4xf32>
+//       CHECK:   return %[[ADD]] : vector<3x4xf32>
+  %all_true = vector.constant_mask [3, 4] : vector<3x4xi1>
+  %0 = vector.mask %all_true { arith.addf %a, %a : vector<3x4xf32> } : vector<3x4xi1> -> vector<3x4xf32>
+  return %0 : vector<3x4xf32>
+}