172 lines · cpp
1//===- Utils.cpp - Utilities to support the Tensor dialect ----------------===//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// This file implements utilities for the Tensor dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Tensor/Utils/Utils.h"14 15#include "mlir/Dialect/Affine/IR/AffineOps.h"16#include "mlir/Dialect/Arith/Utils/Utils.h"17#include "mlir/Dialect/Utils/IndexingUtils.h"18#include "mlir/Interfaces/ValueBoundsOpInterface.h"19 20using namespace mlir;21using namespace mlir::tensor;22 23PadOp mlir::tensor::createPadHighOp(RankedTensorType resType, Value source,24 Value pad, bool nofold, Location loc,25 OpBuilder &b, ValueRange dynOutDims) {26 27 // This assumption simplifies the following logic without limiting what's28 // required _today_. If needed, we can relax it in the future.29 assert(((resType.getNumDynamicDims() == dynOutDims.size()) ||30 dynOutDims.empty()) &&31 "Either none or all output dynamic dims must be specified!");32 33 // Init "low" and "high" padding values ("low" is kept as is, "high" is34 // computed below).35 SmallVector<OpFoldResult> low(resType.getRank(), b.getIndexAttr(0));36 SmallVector<OpFoldResult> high(resType.getRank(), b.getIndexAttr(0));37 38 size_t outDimIdx = 0;39 40 for (const auto [idx, val] : enumerate(resType.getShape())) {41 bool isDimDynamic = ShapedType::isDynamic(val);42 bool updatePadHigh = !isDimDynamic || !dynOutDims.empty();43 44 // Keep the default padding width (i.e. "0") when the output dim is dynamic45 // and no actual output sizes have been provided.46 if (!updatePadHigh)47 continue;48 49 // Compute the padding width: resDim - sourceDim.50 AffineExpr d0, d1;51 bindDims(b.getContext(), d0, d1);52 OpFoldResult sourceDim = tensor::getMixedSize(b, loc, source, idx);53 OpFoldResult outDim = isDimDynamic ? OpFoldResult(dynOutDims[outDimIdx++])54 : OpFoldResult(b.getIndexAttr(val));55 56 high[idx] = affine::makeComposedFoldedAffineApply(b, loc, d0 - d1,57 {outDim, sourceDim});58 }59 return PadOp::create(b, loc, resType, source, low, high, pad, nofold);60}61 62SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,63 Location loc,64 Value rankedTensor) {65 auto tensorTy = cast<RankedTensorType>(rankedTensor.getType());66 SmallVector<Value> dynamicDims;67 for (const auto &en : llvm::enumerate(tensorTy.getShape())) {68 if (en.value() == ShapedType::kDynamic)69 dynamicDims.push_back(70 tensor::DimOp::create(b, loc, rankedTensor, en.index()));71 }72 return dynamicDims;73}74 75FailureOr<RankedTensorType>76mlir::tensor::computeTransposedType(RankedTensorType rankedTensorType,77 ArrayRef<int64_t> transposeVector) {78 if (transposeVector.empty())79 return rankedTensorType;80 81 if (!isPermutationVector(transposeVector) ||82 transposeVector.size() != static_cast<size_t>(rankedTensorType.getRank()))83 return failure();84 85 SmallVector<int64_t> transposedShape(rankedTensorType.getShape());86 applyPermutationToVector(transposedShape, transposeVector);87 88 using RTTBuilder = RankedTensorType::Builder;89 RankedTensorType transposedTensorType =90 RTTBuilder(rankedTensorType).setShape(transposedShape);91 return transposedTensorType;92}93 94CollapseShapeOp95mlir::tensor::dropGivenUnitDims(OpBuilder &b, Location loc, Value src,96 const llvm::SmallBitVector &dropDims) {97 auto srcType = cast<ShapedType>(src.getType());98 int64_t rank = srcType.getRank();99 assert(rank == static_cast<int64_t>(dropDims.size()) &&100 "dropDims dimension does not match src tensor rank");101 assert(llvm::all_of(102 dropDims.set_bits(),103 [&](unsigned dim) { return srcType.getShape()[dim] == 1; }) &&104 "Dropping non unit dimension");105 // Computed reassociation map for the corresponding tensor.collapse_shape.106 SmallVector<ReassociationIndices, 2> reassocMaps;107 // Current reassociation group to add dropped dimension to.108 109 int64_t nextDimToGroup = 0;110 llvm::SmallBitVector keptDims(dropDims);111 keptDims.flip();112 int64_t lastSetBit = keptDims.find_last();113 for (int64_t setBit : keptDims.set_bits()) {114 // Group consecutive dropped dimension with the next non-dropped dimension.115 // If this is the last set dimension, also group all subsequent dropped116 // dimension, if any.117 int64_t upTo = setBit == lastSetBit ? rank - 1 : setBit;118 auto seq = llvm::seq_inclusive(nextDimToGroup, upTo);119 reassocMaps.emplace_back(llvm::make_range(seq.begin(), seq.end()));120 nextDimToGroup = setBit + 1;121 }122 return tensor::CollapseShapeOp::create(b, loc, src, reassocMaps);123}124 125bool mlir::tensor::isCastLikeInsertSliceOp(InsertSliceOp op) {126 llvm::SmallBitVector droppedDims = op.getDroppedDims();127 int64_t srcDim = 0;128 RankedTensorType resultType = op.getDestType();129 // Source dims and destination dims (apart from dropped dims) must have the130 // same size.131 for (int64_t resultDim = 0; resultDim < resultType.getRank(); ++resultDim) {132 if (droppedDims.test(resultDim)) {133 // InsertSlice may expand unit dimensions that result from inserting a134 // size-1 slice into a non-size-1 result dimension.135 if (resultType.getDimSize(resultDim) != 1)136 return false;137 continue;138 }139 FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(140 {op.getSource(), srcDim}, {op.getResult(), resultDim});141 if (failed(equalDimSize) || !*equalDimSize)142 return false;143 ++srcDim;144 }145 146 return true;147}148 149bool mlir::tensor::isCastLikeExtractSliceOp(ExtractSliceOp op) {150 llvm::SmallBitVector droppedDims = op.getDroppedDims();151 int64_t resultDim = 0;152 // Source dims and result dims (apart from dropped dims) must have the same153 // size.154 RankedTensorType sourceType = op.getSourceType();155 for (int64_t dim = 0, e = sourceType.getRank(); dim < e; ++dim) {156 if (droppedDims.test(dim)) {157 // ExtractSlice may drop unit dimensions that result from taking a size-1158 // slice from a non-size-1 source dimension.159 if (sourceType.getDimSize(dim) != 1)160 return false;161 continue;162 }163 FailureOr<bool> equalDimSize = ValueBoundsConstraintSet::areEqual(164 {op.getSource(), dim}, {op.getResult(), resultDim});165 if (failed(equalDimSize) || !*equalDimSize)166 return false;167 ++resultDim;168 }169 170 return true;171}172