MLIR  20.0.0git
ViewLikeInterfaceUtils.cpp
Go to the documentation of this file.
1 //===- ViewLikeInterfaceUtils.cpp -----------------------------------------===//
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 
12 #include "mlir/IR/PatternMatch.h"
13 
14 using namespace mlir;
15 using namespace affine;
16 
18  OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
19  ArrayRef<OpFoldResult> producerSizes,
20  ArrayRef<OpFoldResult> producerStrides,
21  const llvm::SmallBitVector &droppedProducerDims,
22  ArrayRef<OpFoldResult> consumerOffsets,
23  ArrayRef<OpFoldResult> consumerSizes,
24  ArrayRef<OpFoldResult> consumerStrides,
25  SmallVector<OpFoldResult> &combinedOffsets,
26  SmallVector<OpFoldResult> &combinedSizes,
27  SmallVector<OpFoldResult> &combinedStrides) {
28  combinedOffsets.resize(producerOffsets.size());
29  combinedSizes.resize(producerOffsets.size());
30  combinedStrides.resize(producerOffsets.size());
31 
32  AffineExpr s0, s1, s2;
33  bindSymbols(builder.getContext(), s0, s1, s2);
34 
35  unsigned consumerPos = 0;
36  for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
37  if (droppedProducerDims.test(i)) {
38  // For dropped dims, get the values from the producer.
39  combinedOffsets[i] = producerOffsets[i];
40  combinedSizes[i] = producerSizes[i];
41  combinedStrides[i] = producerStrides[i];
42  continue;
43  }
44  SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
45  // The combined offset is computed as
46  // producer_offset + consumer_offset * producer_strides.
47  combinedOffsets[i] = makeComposedFoldedAffineApply(
48  builder, loc, s0 * s1 + s2,
49  {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
50  combinedSizes[i] = consumerSizes[consumerPos];
51  // The combined stride is computed as
52  // consumer_stride * producer_stride.
53  combinedStrides[i] = makeComposedFoldedAffineApply(
54  builder, loc, s0 * s1,
55  {consumerStrides[consumerPos], producerStrides[i]});
56 
57  consumerPos++;
58  }
59  return success();
60 }
61 
63  OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
64  OffsetSizeAndStrideOpInterface consumer,
65  const llvm::SmallBitVector &droppedProducerDims,
66  SmallVector<OpFoldResult> &combinedOffsets,
67  SmallVector<OpFoldResult> &combinedSizes,
68  SmallVector<OpFoldResult> &combinedStrides) {
69  SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
70  SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
71  SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
72  SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
73  SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
74  SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
76  builder, loc, producerOffsets, producerSizes, producerStrides,
77  droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
78  combinedOffsets, combinedSizes, combinedStrides);
79 }
80 
82  RewriterBase &rewriter, Location loc,
83  ArrayRef<OpFoldResult> mixedSourceOffsets,
84  ArrayRef<OpFoldResult> mixedSourceStrides,
85  const llvm::SmallBitVector &rankReducedDims,
86  ArrayRef<OpFoldResult> consumerIndices,
87  SmallVectorImpl<Value> &resolvedIndices) {
88  OpFoldResult zero = rewriter.getIndexAttr(0);
89 
90  // For each dimension that is rank-reduced, add a zero to the indices.
91  int64_t indicesDim = 0;
93  for (auto dim : llvm::seq<int64_t>(0, mixedSourceOffsets.size())) {
94  OpFoldResult ofr =
95  (rankReducedDims.test(dim)) ? zero : consumerIndices[indicesDim++];
96  indices.push_back(ofr);
97  }
98 
99  resolvedIndices.resize(indices.size());
100  resolvedIndices.clear();
101  for (auto [offset, index, stride] :
102  llvm::zip_equal(mixedSourceOffsets, indices, mixedSourceStrides)) {
103  AffineExpr off, idx, str;
104  bindSymbols(rewriter.getContext(), off, idx, str);
106  rewriter, loc, AffineMap::get(0, 3, off + idx * str),
107  {offset, index, stride});
108  resolvedIndices.push_back(
109  getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
110  }
111 }
112 
114  ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes,
115  const llvm::SmallBitVector &rankReducedSourceDims,
116  SmallVectorImpl<OpFoldResult> &resolvedSizes) {
117  int64_t dim = 0;
118  int64_t srcRank = sourceSizes.size();
119  for (int64_t srcDim = 0; srcDim < srcRank; ++srcDim) {
120  if (rankReducedSourceDims[srcDim]) {
121  resolvedSizes.push_back(sourceSizes[srcDim]);
122  continue;
123  }
124  resolvedSizes.push_back(destSizes[dim++]);
125  }
126 }
Base type for affine expression.
Definition: AffineExpr.h:68
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
IntegerAttr getIndexAttr(int64_t value)
Definition: Builders.cpp:148
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:66
This class helps build Operations.
Definition: Builders.h:215
This class represents a single result from folding an operation.
Definition: OpDefinition.h:268
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
Definition: PatternMatch.h:400
void resolveSizesIntoOpWithSizes(ArrayRef< OpFoldResult > sourceSizes, ArrayRef< OpFoldResult > destSizes, const llvm::SmallBitVector &rankReducedSourceDims, SmallVectorImpl< OpFoldResult > &resolvedSizes)
Given sourceSizes, destSizes and information about which dimensions are dropped by the source: rankRe...
void resolveIndicesIntoOpWithOffsetsAndStrides(RewriterBase &rewriter, Location loc, ArrayRef< OpFoldResult > mixedSourceOffsets, ArrayRef< OpFoldResult > mixedSourceStrides, const llvm::SmallBitVector &rankReducedDims, ArrayRef< OpFoldResult > consumerIndices, SmallVectorImpl< Value > &resolvedIndices)
Given the 'consumerIndices' of a load/store operation operating on an op with offsets and strides,...
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
LogicalResult mergeOffsetsSizesAndStrides(OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > producerOffsets, ArrayRef< OpFoldResult > producerSizes, ArrayRef< OpFoldResult > producerStrides, const llvm::SmallBitVector &droppedProducerDims, ArrayRef< OpFoldResult > consumerOffsets, ArrayRef< OpFoldResult > consumerSizes, ArrayRef< OpFoldResult > consumerStrides, SmallVector< OpFoldResult > &combinedOffsets, SmallVector< OpFoldResult > &combinedSizes, SmallVector< OpFoldResult > &combinedStrides)
Fills the combinedOffsets, combinedSizes and combinedStrides to use when combining a producer slice i...
Include the generated interface declarations.
void bindSymbols(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to SymbolExpr at positions: [0 .
Definition: AffineExpr.h:362
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
Definition: Utils.cpp:112