[Mlir-commits] [mlir] 5f5b3bb - [mlir][ArmSME] Add rewrites to swap extract of extend (#80407)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Feb 5 06:13:57 PST 2024
Author: Cullen Rhodes
Date: 2024-02-05T14:13:53Z
New Revision: 5f5b3bb22b2e4ffcd14a8fc8a5edc14bc098a47e
URL: https://github.com/llvm/llvm-project/commit/5f5b3bb22b2e4ffcd14a8fc8a5edc14bc098a47e
DIFF: https://github.com/llvm/llvm-project/commit/5f5b3bb22b2e4ffcd14a8fc8a5edc14bc098a47e.diff
LOG: [mlir][ArmSME] Add rewrites to swap extract of extend (#80407)
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.
Added:
Modified:
mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp
mlir/test/Dialect/ArmSME/outer-product-fusion.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp b/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp
index fa55f9b7b31e1..7dc2aacd7e5f3 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/OuterProductFusion.cpp
@@ -261,6 +261,104 @@ class OuterProductFusion2Way
}
};
+// 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>
+// %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);
+
+ return success();
+ }
+};
+
+// Same as above, but for vector.scalable.extract.
+//
+// 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 =
+ resultType.clone(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);
+
+ return success();
+ }
+};
+
struct OuterProductFusionPass
: public arm_sme::impl::OuterProductFusionBase<OuterProductFusionPass> {
@@ -278,7 +376,11 @@ struct OuterProductFusionPass
void mlir::arm_sme::populateOuterProductFusionPatterns(
RewritePatternSet &patterns) {
- patterns.add<OuterProductFusion2Way>(patterns.getContext());
+ MLIRContext *context = patterns.getContext();
+ // Note: High benefit to ensure extract(extend) are swapped first.
+ patterns.add<SwapVectorExtractOfArithExtend,
+ SwapVectorScalableExtractOfArithExtend>(context, 1024);
+ patterns.add<OuterProductFusion2Way>(context);
}
std::unique_ptr<Pass> mlir::arm_sme::createOuterProductFusionPass() {
diff --git a/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir b/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir
index 6a200848c2d8e..f24943cac4f76 100644
--- a/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir
+++ b/mlir/test/Dialect/ArmSME/outer-product-fusion.mlir
@@ -213,6 +213,48 @@ func.func @outerproduct_sub_widening_2way_unsigned_i16i16i32(
return %1 : vector<[4]x[4]xi32>
}
+/// Tests for related patterns.
+
+// -----
+
+// 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.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.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>
+}
+
+// -----
+
+// 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>
+// 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]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>
+}
+
/// Negative tests
// -----
@@ -362,3 +404,55 @@ func.func @outerproduct_widening_2way__bad_defining_op(
return %1 : vector<[4]x[4]xf32>
}
+
+/// Negative tests for related patterns.
+
+// -----
+
+/// 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>
+}
+
+// -----
+
+/// 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