[Mlir-commits] [mlir] [mlir][ArmSME] Fold MoveTileSliceToVector + TransferWrite to StoreTileSlice (PR #95907)
Benjamin Maxwell
llvmlistbot at llvm.org
Tue Jun 18 10:08:15 PDT 2024
https://github.com/MacDue updated https://github.com/llvm/llvm-project/pull/95907
>From 7be8e7052990f859803dea972231a7671ea09248 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 18 Jun 2024 11:09:02 +0000
Subject: [PATCH 1/2] [mlir][ArmSME] Fold MoveTileSliceToVector + TransferWrite
to StoreTileSlice
---
.../VectorToArmSME/VectorToArmSME.cpp | 62 +++++++++++++++++--
.../VectorToArmSME/vector-to-arm-sme.mlir | 46 ++++++++++++++
2 files changed, 102 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp b/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
index c2f1584e43bac..157dd73c85146 100644
--- a/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
+++ b/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
@@ -666,14 +666,64 @@ struct VectorPrintToArmSMELowering : public OpRewritePattern<vector::PrintOp> {
}
};
+/// Folds a MoveTileSliceToVectorOp + TransferWriteOp to a StoreTileSliceOp.
+///
+/// BEFORE:
+/// ```mlir
+/// %slice = arm_sme.move_tile_slice_to_vector %tile[%index]
+/// : vector<[4]xf32> from vector<[4]x[4]xf32>
+/// vector.transfer_write %slice, %memref[%i, %j], %mask {in_bounds = [true]}
+/// : vector<[4]xf32>, memref<?x?xf32>
+/// ```
+/// AFTER:
+/// ```mlir
+/// arm_sme.store_tile_slice %tile, %index, %mask, %memref[%i, %j]
+/// : memref<?x?xf32>, vector<[4]xi1>, vector<[4]x[4]xf32>
+/// ```
+struct FoldTransferWriteOfExtractTileSlice
+ : public OpRewritePattern<vector::TransferWriteOp> {
+ using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
+ PatternRewriter &rewriter) const final {
+ if (!isa<MemRefType>(writeOp.getSource().getType()))
+ return failure();
+
+ auto moveTileSlice =
+ writeOp.getVector().getDefiningOp<arm_sme::MoveTileSliceToVectorOp>();
+ if (!moveTileSlice)
+ return failure();
+
+ AffineMap map = writeOp.getPermutationMap();
+ if (!map.isMinorIdentity())
+ return rewriter.notifyMatchFailure(writeOp,
+ "unsupported permutation map");
+
+ Value mask = writeOp.getMask();
+ if (!mask) {
+ auto maskType = writeOp.getVectorType().clone(rewriter.getI1Type());
+ mask = rewriter.create<arith::ConstantOp>(
+ writeOp.getLoc(), maskType, DenseElementsAttr::get(maskType, true));
+ }
+
+ rewriter.replaceOpWithNewOp<arm_sme::StoreTileSliceOp>(
+ writeOp, moveTileSlice.getTile(), moveTileSlice.getTileSliceIndex(),
+ mask, writeOp.getSource(), writeOp.getIndices(),
+ moveTileSlice.getLayout());
+ return success();
+ }
+};
+
} // namespace
void mlir::populateVectorToArmSMEPatterns(RewritePatternSet &patterns,
MLIRContext &ctx) {
- patterns.add<BroadcastOpToArmSMELowering, SplatOpToArmSMELowering,
- TransferReadToArmSMELowering, TransferWriteToArmSMELowering,
- TransposeOpToArmSMELowering, VectorLoadToArmSMELowering,
- VectorStoreToArmSMELowering, VectorOuterProductToArmSMELowering,
- VectorExtractToArmSMELowering, VectorInsertToArmSMELowering,
- VectorPrintToArmSMELowering>(&ctx);
+ patterns
+ .add<BroadcastOpToArmSMELowering, SplatOpToArmSMELowering,
+ TransferReadToArmSMELowering, TransferWriteToArmSMELowering,
+ TransposeOpToArmSMELowering, VectorLoadToArmSMELowering,
+ VectorStoreToArmSMELowering, VectorOuterProductToArmSMELowering,
+ VectorExtractToArmSMELowering, VectorInsertToArmSMELowering,
+ VectorPrintToArmSMELowering, FoldTransferWriteOfExtractTileSlice>(
+ &ctx);
}
diff --git a/mlir/test/Conversion/VectorToArmSME/vector-to-arm-sme.mlir b/mlir/test/Conversion/VectorToArmSME/vector-to-arm-sme.mlir
index f22b6de52f367..548dfcc305296 100644
--- a/mlir/test/Conversion/VectorToArmSME/vector-to-arm-sme.mlir
+++ b/mlir/test/Conversion/VectorToArmSME/vector-to-arm-sme.mlir
@@ -334,6 +334,52 @@ func.func @transfer_write_2d_transpose_with_mask_bf16(%vector : vector<[8]x[8]xb
return
}
+// -----
+
+// CHECK-LABEL: func.func @transfer_write_slice(
+// CHECK-SAME: %[[VECTOR:.*]]: vector<[4]x[4]xf32>,
+// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>,
+// CHECK-SAME: %[[INDEX:.*]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<[4]xi1>
+// CHECK: arm_sme.store_tile_slice %[[VECTOR]], %[[INDEX]], %[[MASK]], %[[DEST]][%[[INDEX]], %[[C0]]] : memref<?x?xf32>, vector<[4]xi1>, vector<[4]x[4]xf32>
+func.func @transfer_write_slice(%vector: vector<[4]x[4]xf32>, %dest : memref<?x?xf32>, %slice_index: index) {
+ %c0 = arith.constant 0 : index
+ %slice = vector.extract %vector[%slice_index] : vector<[4]xf32> from vector<[4]x[4]xf32>
+ vector.transfer_write %slice, %dest[%slice_index, %c0] { in_bounds = [true] }: vector<[4]xf32>, memref<?x?xf32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @transfer_write_slice_with_mask(
+// CHECK-SAME: %[[VECTOR:.*]]: vector<[4]x[4]xf32>,
+// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>,
+// CHECK-SAME: %[[MASK:.*]]: vector<[4]xi1>,
+// CHECK-SAME: %[[INDEX:.*]]: index) {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: arm_sme.store_tile_slice %[[VECTOR]], %[[INDEX]], %[[MASK]], %[[DEST]][%[[INDEX]], %[[C0]]] : memref<?x?xf32>, vector<[4]xi1>, vector<[4]x[4]xf32>
+func.func @transfer_write_slice_with_mask(%vector: vector<[4]x[4]xf32>, %dest : memref<?x?xf32>, %mask: vector<[4]xi1>, %slice_index: index) {
+ %c0 = arith.constant 0 : index
+ %slice = vector.extract %vector[%slice_index] : vector<[4]xf32> from vector<[4]x[4]xf32>
+ vector.transfer_write %slice, %dest[%slice_index, %c0], %mask { in_bounds = [true] }: vector<[4]xf32>, memref<?x?xf32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @transfer_write_vertical_slice
+// CHECK: arm_sme.store_tile_slice {{.*}} layout<vertical>
+func.func @transfer_write_vertical_slice(%vector: vector<[4]x[4]xf32>, %dest : memref<?x?xf32>, %slice_index: index) {
+ %c0 = arith.constant 0 : index
+ %slice = arm_sme.move_tile_slice_to_vector %vector[%slice_index] layout<vertical>
+ : vector<[4]xf32> from vector<[4]x[4]xf32>
+ vector.transfer_write %slice, %dest[%slice_index, %c0] { in_bounds = [true] }: vector<[4]xf32>, memref<?x?xf32>
+ return
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// vector.broadcast
//===----------------------------------------------------------------------===//
>From fb7ddac678b485d2f91d4e71b7920f1f3357e21c Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 18 Jun 2024 17:05:42 +0000
Subject: [PATCH 2/2] Fixups
---
mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp | 9 +++++++--
1 file changed, 7 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp b/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
index 157dd73c85146..56ae46a6098ee 100644
--- a/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
+++ b/mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp
@@ -687,12 +687,17 @@ struct FoldTransferWriteOfExtractTileSlice
LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
PatternRewriter &rewriter) const final {
if (!isa<MemRefType>(writeOp.getSource().getType()))
- return failure();
+ return rewriter.notifyMatchFailure(writeOp, "destination not a memref");
+
+ if (writeOp.hasOutOfBoundsDim())
+ return rewriter.notifyMatchFailure(writeOp,
+ "not inbounds transfer write");
auto moveTileSlice =
writeOp.getVector().getDefiningOp<arm_sme::MoveTileSliceToVectorOp>();
if (!moveTileSlice)
- return failure();
+ return rewriter.notifyMatchFailure(
+ writeOp, "vector to store not from MoveTileSliceToVectorOp");
AffineMap map = writeOp.getPermutationMap();
if (!map.isMinorIdentity())
More information about the Mlir-commits
mailing list