e377520a47
tosa.concat is lowered to tensor.insert_slice thus it should be in TosaToTensor rather than in TosaToLinalg. Reviewed By: rsuderman Differential Revision: https://reviews.llvm.org/D145952
427 lines
16 KiB
C++
427 lines
16 KiB
C++
//===- TosaToTensor.cpp - Lowering Tosa to Tensor Dialect -------------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// These rewriters lower from the Tosa to the Tensor dialect.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
using namespace mlir;
|
|
using namespace tosa;
|
|
|
|
static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
|
|
ArrayRef<int64_t> rhsShape,
|
|
SmallVector<int64_t> &intermediateShape,
|
|
bool isDynamic) {
|
|
if (isDynamic) {
|
|
// TODO (natashaknk): Make dynamic intermediate shape not always be rank-1
|
|
intermediateShape = {ShapedType::kDynamic};
|
|
return true;
|
|
}
|
|
|
|
if (lhsShape.empty() || rhsShape.empty()) {
|
|
intermediateShape = {};
|
|
return true;
|
|
}
|
|
|
|
unsigned currLhsDim = 0, currRhsDim = 0;
|
|
while (currLhsDim < lhsShape.size() && currRhsDim < rhsShape.size()) {
|
|
int64_t rhsSize = rhsShape[currRhsDim];
|
|
int64_t lhsSize = lhsShape[currLhsDim];
|
|
while (lhsSize != rhsSize && currLhsDim < lhsShape.size() &&
|
|
currRhsDim < rhsShape.size()) {
|
|
if (lhsSize < rhsSize) {
|
|
currLhsDim++;
|
|
if (currLhsDim < lhsShape.size()) {
|
|
lhsSize *= lhsShape[currLhsDim];
|
|
}
|
|
} else {
|
|
currRhsDim++;
|
|
if (currRhsDim < rhsShape.size()) {
|
|
rhsSize *= rhsShape[currRhsDim];
|
|
}
|
|
}
|
|
}
|
|
if (lhsSize == rhsSize) {
|
|
intermediateShape.push_back(lhsSize);
|
|
}
|
|
currRhsDim++;
|
|
currLhsDim++;
|
|
}
|
|
|
|
// If the iterators didn't reach the end and their leftover dimensions are not
|
|
// equal to 1 an intermediate shape was not found.
|
|
while (currLhsDim < lhsShape.size()) {
|
|
if (lhsShape[currLhsDim++] != 1) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
while (currRhsDim < rhsShape.size()) {
|
|
if (rhsShape[currRhsDim++] != 1) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
static bool createReassociationMapsForCollapse(
|
|
PatternRewriter &rewriter, ArrayRef<int64_t> srcShape,
|
|
ArrayRef<int64_t> dstShape,
|
|
SmallVector<ReassociationExprs, 4> &reassociationMap, bool isDynamic) {
|
|
|
|
// If the shape is dynamic, create a map for collapsing into one dimension.
|
|
if (isDynamic) {
|
|
SmallVector<AffineExpr, 2> exprs;
|
|
for (int i = 0, s = srcShape.size(); i < s; ++i)
|
|
exprs.push_back(rewriter.getAffineDimExpr(i));
|
|
reassociationMap = {exprs};
|
|
return true;
|
|
}
|
|
|
|
if (dstShape.empty()) {
|
|
reassociationMap = {};
|
|
return true;
|
|
}
|
|
|
|
reassociationMap.resize(dstShape.size());
|
|
unsigned currSrcDim = 0, currDstDim = 0;
|
|
while (currSrcDim < srcShape.size() && currDstDim < dstShape.size()) {
|
|
int64_t dstSize = dstShape[currDstDim];
|
|
int64_t srcSize = srcShape[currSrcDim];
|
|
while (srcSize < dstSize && currSrcDim < srcShape.size()) {
|
|
reassociationMap[currDstDim].push_back(
|
|
rewriter.getAffineDimExpr(currSrcDim++));
|
|
srcSize *= srcShape[currSrcDim];
|
|
}
|
|
if (srcSize == dstSize) {
|
|
reassociationMap[currDstDim].push_back(
|
|
rewriter.getAffineDimExpr(currSrcDim++));
|
|
// If the next dim in collapsedShape is not 1, treat subsequent dims in
|
|
// expandedShape which are 1 to be collapsed.
|
|
if (currDstDim == dstShape.size() - 1 || dstShape[currDstDim + 1] != 1) {
|
|
while (currSrcDim < srcShape.size() && srcShape[currSrcDim] == 1) {
|
|
reassociationMap[currDstDim].push_back(
|
|
rewriter.getAffineDimExpr(currSrcDim++));
|
|
}
|
|
}
|
|
}
|
|
currDstDim++;
|
|
}
|
|
|
|
// If both iterators didn't reach the end, we have leftover dimentions which
|
|
// implies that we have a mismatch in shape.
|
|
return currSrcDim == srcShape.size() && currDstDim == dstShape.size();
|
|
}
|
|
|
|
namespace {
|
|
class ReshapeConverterCollapse : public OpConversionPattern<tosa::ReshapeOp> {
|
|
public:
|
|
using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const final {
|
|
ShapedType operandTy = adaptor.getInput1().getType().cast<ShapedType>();
|
|
ShapedType resultTy = reshape.getType().template cast<ShapedType>();
|
|
bool isDynamic = !operandTy.hasStaticShape();
|
|
|
|
if (isDynamic && resultTy.getRank() != 1) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape, "Cannot collapse dynamic dims to more than one dimension");
|
|
}
|
|
|
|
SmallVector<ReassociationExprs, 4> reassociationMap;
|
|
if (!createReassociationMapsForCollapse(rewriter, operandTy.getShape(),
|
|
resultTy.getShape(),
|
|
reassociationMap, isDynamic)) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape,
|
|
"tosa.reshape Attempting to collapse into an incompatible shape");
|
|
}
|
|
|
|
SmallVector<int64_t> intermediateShape;
|
|
if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
|
|
intermediateShape, isDynamic)) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape, "tosa.reshape Cannot collapse into given shape");
|
|
}
|
|
|
|
rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(
|
|
reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
class ReshapeConverterExpand : public OpConversionPattern<tosa::ReshapeOp> {
|
|
public:
|
|
using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const final {
|
|
ShapedType operandTy = adaptor.getInput1().getType().cast<ShapedType>();
|
|
ShapedType resultTy = reshape.getType().template cast<ShapedType>();
|
|
bool isDynamic = !operandTy.hasStaticShape();
|
|
|
|
if (isDynamic && operandTy.getRank() != 1) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape, "Cannot expand dynamic dims from more than one dimension");
|
|
}
|
|
|
|
SmallVector<ReassociationExprs, 4> reassociationMap;
|
|
if (!createReassociationMapsForCollapse(rewriter, resultTy.getShape(),
|
|
operandTy.getShape(),
|
|
reassociationMap, isDynamic)) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape,
|
|
"tosa.reshape Attempting to expand into an incompatible shape");
|
|
}
|
|
|
|
SmallVector<int64_t> intermediateShape;
|
|
if (!findIntermediateShape(operandTy.getShape(), resultTy.getShape(),
|
|
intermediateShape, isDynamic) ||
|
|
intermediateShape != operandTy.getShape()) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape, "tosa.reshape Cannot expand into given shape");
|
|
}
|
|
rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
|
|
reshape, resultTy, adaptor.getOperands()[0], reassociationMap);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
class ReshapeConverterCollapseExpand
|
|
: public OpConversionPattern<tosa::ReshapeOp> {
|
|
public:
|
|
using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const final {
|
|
ShapedType operandTy = adaptor.getInput1().getType().cast<ShapedType>();
|
|
ShapedType resultTy = reshape.getType().template cast<ShapedType>();
|
|
bool isDynamic = !operandTy.hasStaticShape();
|
|
|
|
SmallVector<int64_t> intermediateShape;
|
|
if (!findIntermediateShape(resultTy.getShape(), operandTy.getShape(),
|
|
intermediateShape, isDynamic)) {
|
|
return rewriter.notifyMatchFailure(
|
|
reshape, "tosa.reshape Cannot identify an intermediate shape between "
|
|
"the given two shapes");
|
|
}
|
|
|
|
Value collapse = rewriter.create<tosa::ReshapeOp>(
|
|
reshape.getLoc(),
|
|
RankedTensorType::get(intermediateShape,
|
|
reshape.getType().getElementType()),
|
|
adaptor.getInput1());
|
|
Value expand =
|
|
rewriter.create<tosa::ReshapeOp>(reshape.getLoc(), resultTy, collapse);
|
|
rewriter.replaceOp(reshape, expand);
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
class SliceConverter : public OpConversionPattern<tosa::SliceOp> {
|
|
public:
|
|
using OpConversionPattern<tosa::SliceOp>::OpConversionPattern;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(tosa::SliceOp sliceOp, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const final {
|
|
Location loc = sliceOp.getLoc();
|
|
Value input = adaptor.getInput();
|
|
SmallVector<int64_t> strides, sizes;
|
|
ArrayRef<int64_t> starts = sliceOp.getStart();
|
|
strides.resize(sliceOp.getType().template cast<ShapedType>().getRank(), 1);
|
|
|
|
SmallVector<Value> dynSizes;
|
|
for (const auto &i : llvm::enumerate(sliceOp.getSize())) {
|
|
int64_t size = i.value();
|
|
size_t index = i.index();
|
|
sizes.push_back(size == -1 ? ShapedType::kDynamic : size);
|
|
if (!ShapedType::isDynamic(sizes.back()))
|
|
continue;
|
|
|
|
auto dim = rewriter.create<tensor::DimOp>(loc, input, index);
|
|
auto offset = rewriter.create<arith::ConstantOp>(
|
|
loc, rewriter.getIndexAttr(starts[index]));
|
|
dynSizes.push_back(rewriter.create<arith::SubIOp>(loc, dim, offset));
|
|
}
|
|
|
|
auto newSliceOp = rewriter.create<tensor::ExtractSliceOp>(
|
|
sliceOp.getLoc(), sliceOp.getType(), input, ValueRange({}), dynSizes,
|
|
ValueRange({}), rewriter.getDenseI64ArrayAttr(starts),
|
|
rewriter.getDenseI64ArrayAttr(sizes),
|
|
rewriter.getDenseI64ArrayAttr(strides));
|
|
|
|
rewriter.replaceOp(sliceOp, newSliceOp.getResult());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
class PadConverter : public OpRewritePattern<tosa::PadOp> {
|
|
public:
|
|
using OpRewritePattern<tosa::PadOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::PadOp padOp,
|
|
PatternRewriter &rewriter) const final {
|
|
auto loc = padOp.getLoc();
|
|
auto input = padOp.getInput1();
|
|
auto padding = padOp.getPadding();
|
|
|
|
ShapedType inputTy = input.getType().cast<ShapedType>();
|
|
Type elementTy = inputTy.getElementType();
|
|
int64_t rank = inputTy.getRank();
|
|
|
|
// Setup the default constantAttr.
|
|
|
|
Value padConstant;
|
|
|
|
if (padOp.getPadConst()) {
|
|
padConstant = rewriter.createOrFold<tensor::ExtractOp>(
|
|
loc, padOp.getPadConst(), ValueRange({}));
|
|
} else {
|
|
Attribute constantAttr;
|
|
if (elementTy.isa<FloatType>()) {
|
|
constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
|
|
} else if (elementTy.isa<IntegerType>() && !padOp.getQuantizationInfo()) {
|
|
constantAttr = rewriter.getIntegerAttr(elementTy, 0);
|
|
} else if (elementTy.isa<IntegerType>() && padOp.getQuantizationInfo()) {
|
|
int64_t value = padOp.getQuantizationInfo()->getInputZp();
|
|
constantAttr = rewriter.getIntegerAttr(elementTy, value);
|
|
}
|
|
if (constantAttr)
|
|
padConstant = rewriter.create<arith::ConstantOp>(loc, constantAttr);
|
|
}
|
|
|
|
if (!padConstant) {
|
|
return rewriter.notifyMatchFailure(
|
|
padOp, "tosa.pad was unable to determine the pad constant value.");
|
|
}
|
|
|
|
Value lowIndex =
|
|
rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(0));
|
|
Value highIndex =
|
|
rewriter.create<arith::ConstantOp>(loc, rewriter.getIndexAttr(1));
|
|
|
|
SmallVector<OpFoldResult, 3> lowValues;
|
|
SmallVector<OpFoldResult, 3> highValues;
|
|
|
|
lowValues.reserve(rank);
|
|
highValues.reserve(rank);
|
|
|
|
for (int i = 0; i < rank; i++) {
|
|
Value inputIndex = rewriter.createOrFold<arith::ConstantIndexOp>(loc, i);
|
|
Value lowVal = rewriter.createOrFold<tensor::ExtractOp>(
|
|
loc, padding, ValueRange({inputIndex, lowIndex}));
|
|
Value highVal = rewriter.createOrFold<tensor::ExtractOp>(
|
|
loc, padding, ValueRange({inputIndex, highIndex}));
|
|
|
|
lowVal = rewriter.createOrFold<arith::IndexCastOp>(
|
|
loc, rewriter.getIndexType(), lowVal);
|
|
highVal = rewriter.createOrFold<arith::IndexCastOp>(
|
|
loc, rewriter.getIndexType(), highVal);
|
|
|
|
lowValues.push_back(lowVal);
|
|
highValues.push_back(highVal);
|
|
}
|
|
|
|
auto newPadOp = rewriter.create<tensor::PadOp>(
|
|
loc, padOp.getType(), input, lowValues, highValues, padConstant);
|
|
|
|
rewriter.replaceOp(padOp, newPadOp.getResult());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct ConcatConverter : public OpConversionPattern<tosa::ConcatOp> {
|
|
using OpConversionPattern<tosa::ConcatOp>::OpConversionPattern;
|
|
|
|
LogicalResult
|
|
matchAndRewrite(tosa::ConcatOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
auto inputType = op.getOperand(0).getType().template cast<ShapedType>();
|
|
auto resultType = op.getType().dyn_cast<RankedTensorType>();
|
|
|
|
Location loc = op.getLoc();
|
|
int axis = op.getAxis();
|
|
Value axisValue = rewriter.createOrFold<arith::ConstantOp>(
|
|
loc, rewriter.getIndexAttr(axis));
|
|
int rank = resultType.getRank();
|
|
SmallVector<Value, 3> offsets, sizes, strides;
|
|
sizes.reserve(rank);
|
|
strides.resize(rank, rewriter.create<arith::ConstantIndexOp>(loc, 1));
|
|
offsets.resize(rank, rewriter.create<arith::ConstantIndexOp>(loc, 0));
|
|
|
|
SmallVector<Value> dynDims;
|
|
for (int i = 0; i < rank; ++i) {
|
|
sizes.push_back(rewriter.createOrFold<tensor::DimOp>(
|
|
loc, adaptor.getOperands()[0], i));
|
|
if (inputType.isDynamicDim(i)) {
|
|
dynDims.push_back(
|
|
rewriter.create<tensor::DimOp>(loc, op.getOperand(0), i));
|
|
}
|
|
}
|
|
|
|
Value resultDimSize = sizes[axis];
|
|
for (auto arg : adaptor.getOperands().drop_front()) {
|
|
auto size = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue);
|
|
resultDimSize =
|
|
rewriter.createOrFold<arith::AddIOp>(loc, resultDimSize, size);
|
|
}
|
|
sizes[axis] = resultDimSize;
|
|
|
|
Value emptyTensor = rewriter.create<tensor::EmptyOp>(
|
|
loc, resultType.getShape(), resultType.getElementType(), dynDims);
|
|
|
|
auto toOpFoldResult = [](Value v) -> OpFoldResult {
|
|
auto op = v.getDefiningOp<arith::ConstantIndexOp>();
|
|
if (!op)
|
|
return v;
|
|
return op.getValue();
|
|
};
|
|
Value result = emptyTensor;
|
|
for (auto arg : adaptor.getOperands()) {
|
|
sizes[axis] = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue);
|
|
result = rewriter.createOrFold<tensor::InsertSliceOp>(
|
|
loc, arg, result,
|
|
llvm::to_vector(llvm::map_range(offsets, toOpFoldResult)),
|
|
llvm::to_vector(llvm::map_range(sizes, toOpFoldResult)),
|
|
llvm::to_vector(llvm::map_range(strides, toOpFoldResult)));
|
|
offsets[axis] =
|
|
rewriter.createOrFold<arith::AddIOp>(loc, offsets[axis], sizes[axis]);
|
|
}
|
|
rewriter.replaceOp(op, result);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void mlir::tosa::populateTosaToTensorConversionPatterns(
|
|
RewritePatternSet *patterns) {
|
|
patterns->add<SliceConverter, PadConverter, ConcatConverter>(
|
|
patterns->getContext());
|
|
patterns->add<ReshapeConverterCollapse>(patterns->getContext(),
|
|
/*benefit=*/100);
|
|
patterns->add<ReshapeConverterExpand>(patterns->getContext(),
|
|
/*benefit=*/200);
|
|
patterns->add<ReshapeConverterCollapseExpand>(patterns->getContext(),
|
|
/*benefit=*/300);
|
|
}
|