[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

323

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

def ApplyLowerMasksPatternsOp : Op<Transform_Dialect,
 
   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.
+    Apply opt-in patterns that lower vector.mask operations surrounding 
+    side-effecting ops:
+      - MaskedTransferReadOpPattern
+      - MaskedTransferWriteOpPattern
+      - MaskedGatherOpPattern
 
     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 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 = [{
+    Apply opt-in patterns that convert reductions to contract:
+      - MultiReduceToContract
+      - CombineContractBroadcast
+      - CombineContractABTranspose
+      - CombineContractResultTranspose
+      - ReorderCastOpsOnBroadcast
+      - ReorderElementwiseOpsOnTranspose
+
+    These patterns have the effect of rewriting a vector.multi_reduce into a 
+    vector.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);

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_masked_transfers
+    transform.apply_patterns.vector.transfer_permutation_patterns
+    transform.apply_patterns.vector.reduction_to_contract
+  } : !transform.any_op
+}