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-exception 6 // 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/Dialect/Vector/IR/VectorOps.h" 18 #include "mlir/Interfaces/ViewLikeInterface.h" 19 #include "llvm/ADT/STLExtras.h" 20 21 namespace mlir { 22 namespace memref { 23 24 bool isStaticShapeAndContiguousRowMajor(MemRefType type) { 25 if (!type.hasStaticShape()) 26 return false; 27 28 SmallVector<int64_t> strides; 29 int64_t offset; 30 if (failed(type.getStridesAndOffset(strides, offset))) 31 return false; 32 33 // MemRef is contiguous if outer dimensions are size-1 and inner 34 // dimensions have unit strides. 35 int64_t runningStride = 1; 36 int64_t curDim = strides.size() - 1; 37 // Finds all inner dimensions with unit strides. 38 while (curDim >= 0 && strides[curDim] == runningStride) { 39 runningStride *= type.getDimSize(curDim); 40 --curDim; 41 } 42 43 // Check if other dimensions are size-1. 44 while (curDim >= 0 && type.getDimSize(curDim) == 1) { 45 --curDim; 46 } 47 48 // All dims are unit-strided or size-1. 49 return curDim < 0; 50 } 51 52 std::pair<LinearizedMemRefInfo, OpFoldResult> getLinearizedMemRefOffsetAndSize( 53 OpBuilder &builder, Location loc, int srcBits, int dstBits, 54 OpFoldResult offset, ArrayRef<OpFoldResult> sizes, 55 ArrayRef<OpFoldResult> strides, ArrayRef<OpFoldResult> indices) { 56 unsigned sourceRank = sizes.size(); 57 assert(sizes.size() == strides.size() && 58 "expected as many sizes as strides for a memref"); 59 SmallVector<OpFoldResult> indicesVec = llvm::to_vector(indices); 60 if (indices.empty()) 61 indicesVec.resize(sourceRank, builder.getIndexAttr(0)); 62 assert(indicesVec.size() == strides.size() && 63 "expected as many indices as rank of memref"); 64 65 // Create the affine symbols and values for linearization. 66 SmallVector<AffineExpr> symbols(2 * sourceRank); 67 bindSymbolsList(builder.getContext(), MutableArrayRef{symbols}); 68 AffineExpr addMulMap = builder.getAffineConstantExpr(0); 69 AffineExpr mulMap = builder.getAffineConstantExpr(1); 70 71 SmallVector<OpFoldResult> offsetValues(2 * sourceRank); 72 73 for (unsigned i = 0; i < sourceRank; ++i) { 74 unsigned offsetIdx = 2 * i; 75 addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1]; 76 offsetValues[offsetIdx] = indicesVec[i]; 77 offsetValues[offsetIdx + 1] = strides[i]; 78 79 mulMap = mulMap * symbols[i]; 80 } 81 82 // Adjust linearizedIndices and size by the scale factor (dstBits / srcBits). 83 int64_t scaler = dstBits / srcBits; 84 mulMap = mulMap.floorDiv(scaler); 85 86 OpFoldResult linearizedIndices = affine::makeComposedFoldedAffineApply( 87 builder, loc, addMulMap.floorDiv(scaler), offsetValues); 88 OpFoldResult linearizedSize = 89 affine::makeComposedFoldedAffineApply(builder, loc, mulMap, sizes); 90 91 // Adjust baseOffset by the scale factor (dstBits / srcBits). 92 AffineExpr s0; 93 bindSymbols(builder.getContext(), s0); 94 OpFoldResult adjustBaseOffset = affine::makeComposedFoldedAffineApply( 95 builder, loc, s0.floorDiv(scaler), {offset}); 96 97 OpFoldResult intraVectorOffset = affine::makeComposedFoldedAffineApply( 98 builder, loc, addMulMap % scaler, offsetValues); 99 100 return {{adjustBaseOffset, linearizedSize, intraVectorOffset}, 101 linearizedIndices}; 102 } 103 104 LinearizedMemRefInfo 105 getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits, 106 int dstBits, OpFoldResult offset, 107 ArrayRef<OpFoldResult> sizes) { 108 SmallVector<OpFoldResult> strides(sizes.size()); 109 if (!sizes.empty()) { 110 strides.back() = builder.getIndexAttr(1); 111 AffineExpr s0, s1; 112 bindSymbols(builder.getContext(), s0, s1); 113 for (int index = sizes.size() - 1; index > 0; --index) { 114 strides[index - 1] = affine::makeComposedFoldedAffineApply( 115 builder, loc, s0 * s1, 116 ArrayRef<OpFoldResult>{strides[index], sizes[index]}); 117 } 118 } 119 120 LinearizedMemRefInfo linearizedMemRefInfo; 121 std::tie(linearizedMemRefInfo, std::ignore) = 122 getLinearizedMemRefOffsetAndSize(builder, loc, srcBits, dstBits, offset, 123 sizes, strides); 124 return linearizedMemRefInfo; 125 } 126 127 /// Returns true if all the uses of op are not read/load. 128 /// There can be SubviewOp users as long as all its users are also 129 /// StoreOp/transfer_write. If return true it also fills out the uses, if it 130 /// returns false uses is unchanged. 131 static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) { 132 std::vector<Operation *> opUses; 133 for (OpOperand &use : op->getUses()) { 134 Operation *useOp = use.getOwner(); 135 if (isa<memref::DeallocOp>(useOp) || 136 (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 && 137 !mlir::hasEffect<MemoryEffects::Read>(useOp)) || 138 (isa<memref::SubViewOp>(useOp) && resultIsNotRead(useOp, opUses))) { 139 opUses.push_back(useOp); 140 continue; 141 } 142 return false; 143 } 144 uses.insert(uses.end(), opUses.begin(), opUses.end()); 145 return true; 146 } 147 148 void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp) { 149 std::vector<Operation *> opToErase; 150 parentOp->walk([&](memref::AllocOp op) { 151 std::vector<Operation *> candidates; 152 if (resultIsNotRead(op, candidates)) { 153 opToErase.insert(opToErase.end(), candidates.begin(), candidates.end()); 154 opToErase.push_back(op.getOperation()); 155 } 156 }); 157 for (Operation *op : opToErase) 158 rewriter.eraseOp(op); 159 } 160 161 static SmallVector<OpFoldResult> 162 computeSuffixProductIRBlockImpl(Location loc, OpBuilder &builder, 163 ArrayRef<OpFoldResult> sizes, 164 OpFoldResult unit) { 165 SmallVector<OpFoldResult> strides(sizes.size(), unit); 166 AffineExpr s0, s1; 167 bindSymbols(builder.getContext(), s0, s1); 168 169 for (int64_t r = strides.size() - 1; r > 0; --r) { 170 strides[r - 1] = affine::makeComposedFoldedAffineApply( 171 builder, loc, s0 * s1, {strides[r], sizes[r]}); 172 } 173 return strides; 174 } 175 176 SmallVector<OpFoldResult> 177 computeSuffixProductIRBlock(Location loc, OpBuilder &builder, 178 ArrayRef<OpFoldResult> sizes) { 179 OpFoldResult unit = builder.getIndexAttr(1); 180 return computeSuffixProductIRBlockImpl(loc, builder, sizes, unit); 181 } 182 183 MemrefValue skipFullyAliasingOperations(MemrefValue source) { 184 while (auto op = source.getDefiningOp()) { 185 if (auto subViewOp = dyn_cast<memref::SubViewOp>(op); 186 subViewOp && subViewOp.hasZeroOffset() && subViewOp.hasUnitStride()) { 187 // A `memref.subview` with an all zero offset, and all unit strides, still 188 // points to the same memory. 189 source = cast<MemrefValue>(subViewOp.getSource()); 190 } else if (auto castOp = dyn_cast<memref::CastOp>(op)) { 191 // A `memref.cast` still points to the same memory. 192 source = castOp.getSource(); 193 } else { 194 return source; 195 } 196 } 197 return source; 198 } 199 200 MemrefValue skipViewLikeOps(MemrefValue source) { 201 while (auto op = source.getDefiningOp()) { 202 if (auto viewLike = dyn_cast<ViewLikeOpInterface>(op)) { 203 source = cast<MemrefValue>(viewLike.getViewSource()); 204 continue; 205 } 206 return source; 207 } 208 return source; 209 } 210 211 } // namespace memref 212 } // namespace mlir 213