293 lines · cpp
1//===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements utilities for the MemRef dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Arith/Utils/Utils.h"16#include "mlir/Dialect/MemRef/IR/MemRef.h"17#include "mlir/Interfaces/ViewLikeInterface.h"18#include "llvm/ADT/STLExtras.h"19 20namespace mlir {21namespace memref {22 23bool isStaticShapeAndContiguousRowMajor(MemRefType type) {24 if (!type.hasStaticShape())25 return false;26 27 SmallVector<int64_t> strides;28 int64_t offset;29 if (failed(type.getStridesAndOffset(strides, offset)))30 return false;31 32 // MemRef is contiguous if outer dimensions are size-1 and inner33 // dimensions have unit strides.34 int64_t runningStride = 1;35 int64_t curDim = strides.size() - 1;36 // Finds all inner dimensions with unit strides.37 while (curDim >= 0 && strides[curDim] == runningStride) {38 runningStride *= type.getDimSize(curDim);39 --curDim;40 }41 42 // Check if other dimensions are size-1.43 while (curDim >= 0 && type.getDimSize(curDim) == 1) {44 --curDim;45 }46 47 // All dims are unit-strided or size-1.48 return curDim < 0;49}50 51std::pair<LinearizedMemRefInfo, OpFoldResult> getLinearizedMemRefOffsetAndSize(52 OpBuilder &builder, Location loc, int srcBits, int dstBits,53 OpFoldResult offset, ArrayRef<OpFoldResult> sizes,54 ArrayRef<OpFoldResult> strides, ArrayRef<OpFoldResult> indices) {55 unsigned sourceRank = sizes.size();56 assert(sizes.size() == strides.size() &&57 "expected as many sizes as strides for a memref");58 SmallVector<OpFoldResult> indicesVec = llvm::to_vector(indices);59 if (indices.empty())60 indicesVec.resize(sourceRank, builder.getIndexAttr(0));61 assert(indicesVec.size() == strides.size() &&62 "expected as many indices as rank of memref");63 64 // Create the affine symbols and values for linearization.65 SmallVector<AffineExpr> symbols(2 * sourceRank);66 bindSymbolsList(builder.getContext(), MutableArrayRef{symbols});67 AffineExpr addMulMap = builder.getAffineConstantExpr(0);68 69 SmallVector<OpFoldResult> offsetValues(2 * sourceRank);70 71 for (unsigned i = 0; i < sourceRank; ++i) {72 unsigned offsetIdx = 2 * i;73 addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];74 offsetValues[offsetIdx] = indicesVec[i];75 offsetValues[offsetIdx + 1] = strides[i];76 }77 // Adjust linearizedIndices and size by the scale factor (dstBits / srcBits).78 int64_t scaler = dstBits / srcBits;79 OpFoldResult linearizedIndices = affine::makeComposedFoldedAffineApply(80 builder, loc, addMulMap.floorDiv(scaler), offsetValues);81 82 size_t symbolIndex = 0;83 SmallVector<OpFoldResult> values;84 SmallVector<AffineExpr> productExpressions;85 for (unsigned i = 0; i < sourceRank; ++i) {86 AffineExpr strideExpr = symbols[symbolIndex++];87 values.push_back(strides[i]);88 AffineExpr sizeExpr = symbols[symbolIndex++];89 values.push_back(sizes[i]);90 91 productExpressions.push_back((strideExpr * sizeExpr).floorDiv(scaler));92 }93 AffineMap maxMap = AffineMap::get(94 /*dimCount=*/0, /*symbolCount=*/symbolIndex, productExpressions,95 builder.getContext());96 OpFoldResult linearizedSize =97 affine::makeComposedFoldedAffineMax(builder, loc, maxMap, values);98 99 // Adjust baseOffset by the scale factor (dstBits / srcBits).100 AffineExpr s0;101 bindSymbols(builder.getContext(), s0);102 OpFoldResult adjustBaseOffset = affine::makeComposedFoldedAffineApply(103 builder, loc, s0.floorDiv(scaler), {offset});104 105 OpFoldResult intraVectorOffset = affine::makeComposedFoldedAffineApply(106 builder, loc, addMulMap % scaler, offsetValues);107 108 return {{adjustBaseOffset, linearizedSize, intraVectorOffset},109 linearizedIndices};110}111 112LinearizedMemRefInfo113getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits,114 int dstBits, OpFoldResult offset,115 ArrayRef<OpFoldResult> sizes) {116 SmallVector<OpFoldResult> strides(sizes.size());117 if (!sizes.empty()) {118 strides.back() = builder.getIndexAttr(1);119 AffineExpr s0, s1;120 bindSymbols(builder.getContext(), s0, s1);121 for (int index = sizes.size() - 1; index > 0; --index) {122 strides[index - 1] = affine::makeComposedFoldedAffineApply(123 builder, loc, s0 * s1,124 ArrayRef<OpFoldResult>{strides[index], sizes[index]});125 }126 }127 128 LinearizedMemRefInfo linearizedMemRefInfo;129 std::tie(linearizedMemRefInfo, std::ignore) =130 getLinearizedMemRefOffsetAndSize(builder, loc, srcBits, dstBits, offset,131 sizes, strides);132 return linearizedMemRefInfo;133}134 135/// Returns true if all the uses of op are not read/load.136/// There can be view-like-op users as long as all its users are also137/// StoreOp/transfer_write. If return true it also fills out the uses, if it138/// returns false uses is unchanged.139static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) {140 std::vector<Operation *> opUses;141 for (OpOperand &use : op->getUses()) {142 Operation *useOp = use.getOwner();143 // Use escaped the scope144 if (useOp->mightHaveTrait<OpTrait::IsTerminator>())145 return false;146 if (isa<memref::DeallocOp>(useOp) ||147 (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 &&148 !mlir::hasEffect<MemoryEffects::Read>(useOp)) ||149 (isa<ViewLikeOpInterface>(useOp) && resultIsNotRead(useOp, opUses))) {150 opUses.push_back(useOp);151 continue;152 }153 return false;154 }155 llvm::append_range(uses, opUses);156 return true;157}158 159void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp) {160 std::vector<Operation *> opToErase;161 parentOp->walk([&](Operation *op) {162 std::vector<Operation *> candidates;163 if (isa<memref::AllocOp, memref::AllocaOp>(op) &&164 resultIsNotRead(op, candidates)) {165 llvm::append_range(opToErase, candidates);166 opToErase.push_back(op);167 }168 });169 170 for (Operation *op : opToErase)171 rewriter.eraseOp(op);172}173 174static SmallVector<OpFoldResult>175computeSuffixProductIRBlockImpl(Location loc, OpBuilder &builder,176 ArrayRef<OpFoldResult> sizes,177 OpFoldResult unit) {178 SmallVector<OpFoldResult> strides(sizes.size(), unit);179 AffineExpr s0, s1;180 bindSymbols(builder.getContext(), s0, s1);181 182 for (int64_t r = strides.size() - 1; r > 0; --r) {183 strides[r - 1] = affine::makeComposedFoldedAffineApply(184 builder, loc, s0 * s1, {strides[r], sizes[r]});185 }186 return strides;187}188 189SmallVector<OpFoldResult>190computeSuffixProductIRBlock(Location loc, OpBuilder &builder,191 ArrayRef<OpFoldResult> sizes) {192 OpFoldResult unit = builder.getIndexAttr(1);193 return computeSuffixProductIRBlockImpl(loc, builder, sizes, unit);194}195 196MemrefValue skipFullyAliasingOperations(MemrefValue source) {197 while (auto *op = source.getDefiningOp()) {198 if (auto subViewOp = dyn_cast<memref::SubViewOp>(op);199 subViewOp && subViewOp.hasZeroOffset() && subViewOp.hasUnitStride()) {200 // A `memref.subview` with an all zero offset, and all unit strides, still201 // points to the same memory.202 source = cast<MemrefValue>(subViewOp.getSource());203 } else if (auto castOp = dyn_cast<memref::CastOp>(op)) {204 // A `memref.cast` still points to the same memory.205 source = castOp.getSource();206 } else {207 return source;208 }209 }210 return source;211}212 213MemrefValue skipViewLikeOps(MemrefValue source) {214 while (auto *op = source.getDefiningOp()) {215 if (auto viewLike = dyn_cast<ViewLikeOpInterface>(op)) {216 if (source == viewLike.getViewDest()) {217 source = cast<MemrefValue>(viewLike.getViewSource());218 continue;219 }220 }221 return source;222 }223 return source;224}225 226LogicalResult resolveSourceIndicesExpandShape(227 Location loc, PatternRewriter &rewriter,228 memref::ExpandShapeOp expandShapeOp, ValueRange indices,229 SmallVectorImpl<Value> &sourceIndices, bool startsInbounds) {230 SmallVector<OpFoldResult> destShape = expandShapeOp.getMixedOutputShape();231 232 // Traverse all reassociation groups to determine the appropriate indices233 // corresponding to each one of them post op folding.234 for (ArrayRef<int64_t> group : expandShapeOp.getReassociationIndices()) {235 assert(!group.empty() && "association indices groups cannot be empty");236 int64_t groupSize = group.size();237 if (groupSize == 1) {238 sourceIndices.push_back(indices[group[0]]);239 continue;240 }241 SmallVector<OpFoldResult> groupBasis =242 llvm::map_to_vector(group, [&](int64_t d) { return destShape[d]; });243 SmallVector<Value> groupIndices =244 llvm::map_to_vector(group, [&](int64_t d) { return indices[d]; });245 Value collapsedIndex = affine::AffineLinearizeIndexOp::create(246 rewriter, loc, groupIndices, groupBasis, /*disjoint=*/startsInbounds);247 sourceIndices.push_back(collapsedIndex);248 }249 return success();250}251 252LogicalResult253resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,254 memref::CollapseShapeOp collapseShapeOp,255 ValueRange indices,256 SmallVectorImpl<Value> &sourceIndices) {257 // Note: collapse_shape requires a strided memref, we can do this.258 auto metadata = memref::ExtractStridedMetadataOp::create(259 rewriter, loc, collapseShapeOp.getSrc());260 SmallVector<OpFoldResult> sourceSizes = metadata.getConstifiedMixedSizes();261 for (auto [index, group] :262 llvm::zip(indices, collapseShapeOp.getReassociationIndices())) {263 assert(!group.empty() && "association indices groups cannot be empty");264 int64_t groupSize = group.size();265 266 if (groupSize == 1) {267 sourceIndices.push_back(index);268 continue;269 }270 271 SmallVector<OpFoldResult> basis =272 llvm::map_to_vector(group, [&](int64_t d) { return sourceSizes[d]; });273 auto delinearize = affine::AffineDelinearizeIndexOp::create(274 rewriter, loc, index, basis, /*hasOuterBound=*/true);275 llvm::append_range(sourceIndices, delinearize.getResults());276 }277 if (collapseShapeOp.getReassociationIndices().empty()) {278 auto zeroAffineMap = rewriter.getConstantAffineMap(0);279 int64_t srcRank =280 cast<MemRefType>(collapseShapeOp.getViewSource().getType()).getRank();281 OpFoldResult ofr = affine::makeComposedFoldedAffineApply(282 rewriter, loc, zeroAffineMap, ArrayRef<OpFoldResult>{});283 for (int64_t i = 0; i < srcRank; i++) {284 sourceIndices.push_back(285 getValueOrCreateConstantIndexOp(rewriter, loc, ofr));286 }287 }288 return success();289}290 291} // namespace memref292} // namespace mlir293