[Mlir-commits] [mlir] [mlir][ArmSME] Add rewrites to swap extract of extend (PR #80407)
Cullen Rhodes
llvmlistbot at llvm.org
Mon Feb 5 04:20:29 PST 2024
https://github.com/c-rhodes updated https://github.com/llvm/llvm-project/pull/80407
>From 8c42da9817d69e9cd161d5743fae4f14d5d9ff05 Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Thu, 1 Feb 2024 15:04:46 +0000
Subject: [PATCH 1/3] [mlir][ArmSME] Add rewrites to swap extract of extend
In mixed matmul lowering (e.g., i8 to i32) we're seeing the following
sequence:
%0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
%1 = vector.extract %0[0] : vector<[8]xi32> from vector<4x[8]xi32>
%lhs = vector.scalable.extract %1[0] : vector<[4]xi32> from vector<[8]xi32>
... (same for rhs)
%2 = vector.outerproduct %lhs, %rhs, %acc vector<[4]xi32>, vector<[4]xi32>
// x4 chained by accumulator
This chain of 4 outer products can be fused into a single 4-way widening
variant but the pass doesn't match on the IR, as it expects the source
of the inputs to be an extend and it can't look through the extracts.
This patch fixes this with two rewrites that swaps extract(extend) into
extend(extract).
Related to #78975, #79288.
---
.../ArmSME/Transforms/VectorLegalization.cpp | 104 +++++++++++++++++-
.../Dialect/ArmSME/vector-legalization.mlir | 40 +++++++
2 files changed, 143 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
index 14b9d8e34da65..47de4a8a857bf 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
@@ -415,6 +415,105 @@ struct FoldExtractFromVectorOfSMELikeCreateMasks
}
};
+// Shuffles arith extend ops after vector.extract op.
+//
+// This transforms IR like:
+// %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+// %1 = vector.extract %0[0] : vector<[8]xi32> from vector<4x[8]xi32>
+// Into:
+// %0 = vector.extract %src[0] : vector<[8]xi8> from vector<4x[8]xi8>
+// %1 = arith.extsi %0 : vector<[8]xi8> to vector<[8]xi32>
+//
+// This enables outer product fusion in the `-arm-sme-outer-product-fusion`
+// pass when the result is the input to an outer product.
+struct SwapVectorExtractOfArithExtend
+ : public OpRewritePattern<vector::ExtractOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::ExtractOp extractOp,
+ PatternRewriter &rewriter) const override {
+ VectorType resultType = llvm::dyn_cast<VectorType>(extractOp.getType());
+ if (!resultType)
+ return rewriter.notifyMatchFailure(extractOp,
+ "extracted type is not a vector type");
+
+ auto numScalableDims = llvm::count(resultType.getScalableDims(), true);
+ if (numScalableDims != 1)
+ return rewriter.notifyMatchFailure(
+ extractOp, "extracted type is not a 1-D scalable vector type");
+
+ auto *extendOp = extractOp.getVector().getDefiningOp();
+ if (!isa_and_present<arith::ExtSIOp, arith::ExtUIOp, arith::ExtFOp>(
+ extendOp))
+ return rewriter.notifyMatchFailure(extractOp,
+ "extract not from extend op");
+
+ auto loc = extractOp.getLoc();
+ StringAttr extendOpName = extendOp->getName().getIdentifier();
+ Value extendSource = extendOp->getOperand(0);
+
+ // Create new extract from source of extend.
+ Value newExtract = rewriter.create<vector::ExtractOp>(
+ loc, extendSource, extractOp.getMixedPosition());
+
+ // Extend new extract to original result type.
+ Operation *newExtend =
+ rewriter.create(loc, extendOpName, Value(newExtract), resultType);
+
+ rewriter.replaceOp(extractOp, newExtend->getResult(0));
+
+ return success();
+ }
+};
+
+// Shuffles arith extend ops after vector.scalable.extract op.
+//
+// This transforms IR like:
+// %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
+// %1 = vector.scalable.extract %0[0] : vector<[4]xi32> from vector<[8]xi32>
+// Into:
+// %0 = vector.scalable.extract %src[0] : vector<[4]xi8> from vector<[8]xi8>
+// %1 = arith.extsi %0 : vector<[4]xi8> to vector<[4]xi32>
+//
+// This enables outer product fusion in the `-arm-sme-outer-product-fusion`
+// pass when the result is the input to an outer product.
+struct SwapVectorScalableExtractOfArithExtend
+ : public OpRewritePattern<vector::ScalableExtractOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(vector::ScalableExtractOp extractOp,
+ PatternRewriter &rewriter) const override {
+ auto *extendOp = extractOp.getSource().getDefiningOp();
+ if (!isa_and_present<arith::ExtSIOp, arith::ExtUIOp, arith::ExtFOp>(
+ extendOp))
+ return rewriter.notifyMatchFailure(extractOp,
+ "extract not from extend op");
+
+ auto loc = extractOp.getLoc();
+ VectorType resultType = extractOp.getResultVectorType();
+
+ Value extendSource = extendOp->getOperand(0);
+ StringAttr extendOpName = extendOp->getName().getIdentifier();
+ VectorType extendSourceVectorType =
+ cast<VectorType>(extendSource.getType());
+
+ // Create new extract from source of extend.
+ VectorType extractResultVectorType =
+ VectorType::Builder(resultType)
+ .setElementType(extendSourceVectorType.getElementType());
+ Value newExtract = rewriter.create<vector::ScalableExtractOp>(
+ loc, extractResultVectorType, extendSource, extractOp.getPos());
+
+ // Extend new extract to original result type.
+ Operation *newExtend =
+ rewriter.create(loc, extendOpName, Value(newExtract), resultType);
+
+ rewriter.replaceOp(extractOp, newExtend->getResult(0));
+
+ return success();
+ }
+};
+
struct VectorLegalizationPass
: public arm_sme::impl::VectorLegalizationBase<VectorLegalizationPass> {
void runOnOperation() override {
@@ -434,7 +533,10 @@ struct VectorLegalizationPass
return success();
});
- patterns.add<FoldExtractFromVectorOfSMELikeCreateMasks>(context);
+ patterns.add<FoldExtractFromVectorOfSMELikeCreateMasks,
+ SwapVectorExtractOfArithExtend,
+ SwapVectorScalableExtractOfArithExtend>(context);
+
// Note: High benefit to ensure masked outer products are lowered first.
patterns.add<LegalizeMaskedVectorOuterProductOpsByDecomposition>(
converter, context, 1024);
diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index a2526db9b4831..5ae7be6d1b809 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -302,3 +302,43 @@ func.func @non_constant_extract_from_vector_create_mask_non_constant(%index: ind
%extract = vector.extract %mask[%index] : vector<[4]x[4]xi1> from vector<4x[4]x[4]xi1>
return %extract : vector<[4]x[4]xi1>
}
+
+// -----
+
+// CHECK-LABEL: @extract_from_arith_ext(
+// CHECK-SAME: %[[SRC:.*]]: vector<4x[8]xi8>
+// CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC]][0] : vector<[8]xi8> from vector<4x[8]xi8>
+// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
+// CHECK: return %[[EXTEND]]
+func.func @extract_from_arith_ext(%src: vector<4x[8]xi8>) -> vector<[8]xi32> {
+ %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+ %1 = vector.extract %0[0] : vector<[8]xi32> from vector<4x[8]xi32>
+ return %1 : vector<[8]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @non_constant_extract_from_arith_ext(
+// CHECK-SAME: %[[SRC:[a-z0-9]+]]: vector<4x[8]xi8>,
+// CHECK-SAME: %[[DIM:[a-z0-9]+]]: index
+// CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC]][%[[DIM]]] : vector<[8]xi8> from vector<4x[8]xi8>
+// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
+// CHECK: return %[[EXTEND]]
+func.func @non_constant_extract_from_arith_ext(%src: vector<4x[8]xi8>, %dim: index) -> vector<[8]xi32> {
+ %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+ %1 = vector.extract %0[%dim] : vector<[8]xi32> from vector<4x[8]xi32>
+ return %1 : vector<[8]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @scalable_extract_from_arith_ext(
+// CHECK-SAME: %[[SRC:.*]]: vector<[8]xi8>
+// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xi8> from vector<[8]xi8>
+// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[4]xi8> to vector<[4]xi32>
+// CHECK: return %[[EXTEND]]
+func.func @scalable_extract_from_arith_ext(%src: vector<[8]xi8>) -> vector<[4]xi32> {
+ %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
+ %1 = vector.scalable.extract %0[0] : vector<[4]xi32> from vector<[8]xi32>
+ return %1 : vector<[4]xi32>
+}
>From e26e8362041c2da7917a579d1b2edc04b07c2090 Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Mon, 5 Feb 2024 11:50:58 +0000
Subject: [PATCH 2/3] Address comments
---
.../ArmSME/Transforms/VectorLegalization.cpp | 11 +++++------
.../Dialect/ArmSME/vector-legalization.mlir | 18 +++++++++---------
2 files changed, 14 insertions(+), 15 deletions(-)
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
index 47de4a8a857bf..f598ef5843e47 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
@@ -415,7 +415,7 @@ struct FoldExtractFromVectorOfSMELikeCreateMasks
}
};
-// Shuffles arith extend ops after vector.extract op.
+// Rewrites: vector.extract(arith.extend) -> arith.extend(vector.extract).
//
// This transforms IR like:
// %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
@@ -460,13 +460,13 @@ struct SwapVectorExtractOfArithExtend
Operation *newExtend =
rewriter.create(loc, extendOpName, Value(newExtract), resultType);
- rewriter.replaceOp(extractOp, newExtend->getResult(0));
+ rewriter.replaceOp(extractOp, newExtend);
return success();
}
};
-// Shuffles arith extend ops after vector.scalable.extract op.
+// Same as above, but for vector.scalable.extract.
//
// This transforms IR like:
// %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
@@ -499,8 +499,7 @@ struct SwapVectorScalableExtractOfArithExtend
// Create new extract from source of extend.
VectorType extractResultVectorType =
- VectorType::Builder(resultType)
- .setElementType(extendSourceVectorType.getElementType());
+ resultType.clone(extendSourceVectorType.getElementType());
Value newExtract = rewriter.create<vector::ScalableExtractOp>(
loc, extractResultVectorType, extendSource, extractOp.getPos());
@@ -508,7 +507,7 @@ struct SwapVectorScalableExtractOfArithExtend
Operation *newExtend =
rewriter.create(loc, extendOpName, Value(newExtract), resultType);
- rewriter.replaceOp(extractOp, newExtend->getResult(0));
+ rewriter.replaceOp(extractOp, newExtend);
return success();
}
diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index 5ae7be6d1b809..7009767d95d28 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -322,10 +322,10 @@ func.func @extract_from_arith_ext(%src: vector<4x[8]xi8>) -> vector<[8]xi32> {
// CHECK-SAME: %[[SRC:[a-z0-9]+]]: vector<4x[8]xi8>,
// CHECK-SAME: %[[DIM:[a-z0-9]+]]: index
// CHECK: %[[EXTRACT:.*]] = vector.extract %[[SRC]][%[[DIM]]] : vector<[8]xi8> from vector<4x[8]xi8>
-// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
+// CHECK: %[[EXTEND:.*]] = arith.extui %[[EXTRACT]] : vector<[8]xi8> to vector<[8]xi32>
// CHECK: return %[[EXTEND]]
func.func @non_constant_extract_from_arith_ext(%src: vector<4x[8]xi8>, %dim: index) -> vector<[8]xi32> {
- %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+ %0 = arith.extui %src : vector<4x[8]xi8> to vector<4x[8]xi32>
%1 = vector.extract %0[%dim] : vector<[8]xi32> from vector<4x[8]xi32>
return %1 : vector<[8]xi32>
}
@@ -333,12 +333,12 @@ func.func @non_constant_extract_from_arith_ext(%src: vector<4x[8]xi8>, %dim: ind
// -----
// CHECK-LABEL: @scalable_extract_from_arith_ext(
-// CHECK-SAME: %[[SRC:.*]]: vector<[8]xi8>
-// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xi8> from vector<[8]xi8>
-// CHECK: %[[EXTEND:.*]] = arith.extsi %[[EXTRACT]] : vector<[4]xi8> to vector<[4]xi32>
+// CHECK-SAME: %[[SRC:.*]]: vector<[8]xf16>
+// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xf16> from vector<[8]xf16>
+// CHECK: %[[EXTEND:.*]] = arith.extf %[[EXTRACT]] : vector<[4]xf16> to vector<[4]xf32>
// CHECK: return %[[EXTEND]]
-func.func @scalable_extract_from_arith_ext(%src: vector<[8]xi8>) -> vector<[4]xi32> {
- %0 = arith.extsi %src : vector<[8]xi8> to vector<[8]xi32>
- %1 = vector.scalable.extract %0[0] : vector<[4]xi32> from vector<[8]xi32>
- return %1 : vector<[4]xi32>
+func.func @scalable_extract_from_arith_ext(%src: vector<[8]xf16>) -> vector<[4]xf32> {
+ %0 = arith.extf %src : vector<[8]xf16> to vector<[8]xf32>
+ %1 = vector.scalable.extract %0[0] : vector<[4]xf32> from vector<[8]xf32>
+ return %1 : vector<[4]xf32>
}
>From c75c20aae7d817780340efe0ceed6407f0922db3 Mon Sep 17 00:00:00 2001
From: Cullen Rhodes <cullen.rhodes at arm.com>
Date: Mon, 5 Feb 2024 12:19:24 +0000
Subject: [PATCH 3/3] Address comments.
Add negative tests.
---
.../Dialect/ArmSME/vector-legalization.mlir | 50 +++++++++++++++++++
1 file changed, 50 insertions(+)
diff --git a/mlir/test/Dialect/ArmSME/vector-legalization.mlir b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
index 7009767d95d28..e02a4568eb0c7 100644
--- a/mlir/test/Dialect/ArmSME/vector-legalization.mlir
+++ b/mlir/test/Dialect/ArmSME/vector-legalization.mlir
@@ -332,6 +332,44 @@ func.func @non_constant_extract_from_arith_ext(%src: vector<4x[8]xi8>, %dim: ind
// -----
+/// Non-vector extracts should be ignored.
+
+// CHECK-LABEL: @extract_scalar_from_arith_ext
+// CHECK-NEXT: arith.extsi
+// CHECK-NEXT: vector.extract
+func.func @extract_scalar_from_arith_ext(%src: vector<4x[8]xi8>) -> i32 {
+ %0 = arith.extsi %src : vector<4x[8]xi8> to vector<4x[8]xi32>
+ %1 = vector.extract %0[0, 0] : i32 from vector<4x[8]xi32>
+ return %1 : i32
+}
+
+// -----
+
+/// Extracted type should be a 1-D scalable vector type.
+
+// CHECK-LABEL: @extract_fixed_1d_vec_from_arith_ext
+// CHECK-NEXT: arith.extsi
+// CHECK-NEXT: vector.extract
+func.func @extract_fixed_1d_vec_from_arith_ext(%src: vector<4x8xi8>) -> vector<8xi32> {
+ %0 = arith.extsi %src : vector<4x8xi8> to vector<4x8xi32>
+ %1 = vector.extract %0[0] : vector<8xi32> from vector<4x8xi32>
+ return %1 : vector<8xi32>
+}
+
+// -----
+
+/// Extract must come from an arith extend.
+
+// CHECK-LABEL: @extract_from_non_arith_ext
+// CHECK-NEXT: vector.extract
+// CHECK-NEXT: return
+func.func @extract_from_non_arith_ext(%src: vector<4x[8]xi32>) -> vector<[8]xi32> {
+ %0 = vector.extract %src[0] : vector<[8]xi32> from vector<4x[8]xi32>
+ return %0 : vector<[8]xi32>
+}
+
+// -----
+
// CHECK-LABEL: @scalable_extract_from_arith_ext(
// CHECK-SAME: %[[SRC:.*]]: vector<[8]xf16>
// CHECK: %[[EXTRACT:.*]] = vector.scalable.extract %[[SRC]][0] : vector<[4]xf16> from vector<[8]xf16>
@@ -342,3 +380,15 @@ func.func @scalable_extract_from_arith_ext(%src: vector<[8]xf16>) -> vector<[4]x
%1 = vector.scalable.extract %0[0] : vector<[4]xf32> from vector<[8]xf32>
return %1 : vector<[4]xf32>
}
+
+// -----
+
+/// Scalable extract must come from an arith extend.
+
+// CHECK-LABEL: @scalable_extract_from_non_arith_ext
+// CHECK-NEXT: vector.scalable.extract
+// CHECK-NEXT: return
+func.func @scalable_extract_from_non_arith_ext(%src: vector<[8]xf32>) -> vector<[4]xf32> {
+ %0 = vector.scalable.extract %src[0] : vector<[4]xf32> from vector<[8]xf32>
+ return %0 : vector<[4]xf32>
+}
More information about the Mlir-commits
mailing list