126 lines · cpp
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-exception6//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(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 // The combined offset is computed as45 // producer_offset + consumer_offset * producer_strides.46 combinedOffsets[i] = makeComposedFoldedAffineApply(47 builder, loc, s0 * s1 + s2,48 {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});49 combinedSizes[i] = consumerSizes[consumerPos];50 // The combined stride is computed as51 // consumer_stride * producer_stride.52 combinedStrides[i] = makeComposedFoldedAffineApply(53 builder, loc, s0 * s1,54 {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>(0, mixedSourceOffsets.size())) {93 OpFoldResult ofr =94 (rankReducedDims.test(dim)) ? zero : consumerIndices[indicesDim++];95 indices.push_back(ofr);96 }97 98 resolvedIndices.resize(indices.size());99 resolvedIndices.clear();100 for (auto [offset, index, stride] :101 llvm::zip_equal(mixedSourceOffsets, indices, mixedSourceStrides)) {102 AffineExpr off, idx, str;103 bindSymbols(rewriter.getContext(), off, idx, str);104 OpFoldResult ofr = makeComposedFoldedAffineApply(105 rewriter, loc, AffineMap::get(0, 3, off + idx * str),106 {offset, index, stride});107 resolvedIndices.push_back(108 getValueOrCreateConstantIndexOp(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(sourceSizes[srcDim]);121 continue;122 }123 resolvedSizes.push_back(destSizes[dim++]);124 }125}126