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 @@ -114,6 +114,17 @@ let assemblyFormat = "attr-dict"; } +def ApplyLowerMasksForSideEffectingOpsPatternsOp : Op]> { + let description = [{ + Indicates that vector.mask operations surrounding side-effecting ops should + be lowered. + }]; + + let assemblyFormat = "attr-dict"; +} + def ApplyLowerMaskedTransfersPatternsOp : Op]> { @@ -313,4 +324,14 @@ let assemblyFormat = "attr-dict"; } +def ApplyVectorReductionToContractPatternsOp : Op]> { + let description = [{ + Collect a set of patterns that convert reductions to 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); @@ -66,6 +71,11 @@ populateVectorMaskOpLoweringPatterns(patterns); } +void transform::ApplyLowerMasksForSideEffectingOpsPatternsOp::populatePatterns( + RewritePatternSet &patterns) { + populateVectorMaskLoweringPatternsForSideEffectingOps(patterns); +} + void transform::ApplyLowerMaskedTransfersPatternsOp::populatePatterns( RewritePatternSet &patterns) { populateVectorMaskLoweringPatternsForSideEffectingOps(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_masks_for_side_effecting_ops + transform.apply_patterns.vector.transfer_permutation_patterns + transform.apply_patterns.vector.reduction_to_contract + } : !transform.any_op +}