251 lines · cpp
1//===- ConversionUtils.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// Utility functions for TOSA lowering10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"14#include "mlir/Dialect/Tosa/IR/TosaOps.h"15 16using namespace mlir;17using namespace mlir::tosa;18 19SmallVector<utils::IteratorType>20mlir::tosa::getNParallelLoopsAttrs(unsigned nParallelLoops) {21 return SmallVector<utils::IteratorType>(nParallelLoops,22 utils::IteratorType::parallel);23}24 25SmallVector<Value>26mlir::tosa::condenseValues(const SmallVector<Value> &values) {27 SmallVector<Value> condensedValues;28 for (auto value : values)29 if (value)30 condensedValues.push_back(value);31 return condensedValues;32}33 34Value mlir::tosa::clampFloatHelper(Location loc, Value arg, Value min,35 Value max, OpBuilder &rewriter) {36 Value minValue = arith::MinimumFOp::create(rewriter, loc, arg, max);37 return arith::MaximumFOp::create(rewriter, loc, minValue, min);38}39 40Value mlir::tosa::clampIntHelper(Location loc, Value arg, Value min, Value max,41 OpBuilder &rewriter, bool isUnsigned) {42 if (isUnsigned) {43 auto minOrArg = arith::MaxUIOp::create(rewriter, loc, min, arg);44 return arith::MinUIOp::create(rewriter, loc, max, minOrArg);45 }46 auto minOrArg = arith::MaxSIOp::create(rewriter, loc, min, arg);47 return arith::MinSIOp::create(rewriter, loc, max, minOrArg);48}49 50bool mlir::tosa::validIntegerRange(IntegerType ty, int64_t value) {51 uint64_t bitwidth = ty.getIntOrFloatBitWidth();52 if (ty.getSignedness() == IntegerType::Unsigned) {53 uint64_t uvalue = value;54 APInt intMin = APInt::getMinValue(bitwidth);55 APInt intMax = APInt::getMaxValue(bitwidth);56 return uvalue >= intMin.getZExtValue() && uvalue <= intMax.getZExtValue();57 }58 59 APInt intMin = APInt::getSignedMinValue(bitwidth);60 APInt intMax = APInt::getSignedMaxValue(bitwidth);61 return value >= intMin.getSExtValue() && value <= intMax.getSExtValue();62}63 64namespace {65// Given two tensors of high and low ranks, derive the output shape66// to reshape the lower rank to.67// Examples:68// If lower=[c], higher=[a, b, c], [c] reshaped into [1, 1, c].69// If lower=[b, c], higher=[a, b, c], [b, c] reshaped into [1, b, c].70// If lower=[a], higher=[a, a], [a] reshaped into [1, a].71// If lower=[a], target=[a, b, a], [a] reshaped into [1, 1, a].72// If lower=[], target=[a, b, c], [] reshaped into [1, 1, 1].73// If lower=[c], higher=[?, ?, c], [c] reshaped into [1, 1, c].74// If lower=[?], higher=[?, ?, ?], [?] reshaped into [1, 1, ?].75LogicalResult76computeReshapeOutput(ArrayRef<int64_t> higherRankShape,77 ArrayRef<int64_t> lowerRankShape,78 SmallVectorImpl<int64_t> &reshapeOutputShape) {79 // Initialize new shapes with [1] * higherRank.80 int64_t higherRank = higherRankShape.size();81 int64_t lowerRank = lowerRankShape.size();82 reshapeOutputShape.assign(higherRank, 1);83 84 int64_t higherRankDim;85 int64_t lowerRankDim;86 const int64_t rankDiff = higherRank - lowerRank;87 88 for (int64_t i = lowerRank - 1; i >= 0; i--) {89 higherRankDim = higherRankShape[i + rankDiff];90 lowerRankDim = lowerRankShape[i];91 92 auto isStaticDimAndNotEqualToOne = [](int64_t dim) {93 return dim != 1 && dim != ShapedType::kDynamic;94 };95 96 if (isStaticDimAndNotEqualToOne(lowerRankDim) &&97 isStaticDimAndNotEqualToOne(higherRankDim) &&98 lowerRankDim != higherRankDim)99 return failure();100 101 reshapeOutputShape[i + rankDiff] = lowerRankDim == 1 ? 1 : lowerRankDim;102 }103 return success();104}105} // namespace106 107LogicalResult mlir::tosa::EqualizeRanks(PatternRewriter &rewriter, Location loc,108 Value &input1, Value &input2) {109 ImplicitLocOpBuilder builder(loc, rewriter);110 return EqualizeRanks(builder, input1, input2);111}112 113LogicalResult mlir::tosa::EqualizeRanks(ImplicitLocOpBuilder &builder,114 Value &input1, Value &input2) {115 auto input1Ty = llvm::dyn_cast<RankedTensorType>(input1.getType());116 auto input2Ty = llvm::dyn_cast<RankedTensorType>(input2.getType());117 118 if (!input1Ty || !input2Ty) {119 return failure();120 }121 122 int64_t input1Rank = input1Ty.getRank();123 int64_t input2Rank = input2Ty.getRank();124 125 if (input1Rank == input2Rank)126 return success();127 128 Value higherTensorValue, lowerTensorValue;129 if (input1Rank > input2Rank) {130 higherTensorValue = input1;131 lowerTensorValue = input2;132 } else {133 higherTensorValue = input2;134 lowerTensorValue = input1;135 }136 137 ArrayRef<int64_t> higherRankShape =138 llvm::cast<RankedTensorType>(higherTensorValue.getType()).getShape();139 ArrayRef<int64_t> lowerRankShape =140 llvm::cast<RankedTensorType>(lowerTensorValue.getType()).getShape();141 142 SmallVector<int64_t, 4> reshapeOutputShape;143 144 if (computeReshapeOutput(higherRankShape, lowerRankShape, reshapeOutputShape)145 .failed())146 return failure();147 148 auto reshapeInputType =149 llvm::cast<RankedTensorType>(lowerTensorValue.getType());150 auto reshapeOutputType = RankedTensorType::get(151 ArrayRef<int64_t>(reshapeOutputShape), reshapeInputType.getElementType());152 auto reshapeOutputShapeValue = getTosaConstShape(builder, reshapeOutputShape);153 154 auto reshapeLower = tosa::ReshapeOp::create(155 builder, reshapeOutputType, lowerTensorValue, reshapeOutputShapeValue);156 157 if (input1Rank > input2Rank) {158 input1 = higherTensorValue;159 input2 = reshapeLower.getResult();160 } else {161 input1 = reshapeLower.getResult();162 input2 = higherTensorValue;163 }164 165 return success();166}167 168Value mlir::tosa::getTosaConstShape(ImplicitLocOpBuilder &builder,169 llvm::ArrayRef<int64_t> shape) {170 auto attr = builder.getIndexTensorAttr(convertFromMlirShape(shape));171 auto type = mlir::tosa::shapeType::get(builder.getContext(), shape.size());172 mlir::Operation *mlirOp = tosa::ConstShapeOp::create(builder, type, attr);173 return mlirOp->getResult(0);174}175 176Value mlir::tosa::getTosaConstShape(PatternRewriter &rewriter, Location loc,177 llvm::ArrayRef<int64_t> shape) {178 ImplicitLocOpBuilder builder(loc, rewriter);179 return getTosaConstShape(builder, shape);180}181 182SmallVector<int64_t> mlir::tosa::convertFromMlirShape(ArrayRef<int64_t> shape) {183 return to_vector(llvm::map_range(shape, [](int64_t dim) {184 return ShapedType::isDynamic(dim) ? -1 : dim;185 }));186}187 188bool mlir::tosa::getConstShapeValues(Operation *op,189 llvm::SmallVector<int64_t> &resultShape) {190 if (!op) {191 return false;192 }193 if (auto constOp = mlir::dyn_cast<tosa::ConstShapeOp>(op)) {194 Attribute constOpAttr = constOp->getAttr("values");195 DenseElementsAttr elementsAttr = cast<DenseElementsAttr>(constOpAttr);196 for (int i = 0; i < elementsAttr.size(); i++) {197 int64_t val = elementsAttr.getValues<int64_t>()[i];198 resultShape.push_back(val);199 }200 return true;201 }202 // for undefined op, return false.203 return false;204}205 206// returns a small vector of int64_t values that attr contains207SmallVector<int64_t>208mlir::tosa::convertFromIntAttr(const DenseElementsAttr &attr, const int rank) {209 if (attr.isSplat()) {210 int64_t v = attr.getSplatValue<APInt>().getSExtValue();211 return SmallVector<int64_t>(rank, v);212 }213 214 if (auto intArrayAttr = llvm::dyn_cast<DenseIntElementsAttr>(attr)) {215 SmallVector<int64_t> vec;216 for (APInt val : intArrayAttr.getValues<APInt>()) {217 vec.push_back(val.getSExtValue());218 }219 return vec;220 }221 return {};222}223 224bool mlir::tosa::hasUniqueConstantScatterIndices(225 ShapedType indicesType, DenseIntElementsAttr indicesAttr) {226 const llvm::ArrayRef<int64_t> indicesShape = indicesType.getShape();227 const unsigned int indicesRank = indicesShape.size();228 const unsigned int lastDimSize = indicesShape[indicesRank - 1];229 230 // check each batch of indices from the flat indicesAttr values231 // for duplicates232 auto const indicesValues = indicesAttr.getValues<APInt>();233 assert(234 (indicesValues.size() % lastDimSize == 0) &&235 "Constant indices data length should be a multiple of indicesShape[-1]");236 237 std::vector<APInt> indices(lastDimSize);238 for (auto beg = indicesValues.begin(); beg < indicesValues.end();239 beg += lastDimSize) {240 std::copy(beg, beg + lastDimSize, indices.begin());241 std::sort(indices.begin(), indices.end(),242 [](const APInt &a, const APInt &b) { return a.slt(b); });243 if (std::adjacent_find(indices.begin(), indices.end()) != indices.end()) {244 // found duplicate values in indices in batch245 return false;246 }247 }248 249 return true;250}251