| // Copyright 2024 The IREE Authors |
| // |
| // Licensed 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 |
| |
| #include "iree/compiler/Dialect/Flow/IR/FlowOps.h" |
| #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h" |
| #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h" |
| #include "iree/compiler/DispatchCreation/Passes.h" |
| #include "mlir/Analysis/SliceAnalysis.h" |
| #include "mlir/Analysis/TopologicalSortUtils.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/Linalg/IR/Linalg.h" |
| #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" |
| #include "mlir/Dialect/Linalg/Transforms/Transforms.h" |
| #include "mlir/Dialect/Tensor/IR/Tensor.h" |
| #include "mlir/Dialect/Tensor/Transforms/Transforms.h" |
| #include "mlir/Dialect/Tensor/Utils/Utils.h" |
| #include "mlir/Dialect/Utils/ReshapeOpsUtils.h" |
| #include "mlir/Dialect/Utils/StructuredOpsUtils.h" |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "mlir/IR/Dominance.h" |
| #include "mlir/IR/OpDefinition.h" |
| #include "mlir/IR/OperationSupport.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/Pass/Pass.h" |
| #include "mlir/Support/LLVM.h" |
| #include "mlir/Support/LogicalResult.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| |
| #define DEBUG_TYPE "iree-dispatch-creation-fuse-horizontal-contractions" |
| |
| namespace mlir::iree_compiler::DispatchCreation { |
| |
| #define GEN_PASS_DEF_FUSEHORIZONTALCONTRACTIONSPASS |
| #include "iree/compiler/DispatchCreation/Passes.h.inc" |
| |
| namespace { |
| |
| struct FuseHorizontalContractionsPass final |
| : public impl::FuseHorizontalContractionsPassBase< |
| FuseHorizontalContractionsPass> { |
| using Base::Base; |
| void runOnOperation() override; |
| }; |
| |
| } // namespace |
| |
| /// Structs that captures the ops that are to be fused |
| struct HorizontalFusionGroup { |
| // Contractions op that are to be fused. |
| SmallVector<linalg::LinalgOp> contractionOps; |
| // Optional truncate operations that could be following the contraction op. |
| std::optional<SmallVector<linalg::GenericOp>> truncateOps; |
| }; |
| |
| /// Helper method to check operations equivalence |
| static bool checkOperationEquivalence(Operation *lhsOp, Operation *rhsOp) { |
| // During equivalence check, it would have been easier if `checkEquivalence` |
| // would just use `OpOperands *`. Since it takes `Value`s for now, just |
| // check that the values are the same as operands. This is potentially |
| // making the match too broad, but is an OK work-around for now. |
| // TODO(MaheshRavishankar): Fix upstream `checkEquivalence` signater in |
| // `OperationEquivalence::isEquivalentTo`. |
| llvm::SmallDenseSet<Value, 8> operands; |
| operands.insert(lhsOp->operand_begin(), lhsOp->operand_end()); |
| operands.insert(rhsOp->operand_begin(), rhsOp->operand_end()); |
| |
| llvm::DenseMap<Value, Value> equivalentValues; |
| auto checkEquivalent = [&](Value lhsValue, Value rhsValue) { |
| if (operands.contains(lhsValue) && operands.contains(rhsValue)) { |
| return success(); |
| } |
| return success(equivalentValues.lookup(lhsValue) == rhsValue || |
| equivalentValues.lookup(rhsValue) == lhsValue); |
| }; |
| auto markEquivalent = [&](Value v1, Value v2) { equivalentValues[v1] = v2; }; |
| return OperationEquivalence::isEquivalentTo( |
| lhsOp, rhsOp, checkEquivalent, markEquivalent, |
| /*flags=*/OperationEquivalence::IgnoreLocations); |
| } |
| |
| /// Check that an operation is a `empty -> fill -> contraction` |
| static bool isEmptyFillContractionDAGRootOp( |
| linalg::LinalgOp linalgOp, |
| std::optional<linalg::LinalgOp> seedContractionOp = std::nullopt) { |
| if (!linalg::isaContractionOpInterface(linalgOp)) { |
| return false; |
| } |
| auto fillOp = linalgOp.getDpsInits()[0].getDefiningOp<linalg::FillOp>(); |
| if (!fillOp) { |
| return false; |
| } |
| // For convenience check that the fill value is 0. This is not |
| // a necessity, but easier to handle the rewrite this way. |
| if (!matchPattern(fillOp.getDpsInputOperand(0)->get(), m_AnyZeroFloat()) && |
| !matchPattern(fillOp.getDpsInputOperand(0)->get(), m_Zero())) { |
| return false; |
| } |
| if (!fillOp.getDpsInitOperand(0)->get().getDefiningOp<tensor::EmptyOp>()) { |
| return false; |
| } |
| if (seedContractionOp) { |
| return checkOperationEquivalence(linalgOp, seedContractionOp.value()); |
| } |
| return true; |
| } |
| |
| /// Check that a given operation is "horizontal" to the group. The operation |
| /// is horizontal if the `slice` of the operation does not contain any op |
| /// from the group. |
| static bool isHorizontalToGroup(Operation *op, |
| const llvm::SetVector<Operation *> &currGroup, |
| const DominanceInfo &dominanceInfo, |
| Operation *seedOp) { |
| BackwardSliceOptions options; |
| // Limit the slice to the seed to make sure the slice is small. |
| options.filter = [&](Operation *op) { |
| return !dominanceInfo.properlyDominates(op, seedOp); |
| }; |
| llvm::SetVector<Operation *> slice; |
| getBackwardSlice(op, &slice, options); |
| return !llvm::any_of(currGroup, [&](Operation *groupedOp) { |
| return slice.contains(groupedOp); |
| }); |
| } |
| |
| /// Get user of operation that is a truncate operation. |
| static std::optional<linalg::GenericOp> |
| getTruncateOp(Operation *op, |
| const llvm::SetVector<Operation *> &groupedOperations, |
| const DominanceInfo &dominanceInfo, |
| std::optional<linalg::GenericOp> seedTruncateOp = std::nullopt) { |
| if (!op->hasOneUse()) { |
| return std::nullopt; |
| } |
| Operation *user = *op->user_begin(); |
| // TODO: This test should not be really needed. We should be able to check |
| // for ANY elementwise operation. |
| if (!IREE::LinalgExt::isBitTruncateOp(user)) { |
| return std::nullopt; |
| } |
| auto genericOp = dyn_cast<linalg::GenericOp>(user); |
| if (!genericOp) { |
| return std::nullopt; |
| } |
| if (seedTruncateOp) { |
| if (!checkOperationEquivalence(genericOp, seedTruncateOp.value())) { |
| return std::nullopt; |
| } |
| if (!isHorizontalToGroup(genericOp, groupedOperations, dominanceInfo, |
| seedTruncateOp.value())) { |
| return std::nullopt; |
| } |
| } |
| return genericOp; |
| } |
| |
| /// Find all candidates that can be used for horizontal fusion. For example |
| /// ``` |
| /// %0 = linalg.matmul ins(%arg0, %arg1) |
| /// %1 = linalg.matmul ins(%arg0, %arg2) |
| /// %2 = linalg.matmul ins(%arg0, %arg3) |
| /// ``` |
| /// |
| /// where all matmul share an operand can be combined into |
| /// |
| /// ``` |
| /// %4 = linalg.matmul ins(%arg0, concat(%arg1, %arg2, %arg3)) |
| /// ``` |
| /// |
| /// This method recognizes such patterns. It also accounts for the quantized |
| /// case where individual operations might be have lower-precision operands and |
| /// accumulate in higher precision, followed by a `linalg.generic` that performs |
| /// the `truncf` on the result. |
| static std::optional<HorizontalFusionGroup> getHorizontalFusionGroupMembers( |
| linalg::LinalgOp seedOp, |
| const llvm::SmallDenseSet<linalg::LinalgOp> &groupedOperations, |
| const DominanceInfo &dominanceInfo, int fusionLimit) { |
| |
| Value lhs = seedOp->getOperand(0); |
| auto lhsType = cast<RankedTensorType>(lhs.getType()); |
| Value rhs = seedOp->getOperand(1); |
| auto rhsType = cast<RankedTensorType>(rhs.getType()); |
| Value out = seedOp->getOperand(2); |
| auto outType = cast<RankedTensorType>(out.getType()); |
| |
| if (!lhsType.hasStaticShape() || !rhsType.hasStaticShape() || |
| !outType.hasStaticShape()) { |
| return std::nullopt; |
| } |
| |
| SetVector<Operation *> allOps; |
| SmallVector<linalg::LinalgOp> contractionOps = {seedOp}; |
| std::optional<linalg::GenericOp> seedTruncOp = |
| getTruncateOp(seedOp, allOps, dominanceInfo); |
| std::optional<SmallVector<linalg::GenericOp>> truncateOps; |
| if (seedTruncOp) { |
| truncateOps = {seedTruncOp.value()}; |
| } |
| allOps.insert(seedOp); |
| if (seedTruncOp) { |
| allOps.insert(seedTruncOp.value()); |
| } |
| |
| auto canBeGrouped = [&](linalg::LinalgOp linalgOp) -> bool { |
| if (linalgOp->getParentOp() != seedOp->getParentOp()) { |
| return false; |
| } |
| |
| // Constraints of the operation itself. |
| if (!isEmptyFillContractionDAGRootOp(linalgOp, seedOp)) { |
| return false; |
| } |
| if (linalgOp->getOperand(0).getType() != lhsType || |
| linalgOp->getOperand(1).getType() != rhsType || |
| linalgOp->getOperand(2).getType() != outType) { |
| return false; |
| } |
| if (groupedOperations.contains(linalgOp)) { |
| return false; |
| } |
| |
| // Structural constraints related to being able to fuse the operations. |
| if (!dominanceInfo.properlyDominates(seedOp, linalgOp)) { |
| return false; |
| } |
| if (!isHorizontalToGroup(linalgOp, allOps, dominanceInfo, seedOp)) { |
| return false; |
| } |
| return true; |
| }; |
| |
| // Iterate over users of LHS to find ops that can be grouped with the seed. |
| SmallVector<Operation *> lhsUsers; |
| for (Operation *lhsUser : lhs.getUsers()) { |
| if (lhsUser->getBlock() != seedOp->getBlock() || lhsUser == seedOp) { |
| continue; |
| } |
| |
| auto linalgUser = dyn_cast<linalg::LinalgOp>(lhsUser); |
| if (!linalgUser || !canBeGrouped(linalgUser)) { |
| continue; |
| } |
| lhsUsers.push_back(lhsUser); |
| } |
| |
| // Sort the users so that the order is deterministic |
| llvm::sort(lhsUsers, [&](Operation *lhs, Operation *rhs) { |
| return dominanceInfo.properlyDominates(lhs, rhs); |
| }); |
| |
| // Collect all contraction op users of lhs. |
| for (Operation *lhsUser : lhsUsers) { |
| auto linalgUser = dyn_cast<linalg::LinalgOp>(lhsUser); |
| if (!linalgUser) { |
| continue; |
| } |
| |
| std::optional<linalg::GenericOp> userTruncOp = |
| getTruncateOp(linalgUser, allOps, dominanceInfo, seedTruncOp); |
| // If there are truncate ops to fuse and current contraction op |
| // does not have a compatible truncate op to fuse as well, ignore |
| // the op for horizontal fusion. |
| if (truncateOps && !userTruncOp) { |
| continue; |
| } |
| |
| contractionOps.push_back(linalgUser); |
| allOps.insert(linalgUser); |
| if (truncateOps) { |
| truncateOps.value().push_back(userTruncOp.value()); |
| allOps.insert(userTruncOp.value()); |
| } |
| if (contractionOps.size() >= fusionLimit) { |
| break; |
| } |
| } |
| |
| if (contractionOps.size() == 1) { |
| return std::nullopt; |
| } |
| |
| return HorizontalFusionGroup{contractionOps, truncateOps}; |
| } |
| |
| /// Concatenate the given tensor `values`. The assumption here |
| /// is that all the `values` are the same type. These are concatanted |
| /// by adding a extra outer dimension to each value and concatenating |
| /// along the outer-most dim. |
| static Value concatenateValues(RewriterBase &rewriter, Location loc, |
| ArrayRef<Value> values) { |
| assert((values.size() >= 2) && "Invalid number of operands to concatenate"); |
| auto valueType = cast<RankedTensorType>(values[0].getType()); |
| |
| SmallVector<Value> concatOperands; |
| for (auto v : values) { |
| auto t = cast<RankedTensorType>(v.getType()); |
| SmallVector<int64_t> expandedTypeShape = {1}; |
| expandedTypeShape.append(t.getShape().begin(), t.getShape().end()); |
| auto expandedType = |
| RankedTensorType::get(expandedTypeShape, t.getElementType()); |
| SmallVector<OpFoldResult> expandedShape = {rewriter.getIndexAttr(1)}; |
| auto mixedSizes = tensor::getMixedSizes(rewriter, loc, v); |
| expandedShape.append(mixedSizes.begin(), mixedSizes.end()); |
| |
| SmallVector<ReassociationIndices> reassoc; |
| if (t.getRank() != 0) { |
| reassoc.push_back({0, 1}); |
| for (int i = 0, e = valueType.getRank() - 1; i < e; ++i) { |
| reassoc.push_back({i + 2}); |
| } |
| } |
| |
| Value expanded = rewriter.create<tensor::ExpandShapeOp>( |
| loc, expandedType, v, reassoc, expandedShape); |
| concatOperands.push_back(expanded); |
| } |
| |
| Value concatedVal = |
| rewriter.create<tensor::ConcatOp>(loc, /*dim=*/0, concatOperands); |
| return concatedVal; |
| } |
| |
| /// Compute the indexing map used in the concatenated operation. |
| /// The indexing map is either |
| /// 1) when shiftOnly = false, adds an extra outermost dimension to the indexing |
| /// map and adding that dimension as the outermost dimension in the range. |
| /// This is used for case where the original operands of the operations are |
| /// concatanated as well to get the operand for the horizontally-fused |
| /// operation. |
| /// 2) when shiftOnly = true, adds an extra outermost dimension to the indexing |
| /// map without adding that dimension as the outermost dimension in the |
| /// range. This is used for case where the same value is used as an operand |
| /// for all the concatenated operations. In such cases the original operand |
| /// can just be broadcasted along the concatenated dimension in the |
| /// horizontally-fused operation. |
| static AffineMap getConcatenatedIndexingMap(RewriterBase &rewriter, |
| AffineMap origIndexingMap, |
| bool shiftOnly = false) { |
| AffineMap newIndexingMap = origIndexingMap.shiftDims(1); |
| if (shiftOnly) { |
| return newIndexingMap; |
| } |
| return newIndexingMap.insertResult(rewriter.getAffineDimExpr(0), 0); |
| } |
| |
| /// During horizontal fusion, there might be operands of the fused operations |
| /// whose definitions are interspersed between the fused operations. For groups |
| /// chosen to fuse horizontally, such operations can be moved before the |
| /// seed contraction operation (where the fused operation is generated). |
| template <typename T> |
| static LogicalResult |
| moveOperandDefs(RewriterBase &rewriter, ArrayRef<T> operations, |
| Operation *insertionPoint, DominanceInfo &dominanceInfo, |
| ArrayRef<linalg::LinalgOp> ignoreOperations = {}) { |
| BackwardSliceOptions options; |
| llvm::DenseSet<Operation *> ignoreOperationsSet; |
| ignoreOperationsSet.insert(ignoreOperations.begin(), ignoreOperations.end()); |
| options.filter = [&](Operation *op) { |
| return !dominanceInfo.properlyDominates(op, insertionPoint) && |
| !ignoreOperationsSet.contains(op); |
| }; |
| // Set inclusive to true cause the slice is computed from the operand, and |
| // we want to include the defining op (which is the point here) |
| options.inclusive = true; |
| |
| llvm::SetVector<Operation *> slice; |
| for (auto op : operations) { |
| for (auto operand : op->getOperands()) { |
| getBackwardSlice(operand, &slice, options); |
| } |
| } |
| |
| mlir::topologicalSort(slice); |
| for (auto op : slice) { |
| rewriter.moveOpBefore(op, insertionPoint); |
| } |
| return success(); |
| } |
| |
| /// On finding this pattern |
| /// ``` |
| /// %0 = linalg.matmul ins(%arg0, %arg1) |
| /// %1 = linalg.matmul ins(%arg0, %arg2) |
| /// %2 = linalg.matmul ins(%arg0, %arg3) |
| /// ``` |
| /// |
| /// where all matmul share an operand can be combined into |
| /// rewrite to |
| /// |
| /// ``` |
| /// %arg1_r = tensor.expand_shape %arg1 [[0, 1], ...] : tensor<?x?xf32> to |
| /// tensor<1x?x?xf32> |
| /// %arg2_r = tensor.expand_shape %arg2 [[0, 1], ...] : tensor<?x?xf32> to |
| /// tensor<1x?x?xf32> |
| /// %arg3_r = tensor.expand_shape %arg3 [[0, 1], ...] : tensor<?x?xf32> to |
| /// tensor<1x?x?xf32> |
| /// %rhs = tensor.concat(%arg1_r, %arg2_r, %arg3_r) |
| /// %fused = linalg.generic { |
| /// indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d3)>, |
| /// affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>, |
| /// affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>}], |
| /// iterator_types = ["parallel", "parallel", "parallel", "reduction"]} |
| /// ins(%arg0, %rhs) ... { ... } |
| /// %0 = tensor.extract_slice %fused [0, 0, 0] ... : tensor<1x?x?xf32> to |
| /// tensor<?x?xf32> |
| /// %1 = tensor.extract_slice %fused [1, 0, 0] ... : tensor<1x?x?xf32> to |
| /// tensor<?x?xf32> |
| /// %2 = tensor.extract_slice %fused [2, 0, 0] ... : tensor<1x?x?xf32> to |
| /// tensor<?x?xf32> |
| /// ``` |
| /// |
| /// Also accounts for quantized cases where inputs are at lower precision and |
| /// accumulate is in higher-precision with truncate getting back to the |
| /// quantized sizes. |
| static LogicalResult fuseGroup(RewriterBase &rewriter, |
| HorizontalFusionGroup &fusionGroup, |
| DominanceInfo &dominanceInfo) { |
| linalg::LinalgOp baseContractOp = fusionGroup.contractionOps.front(); |
| Location loc = baseContractOp.getLoc(); |
| OpBuilder::InsertionGuard g(rewriter); |
| rewriter.setInsertionPoint(baseContractOp); |
| |
| if (failed(moveOperandDefs( |
| rewriter, ArrayRef<linalg::LinalgOp>(fusionGroup.contractionOps), |
| baseContractOp, dominanceInfo))) { |
| return baseContractOp.emitOpError("failed to re-order operand definitions"); |
| } |
| |
| SmallVector<Value> rhsValues; |
| SmallVector<Value> initValues; |
| for (auto op : fusionGroup.contractionOps) { |
| Value rhs = op.getDpsInputOperand(1)->get(); |
| Value init = op.getDpsInitOperand(0)->get(); |
| rhsValues.push_back(rhs); |
| initValues.push_back(init); |
| } |
| Value newContractRhs = concatenateValues(rewriter, loc, rhsValues); |
| Value newContractInit = concatenateValues(rewriter, loc, initValues); |
| |
| auto baseContractResultType = |
| cast<RankedTensorType>(baseContractOp->getResult(0).getType()); |
| SmallVector<int64_t> newContractResultShape = { |
| static_cast<int64_t>(rhsValues.size())}; |
| newContractResultShape.append(baseContractResultType.getShape().begin(), |
| baseContractResultType.getShape().end()); |
| auto newContractResultType = RankedTensorType::get( |
| newContractResultShape, baseContractResultType.getElementType()); |
| |
| Value lhs = baseContractOp->getOperand(0); |
| |
| SmallVector<utils::IteratorType> newContractIteratorTypes = { |
| utils::IteratorType::parallel}; |
| newContractIteratorTypes.append(baseContractOp.getIteratorTypesArray()); |
| |
| SmallVector<AffineMap> newContractIndexingMaps = |
| baseContractOp.getIndexingMapsArray(); |
| newContractIndexingMaps[0] = getConcatenatedIndexingMap( |
| rewriter, newContractIndexingMaps[0], /*shiftOnly=*/true); |
| newContractIndexingMaps[1] = |
| getConcatenatedIndexingMap(rewriter, newContractIndexingMaps[1]); |
| newContractIndexingMaps[2] = |
| getConcatenatedIndexingMap(rewriter, newContractIndexingMaps[2]); |
| |
| linalg::GenericOp newContractOp = rewriter.create<linalg::GenericOp>( |
| loc, newContractResultType, ValueRange{lhs, newContractRhs}, |
| newContractInit, newContractIndexingMaps, newContractIteratorTypes); |
| rewriter.cloneRegionBefore(baseContractOp->getRegion(0), |
| newContractOp.getRegion(), |
| newContractOp.getRegion().begin()); |
| |
| linalg::LinalgOp concatResultOp = newContractOp; |
| if (fusionGroup.truncateOps) { |
| SmallVector<Value> newTruncOperands; |
| SmallVector<AffineMap> newTruncIndexingMaps; |
| linalg::GenericOp baseTruncOp = fusionGroup.truncateOps->front(); |
| SmallVector<AffineMap> baseTruncOpIndexingMaps = |
| baseTruncOp.getIndexingMapsArray(); |
| |
| rewriter.setInsertionPoint(baseTruncOp); |
| if (failed(moveOperandDefs( |
| rewriter, |
| ArrayRef<linalg::GenericOp>(fusionGroup.truncateOps.value()), |
| baseTruncOp, dominanceInfo, fusionGroup.contractionOps))) { |
| return baseTruncOp.emitOpError( |
| "failed to move operand defs for truncate operations"); |
| } |
| |
| for (auto [operandIndex, baseTruncOperand, baseIndexingMap] : |
| llvm::enumerate(baseTruncOp->getOperands(), baseTruncOpIndexingMaps)) { |
| // Collect all the operands for the trunc operation. |
| SmallVector<Value> truncOperands; |
| for (auto truncOp : fusionGroup.truncateOps.value()) { |
| truncOperands.push_back(truncOp.getOperand(operandIndex)); |
| } |
| |
| // Three cases to handle here. |
| // Case 1. the operand is the contraction op. |
| if (llvm::all_of(llvm::zip(truncOperands, fusionGroup.contractionOps), |
| [](auto it) { |
| Value operand = std::get<0>(it); |
| return operand.getDefiningOp<linalg::LinalgOp>() == |
| std::get<1>(it); |
| })) { |
| // Use the result of the concatanted generic op |
| newTruncOperands.push_back(newContractOp.getResult(0)); |
| newTruncIndexingMaps.push_back( |
| getConcatenatedIndexingMap(rewriter, baseIndexingMap)); |
| continue; |
| } |
| |
| // Case 2. all the operands are the same. |
| if (operandIndex < baseTruncOp.getNumDpsInputs() && |
| llvm::all_equal(truncOperands)) { |
| newTruncOperands.push_back(truncOperands.front()); |
| newTruncIndexingMaps.push_back(getConcatenatedIndexingMap( |
| rewriter, baseIndexingMap, /*shiftOnly=*/true)); |
| continue; |
| } |
| |
| // Case 3. Concatenate all the operands. |
| newTruncOperands.push_back( |
| concatenateValues(rewriter, loc, truncOperands)); |
| newTruncIndexingMaps.push_back( |
| getConcatenatedIndexingMap(rewriter, baseIndexingMap)); |
| } |
| |
| // Insert truncate operator. |
| auto baseTruncType = |
| cast<RankedTensorType>(baseTruncOp.getResult(0).getType()); |
| SmallVector<int64_t> newTruncShape = { |
| static_cast<int64_t>(rhsValues.size())}; |
| newTruncShape.append(baseTruncType.getShape().begin(), |
| baseTruncType.getShape().end()); |
| auto newTruncType = |
| RankedTensorType::get(newTruncShape, baseTruncType.getElementType()); |
| SmallVector<utils::IteratorType> newTruncIteratorTypes = { |
| utils::IteratorType::parallel}; |
| newTruncIteratorTypes.append(baseTruncOp.getIteratorTypesArray()); |
| |
| ArrayRef newTruncOperandsRef(newTruncOperands); |
| linalg::GenericOp newTruncOp = rewriter.create<linalg::GenericOp>( |
| loc, newTruncType, |
| newTruncOperandsRef.take_front(baseTruncOp.getNumDpsInputs()), |
| newTruncOperandsRef.take_back(baseTruncOp.getNumDpsInits()), |
| newTruncIndexingMaps, newTruncIteratorTypes); |
| |
| rewriter.cloneRegionBefore(baseTruncOp->getRegion(0), |
| newTruncOp.getRegion(), |
| newTruncOp.getRegion().begin()); |
| |
| concatResultOp = cast<linalg::LinalgOp>(newTruncOp.getOperation()); |
| } |
| |
| SmallVector<SmallVector<OpFoldResult>> concatResultShape; |
| if (failed(concatResultOp.reifyResultShapes(rewriter, concatResultShape))) { |
| return baseContractOp.emitOpError( |
| "failed to get shape of concatenated result op"); |
| } |
| Value concatResult = concatResultOp->getResult(0); |
| MutableArrayRef<OpFoldResult> extractSizes(concatResultShape[0]); |
| extractSizes[0] = rewriter.getIndexAttr(1); |
| auto concatResultType = cast<RankedTensorType>(concatResult.getType()); |
| |
| SmallVector<OpFoldResult> extractOffsets(extractSizes.size(), |
| rewriter.getIndexAttr(0)); |
| SmallVector<OpFoldResult> extractStrides(extractSizes.size(), |
| rewriter.getIndexAttr(1)); |
| |
| auto concatResultTypeShape = |
| llvm::map_to_vector(concatResultType.getShape(), |
| [](size_t s) { return static_cast<int64_t>(s); }); |
| auto resultOutType = |
| RankedTensorType::get(ArrayRef(concatResultTypeShape).drop_front(), |
| concatResultType.getElementType()); |
| |
| SmallVector<Value> replacements; |
| for (auto i : llvm::seq<size_t>(0, rhsValues.size())) { |
| extractOffsets[0] = rewriter.getIndexAttr(static_cast<int64_t>(i)); |
| replacements.push_back(rewriter.create<tensor::ExtractSliceOp>( |
| loc, resultOutType, concatResult, extractOffsets, extractSizes, |
| extractStrides)); |
| } |
| |
| for (auto [index, op, replacement] : |
| llvm::enumerate(fusionGroup.contractionOps, replacements)) { |
| Operation *replacedOp = op; |
| if (fusionGroup.truncateOps) { |
| replacedOp = fusionGroup.truncateOps.value()[index]; |
| } |
| rewriter.replaceOp(replacedOp, replacement); |
| } |
| return success(); |
| } |
| |
| void FuseHorizontalContractionsPass::runOnOperation() { |
| MLIRContext *context = &getContext(); |
| DominanceInfo dominanceInfo(getOperation()); |
| |
| SmallVector<HorizontalFusionGroup> horizontalFusionGroups; |
| llvm::SmallDenseSet<linalg::LinalgOp> groupedOperations; |
| |
| getOperation()->walk([&](linalg::LinalgOp linalgOp) { |
| if (!isEmptyFillContractionDAGRootOp(linalgOp)) { |
| return; |
| } |
| // Avoid already grouped operations; |
| if (groupedOperations.contains(linalgOp)) { |
| return; |
| } |
| |
| std::optional<HorizontalFusionGroup> fusionGroup = |
| getHorizontalFusionGroupMembers(linalgOp, groupedOperations, |
| dominanceInfo, fusionLimit); |
| |
| if (!fusionGroup) { |
| return; |
| } |
| |
| // Update statistics. |
| numFusionGroups++; |
| switch (fusionGroup->contractionOps.size()) { |
| case 2: |
| numSize2FusionGroups++; |
| break; |
| case 3: |
| numSize3FusionGroups++; |
| break; |
| default: |
| break; |
| } |
| |
| groupedOperations.insert(fusionGroup->contractionOps.begin(), |
| fusionGroup->contractionOps.end()); |
| horizontalFusionGroups.emplace_back(std::move(fusionGroup.value())); |
| }); |
| |
| if (horizontalFusionGroups.empty()) { |
| return; |
| } |
| |
| IRRewriter rewriter(context); |
| for (auto &fusionGroup : horizontalFusionGroups) { |
| if (failed(fuseGroup(rewriter, fusionGroup, dominanceInfo))) { |
| return signalPassFailure(); |
| } |
| } |
| |
| { |
| RewritePatternSet foldReshapePatterns(context); |
| tensor::populateFoldTensorEmptyPatterns(foldReshapePatterns); |
| linalg::FillOp::getCanonicalizationPatterns(foldReshapePatterns, context); |
| if (failed(applyPatternsGreedily(getOperation(), |
| std::move(foldReshapePatterns)))) { |
| getOperation()->emitOpError("failed during reshape folding patterns"); |
| return signalPassFailure(); |
| } |
| |
| RewritePatternSet foldPatterns(context); |
| tensor::populateFoldTensorEmptyPatterns(foldPatterns); |
| linalg::FillOp::getCanonicalizationPatterns(foldPatterns, context); |
| if (failed( |
| applyPatternsGreedily(getOperation(), std::move(foldPatterns)))) { |
| getOperation()->emitOpError("failed to fold empty/fill with concats"); |
| return signalPassFailure(); |
| } |
| } |
| |
| // Note: Currently these patterns are required due to early lowering of |
| // tensor.concat. When we choose to move the lowering of tensor.concat later, |
| // these patterns should be dropped. |
| RewritePatternSet patterns(context); |
| tensor::populateDecomposeTensorConcatPatterns(patterns); |
| if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) { |
| return signalPassFailure(); |
| } |
| } |
| } // namespace mlir::iree_compiler::DispatchCreation |
