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
14using namespace mlir;
15using namespace affine;
16
17LogicalResult 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(N: producerOffsets.size());
29 combinedSizes.resize(N: producerOffsets.size());
30 combinedStrides.resize(N: producerOffsets.size());
31
32 AffineExpr s0, s1, s2;
33 bindSymbols(ctx: builder.getContext(), exprs&: s0, exprs&: s1, exprs&: s2);
34
35 unsigned consumerPos = 0;
36 for (auto i : llvm::seq<unsigned>(Begin: 0, End: producerOffsets.size())) {
37 if (droppedProducerDims.test(Idx: 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 // The combined offset is computed as
45 // producer_offset + consumer_offset * producer_strides.
46 combinedOffsets[i] = makeComposedFoldedAffineApply(
47 b&: builder, loc, expr: s0 * s1 + s2,
48 operands: {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
49 combinedSizes[i] = consumerSizes[consumerPos];
50 // The combined stride is computed as
51 // consumer_stride * producer_stride.
52 combinedStrides[i] = makeComposedFoldedAffineApply(
53 b&: builder, loc, expr: s0 * s1,
54 operands: {consumerStrides[consumerPos], producerStrides[i]});
55
56 consumerPos++;
57 }
58 return success();
59}
60
61LogicalResult mlir::affine::mergeOffsetsSizesAndStrides(
62 OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
63 OffsetSizeAndStrideOpInterface consumer,
64 const llvm::SmallBitVector &droppedProducerDims,
65 SmallVector<OpFoldResult> &combinedOffsets,
66 SmallVector<OpFoldResult> &combinedSizes,
67 SmallVector<OpFoldResult> &combinedStrides) {
68 SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
69 SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
70 SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
71 SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
72 SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
73 SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
74 return mergeOffsetsSizesAndStrides(
75 builder, loc, producerOffsets, producerSizes, producerStrides,
76 droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
77 combinedOffsets, combinedSizes, combinedStrides);
78}
79
80void mlir::affine::resolveIndicesIntoOpWithOffsetsAndStrides(
81 RewriterBase &rewriter, Location loc,
82 ArrayRef<OpFoldResult> mixedSourceOffsets,
83 ArrayRef<OpFoldResult> mixedSourceStrides,
84 const llvm::SmallBitVector &rankReducedDims,
85 ArrayRef<OpFoldResult> consumerIndices,
86 SmallVectorImpl<Value> &resolvedIndices) {
87 OpFoldResult zero = rewriter.getIndexAttr(0);
88
89 // For each dimension that is rank-reduced, add a zero to the indices.
90 int64_t indicesDim = 0;
91 SmallVector<OpFoldResult> indices;
92 for (auto dim : llvm::seq<int64_t>(Begin: 0, End: mixedSourceOffsets.size())) {
93 OpFoldResult ofr =
94 (rankReducedDims.test(Idx: dim)) ? zero : consumerIndices[indicesDim++];
95 indices.push_back(Elt: ofr);
96 }
97
98 resolvedIndices.resize(N: indices.size());
99 resolvedIndices.clear();
100 for (auto [offset, index, stride] :
101 llvm::zip_equal(t&: mixedSourceOffsets, u&: indices, args&: mixedSourceStrides)) {
102 AffineExpr off, idx, str;
103 bindSymbols(ctx: rewriter.getContext(), exprs&: off, exprs&: idx, exprs&: str);
104 OpFoldResult ofr = makeComposedFoldedAffineApply(
105 b&: rewriter, loc, map: AffineMap::get(dimCount: 0, symbolCount: 3, result: off + idx * str),
106 operands: {offset, index, stride});
107 resolvedIndices.push_back(
108 Elt: getValueOrCreateConstantIndexOp(b&: rewriter, loc, ofr));
109 }
110}
111
112void mlir::affine::resolveSizesIntoOpWithSizes(
113 ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes,
114 const llvm::SmallBitVector &rankReducedSourceDims,
115 SmallVectorImpl<OpFoldResult> &resolvedSizes) {
116 int64_t dim = 0;
117 int64_t srcRank = sourceSizes.size();
118 for (int64_t srcDim = 0; srcDim < srcRank; ++srcDim) {
119 if (rankReducedSourceDims[srcDim]) {
120 resolvedSizes.push_back(Elt: sourceSizes[srcDim]);
121 continue;
122 }
123 resolvedSizes.push_back(Elt: destSizes[dim++]);
124 }
125}
126

source code of mlir/lib/Dialect/Affine/Utils/ViewLikeInterfaceUtils.cpp