diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td --- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td +++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td @@ -118,8 +118,11 @@ "apply_patterns.vector.lower_masked_transfers", [DeclareOpInterfaceMethods]> { 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. @@ -313,4 +316,23 @@ let assemblyFormat = "attr-dict"; } +def ApplyVectorReductionToContractPatternsOp : Op]> { + 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 diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp --- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp +++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp @@ -37,6 +37,11 @@ vector::populateFoldArithExtensionPatterns(patterns); } +void transform::ApplyVectorReductionToContractPatternsOp::populatePatterns( + RewritePatternSet &patterns) { + vector::populateVectorReductionToContractPatterns(patterns); +} + void transform::ApplyRankReducingSubviewPatternsOp::populatePatterns( RewritePatternSet &patterns) { vector::populateVectorTransferDropUnitDimsPatterns(patterns); diff --git a/mlir/test/Dialect/Linalg/masked_vectorization.mlir b/mlir/test/Dialect/Linalg/masked_vectorization.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/Linalg/masked_vectorization.mlir @@ -0,0 +1,39 @@ +// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s + +// CHECK-LABEL: masked_matmul +func.func @masked_matmul(%module: memref, %arg1: memref, %arg2: memref) { + + // CHECK: %[[MLHS:.*]] = vector.create_mask {{.*}} : vector<8x8xi1> + // CHECK: %[[LHS:.*]] = vector.transfer_read %{{.*}}, %[[MLHS]] {in_bounds = [true, true]} : memref>, vector<8x8xf32> + // CHECK: %[[MRHS:.*]] = vector.create_mask {{.*}} : vector<8x8xi1> + // CHECK: %[[RHS:.*]] = vector.transfer_read %{{.*}}, %[[MRHS]] {in_bounds = [true, true]} : memref>, vector<8x8xf32> + // CHECK: %[[MACC:.*]] = vector.create_mask {{.*}} : vector<8x8xi1> + // CHECK: %[[ACC:.*]] = vector.transfer_read {{.*}}, %[[MACC]] {in_bounds = [true, true]} : memref>, 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> + linalg.matmul ins(%module, %arg1 : memref, memref) outs(%arg2 : memref) + 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 +}