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
9 #include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
10 #include "mlir/Dialect/Affine/IR/AffineOps.h"
11 #include "mlir/Dialect/Arith/Utils/Utils.h"
12 #include "mlir/IR/PatternMatch.h"
13
14 using namespace mlir;
15 using namespace affine;
16
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)17 LogicalResult mlir::affine::mergeOffsetsSizesAndStrides(
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
mergeOffsetsSizesAndStrides(OpBuilder & builder,Location loc,OffsetSizeAndStrideOpInterface producer,OffsetSizeAndStrideOpInterface consumer,const llvm::SmallBitVector & droppedProducerDims,SmallVector<OpFoldResult> & combinedOffsets,SmallVector<OpFoldResult> & combinedSizes,SmallVector<OpFoldResult> & combinedStrides)62 LogicalResult mlir::affine::mergeOffsetsSizesAndStrides(
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();
75 return mergeOffsetsSizesAndStrides(
76 builder, loc, producerOffsets, producerSizes, producerStrides,
77 droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
78 combinedOffsets, combinedSizes, combinedStrides);
79 }
80
resolveIndicesIntoOpWithOffsetsAndStrides(RewriterBase & rewriter,Location loc,ArrayRef<OpFoldResult> mixedSourceOffsets,ArrayRef<OpFoldResult> mixedSourceStrides,const llvm::SmallBitVector & rankReducedDims,ArrayRef<OpFoldResult> consumerIndices,SmallVectorImpl<Value> & resolvedIndices)81 void mlir::affine::resolveIndicesIntoOpWithOffsetsAndStrides(
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;
92 SmallVector<OpFoldResult> indices;
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);
105 OpFoldResult ofr = makeComposedFoldedAffineApply(
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
resolveSizesIntoOpWithSizes(ArrayRef<OpFoldResult> sourceSizes,ArrayRef<OpFoldResult> destSizes,const llvm::SmallBitVector & rankReducedSourceDims,SmallVectorImpl<OpFoldResult> & resolvedSizes)113 void mlir::affine::resolveSizesIntoOpWithSizes(
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 }
127