126 lines
5.1 KiB
C++
126 lines
5.1 KiB
C++
//===- ViewLikeInterfaceUtils.cpp -----------------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
|
|
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
|
#include "mlir/Dialect/Arith/Utils/Utils.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
|
|
using namespace mlir;
|
|
using namespace affine;
|
|
|
|
LogicalResult mlir::affine::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) {
|
|
combinedOffsets.resize(producerOffsets.size());
|
|
combinedSizes.resize(producerOffsets.size());
|
|
combinedStrides.resize(producerOffsets.size());
|
|
|
|
AffineExpr s0, s1, s2;
|
|
bindSymbols(builder.getContext(), s0, s1, s2);
|
|
|
|
unsigned consumerPos = 0;
|
|
for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
|
|
if (droppedProducerDims.test(i)) {
|
|
// For dropped dims, get the values from the producer.
|
|
combinedOffsets[i] = producerOffsets[i];
|
|
combinedSizes[i] = producerSizes[i];
|
|
combinedStrides[i] = producerStrides[i];
|
|
continue;
|
|
}
|
|
SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
|
|
// The combined offset is computed as
|
|
// producer_offset + consumer_offset * producer_strides.
|
|
combinedOffsets[i] = makeComposedFoldedAffineApply(
|
|
builder, loc, s0 * s1 + s2,
|
|
{consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
|
|
combinedSizes[i] = consumerSizes[consumerPos];
|
|
// The combined stride is computed as
|
|
// consumer_stride * producer_stride.
|
|
combinedStrides[i] = makeComposedFoldedAffineApply(
|
|
builder, loc, s0 * s1,
|
|
{consumerStrides[consumerPos], producerStrides[i]});
|
|
|
|
consumerPos++;
|
|
}
|
|
return success();
|
|
}
|
|
|
|
LogicalResult mlir::affine::mergeOffsetsSizesAndStrides(
|
|
OpBuilder &builder, Location loc, OffsetSizeAndStrideOpInterface producer,
|
|
OffsetSizeAndStrideOpInterface consumer,
|
|
const llvm::SmallBitVector &droppedProducerDims,
|
|
SmallVector<OpFoldResult> &combinedOffsets,
|
|
SmallVector<OpFoldResult> &combinedSizes,
|
|
SmallVector<OpFoldResult> &combinedStrides) {
|
|
SmallVector<OpFoldResult> consumerOffsets = consumer.getMixedOffsets();
|
|
SmallVector<OpFoldResult> consumerSizes = consumer.getMixedSizes();
|
|
SmallVector<OpFoldResult> consumerStrides = consumer.getMixedStrides();
|
|
SmallVector<OpFoldResult> producerOffsets = producer.getMixedOffsets();
|
|
SmallVector<OpFoldResult> producerSizes = producer.getMixedSizes();
|
|
SmallVector<OpFoldResult> producerStrides = producer.getMixedStrides();
|
|
return mergeOffsetsSizesAndStrides(
|
|
builder, loc, producerOffsets, producerSizes, producerStrides,
|
|
droppedProducerDims, consumerOffsets, consumerSizes, consumerStrides,
|
|
combinedOffsets, combinedSizes, combinedStrides);
|
|
}
|
|
|
|
void mlir::affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
RewriterBase &rewriter, Location loc,
|
|
ArrayRef<OpFoldResult> mixedSourceOffsets,
|
|
ArrayRef<OpFoldResult> mixedSourceStrides,
|
|
const llvm::SmallBitVector &rankReducedDims,
|
|
ArrayRef<OpFoldResult> consumerIndices,
|
|
SmallVectorImpl<Value> &resolvedIndices) {
|
|
OpFoldResult zero = rewriter.getIndexAttr(0);
|
|
|
|
// For each dimension that is rank-reduced, add a zero to the indices.
|
|
int64_t indicesDim = 0;
|
|
SmallVector<OpFoldResult> indices;
|
|
for (auto dim : llvm::seq<int64_t>(0, mixedSourceOffsets.size())) {
|
|
OpFoldResult ofr =
|
|
(rankReducedDims.test(dim)) ? zero : consumerIndices[indicesDim++];
|
|
indices.push_back(ofr);
|
|
}
|
|
|
|
resolvedIndices.resize(indices.size());
|
|
resolvedIndices.clear();
|
|
for (auto [offset, index, stride] :
|
|
llvm::zip_equal(mixedSourceOffsets, indices, mixedSourceStrides)) {
|
|
AffineExpr off, idx, str;
|
|
bindSymbols(rewriter.getContext(), off, idx, str);
|
|
OpFoldResult ofr = makeComposedFoldedAffineApply(
|
|
rewriter, loc, AffineMap::get(0, 3, off + idx * str),
|
|
{offset, index, stride});
|
|
resolvedIndices.push_back(
|
|
getValueOrCreateConstantIndexOp(rewriter, loc, ofr));
|
|
}
|
|
}
|
|
|
|
void mlir::affine::resolveSizesIntoOpWithSizes(
|
|
ArrayRef<OpFoldResult> sourceSizes, ArrayRef<OpFoldResult> destSizes,
|
|
const llvm::SmallBitVector &rankReducedSourceDims,
|
|
SmallVectorImpl<OpFoldResult> &resolvedSizes) {
|
|
int64_t dim = 0;
|
|
int64_t srcRank = sourceSizes.size();
|
|
for (int64_t srcDim = 0; srcDim < srcRank; ++srcDim) {
|
|
if (rankReducedSourceDims[srcDim]) {
|
|
resolvedSizes.push_back(sourceSizes[srcDim]);
|
|
continue;
|
|
}
|
|
resolvedSizes.push_back(destSizes[dim++]);
|
|
}
|
|
}
|