| //===- OuterProductFusion.cpp - Fuse 'arm_sme.outerproduct' ops -----------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file implements rewrites that fuse 'arm_sme.outerproduct' operations |
| // into the 2-way or 4-way widening outerproduct operations. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/ArmSME/IR/ArmSME.h" |
| #include "mlir/Dialect/ArmSME/Transforms/Passes.h" |
| #include "mlir/Dialect/ArmSME/Transforms/Transforms.h" |
| #include "mlir/Dialect/Func/IR/FuncOps.h" |
| #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| #include "llvm/ADT/TypeSwitch.h" |
| |
| #define DEBUG_TYPE "arm-sme-outerproduct-fusion" |
| |
| namespace mlir::arm_sme { |
| #define GEN_PASS_DEF_OUTERPRODUCTFUSION |
| #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc" |
| } // namespace mlir::arm_sme |
| |
| using namespace mlir; |
| using namespace mlir::arm_sme; |
| |
| namespace { |
| |
| // Common match failure reasons. |
| static constexpr StringLiteral |
| kMatchFailureNoAccumulator("no accumulator operand"); |
| static constexpr StringLiteral kMatchFailureExpectedOuterProductDefOp( |
| "defining op of accumulator must be 'arm_sme.outerproduct'"); |
| static constexpr StringLiteral kMatchFailureInconsistentCombiningKind( |
| "combining kind (add or sub) of outer products must match"); |
| static constexpr StringLiteral kMatchFailureInconsistentMasking( |
| "unsupported masking, either both outerproducts are masked " |
| "or neither"); |
| static constexpr StringLiteral kMatchFailureOuterProductNotSingleUse( |
| "outer product(s) not single use and cannot be removed, no benefit to " |
| "fusing"); |
| |
| // An outer product is compatible if all of the following are true: |
| // - the result type matches `resultType`. |
| // - the defining operation of LHS is of the type `LhsExtOp`. |
| // - the defining operation of RHS is of the type `RhsExtOp`. |
| // - the input types of the defining operations are identical and match |
| // `inputType`. |
| template <typename LhsExtOp, typename RhsExtOp = LhsExtOp> |
| static LogicalResult isCompatible(PatternRewriter &rewriter, |
| arm_sme::OuterProductOp op, |
| VectorType resultType, VectorType inputType) { |
| if (op.getResultType() != resultType) |
| return rewriter.notifyMatchFailure(op.getLoc(), [&](Diagnostic &diag) { |
| diag << "unsupported result type, expected " << resultType; |
| }); |
| |
| auto lhsDefOp = op.getLhs().getDefiningOp<LhsExtOp>(); |
| auto rhsDefOp = op.getRhs().getDefiningOp<RhsExtOp>(); |
| |
| if (!lhsDefOp || !rhsDefOp) |
| return rewriter.notifyMatchFailure( |
| op, "defining op of outerproduct operands must be one of: " |
| "'arith.extf' or 'arith.extsi' or 'arith.extui'"); |
| |
| auto lhsInType = cast<VectorType>(lhsDefOp.getIn().getType()); |
| auto rhsInType = cast<VectorType>(rhsDefOp.getIn().getType()); |
| |
| if (lhsInType != inputType || rhsInType != inputType) |
| return rewriter.notifyMatchFailure(op.getLoc(), [&](Diagnostic &diag) { |
| diag << "unsupported input type, expected " << inputType; |
| }); |
| |
| return success(); |
| } |
| |
| // Create 'llvm.experimental.vector.interleave2' intrinsic from `lhs` and `rhs`. |
| static Value createInterleave2Intrinsic(RewriterBase &rewriter, Location loc, |
| Value lhs, Value rhs) { |
| auto inputType = cast<VectorType>(lhs.getType()); |
| VectorType inputTypeX2 = |
| VectorType::Builder(inputType).setDim(0, inputType.getShape()[0] * 2); |
| return rewriter.create<LLVM::vector_interleave2>(loc, inputTypeX2, lhs, rhs); |
| } |
| |
| // Fuse two 'arm_sme.outerproduct' operations that are chained via the |
| // accumulator into 2-way outer product operation. |
| // |
| // For example: |
| // |
| // %a0_ext = arith.extf %a0 : vector<[4]xf16> to vector<[4]xf32> |
| // %b0_ext = arith.extf %b0 : vector<[4]xf16> to vector<[4]xf32> |
| // %0 = arm_sme.outerproduct %a0_ext, %b0_ext : vector<[4]xf32>, |
| // vector<[4]xf32> |
| // |
| // %a1_ext = arith.extf %a1 : vector<[4]xf16> to vector<[4]xf32> |
| // %b1_ext = arith.extf %b1 : vector<[4]xf16> to vector<[4]xf32> |
| // %1 = arm_sme.outerproduct %a1_ext, %b1_ext, %0 : vector<[4]xf32>, |
| // vector<[4]xf32> |
| // |
| // Becomes: |
| // |
| // %a_packed = "llvm.intr.experimental.vector.interleave2"(%a0, %a1) |
| // : (vector<[4]xf16>, vector<[4]xf16>) -> vector<[8]xf16> |
| // %b_packed = "llvm.intr.experimental.vector.interleave2"(%b0, %b1) |
| // : (vector<[4]xf16>, vector<[4]xf16>) -> vector<[8]xf16> |
| // %0 = arm_sme.fmopa_2way %a_packed, %b_packed |
| // : vector<[8]xf16>, vector<[8]xf16> into vector<[4]x[4]xf32> |
| class OuterProductFusion2Way |
| : public OpRewritePattern<arm_sme::OuterProductOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(arm_sme::OuterProductOp op, |
| PatternRewriter &rewriter) const override { |
| Value acc = op.getAcc(); |
| if (!acc) |
| return rewriter.notifyMatchFailure(op, kMatchFailureNoAccumulator); |
| |
| arm_sme::OuterProductOp op1 = acc.getDefiningOp<arm_sme::OuterProductOp>(); |
| arm_sme::OuterProductOp op2 = op; |
| if (!op1) |
| return rewriter.notifyMatchFailure( |
| op, kMatchFailureExpectedOuterProductDefOp); |
| |
| if (op1.getKind() != op2.getKind()) |
| return rewriter.notifyMatchFailure( |
| op, kMatchFailureInconsistentCombiningKind); |
| |
| if (!op1->hasOneUse()) { |
| // If the first outer product has uses other than as the input to another |
| // outer product, it can't be erased after fusion. This is a problem when |
| // it also has an accumulator as this will be used as the root for tile |
| // allocation and since the widening outer product uses the same |
| // accumulator it will get assigned the same tile ID, resulting in 3 |
| // outer products accumulating to the same tile and incorrect results. |
| // |
| // Example: |
| // |
| // %acc = arith.constant dense<0.0> ; root for tile allocation |
| // %0 = arm_sme.outerproduct %a0, %b0 acc(%acc) |
| // vector.print %0 ; intermediary use, can't erase %0 |
| // %1 = arm_sme.outerproduct %a1, %b1 acc(%0) |
| // |
| // After fusion and tile allocation |
| // |
| // %0 = arm_sme.zero {tile_id = 0 : i32} |
| // %1 = arm_sme.outerproduct %a0, %b0 acc(%0) {tile_id = 0 : i32} |
| // vector.print %1 |
| // %2 = arm_sme.fmopa_2way %a, %b acc(%0) {tile_id = 0 : i32} |
| // |
| // No accumulator would be ok, but it's simpler to prevent this |
| // altogether, since it has no benefit. |
| return rewriter.notifyMatchFailure(op, |
| kMatchFailureOuterProductNotSingleUse); |
| } |
| |
| if (bool(op1.getLhsMask()) != bool(op2.getLhsMask())) |
| return rewriter.notifyMatchFailure(op, kMatchFailureInconsistentMasking); |
| |
| if (failed(canFuseOuterProducts(rewriter, op1, op2))) |
| return failure(); |
| |
| auto loc = op.getLoc(); |
| auto packInputs = [&](Value lhs, Value rhs) { |
| return createInterleave2Intrinsic(rewriter, loc, lhs, rhs); |
| }; |
| |
| auto lhs = packInputs(op1.getLhs().getDefiningOp()->getOperand(0), |
| op2.getLhs().getDefiningOp()->getOperand(0)); |
| auto rhs = packInputs(op1.getRhs().getDefiningOp()->getOperand(0), |
| op2.getRhs().getDefiningOp()->getOperand(0)); |
| |
| Value lhsMask, rhsMask; |
| if (op1.getLhsMask() || op2.getLhsMask()) { |
| lhsMask = packInputs(op1.getLhsMask(), op2.getLhsMask()); |
| rhsMask = packInputs(op1.getRhsMask(), op2.getRhsMask()); |
| } |
| |
| auto extOp = op.getLhs().getDefiningOp(); |
| |
| arm_sme::CombiningKind kind = op.getKind(); |
| if (kind == arm_sme::CombiningKind::Add) { |
| TypeSwitch<Operation *>(extOp) |
| .Case<arith::ExtFOp>([&](auto) { |
| rewriter.replaceOpWithNewOp<arm_sme::FMopa2WayOp>( |
| op2, op.getResultType(), lhs, rhs, lhsMask, rhsMask, |
| op1.getAcc()); |
| }) |
| .Case<arith::ExtSIOp>([&](auto) { |
| rewriter.replaceOpWithNewOp<arm_sme::SMopa2WayOp>( |
| op2, op.getResultType(), lhs, rhs, lhsMask, rhsMask, |
| op1.getAcc()); |
| }) |
| .Case<arith::ExtUIOp>([&](auto) { |
| rewriter.replaceOpWithNewOp<arm_sme::UMopa2WayOp>( |
| op2, op.getResultType(), lhs, rhs, lhsMask, rhsMask, |
| op1.getAcc()); |
| }) |
| .Default([&](auto) { llvm_unreachable("unexpected extend op!"); }); |
| } else if (kind == arm_sme::CombiningKind::Sub) { |
| TypeSwitch<Operation *>(extOp) |
| .Case<arith::ExtFOp>([&](auto) { |
| rewriter.replaceOpWithNewOp<arm_sme::FMops2WayOp>( |
| op2, op.getResultType(), lhs, rhs, lhsMask, rhsMask, |
| op1.getAcc()); |
| }) |
| .Case<arith::ExtSIOp>([&](auto) { |
| rewriter.replaceOpWithNewOp<arm_sme::SMops2WayOp>( |
| op2, op.getResultType(), lhs, rhs, lhsMask, rhsMask, |
| op1.getAcc()); |
| }) |
| .Case<arith::ExtUIOp>([&](auto) { |
| rewriter.replaceOpWithNewOp<arm_sme::UMops2WayOp>( |
| op2, op.getResultType(), lhs, rhs, lhsMask, rhsMask, |
| op1.getAcc()); |
| }) |
| .Default([&](auto) { llvm_unreachable("unexpected extend op!"); }); |
| } else { |
| llvm_unreachable("unexpected arm_sme::CombiningKind!"); |
| } |
| |
| rewriter.eraseOp(op1); |
| |
| return success(); |
| } |
| |
| private: |
| // A pair of outer product can be fused if all of the following are true: |
| // - input and result types match. |
| // - the defining operations of the inputs are identical extensions, |
| // specifically either: |
| // - a signed or unsigned extension for integer types. |
| // - a floating-point extension for floating-point types. |
| // - the types and extension are supported, i.e. there's a 2-way operation |
| // they can be fused into. |
| LogicalResult canFuseOuterProducts(PatternRewriter &rewriter, |
| arm_sme::OuterProductOp op1, |
| arm_sme::OuterProductOp op2) const { |
| // Supported result types. |
| auto nxnxv4i32 = |
| VectorType::get({4, 4}, rewriter.getI32Type(), {true, true}); |
| auto nxnxv4f32 = |
| VectorType::get({4, 4}, rewriter.getF32Type(), {true, true}); |
| // Supported input types. |
| // Note: this is before packing so these have half the number of elements |
| // of the input vector types of the 2-way operations. |
| auto nxv4i16 = VectorType::get({4}, rewriter.getI16Type(), true); |
| auto nxv4f16 = VectorType::get({4}, rewriter.getF16Type(), true); |
| auto nxv4bf16 = VectorType::get({4}, rewriter.getBF16Type(), true); |
| if ((failed( |
| isCompatible<arith::ExtFOp>(rewriter, op1, nxnxv4f32, nxv4f16)) || |
| failed( |
| isCompatible<arith::ExtFOp>(rewriter, op2, nxnxv4f32, nxv4f16))) && |
| (failed( |
| isCompatible<arith::ExtFOp>(rewriter, op1, nxnxv4f32, nxv4bf16)) || |
| failed(isCompatible<arith::ExtFOp>(rewriter, op2, nxnxv4f32, |
| nxv4bf16))) && |
| (failed( |
| isCompatible<arith::ExtSIOp>(rewriter, op1, nxnxv4i32, nxv4i16)) || |
| failed(isCompatible<arith::ExtSIOp>(rewriter, op2, nxnxv4i32, |
| nxv4i16))) && |
| (failed( |
| isCompatible<arith::ExtUIOp>(rewriter, op1, nxnxv4i32, nxv4i16)) || |
| failed( |
| isCompatible<arith::ExtUIOp>(rewriter, op2, nxnxv4i32, nxv4i16)))) |
| return failure(); |
| |
| return success(); |
| } |
| }; |
| |
| // Fuse four 'arm_sme.outerproduct' operations that are chained via the |
| // accumulator into 4-way outer product operation. |
| class OuterProductFusion4Way |
| : public OpRewritePattern<arm_sme::OuterProductOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| LogicalResult matchAndRewrite(arm_sme::OuterProductOp op, |
| PatternRewriter &rewriter) const override { |
| SmallVector<arm_sme::OuterProductOp, 4> outerProductChain; |
| outerProductChain.push_back(op); |
| |
| for (int i = 0; i < 3; ++i) { |
| auto currentOp = outerProductChain.back(); |
| auto acc = currentOp.getAcc(); |
| if (!acc) |
| return rewriter.notifyMatchFailure(op, kMatchFailureNoAccumulator); |
| auto previousOp = acc.getDefiningOp<arm_sme::OuterProductOp>(); |
| if (!previousOp) |
| return rewriter.notifyMatchFailure( |
| op, kMatchFailureExpectedOuterProductDefOp); |
| if (!previousOp->hasOneUse()) |
| return rewriter.notifyMatchFailure( |
| op, kMatchFailureOuterProductNotSingleUse); |
| if (previousOp.getKind() != currentOp.getKind()) |
| return rewriter.notifyMatchFailure( |
| op, kMatchFailureInconsistentCombiningKind); |
| if (bool(previousOp.getLhsMask()) != bool(currentOp.getLhsMask())) |
| return rewriter.notifyMatchFailure( |
| op, kMatchFailureInconsistentCombiningKind); |
| outerProductChain.push_back(previousOp); |
| } |
| |
| if (failed(canFuseOuterProducts(rewriter, outerProductChain))) |
| return failure(); |
| |
| arm_sme::OuterProductOp op1 = outerProductChain[3]; |
| arm_sme::OuterProductOp op2 = outerProductChain[2]; |
| arm_sme::OuterProductOp op3 = outerProductChain[1]; |
| arm_sme::OuterProductOp op4 = outerProductChain[0]; |
| |
| auto loc = op.getLoc(); |
| auto packInputs = [&](Value lhs, Value rhs) { |
| return createInterleave2Intrinsic(rewriter, loc, lhs, rhs); |
| }; |
| |
| auto lhs0 = packInputs(op1.getLhs().getDefiningOp()->getOperand(0), |
| op3.getLhs().getDefiningOp()->getOperand(0)); |
| auto lhs1 = packInputs(op2.getLhs().getDefiningOp()->getOperand(0), |
| op4.getLhs().getDefiningOp()->getOperand(0)); |
| auto lhs = packInputs(lhs0, lhs1); |
| |
| auto rhs0 = packInputs(op1.getRhs().getDefiningOp()->getOperand(0), |
| op3.getRhs().getDefiningOp()->getOperand(0)); |
| auto rhs1 = packInputs(op2.getRhs().getDefiningOp()->getOperand(0), |
| op4.getRhs().getDefiningOp()->getOperand(0)); |
| auto rhs = packInputs(rhs0, rhs1); |
| |
| Value lhsMask, rhsMask; |
| if (op1.getLhsMask() || op2.getLhsMask() || op3.getLhsMask() || |
| op4.getLhsMask()) { |
| auto lhs0Mask = packInputs(op1.getLhsMask(), op3.getLhsMask()); |
| auto lhs1Mask = packInputs(op2.getLhsMask(), op4.getLhsMask()); |
| lhsMask = packInputs(lhs0Mask, lhs1Mask); |
| |
| auto rhs0Mask = packInputs(op1.getRhsMask(), op3.getRhsMask()); |
| auto rhs1Mask = packInputs(op2.getRhsMask(), op4.getRhsMask()); |
| rhsMask = packInputs(rhs0Mask, rhs1Mask); |
| } |
| |
| auto lhsExtOp = op.getLhs().getDefiningOp(); |
| auto rhsExtOp = op.getRhs().getDefiningOp(); |
| |
| arm_sme::CombiningKind kind = op.getKind(); |
| if (kind == arm_sme::CombiningKind::Add) { |
| if (isa<arith::ExtSIOp>(lhsExtOp) && isa<arith::ExtSIOp>(rhsExtOp)) { |
| // signed |
| rewriter.replaceOpWithNewOp<arm_sme::SMopa4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else if (isa<arith::ExtUIOp>(lhsExtOp) && |
| isa<arith::ExtUIOp>(rhsExtOp)) { |
| // unsigned |
| rewriter.replaceOpWithNewOp<arm_sme::UMopa4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else if (isa<arith::ExtSIOp>(lhsExtOp) && |
| isa<arith::ExtUIOp>(rhsExtOp)) { |
| // signed by unsigned |
| rewriter.replaceOpWithNewOp<arm_sme::SuMopa4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else if (isa<arith::ExtUIOp>(lhsExtOp) && |
| isa<arith::ExtSIOp>(rhsExtOp)) { |
| // unsigned by signed |
| rewriter.replaceOpWithNewOp<arm_sme::UsMopa4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else { |
| llvm_unreachable("unexpected extend op!"); |
| } |
| } else if (kind == arm_sme::CombiningKind::Sub) { |
| if (isa<arith::ExtSIOp>(lhsExtOp) && isa<arith::ExtSIOp>(rhsExtOp)) { |
| // signed |
| rewriter.replaceOpWithNewOp<arm_sme::SMops4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else if (isa<arith::ExtUIOp>(lhsExtOp) && |
| isa<arith::ExtUIOp>(rhsExtOp)) { |
| // unsigned |
| rewriter.replaceOpWithNewOp<arm_sme::UMops4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else if (isa<arith::ExtSIOp>(lhsExtOp) && |
| isa<arith::ExtUIOp>(rhsExtOp)) { |
| // signed by unsigned |
| rewriter.replaceOpWithNewOp<arm_sme::SuMops4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else if (isa<arith::ExtUIOp>(lhsExtOp) && |
| isa<arith::ExtSIOp>(rhsExtOp)) { |
| // unsigned by signed |
| rewriter.replaceOpWithNewOp<arm_sme::UsMops4WayOp>( |
| op4, op.getResultType(), lhs, rhs, lhsMask, rhsMask, op1.getAcc()); |
| } else { |
| llvm_unreachable("unexpected extend op!"); |
| } |
| } else { |
| llvm_unreachable("unexpected arm_sme::CombiningKind!"); |
| } |
| |
| rewriter.eraseOp(op3); |
| rewriter.eraseOp(op2); |
| rewriter.eraseOp(op1); |
| |
| return success(); |
| } |
| |
| private: |
| // Four outer products can be fused if all of the following are true: |
| // - input and result types match. |
| // - the defining operations of the inputs are identical extensions, |
| // specifically either: |
| // - a signed or unsigned extension for integer types. |
| // - a floating-point extension for floating-point types. |
| // - the types and extension are supported, i.e. there's a 4-way operation |
| // they can be fused into. |
| LogicalResult |
| canFuseOuterProducts(PatternRewriter &rewriter, |
| ArrayRef<arm_sme::OuterProductOp> ops) const { |
| // Supported result types. |
| auto nxnxv4i32 = |
| VectorType::get({4, 4}, rewriter.getI32Type(), {true, true}); |
| auto nxnxv2i64 = |
| VectorType::get({2, 2}, rewriter.getI64Type(), {true, true}); |
| |
| // Supported input types. |
| // Note: this is before packing so these have 1/4 the number of elements |
| // of the input vector types of the 4-way operations. |
| auto nxv4i8 = VectorType::get({4}, rewriter.getI8Type(), true); |
| auto nxv2i16 = VectorType::get({2}, rewriter.getI16Type(), true); |
| |
| auto failedToMatch = [&](VectorType resultType, VectorType inputType, |
| auto lhsExtendOp, auto rhsExtendOp) { |
| using LhsExtendOpTy = decltype(lhsExtendOp); |
| using RhsExtendOpTy = decltype(rhsExtendOp); |
| for (auto op : ops) { |
| if (failed(isCompatible<LhsExtendOpTy, RhsExtendOpTy>( |
| rewriter, op, resultType, inputType))) |
| return true; |
| } |
| return false; |
| }; |
| |
| if (failedToMatch(nxnxv4i32, nxv4i8, arith::ExtSIOp{}, arith::ExtSIOp{}) && |
| failedToMatch(nxnxv4i32, nxv4i8, arith::ExtUIOp{}, arith::ExtUIOp{}) && |
| failedToMatch(nxnxv4i32, nxv4i8, arith::ExtSIOp{}, arith::ExtUIOp{}) && |
| failedToMatch(nxnxv4i32, nxv4i8, arith::ExtUIOp{}, arith::ExtSIOp{}) && |
| failedToMatch(nxnxv2i64, nxv2i16, arith::ExtSIOp{}, arith::ExtSIOp{}) && |
| failedToMatch(nxnxv2i64, nxv2i16, arith::ExtUIOp{}, arith::ExtUIOp{}) && |
| failedToMatch(nxnxv2i64, nxv2i16, arith::ExtSIOp{}, arith::ExtUIOp{}) && |
| failedToMatch(nxnxv2i64, nxv2i16, arith::ExtUIOp{}, arith::ExtSIOp{})) |
| return failure(); |
| |
| return success(); |
| } |
| }; |
| |
| // 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> { |
| |
| void runOnOperation() override { |
| RewritePatternSet patterns(&getContext()); |
| populateOuterProductFusionPatterns(patterns); |
| |
| if (failed( |
| applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) |
| signalPassFailure(); |
| } |
| }; |
| |
| } // namespace |
| |
| void mlir::arm_sme::populateOuterProductFusionPatterns( |
| RewritePatternSet &patterns) { |
| MLIRContext *context = patterns.getContext(); |
| // Note: High benefit to ensure extract(extend) are swapped first. |
| patterns.add<SwapVectorExtractOfArithExtend, |
| SwapVectorScalableExtractOfArithExtend>(context, 1024); |
| patterns.add<OuterProductFusion2Way, OuterProductFusion4Way>(context); |
| } |
| |
| std::unique_ptr<Pass> mlir::arm_sme::createOuterProductFusionPass() { |
| return std::make_unique<OuterProductFusionPass>(); |
| } |
