MLIR  19.0.0git
MemRefUtils.cpp
Go to the documentation of this file.
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 
18 #include "llvm/ADT/STLExtras.h"
19 
20 namespace mlir {
21 namespace memref {
22 
23 bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
24  if (!type.hasStaticShape())
25  return false;
26 
27  SmallVector<int64_t> strides;
28  int64_t offset;
29  if (failed(getStridesAndOffset(type, strides, offset)))
30  return false;
31 
32  // MemRef is contiguous if outer dimensions are size-1 and inner
33  // 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 
51 std::pair<LinearizedMemRefInfo, OpFoldResult> getLinearizedMemRefOffsetAndSize(
52  OpBuilder &builder, Location loc, int srcBits, int dstBits,
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  AffineExpr mulMap = builder.getAffineConstantExpr(1);
69 
70  SmallVector<OpFoldResult> offsetValues(2 * sourceRank);
71 
72  for (unsigned i = 0; i < sourceRank; ++i) {
73  unsigned offsetIdx = 2 * i;
74  addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];
75  offsetValues[offsetIdx] = indicesVec[i];
76  offsetValues[offsetIdx + 1] = strides[i];
77 
78  mulMap = mulMap * symbols[i];
79  }
80 
81  // Adjust linearizedIndices and size by the scale factor (dstBits / srcBits).
82  int64_t scaler = dstBits / srcBits;
83  addMulMap = addMulMap.floorDiv(scaler);
84  mulMap = mulMap.floorDiv(scaler);
85 
87  builder, loc, addMulMap, 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);
95  builder, loc, s0.floorDiv(scaler), {offset});
96 
97  return {{adjustBaseOffset, linearizedSize}, linearizedIndices};
98 }
99 
100 LinearizedMemRefInfo
102  int dstBits, OpFoldResult offset,
103  ArrayRef<OpFoldResult> sizes) {
104  SmallVector<OpFoldResult> strides(sizes.size());
105  if (!sizes.empty()) {
106  strides.back() = builder.getIndexAttr(1);
107  AffineExpr s0, s1;
108  bindSymbols(builder.getContext(), s0, s1);
109  for (int index = sizes.size() - 1; index > 0; --index) {
110  strides[index - 1] = affine::makeComposedFoldedAffineApply(
111  builder, loc, s0 * s1,
112  ArrayRef<OpFoldResult>{strides[index], sizes[index]});
113  }
114  }
115 
116  LinearizedMemRefInfo linearizedMemRefInfo;
117  std::tie(linearizedMemRefInfo, std::ignore) =
118  getLinearizedMemRefOffsetAndSize(builder, loc, srcBits, dstBits, offset,
119  sizes, strides);
120  return linearizedMemRefInfo;
121 }
122 
123 /// Returns true if all the uses of op are not read/load.
124 /// There can be SubviewOp users as long as all its users are also
125 /// StoreOp/transfer_write. If return true it also fills out the uses, if it
126 /// returns false uses is unchanged.
127 static bool resultIsNotRead(Operation *op, std::vector<Operation *> &uses) {
128  std::vector<Operation *> opUses;
129  for (OpOperand &use : op->getUses()) {
130  Operation *useOp = use.getOwner();
131  if (isa<memref::DeallocOp>(useOp) ||
132  (useOp->getNumResults() == 0 && useOp->getNumRegions() == 0 &&
133  !mlir::hasEffect<MemoryEffects::Read>(useOp)) ||
134  (isa<memref::SubViewOp>(useOp) && resultIsNotRead(useOp, opUses))) {
135  opUses.push_back(useOp);
136  continue;
137  }
138  return false;
139  }
140  uses.insert(uses.end(), opUses.begin(), opUses.end());
141  return true;
142 }
143 
144 void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp) {
145  std::vector<Operation *> opToErase;
146  parentOp->walk([&](memref::AllocOp op) {
147  std::vector<Operation *> candidates;
148  if (resultIsNotRead(op, candidates)) {
149  opToErase.insert(opToErase.end(), candidates.begin(), candidates.end());
150  opToErase.push_back(op.getOperation());
151  }
152  });
153  for (Operation *op : opToErase)
154  rewriter.eraseOp(op);
155 }
156 
160  OpFoldResult unit) {
161  SmallVector<OpFoldResult> strides(sizes.size(), unit);
162  AffineExpr s0, s1;
163  bindSymbols(builder.getContext(), s0, s1);
164 
165  for (int64_t r = strides.size() - 1; r > 0; --r) {
166  strides[r - 1] = affine::makeComposedFoldedAffineApply(
167  builder, loc, s0 * s1, {strides[r], sizes[r]});
168  }
169  return strides;
170 }
171 
174  ArrayRef<OpFoldResult> sizes) {
175  OpFoldResult unit = builder.getIndexAttr(1);
176  return computeSuffixProductIRBlockImpl(loc, builder, sizes, unit);
177 }
178 
180  while (auto op = source.getDefiningOp()) {
181  if (auto subViewOp = dyn_cast<memref::SubViewOp>(op);
182  subViewOp && subViewOp.hasZeroOffset() && subViewOp.hasUnitStride()) {
183  // A `memref.subview` with an all zero offset, and all unit strides, still
184  // points to the same memory.
185  source = cast<MemrefValue>(subViewOp.getSource());
186  } else if (auto castOp = dyn_cast<memref::CastOp>(op)) {
187  // A `memref.cast` still points to the same memory.
188  source = castOp.getSource();
189  } else {
190  return source;
191  }
192  }
193  return source;
194 }
195 
197  while (auto op = source.getDefiningOp()) {
198  if (auto subView = dyn_cast<memref::SubViewOp>(op)) {
199  source = cast<MemrefValue>(subView.getSource());
200  } else if (auto cast = dyn_cast<memref::CastOp>(op)) {
201  source = cast.getSource();
202  } else {
203  return source;
204  }
205  }
206  return source;
207 }
208 
209 } // namespace memref
210 } // namespace mlir
Base type for affine expression.
Definition: AffineExpr.h:68
AffineExpr floorDiv(uint64_t v) const
Definition: AffineExpr.cpp:901
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:124
AffineExpr getAffineConstantExpr(int64_t constant)
Definition: Builders.cpp:379
MLIRContext * getContext() const
Definition: Builders.h:55
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
This class helps build Operations.
Definition: Builders.h:209
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
This class represents an operand of an operation.
Definition: Value.h:267
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
Definition: Operation.h:793
unsigned getNumRegions()
Returns the number of regions held by this operation.
Definition: Operation.h:669
unsigned getNumResults()
Return the number of results held by this operation.
Definition: Operation.h:399
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
Definition: AffineOps.cpp:1192
static bool resultIsNotRead(Operation *op, std::vector< Operation * > &uses)
Returns true if all the uses of op are not read/load.
MemrefValue skipFullyAliasingOperations(MemrefValue source)
Walk up the source chain until an operation that changes/defines the view of memory is found (i....
void eraseDeadAllocAndStores(RewriterBase &rewriter, Operation *parentOp)
MemrefValue skipSubViewsAndCasts(MemrefValue source)
Walk up the source chain until something an op other than a memref.subview or memref....
std::pair< LinearizedMemRefInfo, OpFoldResult > getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits, int dstBits, OpFoldResult offset, ArrayRef< OpFoldResult > sizes, ArrayRef< OpFoldResult > strides, ArrayRef< OpFoldResult > indices={})
Definition: MemRefUtils.cpp:51
bool isStaticShapeAndContiguousRowMajor(MemRefType type)
Returns true, if the memref type has static shapes and represents a contiguous chunk of memory.
Definition: MemRefUtils.cpp:23
static SmallVector< OpFoldResult > computeSuffixProductIRBlockImpl(Location loc, OpBuilder &builder, ArrayRef< OpFoldResult > sizes, OpFoldResult unit)
SmallVector< OpFoldResult > computeSuffixProductIRBlock(Location loc, OpBuilder &builder, ArrayRef< OpFoldResult > sizes)
Given a set of sizes, return the suffix product.
Include the generated interface declarations.
LogicalResult getStridesAndOffset(MemRefType t, SmallVectorImpl< int64_t > &strides, int64_t &offset)
Returns the strides of the MemRef if the layout map is in strided form.
TypedValue< BaseMemRefType > MemrefValue
A value with a memref type.
Definition: MemRefUtils.h:26
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Definition: AffineExpr.h:362
void bindSymbolsList(MLIRContext *ctx, MutableArrayRef< AffineExprTy > exprs)
Definition: AffineExpr.h:367
For a memref with offset, sizes and strides, returns the offset and size to use for the linearized me...
Definition: MemRefUtils.h:45