[GlobalOptimization] Expose reusable transpose propagation patterns (#24978) I ran into multiple scenarios in the past (and now again for #24906) where i had to reimplement some form of transpose propagation. To improve this, this PR refactors the transpose propagation pass into reusable patterns. This PR should not cause any functional changes and should be a pure refactor. Assisted by: Claude Signed-off-by: Thomas Ziereis <ziereis@roofline.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel index e971773..f15a9fc 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel +++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -101,11 +101,13 @@ "RaiseSpecialOps.cpp", "RemoveZeroExtentTensors.cpp", "SimplifyPackUnpack.cpp", + "TransposePropagationPatterns.cpp", "Utils.cpp", "WarnOnUninitializedValues.cpp", ], hdrs = [ "Passes.h", + "TransposePropagationPatterns.h", "Utils.h", ], deps = [
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt index e1afc48..068698a 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt +++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -68,6 +68,7 @@ GlobalOptimization HDRS "Passes.h" + "TransposePropagationPatterns.h" "Utils.h" SRCS "CleanupNumericNarrowing.cpp" @@ -94,6 +95,7 @@ "RaiseSpecialOps.cpp" "RemoveZeroExtentTensors.cpp" "SimplifyPackUnpack.cpp" + "TransposePropagationPatterns.cpp" "Utils.cpp" "WarnOnUninitializedValues.cpp" DEPS
diff --git a/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp b/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp index c92fa0d..2541af8 100644 --- a/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp +++ b/compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp
@@ -13,21 +13,17 @@ #include "iree/compiler/Dialect/Flow/Conversion/TensorToFlow/Utils.h" #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h" +#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h" #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h" -#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h" -#include "iree/compiler/Dialect/Util/IR/UtilOps.h" #include "iree/compiler/GlobalOptimization/Passes.h" +#include "iree/compiler/GlobalOptimization/TransposePropagationPatterns.h" #include "llvm/Support/Debug.h" -#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" #include "mlir/Dialect/MemRef/Transforms/Transforms.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Tensor/Transforms/Transforms.h" #include "mlir/Dialect/Utils/IndexingUtils.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -39,54 +35,9 @@ #include "iree/compiler/GlobalOptimization/Passes.h.inc" //===----------------------------------------------------------------------===// -// Transpose permutation helpers +// Propagation control //===----------------------------------------------------------------------===// -static bool isIdentityPermutation(ArrayRef<int64_t> perm) { - for (auto [index, dim] : llvm::enumerate(perm)) { - if (index != dim) { - return false; - } - } - return true; -} - -// Constructs a transpose of the given tensor and permutation. -static Value createTransposeInit(OpBuilder &builder, Value source, - ArrayRef<int64_t> perm) { - SmallVector<OpFoldResult> mixedSizes = - tensor::getMixedSizes(builder, source.getLoc(), source); - applyPermutationToVector(mixedSizes, perm); - Type elemType = cast<RankedTensorType>(source.getType()).getElementType(); - Value empty = - tensor::EmptyOp::create(builder, source.getLoc(), mixedSizes, elemType) - .getResult(); - return empty; -} - -// Constructs a transpose of the given tensor and permutation, -// or produces a transposed version of the producing tensor.empty op. -static Value createTranspose(OpBuilder &builder, Value source, - ArrayRef<int64_t> perm) { - if (auto empty = source.getDefiningOp<tensor::EmptyOp>()) { - Type elementType = empty.getType().getElementType(); - SmallVector<OpFoldResult> mixedSizes = empty.getMixedSizes(); - applyPermutationToVector(mixedSizes, perm); - return tensor::EmptyOp::create(builder, empty.getLoc(), mixedSizes, - elementType); - } - Value empty = createTransposeInit(builder, source, perm); - return linalg::TransposeOp::create(builder, source.getLoc(), source, empty, - perm) - ->getResult(0); -} - -static RankedTensorType getPermutedTensorType(RankedTensorType type, - SmallVector<int64_t> perm) { - SmallVector<int64_t> permutedShape = applyPermutation(type.getShape(), perm); - return RankedTensorType::get(permutedShape, type.getElementType()); -} - static bool isReshapeBlockingFusion(Operation *producer, Operation *consumer) { auto isFusableOp = [](Operation *op) { if (!op) { @@ -99,1032 +50,43 @@ return isFusableOp(producer) && isFusableOp(consumer); } -//===----------------------------------------------------------------------===// -// Transpose specialization -//===----------------------------------------------------------------------===// - -// Specializes linalg.generic op to linalg.transpose if it is transposing a -// single input. -static void specializeGenericTransposeOp(RewriterBase &rewriter, - linalg::GenericOp genericOp) { - if (!IREE::LinalgExt::isaTransposeOpInterface(genericOp)) { - return; - } - - auto mapRange = genericOp.getIndexingMapsArray(); - AffineMap outMap = mapRange.back(); - AffineMap inMap = mapRange.front(); - SmallVector<int64_t> perm; - // To get the permutation, look at each output index and find which - // dimension in the input we're reading from for that index. - for (AffineExpr expr : outMap.getResults()) { - perm.push_back(*inMap.getResultPosition(expr)); - } - rewriter.replaceOpWithNewOp<linalg::TransposeOp>( - genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0], perm); +// Restricts propagation to edges whose ops are both outside of pre-formed +// dispatches. +static bool isOutsideDispatch(OpOperand *operand) { + return IREE::Flow::isNonNullAndOutsideDispatch( + {operand->get().getDefiningOp(), operand->getOwner()}); } -//===----------------------------------------------------------------------===// -// Other pattern helpers -//===----------------------------------------------------------------------===// +// Profitability heuristic for sinking a transpose through a +// tensor.extract_slice: accepts the edge only when the slice, taken from the +// untransposed source instead, is still mappable to Flow. +static bool isUntransposedSliceMappableToFlow(OpOperand *operand) { + if (!isOutsideDispatch(operand)) { + return false; + } + auto transposeOp = cast<linalg::TransposeOp>(operand->get().getDefiningOp()); + auto extractOp = cast<tensor::ExtractSliceOp>(operand->getOwner()); + auto invPerm = invertPermutationVector(transposeOp.getPermutation()); + SmallVector<OpFoldResult> offsets = extractOp.getMixedOffsets(); + SmallVector<OpFoldResult> sizes = extractOp.getMixedSizes(); + SmallVector<OpFoldResult> strides = extractOp.getMixedStrides(); + ArrayRef<int64_t> srcShape = extractOp.getSourceType().getShape(); -/// Returns the `op` if it is a linalg::GenericOp. If it is a named op and -/// `allowGeneralizing` is true, returns the generalized op. Otherwise, returns -/// failure. -/// TODO: Due to fragility around handling of convolutions, convolution -/// propagation is behind a flag. -static FailureOr<linalg::GenericOp> -getAllowedGenericOpOrGeneralizeNamedOp(RewriterBase &rewriter, Operation *op, - bool allowGeneralizing, bool convProp) { - auto linalgOp = dyn_cast<linalg::LinalgOp>(op); - if (!linalgOp) { - return failure(); - } + // Permute the offsets, sizes, and strides to pre-transpose ordering. + applyPermutationToVector(offsets, invPerm); + applyPermutationToVector(sizes, invPerm); + applyPermutationToVector(strides, invPerm); + SmallVector<int64_t> baseShape = applyPermutation(srcShape, invPerm); - if (!convProp && linalg::isaConvolutionOpInterface(linalgOp)) { - return failure(); - } - - if (!isa<linalg::GenericOp>(linalgOp) && - !(allowGeneralizing && linalg::isaContractionOpInterface(linalgOp)) && - !(convProp && linalg::isaConvolutionOpInterface(linalgOp))) { - return failure(); - } - - auto genericOp = dyn_cast<linalg::GenericOp>(op); - if (genericOp) { - return genericOp; - } - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPoint(linalgOp); - FailureOr<linalg::LinalgOp> generalizedOp = - linalg::generalizeNamedOp(rewriter, linalgOp); - if (failed(generalizedOp)) { - return failure(); - } - return cast<linalg::GenericOp>(generalizedOp->getOperation()); + // Check if the resulting offsets, sizes, and strides correspond to a + // contiguous slice and can thus be mappable to a `flow.tensor.update` op. + // This should always be worth doing because this can remove a dispatch for + // the slice, and the transpose is on the slice rather than the full tensor. + return IREE::Flow::isOffsetSizeAndStrideMappableToFlow(offsets, sizes, + strides, baseShape); } //===----------------------------------------------------------------------===// -// Transpose Bubbling Patterns -//===----------------------------------------------------------------------===// - -namespace { - -// Fuses a transpose with the init of a linalg.generic op or contraction op. -// Contraction ops are generalized and then treated as a generic. For example, -// -// linalg.generic { -// indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, -// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] -// ins(%0 : tensor<2x7x5>) outs(%1 : tensor<7x2x5>) -// -// %2 = linalg.transpose ... permutation = [0, 2, 1] : -// tensor<7x2x5> -> tensor<7x5x2> -// Becomes -// -// linalg.generic { -// indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0, d1)>, -// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] -// ins(%0 : tensor<2x7x5>) outs(%3 : tensor<7x5x2>) -class FuseTransposeWithProducerLinalgOp - : public OpRewritePattern<linalg::TransposeOp> { -public: - using Base::Base; - FuseTransposeWithProducerLinalgOp(MLIRContext *ctx, bool aggressiveProp, - bool convProp, PatternBenefit b = 1) - : OpRewritePattern<linalg::TransposeOp>(ctx, b), - allowGeneralizing(aggressiveProp), convProp(convProp) {} - - LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(transposeOp)) { - return failure(); - } - OpResult result = dyn_cast<OpResult>(transposeOp.getInput()); - if (!result) { - return rewriter.notifyMatchFailure( - transposeOp, "transpose input defined by block argument"); - } - if (!result.hasOneUse()) { - return rewriter.notifyMatchFailure(transposeOp, - "multi use transpose input"); - } - auto linalgOp = dyn_cast<linalg::LinalgOp>(result.getOwner()); - if (!linalgOp) { - return rewriter.notifyMatchFailure( - transposeOp, "non-linalg op producer for transpose input"); - } - - int64_t resultIndex = result.getResultNumber(); - auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp( - rewriter, result.getOwner(), allowGeneralizing, convProp); - if (failed(maybeGenericOp)) { - return rewriter.notifyMatchFailure( - transposeOp, "linalg op producer is not generic or contraction"); - } - - auto genericOp = maybeGenericOp.value(); - result = genericOp->getOpResult(resultIndex); - - ArrayRef<int64_t> perm = transposeOp.getPermutation(); - auto invPerm = invertPermutationVector(perm); - - // 1. Get the transposed init of the generic. - Value init = genericOp.getDpsInits()[resultIndex]; - SmallVector<Value> inits = genericOp.getDpsInits(); - Value newInit = createTranspose(rewriter, init, perm); - inits[resultIndex] = newInit; - - SmallVector<Type> resultTypes(genericOp->getResultTypes()); - resultTypes[resultIndex] = newInit.getType(); - - // 2. Update the indexing map of the transposed init operand by permuting - // the results of the map. - SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray(); - AffineMap resultMap = - newIndexingMaps[genericOp.getNumDpsInputs() + resultIndex]; - SmallVector<AffineExpr> newExprs = - applyPermutation(resultMap.getResults(), perm); - AffineMap transposedMap = - AffineMap::get(resultMap.getNumDims(), resultMap.getNumSymbols(), - newExprs, rewriter.getContext()); - newIndexingMaps[genericOp.getNumDpsInputs() + resultIndex] = transposedMap; - - // 3. Create the new generic with the same iteration order. - auto newGenericOp = linalg::GenericOp::create( - rewriter, genericOp.getLoc(), resultTypes, genericOp.getDpsInputs(), - newInit, newIndexingMaps, genericOp.getIteratorTypesArray(), - /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp)); - rewriter.cloneRegionBefore(genericOp.getRegion(), newGenericOp.getRegion(), - newGenericOp.getRegion().begin()); - - // 4. Remap iteration space of the generic to match the dimension order of - // the output. - if (newGenericOp.getNumResults() == 1) { - SmallVector<unsigned int> interchange; - int64_t permIdx = 0; - for (int i = 0, e = transposedMap.getNumDims(); i < e; ++i) { - if (transposedMap.isFunctionOfDim(i)) { - interchange.push_back( - cast<AffineDimExpr>(transposedMap.getResult(permIdx)) - .getPosition()); - permIdx++; - continue; - } - interchange.push_back(i); - } - auto interchangedGenericOp = - linalg::interchangeGenericOp(rewriter, newGenericOp, interchange); - // Interchange only fails if interchangeGenericOpPrecondition fails, which - // only fails if the interchange vector is not invertible or doesn't match - // the number of loops in the generic, both of which are guaranteed by - // the fact that the output map must be a projection in the above - // construction. - assert(succeeded(interchangedGenericOp) && - "failed to interchange transposed generic"); - newGenericOp = *interchangedGenericOp; - } - - // 5. Replace the result of the transpose with the transposed init. - rewriter.replaceOp(transposeOp, newGenericOp->getResult(resultIndex)); - for (auto [oldRes, newRes] : - llvm::zip_equal(genericOp.getResults(), newGenericOp->getResults())) { - if (oldRes.getResultNumber() == resultIndex) { - continue; - } - rewriter.replaceAllUsesWith(oldRes, newRes); - } - return success(); - } - -private: - bool allowGeneralizing = false; - bool convProp = false; -}; - -// Bubbles a transpose through a tensor.collapse_shape. -class BubbleTransposeThroughCollapseShape - : public OpRewritePattern<linalg::TransposeOp> { -public: - using Base::Base; - BubbleTransposeThroughCollapseShape(MLIRContext *ctx, - bool enableEdgeReshapeProp, - PatternBenefit b = 1) - : OpRewritePattern<linalg::TransposeOp>(ctx, b), - enableEdgeReshapePropagation(enableEdgeReshapeProp) {} - - LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(transposeOp)) { - return failure(); - } - Value source = transposeOp.getDpsInputOperand(0)->get(); - auto collapseOp = source.getDefiningOp<tensor::CollapseShapeOp>(); - // Do not propagate through reshapes if the transpose has multiple users, as - // this could end up duplicating the transposes. We should only propagate - // through reshape when it is free to do so. - if (!collapseOp || !collapseOp->hasOneUse()) { - return rewriter.notifyMatchFailure( - transposeOp, "transpose input is not a single-use collapse shape"); - } - - if (!enableEdgeReshapePropagation && - !isReshapeBlockingFusion(transposeOp, - collapseOp.getSrc().getDefiningOp())) { - return rewriter.notifyMatchFailure(transposeOp, - "transpose not blocking fusion"); - } - - SmallVector<ReassociationIndices> reassociations = - collapseOp.getReassociationIndices(); - - // Because we are doing transpose(collapse_shape), all expanded groups are - // transposed together. As a result, to get the permutation of the new - // transpose, we can just flatten the transposed reassociation indices. - // For example, - // - // reassociation_map = [[0, 1, 2], [3], [4, 5]] - // permutation = [1, 2, 0] - // - // Becomes - // - // permutation = [3, 4, 5, 0, 1, 2] - // reassociation_map = [[0], [1, 2], [3, 4, 5]] - applyPermutationToVector(reassociations, transposeOp.getPermutation()); - - SmallVector<int64_t> newPerm; - SmallVector<ReassociationIndices> newReassociations; - int64_t expandedDim = 0; - for (auto reassoc : reassociations) { - ReassociationIndices newReassoc; - for (auto dim : reassoc) { - newPerm.push_back(dim); - newReassoc.push_back(expandedDim++); - } - newReassociations.push_back(newReassoc); - } - - Value newTranspose = - createTranspose(rewriter, collapseOp.getSrc(), newPerm); - Value newReshape = tensor::CollapseShapeOp::create( - rewriter, collapseOp.getLoc(), transposeOp.getResultTypes()[0], - newTranspose, newReassociations); - rewriter.replaceOp(transposeOp, newReshape); - return success(); - } - -private: - bool enableEdgeReshapePropagation = true; -}; - -} // namespace - -//===----------------------------------------------------------------------===// -// Transpose Sinking Patterns -//===----------------------------------------------------------------------===// - -namespace { - -// Combines two transposes into one. This shouldn't be strictly necessary as -// fusion should cancel inverse transposes, but doing this here can open up -// new propagation opportunities and eases the analysis in fusion/later passes. -class ComposeTransposes : public OpRewritePattern<linalg::TransposeOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::TransposeOp consumer, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(consumer)) { - return failure(); - } - Value input = consumer.getInput(); - auto producer = input.getDefiningOp<linalg::TransposeOp>(); - if (!producer) { - return failure(); - } - - ArrayRef<int64_t> producerPerm = producer.getPermutation(); - ArrayRef<int64_t> consumerPerm = consumer.getPermutation(); - SmallVector<int64_t> composedPerm = - applyPermutation(producerPerm, consumerPerm); - - Value transposedSource = producer.getInput(); - if (!isIdentityPermutation(composedPerm)) { - transposedSource = - createTranspose(rewriter, transposedSource, composedPerm); - } - rewriter.replaceOp(consumer, transposedSource); - return success(); - } -}; - -// Sinks a transpose through a tensor.extract_slice iff the transpose turns -// the extracted slice into a contiguous slice. -class SinkTransposeThroughExtractSlice - : public OpRewritePattern<tensor::ExtractSliceOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(extractOp)) { - return failure(); - } - Value source = extractOp.getSource(); - auto transposeOp = source.getDefiningOp<linalg::TransposeOp>(); - if (!transposeOp) { - return failure(); - } - - // Applying `perm` takes a list from the pre-transpose ordering to the - // post-transpose ordering. - ArrayRef<int64_t> perm = transposeOp.getPermutation(); - // Applying `invPerm` takes a list from the post-transpose ordering to the - // pre-transpose ordering. Sinking a transpose through an op is largely a - // matter of rewriting it in pre-transpose space, and thus just applying - // the inverse permutation. - auto invPerm = invertPermutationVector(perm); - - SmallVector<OpFoldResult> offsets = extractOp.getMixedOffsets(); - SmallVector<OpFoldResult> sizes = extractOp.getMixedSizes(); - SmallVector<OpFoldResult> strides = extractOp.getMixedStrides(); - ArrayRef<int64_t> srcShape = extractOp.getSourceType().getShape(); - - // Permute the offsets, sizes, and strides to pre-transpose ordering. - applyPermutationToVector(offsets, invPerm); - applyPermutationToVector(sizes, invPerm); - applyPermutationToVector(strides, invPerm); - SmallVector<int64_t> baseShape = applyPermutation(srcShape, invPerm); - - // Check if the resulting offsets, sizes, and strides correspond to a - // contiguous slice and can thus be mappable to a `flow.tensor.update` op. - // This should always be worth doing because this can remove a dispatch for - // the slice, and the transpose is on the slice rather than the full tensor. - if (!IREE::Flow::isOffsetSizeAndStrideMappableToFlow(offsets, sizes, - strides, baseShape)) { - return rewriter.notifyMatchFailure( - extractOp, "transposed slice not mappable to flow ops"); - } - - ArrayRef<int64_t> staticSizes = extractOp.getStaticSizes(); - ArrayRef<int64_t> sliceShape = extractOp.getResultType().getShape(); - std::optional<llvm::SmallDenseSet<unsigned>> maybeRankReducingMask = - mlir::computeRankReductionMask(staticSizes, sliceShape); - if (!maybeRankReducingMask) { - return rewriter.notifyMatchFailure( - extractOp, "failed to compute rank reducing mask"); - } - llvm::SmallDenseSet<unsigned> rankReducingMask = *maybeRankReducingMask; - - // Find rank reducing map in the pre-transposed domain. - int64_t dim = 0; - llvm::SmallDenseMap<int64_t, int64_t> rankReducedMap; - // Since `dim` is in the pre-transposed domain, and is incrementing each - // iteration, `idx` must also be in the pre-transposed domain. - for (int64_t idx = 0, e = invPerm.size(); idx < e; ++idx) { - // Get index in the transposed domain, since `rankReducingMask` is in - // the transposed domain. - if (!rankReducingMask.contains(invPerm[idx])) { - // Domain of `rankReducedMap` is in pre-transposed domain. - rankReducedMap[idx] = dim++; - } - } - - // Compute the new permutation by dropping all rank-reduced dimensions. - SmallVector<int64_t> rankReducedPerm; - for (int64_t i : perm) { - if (rankReducedMap.contains(i)) { - rankReducedPerm.push_back(rankReducedMap[i]); - } - } - - auto rankReducedInvPerm = invertPermutationVector(rankReducedPerm); - - RankedTensorType sliceType = getPermutedTensorType( - cast<RankedTensorType>(extractOp.getType()), rankReducedInvPerm); - Value slice = tensor::ExtractSliceOp::create( - rewriter, extractOp.getLoc(), sliceType, transposeOp.getInput(), - offsets, sizes, strides); - // Transpose back to the original slice. - if (!isIdentityPermutation(rankReducedPerm)) { - slice = createTranspose(rewriter, slice, rankReducedPerm); - } - rewriter.replaceOp(extractOp, slice); - return success(); - } -}; - -// Sinks a transpose through a tensor.expand_shape. -class SinkTransposeThroughExpandShape - : public OpRewritePattern<tensor::ExpandShapeOp> { -public: - using Base::Base; - SinkTransposeThroughExpandShape(MLIRContext *ctx, bool enableEdgeReshapeProp, - PatternBenefit b = 1) - : OpRewritePattern<tensor::ExpandShapeOp>(ctx, b), - enableEdgeReshapePropagation(enableEdgeReshapeProp) {} - - LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(expandOp)) { - return failure(); - } - Value source = expandOp.getSrc(); - auto transposeOp = source.getDefiningOp<linalg::TransposeOp>(); - // Do not propagate through reshapes if the transpose has multiple users, as - // this could end up duplicating the transposes. We should only propagate - // through reshape when it is free to do so. - if (!transposeOp || !transposeOp->hasOneUse()) { - return rewriter.notifyMatchFailure( - expandOp, "expand shape input is not a single-use transpose"); - } - - if (!enableEdgeReshapePropagation && - llvm::none_of(expandOp->getUsers(), [&](Operation *consumer) { - return isReshapeBlockingFusion(transposeOp, consumer); - })) { - return rewriter.notifyMatchFailure(transposeOp, - "transpose not blocking fusion"); - } - - auto invPerm = invertPermutationVector(transposeOp.getPermutation()); - SmallVector<ReassociationIndices> reassociations = - expandOp.getReassociationIndices(); - - // Because we are doing expand_shape(transpose), all expanded groups are - // transposed together. As a result, to get the permutation of the new - // transpose, we can just flatten the transposed reassociation indices. - // For example, - // - // permutation = [0, 2, 1] - // reassociation_map = [[0, 1, 2], [3], [4, 5]] - // - // Becomes - // - // reassociation_map = [[0, 1, 2], [3, 4], [5]] - // permutation = [0, 1, 2, 4, 5, 3] - applyPermutationToVector(reassociations, invPerm); - - SmallVector<int64_t> newInvPerm; - SmallVector<ReassociationIndices> newReassociations; - int64_t expandedDim = 0; - for (auto reassoc : reassociations) { - ReassociationIndices newReassoc; - for (auto dim : reassoc) { - newInvPerm.push_back(dim); - newReassoc.push_back(expandedDim++); - } - newReassociations.push_back(newReassoc); - } - - auto newPerm = invertPermutationVector(newInvPerm); - - RankedTensorType expandedType = getPermutedTensorType( - cast<RankedTensorType>(expandOp.getType()), newInvPerm); - // Permute the output shape to match the permuted type. - SmallVector<OpFoldResult> permutedOutputShape = - expandOp.getMixedOutputShape(); - applyPermutationToVector(permutedOutputShape, newInvPerm); - Value transposedReshape = tensor::ExpandShapeOp::create( - rewriter, expandOp.getLoc(), expandedType, transposeOp.getInput(), - newReassociations, permutedOutputShape); - Value originalReshape = - createTranspose(rewriter, transposedReshape, newPerm); - rewriter.replaceOp(expandOp, originalReshape); - return success(); - } - -private: - bool enableEdgeReshapePropagation = true; -}; - -// Sinks a transpose through a tensor.pad. -class SinkTransposeThroughPad : public OpRewritePattern<tensor::PadOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(tensor::PadOp padOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(padOp)) { - return failure(); - } - Value source = padOp.getSource(); - auto transposeOp = source.getDefiningOp<linalg::TransposeOp>(); - if (!transposeOp) { - return failure(); - } - - Block &block = padOp.getRegion().front(); - if (llvm::any_of(block.getArguments(), [](BlockArgument blockArg) { - return blockArg.getNumUses(); - })) { - return failure(); - } - - auto invPerm = invertPermutationVector(transposeOp.getPermutation()); - SmallVector<OpFoldResult> lowSizes = padOp.getMixedLowPad(); - SmallVector<OpFoldResult> highSizes = padOp.getMixedHighPad(); - applyPermutationToVector(lowSizes, invPerm); - applyPermutationToVector(highSizes, invPerm); - - RankedTensorType oldPaddedType = cast<RankedTensorType>(padOp.getType()); - RankedTensorType newPaddedType = oldPaddedType.clone( - applyPermutation(oldPaddedType.getShape(), invPerm)); - - auto newPadOp = tensor::PadOp::create( - rewriter, padOp.getLoc(), newPaddedType, transposeOp.getInput(), - lowSizes, highSizes, padOp.getNofold()); - rewriter.cloneRegionBefore(padOp.getRegion(), newPadOp.getRegion(), - newPadOp.getRegion().begin()); - - Value newTransposeOp = - createTranspose(rewriter, newPadOp, transposeOp.getPermutation()); - rewriter.replaceOp(padOp, newTransposeOp); - return success(); - } -}; - -// Fuses a transpose with the input of a linalg.generic op or contraction op. -// Contraction ops are generalized and then treated as a generic. For example, -// -// %0 = linalg.transpose ... permutation = [0, 2, 1] : -// tensor<2x5x7> -> tensor<2x7x5> -// linalg.generic { -// indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, -// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] -// ins(%0 : tensor<2x7x5>) -// -// Becomes -// -// linalg.generic { -// indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, -// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] -// ins(%0 : tensor<2x5x7>) -// -// This is considered just one way to model transpose propagation to generics, -// another option would be to interpret the transpose on the iterators of the -// generic, thus producing a transpose on the output and any other inputs to -// the generic. This has the potential introduce more transposes/data movement -// and isn't the way this pass is modeled. Global data layout transformations -// like that are better suited for pack/unpack propagation rooted on specific -// operations. -// -// TODO: Rewrite this to use elementwise op fusion patterns. -class FuseTransposeWithLinalgOpConsumer - : public OpInterfaceRewritePattern<linalg::LinalgOp> { -public: - using OpInterfaceRewritePattern<linalg::LinalgOp>::OpInterfaceRewritePattern; - FuseTransposeWithLinalgOpConsumer(MLIRContext *ctx, bool aggressiveProp, - bool convProp, PatternBenefit b = 1) - : OpInterfaceRewritePattern<linalg::LinalgOp>(ctx, b), - allowGeneralizing(aggressiveProp), convProp(convProp) {} - - LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(linalgOp)) { - return failure(); - } - OpOperand *transposeOperand = nullptr; - linalg::TransposeOp transposeOp; - for (OpOperand *input : linalgOp.getDpsInputOperands()) { - auto maybeTransposeOp = input->get().getDefiningOp<linalg::TransposeOp>(); - if (maybeTransposeOp && maybeTransposeOp->hasOneUse()) { - transposeOp = maybeTransposeOp; - transposeOperand = input; - break; - } - } - if (!transposeOperand) { - return rewriter.notifyMatchFailure(linalgOp, "no transpose operand"); - } - - int64_t inputIndex = transposeOperand->getOperandNumber(); - ArrayRef<int64_t> perm = transposeOp.getPermutation(); - auto invPerm = invertPermutationVector(perm); - - // To do the fusion, we can simply apply the permutation of the transpose - // to the results of the associated input's indexing map, and then forward - // the input to the transpose to the consumer generic. - auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp( - rewriter, linalgOp, allowGeneralizing, convProp); - if (failed(maybeGenericOp)) { - return failure(); - } - auto genericOp = maybeGenericOp.value(); - transposeOperand = genericOp.getDpsInputOperand(inputIndex); - rewriter.startOpModification(genericOp); - - SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray(); - AffineMap inputMap = genericOp.getMatchingIndexingMap(transposeOperand); - SmallVector<AffineExpr> newExprs = - applyPermutation(inputMap.getResults(), invPerm); - AffineMap transposedMap = - AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(), - newExprs, rewriter.getContext()); - newIndexingMaps[inputIndex] = transposedMap; - genericOp.setIndexingMapsAttr( - rewriter.getAffineMapArrayAttr(newIndexingMaps)); - - genericOp.setOperand(inputIndex, transposeOp.getInput()); - rewriter.finalizeOpModification(genericOp); - return success(); - } - -private: - bool allowGeneralizing = false; - bool convProp = false; -}; - -static bool isIndexingMapAffectedByTransposeMap( - AffineMap indexingMap, ArrayRef<int64_t> iterationSpacePermutation) { - int64_t prevIdx = -1; - for (auto result : indexingMap.getResults()) { - int64_t idx = - iterationSpacePermutation[cast<AffineDimExpr>(result).getPosition()]; - // Verify that the relative ordering of indices in the map remain the same. - // If not, then the transposition affects the access order for the given - // map (and associated operand). - if (idx <= prevIdx) { - return true; - } - prevIdx = idx; - } - return false; -} - -// Finds a single DPS input operand of the given |genericOp| that is affected by -// the |iterationSpacePermutation|. In other words, the permutation changes the -// relative ordering of any of the dimensions of that input operand. -// -// For example, with permutation [1, 0, 2], affine map (d0, d1, d2) -> (d0, d1) -// is affected by the permutation because the first two dimensions are iterated -// in a different order while (d0, d1, d2) -> (d0, d2) is unaffected. -// -// If no such operand is found or there is more than one such operation, nullptr -// is returned. -static OpOperand * -getSingleTransposedInputOperand(linalg::GenericOp genericOp, - ArrayRef<int64_t> iterationSpacePermutation) { - OpOperand *operand = nullptr; - for (auto input : genericOp.getDpsInputOperands()) { - if (!isIndexingMapAffectedByTransposeMap( - genericOp.getMatchingIndexingMap(input), - iterationSpacePermutation)) { - continue; - } - if (operand) { - return nullptr; - } - operand = input; - } - return operand; -} - -// Returns a new list of indexing maps that composes the iteration space -// permutation map |transposeMap| with all indexing maps of |genericOp| except -// for the |transposedInputIdx|'th operand. The unchanged operand is expected -// to have an explicit `linalg.transpose` op constructed for it so its map does -// not need to be updated. -static SmallVector<AffineMap> -getTransposedIndexingMaps(linalg::GenericOp genericOp, - int64_t transposedInputIdx, AffineMap transposeMap) { - SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray(); - for (unsigned i = 0, e = genericOp.getNumDpsInputs(); i < e; ++i) { - if (i == transposedInputIdx) { - continue; - } - indexingMaps[i] = indexingMaps[i].compose(transposeMap); - } - return indexingMaps; -} - -// Sinks a transpose through the input of a elementwise operation where the -// transposition of the iteration space only affects a single input operand. -class SinkTransposeThroughUnaryElementwiseInput - : public OpRewritePattern<linalg::GenericOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::GenericOp genericOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(genericOp)) { - return rewriter.notifyMatchFailure(genericOp, "pre-formed dispatch"); - } - - if (!linalg::isElementwise(genericOp)) { - return rewriter.notifyMatchFailure(genericOp, "non-elementwise generic"); - } - - if (genericOp.hasIndexSemantics()) { - return rewriter.notifyMatchFailure(genericOp, "has index semantics"); - } - - if (genericOp.getNumDpsInits() != 1) { - return rewriter.notifyMatchFailure(genericOp, - "unimplemented: multiple results"); - } - - AffineMap resultMap = - genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0)); - if (!resultMap.isIdentity()) { - return rewriter.notifyMatchFailure( - genericOp, "unimplemented: non-identity result map"); - } - - linalg::TransposeOp transposeOp; - OpOperand *inputOperand; - for (auto input : genericOp.getDpsInputOperands()) { - // Skip broadcasted operands and transposed operands. If the input is - // broadcasted then we would not want to propagate because that would - // do the transpose on larger data, and if transposed we would rather - // simply compose the transposes (handled in a separate pattern). - if (genericOp.getMatchingIndexingMap(input) != resultMap) { - continue; - } - - auto maybeTransposeOp = input->get().getDefiningOp<linalg::TransposeOp>(); - // Skip multi-use transposes. - if (!maybeTransposeOp || !maybeTransposeOp->hasOneUse()) { - continue; - } - - auto transposableInputOperand = getSingleTransposedInputOperand( - genericOp, maybeTransposeOp.getPermutation()); - // Skip if more than one operand is affected by the transpose. - if (transposableInputOperand != input) { - continue; - } - - transposeOp = maybeTransposeOp; - inputOperand = transposableInputOperand; - break; - } - - if (!transposeOp) { - return rewriter.notifyMatchFailure(genericOp, - "no single use transpose operand"); - } - - ArrayRef<int64_t> perm = transposeOp.getPermutation(); - auto invPerm = invertPermutationVector(perm); - - // Create a new empty init for the transposed generic. - Value newInit = - createTransposeInit(rewriter, genericOp.getDpsInits()[0], invPerm); - - // We do not need to update iterator types because this is an elementwise - // op. We just need to update the indexing maps of all other input operands - // by composing the transpose map. - AffineMap transposeMap = - AffineMap::getPermutationMap(perm, rewriter.getContext()); - SmallVector<AffineMap> indexingMaps = getTransposedIndexingMaps( - genericOp, inputOperand->getOperandNumber(), transposeMap); - - SmallVector<Value> newOperands = genericOp->getOperands(); - newOperands[inputOperand->getOperandNumber()] = transposeOp.getInput(); - newOperands[genericOp.getDpsInitOperand(0)->getOperandNumber()] = newInit; - - auto newGenericOp = - mlir::clone(rewriter, genericOp, newInit.getType(), newOperands); - newGenericOp.setIndexingMapsAttr( - rewriter.getAffineMapArrayAttr(indexingMaps)); - rewriter.replaceOp( - genericOp, createTranspose(rewriter, newGenericOp->getResult(0), perm)); - return success(); - } -}; - -// Bubbles a transpose through the init of a elementwise operation where the -// transposition of the iteration space only affects a single input operand. -class BubbleTransposeThroughUnaryElementwiseDpsInit - : public OpRewritePattern<linalg::TransposeOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, - PatternRewriter &rewriter) const override { - auto genericOp = transposeOp.getInput().getDefiningOp<linalg::GenericOp>(); - if (!genericOp) { - return rewriter.notifyMatchFailure(transposeOp, "non-generic producer"); - } - - if (genericOp.getNumDpsInits() != 1) { - return rewriter.notifyMatchFailure(transposeOp, - "unimplemented: multiple results"); - } - - if (!IREE::Flow::isNonNullAndOutsideDispatch({genericOp, transposeOp})) { - return failure(); - } - - if (!linalg::isElementwise(genericOp) || - !genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0)) - .isIdentity()) { - return rewriter.notifyMatchFailure(transposeOp, "not elementwise"); - } - - if (genericOp.hasIndexSemantics()) { - return rewriter.notifyMatchFailure(genericOp, "has index semantics"); - } - - if (!genericOp->hasOneUse()) { - return rewriter.notifyMatchFailure(transposeOp, "not single user"); - } - - ArrayRef<int64_t> perm = transposeOp.getPermutation(); - auto invPerm = invertPermutationVector(perm); - - auto inputOperand = getSingleTransposedInputOperand(genericOp, invPerm); - if (!inputOperand || - !genericOp.getMatchingIndexingMap(inputOperand).isIdentity()) { - return rewriter.notifyMatchFailure( - genericOp, "no single transposable input operand"); - } - - Value newTranspose = createTranspose(rewriter, inputOperand->get(), perm); - - // Create a new empty init for the transposed generic. - Value newInit = - createTransposeInit(rewriter, genericOp.getDpsInits()[0], perm); - - SmallVector<Value> newOperands = genericOp->getOperands(); - newOperands[inputOperand->getOperandNumber()] = newTranspose; - newOperands[genericOp.getDpsInitOperand(0)->getOperandNumber()] = newInit; - - AffineMap transposeMap = - AffineMap::getPermutationMap(invPerm, rewriter.getContext()); - - // We do not need to update iterator types because this is an elementwise - // op. We just need to update the indexing maps of all other input operands - // by composing the transpose map. - SmallVector<AffineMap> indexingMaps = getTransposedIndexingMaps( - genericOp, inputOperand->getOperandNumber(), transposeMap); - - // We do not need to update indexing maps because this is an elementwise - // op where the input and output maps are the same. - // Just replace the operands with transposed variants. - auto newGenericOp = - mlir::clone(rewriter, genericOp, newInit.getType(), newOperands); - newGenericOp.setIndexingMapsAttr( - rewriter.getAffineMapArrayAttr(indexingMaps)); - rewriter.replaceOp(transposeOp, newGenericOp); - return success(); - } -}; - -} // namespace - -//===----------------------------------------------------------------------===// -// Linalg Named Op -> Named Op Conversions -//===----------------------------------------------------------------------===// - -namespace { - -template <typename OpTy, int64_t inputIdx> -class NamedOpConversion : public OpRewritePattern<OpTy> { -public: - using OpRewritePattern<OpTy>::OpRewritePattern; - NamedOpConversion(MLIRContext *ctx, SmallVector<int64_t> perm, - PatternBenefit b = 1) - : OpRewritePattern<OpTy>(ctx, b), permutation(perm) {} - - LogicalResult matchAndRewrite(OpTy namedOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(namedOp)) { - return failure(); - } - - Value input = namedOp.getInputs()[inputIdx]; - auto transpose = input.getDefiningOp<linalg::TransposeOp>(); - if (!transpose) { - return failure(); - } - - SmallVector<int64_t> transPerm(transpose.getPermutation()); - if (transPerm != permutation) { - return rewriter.notifyMatchFailure( - namedOp, "transpose permutation does not match target permutation"); - } - SmallVector<NamedAttribute> attrs = getPrunedAttributeList(namedOp); - SmallVector<Value> newInputs = namedOp.getInputs(); - newInputs[inputIdx] = transpose.getInput(); - - auto replaceOp = [&](auto *typePtr) { - rewriter.replaceOpWithNewOp<std::remove_pointer_t<decltype(typePtr)>>( - namedOp, newInputs, namedOp.getDpsInits(), attrs); - }; - - Operation *op = namedOp.getOperation(); - if (isa<linalg::MatmulTransposeAOp>(op) && inputIdx == 0) { - replaceOp(static_cast<linalg::MatmulOp *>(nullptr)); - } else if (isa<linalg::MatmulTransposeBOp>(op) && inputIdx == 1) { - replaceOp(static_cast<linalg::MatmulOp *>(nullptr)); - } else if (IREE::LinalgExt::isPureMatmul(op) && inputIdx == 0) { - replaceOp(static_cast<linalg::MatmulTransposeAOp *>(nullptr)); - } else if (IREE::LinalgExt::isPureMatmul(op) && inputIdx == 1) { - replaceOp(static_cast<linalg::MatmulTransposeBOp *>(nullptr)); - } else if (isa<linalg::BatchMatmulTransposeAOp>(op) && inputIdx == 0) { - replaceOp(static_cast<linalg::BatchMatmulOp *>(nullptr)); - } else if (isa<linalg::BatchMatmulTransposeBOp>(op) && inputIdx == 1) { - replaceOp(static_cast<linalg::BatchMatmulOp *>(nullptr)); - } else if (IREE::LinalgExt::isPureBatchMatmul(op) && inputIdx == 0) { - replaceOp(static_cast<linalg::BatchMatmulTransposeAOp *>(nullptr)); - } else if (IREE::LinalgExt::isPureBatchMatmul(op) && inputIdx == 1) { - replaceOp(static_cast<linalg::BatchMatmulTransposeBOp *>(nullptr)); - } else { - return failure(); - } - return success(); - } - -private: - // Non-type literal array template parameters are a C++20 feature, so instead - // all the named op patterns pass their permutation explicitly as a - // SmallVector. - SmallVector<int64_t> permutation; -}; - -// Fuses a transpose into a reduction linalg.generic by absorbing it into the -// indexing map. -class FuseTransposeThroughGenericReduction - : public OpRewritePattern<linalg::GenericOp> { -public: - using Base::Base; - - LogicalResult matchAndRewrite(linalg::GenericOp genericOp, - PatternRewriter &rewriter) const override { - if (!IREE::Flow::isNonNullAndOutsideDispatch(genericOp)) { - return failure(); - } - - // Only try to fuse when the generic has at least one reduction dimension. - if (genericOp.getNumParallelLoops() == genericOp.getNumLoops()) { - return rewriter.notifyMatchFailure(genericOp, "not a reduction"); - } - - // All maps must be projected permutations. - if (!llvm::all_of(genericOp.getIndexingMapsArray(), [](AffineMap map) { - return map.isProjectedPermutation(); - })) { - return rewriter.notifyMatchFailure(genericOp, - "not a projected permutation"); - } - - // Look for a transpose on any input. - for (int64_t inputIdx = 0, endIdx = genericOp.getNumDpsInputs(); - inputIdx < endIdx; ++inputIdx) { - auto transpose = genericOp.getDpsInputs()[inputIdx] - .getDefiningOp<linalg::TransposeOp>(); - if (!transpose) { - continue; - } - - // Update the indexing map according to the transpose. - AffineMap inputMap = genericOp.getMatchingIndexingMap( - genericOp.getDpsInputOperand(inputIdx)); - - // Fuse by updating the indexing map to absorb the transpose. - auto invPerm = invertPermutationVector(transpose.getPermutation()); - SmallVector<AffineExpr> newExprs = - applyPermutation(inputMap.getResults(), invPerm); - - // Only fuse if dimension indices are in increasing order to maintain - // efficient memory access patterns. This should offset the fact that the - // transpose may have multiple uses. - int64_t prevDim = -1; - for (AffineExpr expr : newExprs) { - const int64_t dim = cast<AffineDimExpr>(expr).getPosition(); - if (dim < prevDim) { - return rewriter.notifyMatchFailure(genericOp, - "newExprs are not contiguous"); - } - prevDim = dim; - } - - auto transposedMap = - AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(), - newExprs, rewriter.getContext()); - - SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray(); - newIndexingMaps[inputIdx] = transposedMap; - - rewriter.startOpModification(genericOp); - genericOp.setIndexingMapsAttr( - rewriter.getAffineMapArrayAttr(newIndexingMaps)); - genericOp.setOperand(inputIdx, transpose.getInput()); - rewriter.finalizeOpModification(genericOp); - return success(); - } - - return rewriter.notifyMatchFailure(genericOp, - "no matching transpose found"); - } -}; - -} // namespace - -//===----------------------------------------------------------------------===// // Pass definition //===----------------------------------------------------------------------===// @@ -1143,21 +105,11 @@ }; } // namespace -static void populateMatmulSinkingPatterns(MLIRContext *context, - RewritePatternSet &sinkingPatterns) { - sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::MatmulOp, - /*inputIdx=*/1>>( - context, SmallVector<int64_t>{1, 0}); - sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::MatmulOp, - /*inputIdx=*/0>>( - context, SmallVector<int64_t>{1, 0}); - sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::BatchMatmulOp, - /*inputIdx=*/1>>( - context, SmallVector<int64_t>{0, 2, 1}); - sinkingPatterns.insert<NamedOpConversion</*OpType=*/linalg::BatchMatmulOp, - /*inputIdx=*/0>>( - context, SmallVector<int64_t>{0, 2, 1}); - sinkingPatterns.insert<FuseTransposeThroughGenericReduction>(context); +static void populateMatmulSinkingPatterns(RewritePatternSet &sinkingPatterns) { + populateFoldTransposeIntoNamedMatmulPatterns(sinkingPatterns, + isOutsideDispatch); + populateFuseTransposeThroughGenericReductionPattern(sinkingPatterns, + isOutsideDispatch); } static void @@ -1175,6 +127,34 @@ void PropagateLinalgTransposePass::runOnOperation() { MLIRContext *context = &getContext(); mlir::FunctionOpInterface funcOp = getOperation(); + + // Unless edge reshape propagation is enabled, only move transposes through + // reshapes when the transpose blocks fusion of the reshape with its + // neighbors. + ControlTransposePropagationFn bubbleThroughCollapseShapeFn = + [this](OpOperand *operand) { + if (!isOutsideDispatch(operand)) { + return false; + } + auto collapseOp = + cast<tensor::CollapseShapeOp>(operand->get().getDefiningOp()); + return enableEdgeReshapePropagation || + isReshapeBlockingFusion(operand->getOwner(), + collapseOp.getSrc().getDefiningOp()); + }; + ControlTransposePropagationFn sinkThroughExpandShapeFn = + [this](OpOperand *operand) { + if (!isOutsideDispatch(operand)) { + return false; + } + Operation *transposeOp = operand->get().getDefiningOp(); + return enableEdgeReshapePropagation || + llvm::any_of( + operand->getOwner()->getUsers(), [&](Operation *consumer) { + return isReshapeBlockingFusion(transposeOp, consumer); + }); + }; + // First, specialize all transposes to `linalg.transpose`. This dramatically // simplifies all subsequent propagation patterns, both in matching and // rewriting. @@ -1188,7 +168,7 @@ IRRewriter rewriter(&getContext()); for (auto genericOp : genericCandidates) { rewriter.setInsertionPoint(genericOp); - specializeGenericTransposeOp(rewriter, genericOp); + (void)specializeGenericTransposeOp(rewriter, genericOp); } } @@ -1203,13 +183,14 @@ // and it is more canonical if we can fuse the ops into a new named op. if (!testBubblingOnly) { RewritePatternSet sinkingPatterns(context); - sinkingPatterns.insert<SinkTransposeThroughExtractSlice>(context); - sinkingPatterns.insert<SinkTransposeThroughExpandShape>( - context, enableEdgeReshapePropagation); - populateMatmulSinkingPatterns(context, sinkingPatterns); + populateSinkTransposeThroughExtractSlicePattern( + sinkingPatterns, isUntransposedSliceMappableToFlow); + populateSinkTransposeThroughExpandShapePattern(sinkingPatterns, + sinkThroughExpandShapeFn); + populateMatmulSinkingPatterns(sinkingPatterns); populateCommonCanonicalizationPatterns(context, sinkingPatterns); - sinkingPatterns.add<SinkTransposeThroughUnaryElementwiseInput>( - context, /*benefit=*/2); + populateSinkTransposeThroughUnaryElementwiseInputPattern( + sinkingPatterns, isOutsideDispatch, /*benefit=*/2); if (failed(applyPatternsGreedily(funcOp, std::move(sinkingPatterns)))) { funcOp.emitError("Transpose initial sinking patterns failed"); return signalPassFailure(); @@ -1269,24 +250,17 @@ linalg::FillOp::getCanonicalizationPatterns(bubblingPatterns, context); if (enableAttentionVTranspose) { - linalg::ControlFusionFn bubbleTransposeControlFn = - [](OpOperand *fusedOperand) { - Operation *producer = fusedOperand->get().getDefiningOp(); - Operation *consumer = fusedOperand->getOwner(); - - return IREE::Flow::isNonNullAndOutsideDispatch( - {producer, consumer}); - }; IREE::LinalgExt::populateBubbleTransposeFromLinalgExtOps( - bubblingPatterns, bubbleTransposeControlFn); + bubblingPatterns, isOutsideDispatch); } - bubblingPatterns.insert<FuseTransposeWithProducerLinalgOp>( - context, enableAggressivePropagation, enableConvolutionPropagation); - bubblingPatterns.insert<BubbleTransposeThroughCollapseShape>( - context, enableEdgeReshapePropagation); - bubblingPatterns.add<BubbleTransposeThroughUnaryElementwiseDpsInit>( - context, /*benefit=*/2); - bubblingPatterns.insert<ComposeTransposes>(context); + populateFuseTransposeWithProducerLinalgOpPattern( + bubblingPatterns, isOutsideDispatch, enableAggressivePropagation, + enableConvolutionPropagation); + populateBubbleTransposeThroughCollapseShapePattern( + bubblingPatterns, bubbleThroughCollapseShapeFn); + populateBubbleTransposeThroughUnaryElementwiseDpsInitPattern( + bubblingPatterns, isOutsideDispatch, /*benefit=*/2); + populateComposeTransposesPattern(bubblingPatterns, isOutsideDispatch); populateCommonCanonicalizationPatterns(context, bubblingPatterns); GreedyRewriteConfig config; @@ -1344,19 +318,22 @@ }; linalg::populateFoldReshapeOpsByExpansionPatterns(sinkingPatterns, reshapePropagationFn); - sinkingPatterns.insert<SinkTransposeThroughExtractSlice>(context); - sinkingPatterns.insert<SinkTransposeThroughExpandShape>( - context, enableEdgeReshapePropagation); + populateSinkTransposeThroughExtractSlicePattern( + sinkingPatterns, isUntransposedSliceMappableToFlow); + populateSinkTransposeThroughExpandShapePattern(sinkingPatterns, + sinkThroughExpandShapeFn); if (enableSinkTransposeThroughPad) { - sinkingPatterns.insert<SinkTransposeThroughPad>(context); + populateSinkTransposeThroughPadPattern(sinkingPatterns, + isOutsideDispatch); } - sinkingPatterns.insert<FuseTransposeWithLinalgOpConsumer>( - context, enableAggressivePropagation, enableConvolutionPropagation); - sinkingPatterns.insert<ComposeTransposes>(context); - populateMatmulSinkingPatterns(context, sinkingPatterns); + populateFuseTransposeWithLinalgOpConsumerPattern( + sinkingPatterns, isOutsideDispatch, enableAggressivePropagation, + enableConvolutionPropagation); + populateComposeTransposesPattern(sinkingPatterns, isOutsideDispatch); + populateMatmulSinkingPatterns(sinkingPatterns); populateCommonCanonicalizationPatterns(context, sinkingPatterns); - sinkingPatterns.add<SinkTransposeThroughUnaryElementwiseInput>( - context, /*benefit=*/2); + populateSinkTransposeThroughUnaryElementwiseInputPattern( + sinkingPatterns, isOutsideDispatch, /*benefit=*/2); GreedyRewriteConfig config; // TODO: This is inefficient. Consider rewriting this pass to use a // worklist of just the transpose operations.
diff --git a/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp new file mode 100644 index 0000000..fdfb5ba --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.cpp
@@ -0,0 +1,1281 @@ +// Copyright 2023 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 + +//===- TransposePropagationPatterns.cpp - Transpose propagation patterns --===// +// +// Patterns that propagate linalg.transpose operations through a restricted +// set of operations based on local propagation decisions. +// +//===----------------------------------------------------------------------===// + +#include "iree/compiler/GlobalOptimization/TransposePropagationPatterns.h" +#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/Transforms.h" +#include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/IndexingUtils.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/PatternMatch.h" + +namespace mlir::iree_compiler::GlobalOptimization { + +//===----------------------------------------------------------------------===// +// Transpose permutation helpers +//===----------------------------------------------------------------------===// + +static bool isIdentityPermutation(ArrayRef<int64_t> perm) { + for (auto [index, dim] : llvm::enumerate(perm)) { + if (index != dim) { + return false; + } + } + return true; +} + +// Constructs a transpose of the given tensor and permutation. +static Value createTransposeInit(OpBuilder &builder, Value source, + ArrayRef<int64_t> perm) { + SmallVector<OpFoldResult> mixedSizes = + tensor::getMixedSizes(builder, source.getLoc(), source); + applyPermutationToVector(mixedSizes, perm); + Type elemType = cast<RankedTensorType>(source.getType()).getElementType(); + Value empty = + tensor::EmptyOp::create(builder, source.getLoc(), mixedSizes, elemType) + .getResult(); + return empty; +} + +// Constructs a transpose of the given tensor and permutation, +// or produces a transposed version of the producing tensor.empty op. +static Value createTranspose(OpBuilder &builder, Value source, + ArrayRef<int64_t> perm) { + if (auto empty = source.getDefiningOp<tensor::EmptyOp>()) { + Type elementType = empty.getType().getElementType(); + SmallVector<OpFoldResult> mixedSizes = empty.getMixedSizes(); + applyPermutationToVector(mixedSizes, perm); + return tensor::EmptyOp::create(builder, empty.getLoc(), mixedSizes, + elementType); + } + Value empty = createTransposeInit(builder, source, perm); + return linalg::TransposeOp::create(builder, source.getLoc(), source, empty, + perm) + ->getResult(0); +} + +static RankedTensorType getPermutedTensorType(RankedTensorType type, + SmallVector<int64_t> perm) { + SmallVector<int64_t> permutedShape = applyPermutation(type.getShape(), perm); + return RankedTensorType::get(permutedShape, type.getElementType()); +} + +//===----------------------------------------------------------------------===// +// Transpose specialization +//===----------------------------------------------------------------------===// + +// Replaces |genericOp|, which must satisfy isaTransposeOpInterface, with the +// equivalent linalg.transpose. +static void replaceWithTransposeOp(RewriterBase &rewriter, + linalg::GenericOp genericOp) { + auto mapRange = genericOp.getIndexingMapsArray(); + AffineMap outMap = mapRange.back(); + AffineMap inMap = mapRange.front(); + SmallVector<int64_t> perm; + // To get the permutation, look at each output index and find which + // dimension in the input we're reading from for that index. + for (AffineExpr expr : outMap.getResults()) { + perm.push_back(*inMap.getResultPosition(expr)); + } + rewriter.replaceOpWithNewOp<linalg::TransposeOp>( + genericOp, genericOp.getDpsInputs()[0], genericOp.getDpsInits()[0], perm); +} + +LogicalResult specializeGenericTransposeOp(RewriterBase &rewriter, + linalg::GenericOp genericOp) { + if (!IREE::LinalgExt::isaTransposeOpInterface(genericOp)) { + return failure(); + } + replaceWithTransposeOp(rewriter, genericOp); + return success(); +} + +namespace { + +class SpecializeGenericTranspose : public OpRewritePattern<linalg::GenericOp> { +public: + SpecializeGenericTranspose(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::GenericOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::GenericOp genericOp, + PatternRewriter &rewriter) const override { + if (!IREE::LinalgExt::isaTransposeOpInterface(genericOp)) { + return rewriter.notifyMatchFailure(genericOp, "not a transpose"); + } + if (!controlFn(genericOp.getDpsInputOperand(0))) { + return rewriter.notifyMatchFailure(genericOp, + "rejected by control function"); + } + replaceWithTransposeOp(rewriter, genericOp); + return success(); + } + +private: + // Decides whether the generic transpose may be specialized. + ControlTransposePropagationFn controlFn; +}; + +} // namespace + +//===----------------------------------------------------------------------===// +// Other pattern helpers +//===----------------------------------------------------------------------===// + +/// Returns the `op` if it is a linalg::GenericOp. If it is a named op and +/// `allowGeneralizing` is true, returns the generalized op. Otherwise, returns +/// failure. +/// TODO: Due to fragility around handling of convolutions, convolution +/// propagation is behind a flag. +static FailureOr<linalg::GenericOp> +getAllowedGenericOpOrGeneralizeNamedOp(RewriterBase &rewriter, Operation *op, + bool allowGeneralizing, bool convProp) { + auto linalgOp = dyn_cast<linalg::LinalgOp>(op); + if (!linalgOp) { + return failure(); + } + + if (!convProp && linalg::isaConvolutionOpInterface(linalgOp)) { + return failure(); + } + + if (!isa<linalg::GenericOp>(linalgOp) && + !(allowGeneralizing && linalg::isaContractionOpInterface(linalgOp)) && + !(convProp && linalg::isaConvolutionOpInterface(linalgOp))) { + return failure(); + } + + auto genericOp = dyn_cast<linalg::GenericOp>(op); + if (genericOp) { + return genericOp; + } + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(linalgOp); + FailureOr<linalg::LinalgOp> generalizedOp = + linalg::generalizeNamedOp(rewriter, linalgOp); + if (failed(generalizedOp)) { + return failure(); + } + return cast<linalg::GenericOp>(generalizedOp->getOperation()); +} + +//===----------------------------------------------------------------------===// +// Transpose Bubbling Patterns +//===----------------------------------------------------------------------===// + +namespace { + +// Fuses a transpose with the init of a linalg.generic op or contraction op. +// Contraction ops are generalized and then treated as a generic. For example, +// +// linalg.generic { +// indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, +// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// ins(%0 : tensor<2x7x5>) outs(%1 : tensor<7x2x5>) +// +// %2 = linalg.transpose ... permutation = [0, 2, 1] : +// tensor<7x2x5> -> tensor<7x5x2> +// Becomes +// +// linalg.generic { +// indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0, d1)>, +// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// ins(%0 : tensor<2x7x5>) outs(%3 : tensor<7x5x2>) +class FuseTransposeWithProducerLinalgOp + : public OpRewritePattern<linalg::TransposeOp> { +public: + FuseTransposeWithProducerLinalgOp(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + bool aggressiveProp, bool convProp, + PatternBenefit b = 1) + : OpRewritePattern<linalg::TransposeOp>(ctx, b), + controlFn(std::move(controlFn)), allowGeneralizing(aggressiveProp), + convProp(convProp) {} + + LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, + PatternRewriter &rewriter) const override { + OpResult result = dyn_cast<OpResult>(transposeOp.getInput()); + if (!result) { + return rewriter.notifyMatchFailure( + transposeOp, "transpose input defined by block argument"); + } + if (!result.hasOneUse()) { + return rewriter.notifyMatchFailure(transposeOp, + "multi use transpose input"); + } + auto linalgOp = dyn_cast<linalg::LinalgOp>(result.getOwner()); + if (!linalgOp) { + return rewriter.notifyMatchFailure( + transposeOp, "non-linalg op producer for transpose input"); + } + if (!controlFn(&transposeOp.getInputMutable())) { + return rewriter.notifyMatchFailure(transposeOp, + "rejected by control function"); + } + + int64_t resultIndex = result.getResultNumber(); + auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp( + rewriter, result.getOwner(), allowGeneralizing, convProp); + if (failed(maybeGenericOp)) { + return rewriter.notifyMatchFailure( + transposeOp, "linalg op producer is not generic or contraction"); + } + + auto genericOp = maybeGenericOp.value(); + result = genericOp->getOpResult(resultIndex); + + ArrayRef<int64_t> perm = transposeOp.getPermutation(); + auto invPerm = invertPermutationVector(perm); + + // 1. Get the transposed init of the generic. + Value init = genericOp.getDpsInits()[resultIndex]; + SmallVector<Value> inits = genericOp.getDpsInits(); + Value newInit = createTranspose(rewriter, init, perm); + inits[resultIndex] = newInit; + + SmallVector<Type> resultTypes(genericOp->getResultTypes()); + resultTypes[resultIndex] = newInit.getType(); + + // 2. Update the indexing map of the transposed init operand by permuting + // the results of the map. + SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray(); + AffineMap resultMap = + newIndexingMaps[genericOp.getNumDpsInputs() + resultIndex]; + SmallVector<AffineExpr> newExprs = + applyPermutation(resultMap.getResults(), perm); + AffineMap transposedMap = + AffineMap::get(resultMap.getNumDims(), resultMap.getNumSymbols(), + newExprs, rewriter.getContext()); + newIndexingMaps[genericOp.getNumDpsInputs() + resultIndex] = transposedMap; + + // 3. Create the new generic with the same iteration order. + auto newGenericOp = linalg::GenericOp::create( + rewriter, genericOp.getLoc(), resultTypes, genericOp.getDpsInputs(), + newInit, newIndexingMaps, genericOp.getIteratorTypesArray(), + /*bodyBuild=*/nullptr, linalg::getPrunedAttributeList(genericOp)); + rewriter.cloneRegionBefore(genericOp.getRegion(), newGenericOp.getRegion(), + newGenericOp.getRegion().begin()); + + // 4. Remap iteration space of the generic to match the dimension order of + // the output. + if (newGenericOp.getNumResults() == 1) { + SmallVector<unsigned int> interchange; + int64_t permIdx = 0; + for (int i = 0, e = transposedMap.getNumDims(); i < e; ++i) { + if (transposedMap.isFunctionOfDim(i)) { + interchange.push_back( + cast<AffineDimExpr>(transposedMap.getResult(permIdx)) + .getPosition()); + permIdx++; + continue; + } + interchange.push_back(i); + } + auto interchangedGenericOp = + linalg::interchangeGenericOp(rewriter, newGenericOp, interchange); + // Interchange only fails if interchangeGenericOpPrecondition fails, which + // only fails if the interchange vector is not invertible or doesn't match + // the number of loops in the generic, both of which are guaranteed by + // the fact that the output map must be a projection in the above + // construction. + assert(succeeded(interchangedGenericOp) && + "failed to interchange transposed generic"); + newGenericOp = *interchangedGenericOp; + } + + // 5. Replace the result of the transpose with the transposed init. + rewriter.replaceOp(transposeOp, newGenericOp->getResult(resultIndex)); + for (auto [oldRes, newRes] : + llvm::zip_equal(genericOp.getResults(), newGenericOp->getResults())) { + if (oldRes.getResultNumber() == resultIndex) { + continue; + } + rewriter.replaceAllUsesWith(oldRes, newRes); + } + return success(); + } + +private: + // Decides whether the transpose may fuse into its producer. + ControlTransposePropagationFn controlFn; + // Whether named contraction producers may be generalized for fusion. + bool allowGeneralizing = false; + // Whether convolution producers may be generalized for fusion. + bool convProp = false; +}; + +// Bubbles a transpose through a tensor.collapse_shape. +class BubbleTransposeThroughCollapseShape + : public OpRewritePattern<linalg::TransposeOp> { +public: + BubbleTransposeThroughCollapseShape(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::TransposeOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, + PatternRewriter &rewriter) const override { + Value source = transposeOp.getDpsInputOperand(0)->get(); + auto collapseOp = source.getDefiningOp<tensor::CollapseShapeOp>(); + // Do not propagate through reshapes if the transpose has multiple users, as + // this could end up duplicating the transposes. We should only propagate + // through reshape when it is free to do so. + if (!collapseOp || !collapseOp->hasOneUse()) { + return rewriter.notifyMatchFailure( + transposeOp, "transpose input is not a single-use collapse shape"); + } + + if (!controlFn(&transposeOp.getInputMutable())) { + return rewriter.notifyMatchFailure(transposeOp, + "rejected by control function"); + } + + SmallVector<ReassociationIndices> reassociations = + collapseOp.getReassociationIndices(); + + // Because we are doing transpose(collapse_shape), all expanded groups are + // transposed together. As a result, to get the permutation of the new + // transpose, we can just flatten the transposed reassociation indices. + // For example, + // + // reassociation_map = [[0, 1, 2], [3], [4, 5]] + // permutation = [1, 2, 0] + // + // Becomes + // + // permutation = [3, 4, 5, 0, 1, 2] + // reassociation_map = [[0], [1, 2], [3, 4, 5]] + applyPermutationToVector(reassociations, transposeOp.getPermutation()); + + SmallVector<int64_t> newPerm; + SmallVector<ReassociationIndices> newReassociations; + int64_t expandedDim = 0; + for (auto reassoc : reassociations) { + ReassociationIndices newReassoc; + for (auto dim : reassoc) { + newPerm.push_back(dim); + newReassoc.push_back(expandedDim++); + } + newReassociations.push_back(newReassoc); + } + + Value newTranspose = + createTranspose(rewriter, collapseOp.getSrc(), newPerm); + Value newReshape = tensor::CollapseShapeOp::create( + rewriter, collapseOp.getLoc(), transposeOp.getResultTypes()[0], + newTranspose, newReassociations); + rewriter.replaceOp(transposeOp, newReshape); + return success(); + } + +private: + // Decides whether the transpose may bubble through the collapse_shape. + ControlTransposePropagationFn controlFn; +}; + +} // namespace + +//===----------------------------------------------------------------------===// +// Transpose Sinking Patterns +//===----------------------------------------------------------------------===// + +namespace { + +// Combines two transposes into one. This shouldn't be strictly necessary as +// fusion should cancel inverse transposes, but doing this here can open up +// new propagation opportunities and eases the analysis in fusion/later passes. +class ComposeTransposes : public OpRewritePattern<linalg::TransposeOp> { +public: + ComposeTransposes(MLIRContext *ctx, ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::TransposeOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::TransposeOp consumer, + PatternRewriter &rewriter) const override { + Value input = consumer.getInput(); + auto producer = input.getDefiningOp<linalg::TransposeOp>(); + if (!producer) { + return failure(); + } + if (!controlFn(&consumer.getInputMutable())) { + return rewriter.notifyMatchFailure(consumer, + "rejected by control function"); + } + + ArrayRef<int64_t> producerPerm = producer.getPermutation(); + ArrayRef<int64_t> consumerPerm = consumer.getPermutation(); + SmallVector<int64_t> composedPerm = + applyPermutation(producerPerm, consumerPerm); + + Value transposedSource = producer.getInput(); + if (!isIdentityPermutation(composedPerm)) { + transposedSource = + createTranspose(rewriter, transposedSource, composedPerm); + } + rewriter.replaceOp(consumer, transposedSource); + return success(); + } + +private: + // Decides whether the producer transpose may be composed into the consumer. + ControlTransposePropagationFn controlFn; +}; + +// Sinks a transpose through a tensor.extract_slice. +class SinkTransposeThroughExtractSlice + : public OpRewritePattern<tensor::ExtractSliceOp> { +public: + SinkTransposeThroughExtractSlice(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<tensor::ExtractSliceOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractOp, + PatternRewriter &rewriter) const override { + Value source = extractOp.getSource(); + auto transposeOp = source.getDefiningOp<linalg::TransposeOp>(); + if (!transposeOp) { + return failure(); + } + + // Applying `perm` takes a list from the pre-transpose ordering to the + // post-transpose ordering. + ArrayRef<int64_t> perm = transposeOp.getPermutation(); + // Applying `invPerm` takes a list from the post-transpose ordering to the + // pre-transpose ordering. Sinking a transpose through an op is largely a + // matter of rewriting it in pre-transpose space, and thus just applying + // the inverse permutation. + auto invPerm = invertPermutationVector(perm); + + SmallVector<OpFoldResult> offsets = extractOp.getMixedOffsets(); + SmallVector<OpFoldResult> sizes = extractOp.getMixedSizes(); + SmallVector<OpFoldResult> strides = extractOp.getMixedStrides(); + + // Permute the offsets, sizes, and strides to pre-transpose ordering. + applyPermutationToVector(offsets, invPerm); + applyPermutationToVector(sizes, invPerm); + applyPermutationToVector(strides, invPerm); + + if (!controlFn(&extractOp.getSourceMutable())) { + return rewriter.notifyMatchFailure(extractOp, + "rejected by control function"); + } + + ArrayRef<int64_t> staticSizes = extractOp.getStaticSizes(); + ArrayRef<int64_t> sliceShape = extractOp.getResultType().getShape(); + std::optional<llvm::SmallDenseSet<unsigned>> maybeRankReducingMask = + mlir::computeRankReductionMask(staticSizes, sliceShape); + if (!maybeRankReducingMask) { + return rewriter.notifyMatchFailure( + extractOp, "failed to compute rank reducing mask"); + } + llvm::SmallDenseSet<unsigned> rankReducingMask = *maybeRankReducingMask; + + // Find rank reducing map in the pre-transposed domain. + int64_t dim = 0; + llvm::SmallDenseMap<int64_t, int64_t> rankReducedMap; + // Since `dim` is in the pre-transposed domain, and is incrementing each + // iteration, `idx` must also be in the pre-transposed domain. + for (int64_t idx = 0, e = invPerm.size(); idx < e; ++idx) { + // Get index in the transposed domain, since `rankReducingMask` is in + // the transposed domain. + if (!rankReducingMask.contains(invPerm[idx])) { + // Domain of `rankReducedMap` is in pre-transposed domain. + rankReducedMap[idx] = dim++; + } + } + + // Compute the new permutation by dropping all rank-reduced dimensions. + SmallVector<int64_t> rankReducedPerm; + for (int64_t i : perm) { + if (rankReducedMap.contains(i)) { + rankReducedPerm.push_back(rankReducedMap[i]); + } + } + + auto rankReducedInvPerm = invertPermutationVector(rankReducedPerm); + + RankedTensorType sliceType = getPermutedTensorType( + cast<RankedTensorType>(extractOp.getType()), rankReducedInvPerm); + Value slice = tensor::ExtractSliceOp::create( + rewriter, extractOp.getLoc(), sliceType, transposeOp.getInput(), + offsets, sizes, strides); + // Transpose back to the original slice. + if (!isIdentityPermutation(rankReducedPerm)) { + slice = createTranspose(rewriter, slice, rankReducedPerm); + } + rewriter.replaceOp(extractOp, slice); + return success(); + } + +private: + // Decides whether the transpose may sink through the extract_slice. + ControlTransposePropagationFn controlFn; +}; + +// Sinks a transpose through a tensor.expand_shape. +class SinkTransposeThroughExpandShape + : public OpRewritePattern<tensor::ExpandShapeOp> { +public: + SinkTransposeThroughExpandShape(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<tensor::ExpandShapeOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(tensor::ExpandShapeOp expandOp, + PatternRewriter &rewriter) const override { + Value source = expandOp.getSrc(); + auto transposeOp = source.getDefiningOp<linalg::TransposeOp>(); + // Do not propagate through reshapes if the transpose has multiple users, as + // this could end up duplicating the transposes. We should only propagate + // through reshape when it is free to do so. + if (!transposeOp || !transposeOp->hasOneUse()) { + return rewriter.notifyMatchFailure( + expandOp, "expand shape input is not a single-use transpose"); + } + + if (!controlFn(&expandOp.getSrcMutable())) { + return rewriter.notifyMatchFailure(expandOp, + "rejected by control function"); + } + + auto invPerm = invertPermutationVector(transposeOp.getPermutation()); + SmallVector<ReassociationIndices> reassociations = + expandOp.getReassociationIndices(); + + // Because we are doing expand_shape(transpose), all expanded groups are + // transposed together. As a result, to get the permutation of the new + // transpose, we can just flatten the transposed reassociation indices. + // For example, + // + // permutation = [0, 2, 1] + // reassociation_map = [[0, 1, 2], [3], [4, 5]] + // + // Becomes + // + // reassociation_map = [[0, 1, 2], [3, 4], [5]] + // permutation = [0, 1, 2, 4, 5, 3] + applyPermutationToVector(reassociations, invPerm); + + SmallVector<int64_t> newInvPerm; + SmallVector<ReassociationIndices> newReassociations; + int64_t expandedDim = 0; + for (auto reassoc : reassociations) { + ReassociationIndices newReassoc; + for (auto dim : reassoc) { + newInvPerm.push_back(dim); + newReassoc.push_back(expandedDim++); + } + newReassociations.push_back(newReassoc); + } + + auto newPerm = invertPermutationVector(newInvPerm); + + RankedTensorType expandedType = getPermutedTensorType( + cast<RankedTensorType>(expandOp.getType()), newInvPerm); + // Permute the output shape to match the permuted type. + SmallVector<OpFoldResult> permutedOutputShape = + expandOp.getMixedOutputShape(); + applyPermutationToVector(permutedOutputShape, newInvPerm); + Value transposedReshape = tensor::ExpandShapeOp::create( + rewriter, expandOp.getLoc(), expandedType, transposeOp.getInput(), + newReassociations, permutedOutputShape); + Value originalReshape = + createTranspose(rewriter, transposedReshape, newPerm); + rewriter.replaceOp(expandOp, originalReshape); + return success(); + } + +private: + // Decides whether the transpose may sink through the expand_shape. + ControlTransposePropagationFn controlFn; +}; + +// Sinks a transpose through a tensor.pad. +class SinkTransposeThroughPad : public OpRewritePattern<tensor::PadOp> { +public: + SinkTransposeThroughPad(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<tensor::PadOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(tensor::PadOp padOp, + PatternRewriter &rewriter) const override { + Value source = padOp.getSource(); + auto transposeOp = source.getDefiningOp<linalg::TransposeOp>(); + if (!transposeOp) { + return failure(); + } + if (!controlFn(&padOp.getSourceMutable())) { + return rewriter.notifyMatchFailure(padOp, "rejected by control function"); + } + + Block &block = padOp.getRegion().front(); + if (llvm::any_of(block.getArguments(), [](BlockArgument blockArg) { + return blockArg.getNumUses(); + })) { + return failure(); + } + + auto invPerm = invertPermutationVector(transposeOp.getPermutation()); + SmallVector<OpFoldResult> lowSizes = padOp.getMixedLowPad(); + SmallVector<OpFoldResult> highSizes = padOp.getMixedHighPad(); + applyPermutationToVector(lowSizes, invPerm); + applyPermutationToVector(highSizes, invPerm); + + RankedTensorType oldPaddedType = cast<RankedTensorType>(padOp.getType()); + RankedTensorType newPaddedType = oldPaddedType.clone( + applyPermutation(oldPaddedType.getShape(), invPerm)); + + auto newPadOp = tensor::PadOp::create( + rewriter, padOp.getLoc(), newPaddedType, transposeOp.getInput(), + lowSizes, highSizes, padOp.getNofold()); + rewriter.cloneRegionBefore(padOp.getRegion(), newPadOp.getRegion(), + newPadOp.getRegion().begin()); + + Value newTransposeOp = + createTranspose(rewriter, newPadOp, transposeOp.getPermutation()); + rewriter.replaceOp(padOp, newTransposeOp); + return success(); + } + +private: + // Decides whether the transpose may sink through the pad. + ControlTransposePropagationFn controlFn; +}; + +// Fuses a transpose with the input of a linalg.generic op or contraction op. +// Contraction ops are generalized and then treated as a generic. For example, +// +// %0 = linalg.transpose ... permutation = [0, 2, 1] : +// tensor<2x5x7> -> tensor<2x7x5> +// linalg.generic { +// indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, +// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// ins(%0 : tensor<2x7x5>) +// +// Becomes +// +// linalg.generic { +// indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2, d0)>, +// affine_map<(d0, d1, d2) -> (d0, d1, d2)>] +// ins(%0 : tensor<2x5x7>) +// +// This is considered just one way to model transpose propagation to generics, +// another option would be to interpret the transpose on the iterators of the +// generic, thus producing a transpose on the output and any other inputs to +// the generic. This has the potential introduce more transposes/data movement +// and isn't the way this pass is modeled. Global data layout transformations +// like that are better suited for pack/unpack propagation rooted on specific +// operations. +// +// TODO: Rewrite this to use elementwise op fusion patterns. +class FuseTransposeWithLinalgOpConsumer + : public OpInterfaceRewritePattern<linalg::LinalgOp> { +public: + FuseTransposeWithLinalgOpConsumer(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + bool aggressiveProp, bool convProp, + PatternBenefit b = 1) + : OpInterfaceRewritePattern<linalg::LinalgOp>(ctx, b), + controlFn(std::move(controlFn)), allowGeneralizing(aggressiveProp), + convProp(convProp) {} + + LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp, + PatternRewriter &rewriter) const override { + OpOperand *transposeOperand = nullptr; + linalg::TransposeOp transposeOp; + for (OpOperand *input : linalgOp.getDpsInputOperands()) { + auto maybeTransposeOp = input->get().getDefiningOp<linalg::TransposeOp>(); + if (maybeTransposeOp && maybeTransposeOp->hasOneUse()) { + transposeOp = maybeTransposeOp; + transposeOperand = input; + break; + } + } + if (!transposeOperand) { + return rewriter.notifyMatchFailure(linalgOp, "no transpose operand"); + } + if (!controlFn(transposeOperand)) { + return rewriter.notifyMatchFailure(linalgOp, + "rejected by control function"); + } + + int64_t inputIndex = transposeOperand->getOperandNumber(); + ArrayRef<int64_t> perm = transposeOp.getPermutation(); + auto invPerm = invertPermutationVector(perm); + + // To do the fusion, we can simply apply the permutation of the transpose + // to the results of the associated input's indexing map, and then forward + // the input to the transpose to the consumer generic. + auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp( + rewriter, linalgOp, allowGeneralizing, convProp); + if (failed(maybeGenericOp)) { + return failure(); + } + auto genericOp = maybeGenericOp.value(); + transposeOperand = genericOp.getDpsInputOperand(inputIndex); + rewriter.startOpModification(genericOp); + + SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray(); + AffineMap inputMap = genericOp.getMatchingIndexingMap(transposeOperand); + SmallVector<AffineExpr> newExprs = + applyPermutation(inputMap.getResults(), invPerm); + AffineMap transposedMap = + AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(), + newExprs, rewriter.getContext()); + newIndexingMaps[inputIndex] = transposedMap; + genericOp.setIndexingMapsAttr( + rewriter.getAffineMapArrayAttr(newIndexingMaps)); + + genericOp.setOperand(inputIndex, transposeOp.getInput()); + rewriter.finalizeOpModification(genericOp); + return success(); + } + +private: + // Decides whether the transpose may fuse into its consumer. + ControlTransposePropagationFn controlFn; + // Whether named contraction consumers may be generalized for fusion. + bool allowGeneralizing = false; + // Whether convolution consumers may be generalized for fusion. + bool convProp = false; +}; + +static bool isIndexingMapAffectedByTransposeMap( + AffineMap indexingMap, ArrayRef<int64_t> iterationSpacePermutation) { + int64_t prevIdx = -1; + for (auto result : indexingMap.getResults()) { + int64_t idx = + iterationSpacePermutation[cast<AffineDimExpr>(result).getPosition()]; + // Verify that the relative ordering of indices in the map remain the same. + // If not, then the transposition affects the access order for the given + // map (and associated operand). + if (idx <= prevIdx) { + return true; + } + prevIdx = idx; + } + return false; +} + +// Finds a single DPS input operand of the given |genericOp| that is affected by +// the |iterationSpacePermutation|. In other words, the permutation changes the +// relative ordering of any of the dimensions of that input operand. +// +// For example, with permutation [1, 0, 2], affine map (d0, d1, d2) -> (d0, d1) +// is affected by the permutation because the first two dimensions are iterated +// in a different order while (d0, d1, d2) -> (d0, d2) is unaffected. +// +// If no such operand is found or there is more than one such operation, nullptr +// is returned. +static OpOperand * +getSingleTransposedInputOperand(linalg::GenericOp genericOp, + ArrayRef<int64_t> iterationSpacePermutation) { + OpOperand *operand = nullptr; + for (auto input : genericOp.getDpsInputOperands()) { + if (!isIndexingMapAffectedByTransposeMap( + genericOp.getMatchingIndexingMap(input), + iterationSpacePermutation)) { + continue; + } + if (operand) { + return nullptr; + } + operand = input; + } + return operand; +} + +// Returns a new list of indexing maps that composes the iteration space +// permutation map |transposeMap| with all indexing maps of |genericOp| except +// for the |transposedInputIdx|'th operand. The unchanged operand is expected +// to have an explicit `linalg.transpose` op constructed for it so its map does +// not need to be updated. +static SmallVector<AffineMap> +getTransposedIndexingMaps(linalg::GenericOp genericOp, + int64_t transposedInputIdx, AffineMap transposeMap) { + SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray(); + for (unsigned i = 0, e = genericOp.getNumDpsInputs(); i < e; ++i) { + if (i == transposedInputIdx) { + continue; + } + indexingMaps[i] = indexingMaps[i].compose(transposeMap); + } + return indexingMaps; +} + +// Sinks a transpose through the input of a elementwise operation where the +// transposition of the iteration space only affects a single input operand. +class SinkTransposeThroughUnaryElementwiseInput + : public OpRewritePattern<linalg::GenericOp> { +public: + SinkTransposeThroughUnaryElementwiseInput( + MLIRContext *ctx, ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::GenericOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::GenericOp genericOp, + PatternRewriter &rewriter) const override { + if (!linalg::isElementwise(genericOp)) { + return rewriter.notifyMatchFailure(genericOp, "non-elementwise generic"); + } + + if (genericOp.hasIndexSemantics()) { + return rewriter.notifyMatchFailure(genericOp, "has index semantics"); + } + + if (genericOp.getNumDpsInits() != 1) { + return rewriter.notifyMatchFailure(genericOp, + "unimplemented: multiple results"); + } + + AffineMap resultMap = + genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0)); + if (!resultMap.isIdentity()) { + return rewriter.notifyMatchFailure( + genericOp, "unimplemented: non-identity result map"); + } + + linalg::TransposeOp transposeOp; + OpOperand *inputOperand; + for (auto input : genericOp.getDpsInputOperands()) { + // Skip broadcasted operands and transposed operands. If the input is + // broadcasted then we would not want to propagate because that would + // do the transpose on larger data, and if transposed we would rather + // simply compose the transposes (handled in a separate pattern). + if (genericOp.getMatchingIndexingMap(input) != resultMap) { + continue; + } + + auto maybeTransposeOp = input->get().getDefiningOp<linalg::TransposeOp>(); + // Skip multi-use transposes. + if (!maybeTransposeOp || !maybeTransposeOp->hasOneUse()) { + continue; + } + + auto transposableInputOperand = getSingleTransposedInputOperand( + genericOp, maybeTransposeOp.getPermutation()); + // Skip if more than one operand is affected by the transpose. + if (transposableInputOperand != input) { + continue; + } + + transposeOp = maybeTransposeOp; + inputOperand = transposableInputOperand; + break; + } + + if (!transposeOp) { + return rewriter.notifyMatchFailure(genericOp, + "no single use transpose operand"); + } + if (!controlFn(inputOperand)) { + return rewriter.notifyMatchFailure(genericOp, + "rejected by control function"); + } + + ArrayRef<int64_t> perm = transposeOp.getPermutation(); + auto invPerm = invertPermutationVector(perm); + + // Create a new empty init for the transposed generic. + Value newInit = + createTransposeInit(rewriter, genericOp.getDpsInits()[0], invPerm); + + // We do not need to update iterator types because this is an elementwise + // op. We just need to update the indexing maps of all other input operands + // by composing the transpose map. + AffineMap transposeMap = + AffineMap::getPermutationMap(perm, rewriter.getContext()); + SmallVector<AffineMap> indexingMaps = getTransposedIndexingMaps( + genericOp, inputOperand->getOperandNumber(), transposeMap); + + SmallVector<Value> newOperands = genericOp->getOperands(); + newOperands[inputOperand->getOperandNumber()] = transposeOp.getInput(); + newOperands[genericOp.getDpsInitOperand(0)->getOperandNumber()] = newInit; + + auto newGenericOp = + mlir::clone(rewriter, genericOp, newInit.getType(), newOperands); + newGenericOp.setIndexingMapsAttr( + rewriter.getAffineMapArrayAttr(indexingMaps)); + rewriter.replaceOp( + genericOp, createTranspose(rewriter, newGenericOp->getResult(0), perm)); + return success(); + } + +private: + // Decides whether the transpose may sink through the elementwise op. + ControlTransposePropagationFn controlFn; +}; + +// Bubbles a transpose through the init of a elementwise operation where the +// transposition of the iteration space only affects a single input operand. +class BubbleTransposeThroughUnaryElementwiseDpsInit + : public OpRewritePattern<linalg::TransposeOp> { +public: + BubbleTransposeThroughUnaryElementwiseDpsInit( + MLIRContext *ctx, ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::TransposeOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp, + PatternRewriter &rewriter) const override { + auto genericOp = transposeOp.getInput().getDefiningOp<linalg::GenericOp>(); + if (!genericOp) { + return rewriter.notifyMatchFailure(transposeOp, "non-generic producer"); + } + + if (genericOp.getNumDpsInits() != 1) { + return rewriter.notifyMatchFailure(transposeOp, + "unimplemented: multiple results"); + } + + if (!controlFn(&transposeOp.getInputMutable())) { + return rewriter.notifyMatchFailure(transposeOp, + "rejected by control function"); + } + + if (!linalg::isElementwise(genericOp) || + !genericOp.getMatchingIndexingMap(genericOp.getDpsInitOperand(0)) + .isIdentity()) { + return rewriter.notifyMatchFailure(transposeOp, "not elementwise"); + } + + if (genericOp.hasIndexSemantics()) { + return rewriter.notifyMatchFailure(genericOp, "has index semantics"); + } + + if (!genericOp->hasOneUse()) { + return rewriter.notifyMatchFailure(transposeOp, "not single user"); + } + + ArrayRef<int64_t> perm = transposeOp.getPermutation(); + auto invPerm = invertPermutationVector(perm); + + auto inputOperand = getSingleTransposedInputOperand(genericOp, invPerm); + if (!inputOperand || + !genericOp.getMatchingIndexingMap(inputOperand).isIdentity()) { + return rewriter.notifyMatchFailure( + genericOp, "no single transposable input operand"); + } + + Value newTranspose = createTranspose(rewriter, inputOperand->get(), perm); + + // Create a new empty init for the transposed generic. + Value newInit = + createTransposeInit(rewriter, genericOp.getDpsInits()[0], perm); + + SmallVector<Value> newOperands = genericOp->getOperands(); + newOperands[inputOperand->getOperandNumber()] = newTranspose; + newOperands[genericOp.getDpsInitOperand(0)->getOperandNumber()] = newInit; + + AffineMap transposeMap = + AffineMap::getPermutationMap(invPerm, rewriter.getContext()); + + // We do not need to update iterator types because this is an elementwise + // op. We just need to update the indexing maps of all other input operands + // by composing the transpose map. + SmallVector<AffineMap> indexingMaps = getTransposedIndexingMaps( + genericOp, inputOperand->getOperandNumber(), transposeMap); + + // We do not need to update indexing maps because this is an elementwise + // op where the input and output maps are the same. + // Just replace the operands with transposed variants. + auto newGenericOp = + mlir::clone(rewriter, genericOp, newInit.getType(), newOperands); + newGenericOp.setIndexingMapsAttr( + rewriter.getAffineMapArrayAttr(indexingMaps)); + rewriter.replaceOp(transposeOp, newGenericOp); + return success(); + } + +private: + // Decides whether the transpose may bubble through the elementwise op. + ControlTransposePropagationFn controlFn; +}; + +} // namespace + +//===----------------------------------------------------------------------===// +// Linalg Named Op -> Named Op Conversions +//===----------------------------------------------------------------------===// + +namespace { + +template <typename OpTy, int64_t inputIdx> +class NamedOpConversion : public OpRewritePattern<OpTy> { +public: + NamedOpConversion(MLIRContext *ctx, SmallVector<int64_t> perm, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<OpTy>(ctx, b), permutation(perm), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(OpTy namedOp, + PatternRewriter &rewriter) const override { + Value input = namedOp.getInputs()[inputIdx]; + auto transpose = input.getDefiningOp<linalg::TransposeOp>(); + if (!transpose) { + return failure(); + } + if (!controlFn(namedOp.getDpsInputOperand(inputIdx))) { + return rewriter.notifyMatchFailure(namedOp, + "rejected by control function"); + } + + SmallVector<int64_t> transPerm(transpose.getPermutation()); + if (transPerm != permutation) { + return rewriter.notifyMatchFailure( + namedOp, "transpose permutation does not match target permutation"); + } + SmallVector<NamedAttribute> attrs = getPrunedAttributeList(namedOp); + SmallVector<Value> newInputs = namedOp.getInputs(); + newInputs[inputIdx] = transpose.getInput(); + + auto replaceOp = [&](auto *typePtr) { + rewriter.replaceOpWithNewOp<std::remove_pointer_t<decltype(typePtr)>>( + namedOp, newInputs, namedOp.getDpsInits(), attrs); + }; + + Operation *op = namedOp.getOperation(); + if (isa<linalg::MatmulTransposeAOp>(op) && inputIdx == 0) { + replaceOp(static_cast<linalg::MatmulOp *>(nullptr)); + } else if (isa<linalg::MatmulTransposeBOp>(op) && inputIdx == 1) { + replaceOp(static_cast<linalg::MatmulOp *>(nullptr)); + } else if (IREE::LinalgExt::isPureMatmul(op) && inputIdx == 0) { + replaceOp(static_cast<linalg::MatmulTransposeAOp *>(nullptr)); + } else if (IREE::LinalgExt::isPureMatmul(op) && inputIdx == 1) { + replaceOp(static_cast<linalg::MatmulTransposeBOp *>(nullptr)); + } else if (isa<linalg::BatchMatmulTransposeAOp>(op) && inputIdx == 0) { + replaceOp(static_cast<linalg::BatchMatmulOp *>(nullptr)); + } else if (isa<linalg::BatchMatmulTransposeBOp>(op) && inputIdx == 1) { + replaceOp(static_cast<linalg::BatchMatmulOp *>(nullptr)); + } else if (IREE::LinalgExt::isPureBatchMatmul(op) && inputIdx == 0) { + replaceOp(static_cast<linalg::BatchMatmulTransposeAOp *>(nullptr)); + } else if (IREE::LinalgExt::isPureBatchMatmul(op) && inputIdx == 1) { + replaceOp(static_cast<linalg::BatchMatmulTransposeBOp *>(nullptr)); + } else { + return failure(); + } + return success(); + } + +private: + // Non-type literal array template parameters are a C++20 feature, so instead + // all the named op patterns pass their permutation explicitly as a + // SmallVector. + SmallVector<int64_t> permutation; + // Decides whether the transpose may fold into the named op. + ControlTransposePropagationFn controlFn; +}; + +// Fuses a transpose into a reduction linalg.generic by absorbing it into the +// indexing map. +class FuseTransposeThroughGenericReduction + : public OpRewritePattern<linalg::GenericOp> { +public: + FuseTransposeThroughGenericReduction(MLIRContext *ctx, + ControlTransposePropagationFn controlFn, + PatternBenefit b = 1) + : OpRewritePattern<linalg::GenericOp>(ctx, b), + controlFn(std::move(controlFn)) {} + + LogicalResult matchAndRewrite(linalg::GenericOp genericOp, + PatternRewriter &rewriter) const override { + // Only try to fuse when the generic has at least one reduction dimension. + if (genericOp.getNumParallelLoops() == genericOp.getNumLoops()) { + return rewriter.notifyMatchFailure(genericOp, "not a reduction"); + } + + // All maps must be projected permutations. + if (!llvm::all_of(genericOp.getIndexingMapsArray(), [](AffineMap map) { + return map.isProjectedPermutation(); + })) { + return rewriter.notifyMatchFailure(genericOp, + "not a projected permutation"); + } + + // Look for a transpose on any input. + for (int64_t inputIdx = 0, endIdx = genericOp.getNumDpsInputs(); + inputIdx < endIdx; ++inputIdx) { + auto transpose = genericOp.getDpsInputs()[inputIdx] + .getDefiningOp<linalg::TransposeOp>(); + if (!transpose) { + continue; + } + if (!controlFn(genericOp.getDpsInputOperand(inputIdx))) { + return rewriter.notifyMatchFailure(genericOp, + "rejected by control function"); + } + + // Update the indexing map according to the transpose. + AffineMap inputMap = genericOp.getMatchingIndexingMap( + genericOp.getDpsInputOperand(inputIdx)); + + // Fuse by updating the indexing map to absorb the transpose. + auto invPerm = invertPermutationVector(transpose.getPermutation()); + SmallVector<AffineExpr> newExprs = + applyPermutation(inputMap.getResults(), invPerm); + + // Only fuse if dimension indices are in increasing order to maintain + // efficient memory access patterns. This should offset the fact that the + // transpose may have multiple uses. + int64_t prevDim = -1; + for (AffineExpr expr : newExprs) { + const int64_t dim = cast<AffineDimExpr>(expr).getPosition(); + if (dim < prevDim) { + return rewriter.notifyMatchFailure(genericOp, + "newExprs are not contiguous"); + } + prevDim = dim; + } + + auto transposedMap = + AffineMap::get(inputMap.getNumDims(), inputMap.getNumSymbols(), + newExprs, rewriter.getContext()); + + SmallVector<AffineMap> newIndexingMaps = genericOp.getIndexingMapsArray(); + newIndexingMaps[inputIdx] = transposedMap; + + rewriter.startOpModification(genericOp); + genericOp.setIndexingMapsAttr( + rewriter.getAffineMapArrayAttr(newIndexingMaps)); + genericOp.setOperand(inputIdx, transpose.getInput()); + rewriter.finalizeOpModification(genericOp); + return success(); + } + + return rewriter.notifyMatchFailure(genericOp, + "no matching transpose found"); + } + +private: + // Decides whether the transpose may fuse into the reduction. + ControlTransposePropagationFn controlFn; +}; + +} // namespace + +//===----------------------------------------------------------------------===// +// Pattern population +//===----------------------------------------------------------------------===// + +void populateSpecializeGenericTransposePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<SpecializeGenericTranspose>(patterns.getContext(), + std::move(controlFn), benefit); +} + +void populateFuseTransposeWithProducerLinalgOpPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + bool allowGeneralizing, bool allowConvolution, PatternBenefit benefit) { + patterns.add<FuseTransposeWithProducerLinalgOp>( + patterns.getContext(), std::move(controlFn), allowGeneralizing, + allowConvolution, benefit); +} + +void populateBubbleTransposeThroughCollapseShapePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<BubbleTransposeThroughCollapseShape>( + patterns.getContext(), std::move(controlFn), benefit); +} + +void populateBubbleTransposeThroughUnaryElementwiseDpsInitPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<BubbleTransposeThroughUnaryElementwiseDpsInit>( + patterns.getContext(), std::move(controlFn), benefit); +} + +void populateComposeTransposesPattern(RewritePatternSet &patterns, + ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<ComposeTransposes>(patterns.getContext(), std::move(controlFn), + benefit); +} + +void populateSinkTransposeThroughExtractSlicePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<SinkTransposeThroughExtractSlice>(patterns.getContext(), + std::move(controlFn), benefit); +} + +void populateSinkTransposeThroughExpandShapePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<SinkTransposeThroughExpandShape>(patterns.getContext(), + std::move(controlFn), benefit); +} + +void populateSinkTransposeThroughPadPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<SinkTransposeThroughPad>(patterns.getContext(), + std::move(controlFn), benefit); +} + +void populateFuseTransposeWithLinalgOpConsumerPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + bool allowGeneralizing, bool allowConvolution, PatternBenefit benefit) { + patterns.add<FuseTransposeWithLinalgOpConsumer>( + patterns.getContext(), std::move(controlFn), allowGeneralizing, + allowConvolution, benefit); +} + +void populateSinkTransposeThroughUnaryElementwiseInputPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<SinkTransposeThroughUnaryElementwiseInput>( + patterns.getContext(), std::move(controlFn), benefit); +} + +void populateFoldTransposeIntoNamedMatmulPatterns( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + MLIRContext *context = patterns.getContext(); + patterns.add<NamedOpConversion</*OpType=*/linalg::MatmulOp, + /*inputIdx=*/1>>( + context, SmallVector<int64_t>{1, 0}, controlFn, benefit); + patterns.add<NamedOpConversion</*OpType=*/linalg::MatmulOp, + /*inputIdx=*/0>>( + context, SmallVector<int64_t>{1, 0}, controlFn, benefit); + patterns.add<NamedOpConversion</*OpType=*/linalg::BatchMatmulOp, + /*inputIdx=*/1>>( + context, SmallVector<int64_t>{0, 2, 1}, controlFn, benefit); + patterns.add<NamedOpConversion</*OpType=*/linalg::BatchMatmulOp, + /*inputIdx=*/0>>( + context, SmallVector<int64_t>{0, 2, 1}, controlFn, benefit); +} + +void populateFuseTransposeThroughGenericReductionPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit) { + patterns.add<FuseTransposeThroughGenericReduction>( + patterns.getContext(), std::move(controlFn), benefit); +} + +} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h new file mode 100644 index 0000000..abfd656 --- /dev/null +++ b/compiler/src/iree/compiler/GlobalOptimization/TransposePropagationPatterns.h
@@ -0,0 +1,112 @@ +// Copyright 2026 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 + +#ifndef IREE_COMPILER_GLOBALOPTIMIZATION_TRANSPOSEPROPAGATIONPATTERNS_H_ +#define IREE_COMPILER_GLOBALOPTIMIZATION_TRANSPOSEPROPAGATIONPATTERNS_H_ + +#include <functional> + +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/IR/PatternMatch.h" + +namespace mlir::iree_compiler::GlobalOptimization { + +// Decides whether a transpose may be moved across the use edge |operand|. +// Patterns call it once, after the structural match, so the edge always +// connects the matched ops: the owner of |operand| consumes the value produced +// by the op on the other side of the edge. When a transpose is sunk, the +// transpose defines the operand; when a transpose is bubbled, the transpose +// owns the operand. +using ControlTransposePropagationFn = std::function<bool(OpOperand *operand)>; + +// Rewrites |genericOp| into a linalg.transpose if it only transposes its single +// input. Returns failure without modifying the IR otherwise. +LogicalResult specializeGenericTransposeOp(RewriterBase &rewriter, + linalg::GenericOp genericOp); + +// Specializes linalg.generic transposes into linalg.transpose. |controlFn| +// receives the input operand of the generic transpose. +void populateSpecializeGenericTransposePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +//===----------------------------------------------------------------------===// +// Transpose bubbling patterns +//===----------------------------------------------------------------------===// + +// Fuses a transpose into the init of its producer linalg.generic. Named +// contraction producers are generalized first when |allowGeneralizing| is set +// and convolution producers only when |allowConvolution| is set. +void populateFuseTransposeWithProducerLinalgOpPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + bool allowGeneralizing, bool allowConvolution, PatternBenefit benefit = 1); + +// Bubbles a transpose above its single-use tensor.collapse_shape producer. +void populateBubbleTransposeThroughCollapseShapePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Bubbles a transpose above an elementwise producer whose iteration-space +// permutation only affects a single input. +void populateBubbleTransposeThroughUnaryElementwiseDpsInitPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +//===----------------------------------------------------------------------===// +// Transpose sinking patterns +//===----------------------------------------------------------------------===// + +// Composes a transpose of a transpose into a single transpose, or removes both +// when they cancel. +void populateComposeTransposesPattern(RewritePatternSet &patterns, + ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Sinks a transpose below a tensor.extract_slice consumer. +void populateSinkTransposeThroughExtractSlicePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Sinks a single-use transpose below a tensor.expand_shape consumer. +void populateSinkTransposeThroughExpandShapePattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Sinks a transpose below a tensor.pad consumer whose padding value does not +// depend on the padded index. +void populateSinkTransposeThroughPadPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Fuses a single-use transpose into the indexing map of a consumer +// linalg.generic. Named contraction consumers are generalized first when +// |allowGeneralizing| is set and convolution consumers only when +// |allowConvolution| is set. +void populateFuseTransposeWithLinalgOpConsumerPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + bool allowGeneralizing, bool allowConvolution, PatternBenefit benefit = 1); + +// Sinks a transpose below an elementwise consumer whose iteration-space +// permutation only affects the transposed input. +void populateSinkTransposeThroughUnaryElementwiseInputPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Folds a transpose of a matmul or batch_matmul operand into the transposed +// named op variant, and a transposed named op input back into the plain op. +void populateFoldTransposeIntoNamedMatmulPatterns( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +// Fuses a transpose into the indexing map of a reduction linalg.generic when +// the resulting access order stays increasing. +void populateFuseTransposeThroughGenericReductionPattern( + RewritePatternSet &patterns, ControlTransposePropagationFn controlFn, + PatternBenefit benefit = 1); + +} // namespace mlir::iree_compiler::GlobalOptimization + +#endif // IREE_COMPILER_GLOBALOPTIMIZATION_TRANSPOSEPROPAGATIONPATTERNS_H_