| //===----- FlattenMemRefs.cpp - MemRef ops flattener pass ----------------===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file contains patterns for flattening an multi-rank memref-related |
| // ops into 1-d memref ops. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/Affine/IR/AffineOps.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/MemRef/IR/MemRef.h" |
| #include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h" |
| #include "mlir/Dialect/MemRef/Transforms/Passes.h" |
| #include "mlir/Dialect/MemRef/Transforms/Transforms.h" |
| #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h" |
| #include "mlir/Dialect/Utils/IndexingUtils.h" |
| #include "mlir/Dialect/Utils/StaticValueUtils.h" |
| #include "mlir/Dialect/Vector/IR/VectorOps.h" |
| #include "mlir/IR/Attributes.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "mlir/IR/OpDefinition.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include <algorithm> |
| |
| namespace mlir { |
| namespace memref { |
| #define GEN_PASS_DEF_FLATTENMEMREFSPASS |
| #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc" |
| } // namespace memref |
| } // namespace mlir |
| |
| using namespace mlir; |
| |
| static Value getValueFromOpFoldResult(OpBuilder &rewriter, Location loc, |
| OpFoldResult in) { |
| if (Attribute offsetAttr = dyn_cast<Attribute>(in)) { |
| return arith::ConstantIndexOp::create( |
| rewriter, loc, cast<IntegerAttr>(offsetAttr).getInt()); |
| } |
| return cast<Value>(in); |
| } |
| |
| /// Returns a collapsed memref and the linearized index to access the element |
| /// at the specified indices. |
| static std::pair<Value, Value> getFlattenMemrefAndOffset(OpBuilder &rewriter, |
| Location loc, |
| Value source, |
| ValueRange indices) { |
| int64_t sourceOffset; |
| SmallVector<int64_t, 4> sourceStrides; |
| auto sourceType = cast<MemRefType>(source.getType()); |
| if (failed(sourceType.getStridesAndOffset(sourceStrides, sourceOffset))) { |
| assert(false); |
| } |
| |
| memref::ExtractStridedMetadataOp stridedMetadata = |
| memref::ExtractStridedMetadataOp::create(rewriter, loc, source); |
| |
| auto typeBit = sourceType.getElementType().getIntOrFloatBitWidth(); |
| OpFoldResult linearizedIndices; |
| memref::LinearizedMemRefInfo linearizedInfo; |
| std::tie(linearizedInfo, linearizedIndices) = |
| memref::getLinearizedMemRefOffsetAndSize( |
| rewriter, loc, typeBit, typeBit, |
| stridedMetadata.getConstifiedMixedOffset(), |
| stridedMetadata.getConstifiedMixedSizes(), |
| stridedMetadata.getConstifiedMixedStrides(), |
| getAsOpFoldResult(indices)); |
| |
| return std::make_pair( |
| memref::ReinterpretCastOp::create( |
| rewriter, loc, source, |
| /*offset=*/linearizedInfo.linearizedOffset, |
| /*sizes=*/ |
| ArrayRef<OpFoldResult>{linearizedInfo.linearizedSize}, |
| /*strides=*/ |
| ArrayRef<OpFoldResult>{rewriter.getIndexAttr(1)}), |
| getValueFromOpFoldResult(rewriter, loc, linearizedIndices)); |
| } |
| |
| static bool needFlattening(Value val) { |
| auto type = cast<MemRefType>(val.getType()); |
| return type.getRank() > 1; |
| } |
| |
| static bool checkLayout(Value val) { |
| auto type = cast<MemRefType>(val.getType()); |
| return type.getLayout().isIdentity() || |
| isa<StridedLayoutAttr>(type.getLayout()); |
| } |
| |
| namespace { |
| static bool hasSupportedElementType(Value memref) { |
| auto type = cast<MemRefType>(memref.getType()); |
| return type.getElementType().isIntOrFloat(); |
| } |
| |
| /// Compute the type that will be used to linearize the memref. |
| /// Used so we don't create IR like `getLinearizedMemRefOffsetAndSize` would. |
| static FailureOr<MemRefType> getFlattenedMemRefType(MemRefType sourceType) { |
| int64_t sourceOffset; |
| SmallVector<int64_t> sourceStrides; |
| if (failed(sourceType.getStridesAndOffset(sourceStrides, sourceOffset))) |
| return failure(); |
| |
| auto flatDimSize = SaturatedInteger::wrap(0); |
| for (auto [size, stride] : |
| llvm::zip_equal(sourceType.getShape(), sourceStrides)) { |
| auto dimSize = |
| SaturatedInteger::wrap(size) * SaturatedInteger::wrap(stride); |
| flatDimSize = flatDimSize.smax(dimSize); |
| if (flatDimSize.isSaturated()) |
| break; |
| } |
| |
| if (sourceType.getLayout().isIdentity()) |
| return MemRefType::get( |
| {flatDimSize.asInteger()}, sourceType.getElementType(), |
| MemRefLayoutAttrInterface{}, sourceType.getMemorySpace()); |
| |
| return MemRefType::get( |
| {flatDimSize.asInteger()}, sourceType.getElementType(), |
| StridedLayoutAttr::get(sourceType.getContext(), sourceOffset, {1}), |
| sourceType.getMemorySpace()); |
| } |
| |
| /// Return whether `memref` has the basic properties needed for linearizing it |
| /// into a 1-D reinterpret_cast. |
| static LogicalResult checkFlattenableMemref(Operation *op, Value memref, |
| PatternRewriter &rewriter) { |
| if (!needFlattening(memref)) |
| return rewriter.notifyMatchFailure(op, "memref does not need flattening"); |
| if (!checkLayout(memref)) |
| return rewriter.notifyMatchFailure(op, "unsupported memref layout"); |
| if (!hasSupportedElementType(memref)) |
| return rewriter.notifyMatchFailure(op, "unsupported element type"); |
| return success(); |
| } |
| |
| /// Wrapeer around checking if the last memref dimension is contiguous that |
| /// provides nice failures message. |
| static LogicalResult hasUnitTrailingStride(Operation *op, |
| TypedValue<MemRefType> memref, |
| PatternRewriter &rewriter) { |
| if (!memref.getType().areTrailingDimsContiguous(1)) |
| return rewriter.notifyMatchFailure( |
| op, "cannot preserve non-unit trailing access stride"); |
| |
| return success(); |
| } |
| |
| static LogicalResult |
| canLinearizeAccessedShape(memref::IndexedAccessOpInterface op, |
| TypedValue<MemRefType> memref, |
| PatternRewriter &rewriter) { |
| SmallVector<int64_t> accessedShape = op.getAccessedShape(); |
| if (accessedShape.empty()) |
| return success(); |
| if (accessedShape.size() > 1) |
| return rewriter.notifyMatchFailure( |
| op, "cannot preserve multi-dimensional accessed shape"); |
| |
| return hasUnitTrailingStride(op, memref, rewriter); |
| } |
| |
| static LogicalResult canFlattenTransferOp(VectorTransferOpInterface op, |
| TypedValue<MemRefType> memref, |
| PatternRewriter &rewriter) { |
| // For vector.transfer_read/write, must make sure: |
| // 1. all accesses are inbounds, |
| // 2. has a minor identity permutation map, and |
| // 3. has at most one transfer dimension. |
| AffineMap permutationMap = op.getPermutationMap(); |
| if (!permutationMap.isMinorIdentity()) |
| return rewriter.notifyMatchFailure( |
| op, "only identity or minor identity permutation map is supported"); |
| |
| if (op.hasOutOfBoundsDim()) |
| return rewriter.notifyMatchFailure(op, "only inbounds are supported"); |
| |
| if (op.getTransferRank() > 1) |
| return rewriter.notifyMatchFailure( |
| op, "cannot flatten multi-dimensional vector transfer"); |
| |
| if (op.getTransferRank() > 0 && |
| failed(hasUnitTrailingStride(op, memref, rewriter))) |
| return failure(); |
| |
| return success(); |
| } |
| |
| // Pattern for memref::AllocOp and memref::AllocaOp. |
| // |
| // The "source" memref for these ops is the op's own result, so the generic |
| // indexed access pattern cannot be used: getFlattenMemrefAndOffset would |
| // insert ExtractStridedMetadataOp and ReinterpretCastOp that use op.result |
| // before this op in the block. After replaceOpWithNewOp the original result is |
| // RAUW'd to the new ReinterpretCastOp, leaving the earlier ops with forward |
| // references (domination violations) caught by |
| // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS. |
| // |
| // Instead, sizes and strides are computed from the op's operands and type |
| // (which all dominate the op), avoiding any reference to op.result until the |
| // final replaceOpWithNewOp. |
| template <typename AllocLikeOp> |
| struct AllocLikeFlattenPattern final : public OpRewritePattern<AllocLikeOp> { |
| using Base = OpRewritePattern<AllocLikeOp>; |
| using Base::Base; |
| |
| LogicalResult matchAndRewrite(AllocLikeOp op, |
| PatternRewriter &rewriter) const override { |
| if (!needFlattening(op.getMemref()) || !checkLayout(op.getMemref())) |
| return failure(); |
| |
| Location loc = op->getLoc(); |
| auto memrefType = cast<MemRefType>(op.getType()); |
| auto elemType = memrefType.getElementType(); |
| if (!elemType.isIntOrFloat()) |
| return failure(); |
| unsigned elemBitWidth = elemType.getIntOrFloatBitWidth(); |
| |
| SmallVector<OpFoldResult> sizes = op.getMixedSizes(); |
| |
| int64_t staticOffset; |
| SmallVector<int64_t> staticStrides; |
| if (failed(memrefType.getStridesAndOffset(staticStrides, staticOffset))) |
| return failure(); |
| if (staticOffset == ShapedType::kDynamic) |
| return rewriter.notifyMatchFailure(op, "dynamic offset not supported"); |
| SmallVector<OpFoldResult> strides; |
| strides.reserve(staticStrides.size()); |
| for (int64_t stride : staticStrides) { |
| if (stride == ShapedType::kDynamic) |
| return rewriter.notifyMatchFailure(op, |
| "dynamic stride cannot be computed"); |
| strides.push_back(rewriter.getIndexAttr(stride)); |
| } |
| |
| // Compute the linearized flat extent from sizes and strides (no SSA ops |
| // referencing op.result are created here). |
| memref::LinearizedMemRefInfo linearizedInfo; |
| OpFoldResult linearizedOffset; |
| std::tie(linearizedInfo, linearizedOffset) = |
| memref::getLinearizedMemRefOffsetAndSize( |
| rewriter, loc, elemBitWidth, elemBitWidth, rewriter.getIndexAttr(0), |
| sizes, strides); |
| (void)linearizedOffset; |
| |
| // The total allocation must cover [0, staticOffset + linearizedExtent). |
| // When the offset is non-zero, add it to the computed extent so that the |
| // buffer is large enough for elements accessed at positions |
| // [staticOffset, staticOffset + linearizedExtent). |
| OpFoldResult flatSizeOfr = linearizedInfo.linearizedSize; |
| if (staticOffset != 0) { |
| AffineExpr s0; |
| bindSymbols(rewriter.getContext(), s0); |
| flatSizeOfr = affine::makeComposedFoldedAffineApply( |
| rewriter, loc, s0 + staticOffset, {flatSizeOfr}); |
| } |
| |
| // Build the flat 1-D MemRefType. The linearized size may be static or |
| // dynamic (OpFoldResult of either IntegerAttr or a Value). |
| int64_t flatDimSize = ShapedType::kDynamic; |
| if (auto attr = dyn_cast<Attribute>(flatSizeOfr)) |
| if (auto intAttr = dyn_cast<IntegerAttr>(attr)) |
| flatDimSize = intAttr.getInt(); |
| |
| auto flatMemrefType = |
| MemRefType::get({flatDimSize}, memrefType.getElementType(), |
| StridedLayoutAttr::get(rewriter.getContext(), 0, {1}), |
| memrefType.getMemorySpace()); |
| |
| // Collect the flat dynamic-size operand (empty for fully-static case). |
| SmallVector<Value, 1> dynSizes; |
| if (flatDimSize == ShapedType::kDynamic) |
| dynSizes.push_back(getValueFromOpFoldResult(rewriter, loc, flatSizeOfr)); |
| |
| auto newOp = AllocLikeOp::create(rewriter, loc, flatMemrefType, dynSizes, |
| op.getAlignmentAttr()); |
| rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>( |
| op, cast<MemRefType>(op.getType()), newOp, |
| rewriter.getIndexAttr(staticOffset), sizes, strides); |
| return success(); |
| } |
| }; |
| |
| /// Pattern that flattens any IndexedAccessOpInterface op. |
| struct IndexedAccessOpFlattenPattern final |
| : public OpInterfaceRewritePattern<memref::IndexedAccessOpInterface> { |
| using Base::Base; |
| |
| LogicalResult matchAndRewrite(memref::IndexedAccessOpInterface op, |
| PatternRewriter &rewriter) const override { |
| TypedValue<MemRefType> memref = op.getAccessedMemref(); |
| if (!memref) |
| return rewriter.notifyMatchFailure(op, "not accessing a memref"); |
| if (failed(checkFlattenableMemref(op, memref, rewriter))) |
| return failure(); |
| if (failed(canLinearizeAccessedShape(op, memref, rewriter))) |
| return failure(); |
| |
| auto [flatMemref, offset] = getFlattenMemrefAndOffset( |
| rewriter, op->getLoc(), memref, op.getIndices()); |
| std::optional<SmallVector<Value>> replacementValues = |
| op.updateMemrefAndIndices(rewriter, flatMemref, ValueRange{offset}); |
| if (replacementValues) |
| rewriter.replaceOp(op, *replacementValues); |
| return success(); |
| } |
| }; |
| |
| /// Flatten operations that use VectorTransferOpInterface. Transfer ops have |
| /// permutation-map and in_bounds semantics that are separate from |
| /// IndexedAccessOpInterface, so use updateStartingPosition directly. |
| struct VectorTransferOpFlattenPattern final |
| : public OpInterfaceRewritePattern<VectorTransferOpInterface> { |
| using Base::Base; |
| |
| LogicalResult matchAndRewrite(VectorTransferOpInterface op, |
| PatternRewriter &rewriter) const override { |
| auto memref = dyn_cast<TypedValue<MemRefType>>(op.getBase()); |
| if (!memref) |
| return rewriter.notifyMatchFailure(op, "not accessing a memref"); |
| if (failed(checkFlattenableMemref(op, memref, rewriter))) |
| return failure(); |
| if (failed(canFlattenTransferOp(op, memref, rewriter))) |
| return failure(); |
| |
| FailureOr<MemRefType> flatMemrefType = |
| getFlattenedMemRefType(memref.getType()); |
| if (failed(flatMemrefType)) |
| return failure(); |
| AffineMap newPermutationMap = AffineMap::getMinorIdentityMap( |
| /*dims=*/1, op.getTransferRank(), op.getContext()); |
| if (failed( |
| op.mayUpdateStartingPosition(*flatMemrefType, newPermutationMap))) |
| return rewriter.notifyMatchFailure(op, |
| "failed op-specific preconditions"); |
| |
| auto [flatMemref, offset] = getFlattenMemrefAndOffset( |
| rewriter, op->getLoc(), memref, op.getIndices()); |
| op.updateStartingPosition(rewriter, flatMemref, ValueRange{offset}, |
| AffineMapAttr::get(newPermutationMap)); |
| return success(); |
| } |
| }; |
| |
| /// Flatten the source and destination memref/index pairs of indexed memcpy-like |
| /// operations such as memref.dma_start. |
| struct FlattenedMemrefAccess { |
| Value memref; |
| Value index; |
| }; |
| |
| /// Flatten all IndexedMemCopyOpInterface operations. |
| struct IndexedMemCopyOpFlattenPattern final |
| : public OpInterfaceRewritePattern<memref::IndexedMemCopyOpInterface> { |
| using Base::Base; |
| |
| LogicalResult matchAndRewrite(memref::IndexedMemCopyOpInterface op, |
| PatternRewriter &rewriter) const override { |
| TypedValue<MemRefType> src = op.getSrc(); |
| TypedValue<MemRefType> dst = op.getDst(); |
| if (!src && !dst) |
| return rewriter.notifyMatchFailure(op, "not copying between memrefs"); |
| |
| auto tryFlatten = |
| [&](TypedValue<MemRefType> memref, |
| ValueRange indices) -> std::optional<FlattenedMemrefAccess> { |
| if (!memref || !needFlattening(memref)) |
| return std::nullopt; |
| if (failed(checkFlattenableMemref(op, memref, rewriter))) |
| return std::nullopt; |
| |
| auto [flatMemref, offset] = |
| getFlattenMemrefAndOffset(rewriter, op->getLoc(), memref, indices); |
| return FlattenedMemrefAccess{flatMemref, offset}; |
| }; |
| |
| std::optional<FlattenedMemrefAccess> newSrc = |
| tryFlatten(src, op.getSrcIndices()); |
| std::optional<FlattenedMemrefAccess> newDst = |
| tryFlatten(dst, op.getDstIndices()); |
| if (!newSrc && !newDst) |
| return rewriter.notifyMatchFailure( |
| op, "no source or destination memref needed flattening"); |
| |
| Value srcMemref = src; |
| ValueRange srcIndices = op.getSrcIndices(); |
| if (newSrc) { |
| srcMemref = newSrc->memref; |
| srcIndices = ValueRange(newSrc->index); |
| } |
| |
| Value dstMemref = dst; |
| ValueRange dstIndices = op.getDstIndices(); |
| if (newDst) { |
| dstMemref = newDst->memref; |
| dstIndices = ValueRange(newDst->index); |
| } |
| |
| op.setMemrefsAndIndices(rewriter, srcMemref, srcIndices, dstMemref, |
| dstIndices); |
| return success(); |
| } |
| }; |
| |
| struct FlattenMemrefsPass |
| : public mlir::memref::impl::FlattenMemrefsPassBase<FlattenMemrefsPass> { |
| using Base::Base; |
| |
| void getDependentDialects(DialectRegistry ®istry) const override { |
| registry.insert<affine::AffineDialect, arith::ArithDialect, |
| memref::MemRefDialect>(); |
| } |
| |
| void runOnOperation() override { |
| RewritePatternSet patterns(&getContext()); |
| |
| memref::populateFlattenMemrefsPatterns(patterns); |
| |
| if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) |
| return signalPassFailure(); |
| } |
| }; |
| |
| } // namespace |
| |
| void memref::populateFlattenMemrefsPatterns(RewritePatternSet &patterns) { |
| patterns.insert<IndexedAccessOpFlattenPattern, IndexedMemCopyOpFlattenPattern, |
| VectorTransferOpFlattenPattern, |
| AllocLikeFlattenPattern<memref::AllocOp>, |
| AllocLikeFlattenPattern<memref::AllocaOp>>( |
| patterns.getContext()); |
| } |
