[mlir][transform] Expose vector patterns useful for cleaning up masked vectorization and add test to showcase composition

[ftynse]

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

331

Can this be made more similar to the other descriptions? It's unclear where are the patterns collected.

Ideally, we'd have some more detailed description for all of these, potentially with an example or reference to code that contains an example. But it's the issue with everything in this file, so not for this commit.

[nicolasvasilache]

Address

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

let description = [{
 
     This is usually a late step that is run after bufferization as part of the
     process of lowering to e.g. LLVM or NVVM.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyLowerMasksForSideEffectingOpsPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.vector.lower_masks_for_side_effecting_ops",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Indicates that vector.mask operations surrounding side-effecting ops should
+    be lowered.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 def ApplyLowerMaskedTransfersPatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.lower_masked_transfers",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that masked vector.transfer and vector.gather operations should
     be lowered to finer-grained vector primitives.
 
     This is usually a late step that is run after bufferization as part of the
def ApplyFoldArithExtensionPatternsOp : Op<Transform_Dialect,
   let description = [{
     Collect a set of patterns that fold arithmetic extension on floating point
     into vector contract for the backends with native support.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyVectorReductionToContractPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.vector.reduction_to_contract",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Collect a set of patterns that convert reductions to contract.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 #endif // VECTOR_TRANSFORM_OPS

mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp

void transform::ApplyCastAwayVectorLeadingOneDimPatternsOp::populatePatterns(
   vector::populateCastAwayVectorLeadingOneDimPatterns(patterns);
 }
 
 void transform::ApplyFoldArithExtensionPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   vector::populateFoldArithExtensionPatterns(patterns);
 }
 
+void transform::ApplyVectorReductionToContractPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  vector::populateVectorReductionToContractPatterns(patterns);
+}
+
 void transform::ApplyRankReducingSubviewPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   vector::populateVectorTransferDropUnitDimsPatterns(patterns);
 }
 
 void transform::ApplyTransferPermutationPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
populateVectorContractLoweringPatterns(patterns, vectorTransformOptions,
                                          /*disableOuterProductLowering=*/true);
 }
 
 void transform::ApplyLowerMasksPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   populateVectorMaskOpLoweringPatterns(patterns);
 }
 
+void transform::ApplyLowerMasksForSideEffectingOpsPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  populateVectorMaskLoweringPatternsForSideEffectingOps(patterns);
+}
+
 void transform::ApplyLowerMaskedTransfersPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   populateVectorMaskLoweringPatternsForSideEffectingOps(patterns);
 }
 
 void transform::ApplyMaterializeMasksPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   populateVectorMaskMaterializationPatterns(patterns,

mlir/test/Dialect/Linalg/masked_vectorization.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
+
+//      CHECK-LABEL: masked_matmul
+func.func @masked_matmul(%module: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
+
+  //      CHECK: %[[MLHS:.*]] = vector.create_mask {{.*}} : vector<8x8xi1>
+  //      CHECK: %[[LHS:.*]] = vector.transfer_read %{{.*}}, %[[MLHS]] {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, 1], offset: ?>>, vector<8x8xf32> 
+  //      CHECK: %[[MRHS:.*]] = vector.create_mask {{.*}} : vector<8x8xi1> 
+  //      CHECK: %[[RHS:.*]] = vector.transfer_read %{{.*}}, %[[MRHS]] {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, 1], offset: ?>>, vector<8x8xf32>
+  //      CHECK: %[[MACC:.*]] = vector.create_mask {{.*}} : vector<8x8xi1>
+  //      CHECK: %[[ACC:.*]] = vector.transfer_read {{.*}}, %[[MACC]] {in_bounds = [true, true]} : memref<?x?xf32, strided<[?, 1], offset: ?>>, vector<8x8xf32> 
+  //      CHECK: %[[MRES:.*]] = vector.create_mask {{.*}} : vector<8x8x8xi1>
+  //      CHECK: %[[RES:.*]] = vector.mask %[[MRES]] { vector.contract
+  // CHECK-SAME:   : vector<8x8xf32>, vector<8x8xf32> into vector<8x8xf32>
+  // CHECK-SAME:   : vector<8x8x8xi1> -> vector<8x8xf32>
+  //      CHECK: vector.transfer_write %[[RES]], %{{.*}}, %[[MACC]] {in_bounds = [true, true]} : vector<8x8xf32>, memref<?x?xf32, strided<[?, 1], offset: ?>> 
+  linalg.matmul ins(%module, %arg1 : memref<?x?xf32>, memref<?x?xf32>) outs(%arg2 : memref<?x?xf32>)
+  return
+}
+
+transform.sequence  failures(propagate) {
+^bb0(%module: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %module 
+    : (!transform.any_op) -> !transform.any_op
+  %tiled_linalg_op, %loops:3 = transform.structured.tile_to_scf_for %0[64, 128, 256] 
+    : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+  %tiled_linalg_op_0, %loops_1:3 = transform.structured.tile_to_scf_for %tiled_linalg_op[8, 8, 8] 
+    : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+  transform.structured.masked_vectorize %tiled_linalg_op_0 vector_sizes [8, 8, 8] 
+    : !transform.any_op
+
+  %func = transform.structured.match ops{["func.func"]} in %module 
+    : (!transform.any_op) -> !transform.any_op
+  apply_patterns to %func {
+    transform.apply_patterns.vector.lower_masks_for_side_effecting_ops
+    transform.apply_patterns.vector.transfer_permutation_patterns
+    transform.apply_patterns.vector.reduction_to_contract
+  } : !transform.any_op
+}