derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
2471 lines
88 KiB
C++
2471 lines
88 KiB
C++
//===- TosaCanonicalizations.cpp - Canonicalization patterns & folders ----===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// \file
|
|
// TOSA canonicalization patterns and folders.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Quant/IR/Quant.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tosa/IR/TargetEnv.h"
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
|
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
|
|
#include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
|
|
#include "mlir/Dialect/Traits.h"
|
|
#include "mlir/IR/BuiltinTypeInterfaces.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Transforms/FoldUtils.h"
|
|
#include "mlir/Transforms/InliningUtils.h"
|
|
#include "llvm/ADT/APFloat.h"
|
|
#include "llvm/ADT/APInt.h"
|
|
|
|
#include <functional>
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::tosa;
|
|
|
|
namespace {
|
|
OpFoldResult foldToInputIfTypeMatches(Type typeRef, Value input) {
|
|
return input.getType() == typeRef ? OpFoldResult(input) : OpFoldResult{};
|
|
}
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Operator Canonicalizers.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Tensor Data Engine Operators.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// Check that the zero point of the tensor and padding operations are aligned.
|
|
static bool checkMatchingPadConstAndZp(Value padConst, Value zp) {
|
|
// Check that padConst is a constant value and a scalar tensor
|
|
DenseElementsAttr padConstAttr;
|
|
if (!matchPattern(padConst, m_Constant(&padConstAttr)) ||
|
|
(padConstAttr.size() != 1)) {
|
|
return false;
|
|
}
|
|
|
|
// Check that floating point pad is zero
|
|
if (auto padConstFpAttr = mlir::dyn_cast<DenseFPElementsAttr>(padConstAttr)) {
|
|
float padConstVal = (*padConstFpAttr.begin()).convertToFloat();
|
|
return padConstVal == 0.0f;
|
|
}
|
|
|
|
// Check that the zp and padConst align for the integer (quantized) case
|
|
if (auto padConstIntAttr =
|
|
mlir::dyn_cast<DenseIntElementsAttr>(padConstAttr)) {
|
|
DenseIntElementsAttr zpAttr;
|
|
// Check that zp is a constant value and a scalar tensor
|
|
if (!matchPattern(zp, m_Constant(&zpAttr)) || (padConstAttr.size() != 1)) {
|
|
return false;
|
|
}
|
|
|
|
// Check equality
|
|
int64_t zpVal = (*zpAttr.begin()).getSExtValue();
|
|
int64_t padConstVal = (*padConstIntAttr.begin()).getSExtValue();
|
|
return zpVal == padConstVal;
|
|
}
|
|
|
|
// Bail-out on unsupported type
|
|
return false;
|
|
}
|
|
|
|
namespace {
|
|
template <typename OpTy>
|
|
struct PoolPadFoldAdaptor;
|
|
|
|
template <>
|
|
struct PoolPadFoldAdaptor<tosa::MaxPool2dOp> {
|
|
using OpTy = tosa::MaxPool2dOp;
|
|
static bool checkKernelCompliance(OpTy op, const ArrayRef<int64_t> newPad) {
|
|
const llvm::ArrayRef<int64_t> kernel = op.getKernel();
|
|
if (newPad[2] >= kernel[1] || newPad[3] >= kernel[1] ||
|
|
newPad[0] >= kernel[0] || newPad[1] >= kernel[0])
|
|
return false;
|
|
return true;
|
|
}
|
|
static bool checkPadConstCompliance(OpTy, Value padConst) {
|
|
// Check that padConst is a constant value and a scalar tensor
|
|
DenseElementsAttr padConstAttr;
|
|
if (!matchPattern(padConst, m_Constant(&padConstAttr)) ||
|
|
padConstAttr.size() != 1) {
|
|
return false;
|
|
}
|
|
|
|
// Pad needs to be in the minimum value to be able to merge
|
|
if (auto padConstFpAttr =
|
|
mlir::dyn_cast<DenseFPElementsAttr>(padConstAttr)) {
|
|
const APFloat padConstVal = *padConstFpAttr.begin();
|
|
const APFloat lowestVal =
|
|
APFloat::getLargest(padConstVal.getSemantics(), true);
|
|
return padConstVal == lowestVal;
|
|
}
|
|
if (auto padConstIntAttr =
|
|
mlir::dyn_cast<DenseIntElementsAttr>(padConstAttr)) {
|
|
const APInt padConstVal = *padConstIntAttr.begin();
|
|
const unsigned int bitWidth = padConstVal.getBitWidth();
|
|
const APInt lowestVal =
|
|
padConstIntAttr.getElementType().isUnsignedInteger()
|
|
? APInt::getZero(bitWidth)
|
|
: APInt::getSignedMinValue(bitWidth);
|
|
return padConstVal == lowestVal;
|
|
}
|
|
|
|
// Bail-out on unsupported type
|
|
return false;
|
|
}
|
|
static void replaceOpWithNewPad(PatternRewriter &rewriter, OpTy op,
|
|
Value padInput, ArrayRef<int64_t> newPad) {
|
|
rewriter.replaceOpWithNewOp<tosa::MaxPool2dOp>(
|
|
op, op.getType(), padInput, op.getKernel(), op.getStride(),
|
|
rewriter.getDenseI64ArrayAttr(newPad), op.getNanMode());
|
|
}
|
|
};
|
|
|
|
template <typename OpTy>
|
|
struct ConvPadFoldAdaptor {
|
|
static bool checkKernelCompliance(OpTy, const ArrayRef<int64_t>) {
|
|
return true;
|
|
}
|
|
static bool checkPadConstCompliance(OpTy op, Value padConst) {
|
|
return checkMatchingPadConstAndZp(padConst, op.getInputZp());
|
|
}
|
|
static void replaceOpWithNewPad(PatternRewriter &rewriter, OpTy op,
|
|
Value padInput, ArrayRef<int64_t> newPad) {
|
|
rewriter.replaceOpWithNewOp<OpTy>(
|
|
op, op.getResult().getType(), padInput, op.getWeight(), op.getBias(),
|
|
op.getInputZp(), op.getWeightZp(), newPad, op.getStrideAttr(),
|
|
op.getDilationAttr(), op.getAccType(), op.getLocalBound());
|
|
}
|
|
};
|
|
|
|
// Pattern attempts to fold a `tosa.pad` operator to a following tensor
|
|
// operation like `tosa.conv2d` by merging the padding associated with the
|
|
// pad operator directly to the implicit padding of the tensor operation.
|
|
// This helps eliminate the explicit padding operator if unused.
|
|
template <typename OpTy, typename AdaptorTy>
|
|
struct FoldPadToTensorOp : public OpRewritePattern<OpTy> {
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy tensorOp,
|
|
PatternRewriter &rewriter) const override {
|
|
// Check producer is a tosa::PadOp
|
|
auto padOp = tensorOp.getInput().template getDefiningOp<tosa::PadOp>();
|
|
if (!padOp)
|
|
return rewriter.notifyMatchFailure(tensorOp,
|
|
"Producer must be a tosa::PadOp.");
|
|
|
|
// Validate that tensor operation has sane padding
|
|
const std::vector<int64_t> &tensorOpPad = tensorOp.getPad().vec();
|
|
if (tensorOpPad.size() != 4) // pad_top, pad_bottom, pad_left, pad_right
|
|
return rewriter.notifyMatchFailure(
|
|
tensorOp, "Tensor operation padding shall have 4 elements.");
|
|
|
|
// Validate tosa::PadOp padding
|
|
DenseIntElementsAttr padOpPadding;
|
|
if (!matchPattern(padOp.getPadding(), m_Constant(&padOpPadding))) {
|
|
return rewriter.notifyMatchFailure(
|
|
tensorOp,
|
|
"The `padding` input specified on the tosa::PadOp must be constant.");
|
|
}
|
|
// N_before, N_after, H_before, H_after, W_before, W_after, C_before,
|
|
// C_after
|
|
if (padOpPadding.size() != 8)
|
|
return rewriter.notifyMatchFailure(tensorOp,
|
|
"Pad padding should have 8 elements.");
|
|
int64_t padNBefore = (*(padOpPadding.begin() + 0)).getLimitedValue();
|
|
int64_t padNAfter = (*(padOpPadding.begin() + 1)).getLimitedValue();
|
|
int64_t padHBefore = (*(padOpPadding.begin() + 2)).getLimitedValue();
|
|
int64_t padHAfter = (*(padOpPadding.begin() + 3)).getLimitedValue();
|
|
int64_t padWBefore = (*(padOpPadding.begin() + 4)).getLimitedValue();
|
|
int64_t padWAfter = (*(padOpPadding.begin() + 5)).getLimitedValue();
|
|
int64_t padCBefore = (*(padOpPadding.begin() + 6)).getLimitedValue();
|
|
int64_t padCAfter = (*(padOpPadding.begin() + 7)).getLimitedValue();
|
|
|
|
if (padNBefore != 0 || padNAfter != 0 || padCBefore != 0 || padCAfter != 0)
|
|
return rewriter.notifyMatchFailure(
|
|
tensorOp, "Folding padding in N or C dimensions is not supported.");
|
|
|
|
// Fold padding from Pad into the tensor operation
|
|
// 4 elements - pad_top, pad_bottom, pad_left, pad_right
|
|
SmallVector<int64_t> foldedPad(tensorOpPad.size());
|
|
foldedPad[0] = padHBefore + tensorOpPad[0];
|
|
foldedPad[1] = padHAfter + tensorOpPad[1];
|
|
foldedPad[2] = padWBefore + tensorOpPad[2];
|
|
foldedPad[3] = padWAfter + tensorOpPad[3];
|
|
|
|
// Check kernel related restrictions
|
|
if (!AdaptorTy::checkKernelCompliance(tensorOp, foldedPad)) {
|
|
return rewriter.notifyMatchFailure(
|
|
tensorOp, "Padding size not aligned with kernel restrictions.");
|
|
}
|
|
|
|
// Check padding constant restrictions
|
|
if (!AdaptorTy::checkPadConstCompliance(tensorOp, padOp.getPadConst())) {
|
|
return rewriter.notifyMatchFailure(
|
|
tensorOp,
|
|
"Padding constant is not aligned with operator zero-point.");
|
|
}
|
|
|
|
// Check that padding doesn't grow more than 8K level (8192) for now
|
|
if (llvm::any_of(foldedPad, [](int64_t padVal) { return padVal > 8192; })) {
|
|
return rewriter.notifyMatchFailure(
|
|
tensorOp, "Padding size more than the 8K level limit.");
|
|
}
|
|
|
|
// Create operator
|
|
AdaptorTy::replaceOpWithNewPad(rewriter, tensorOp, padOp.getInput1(),
|
|
foldedPad);
|
|
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void Conv2DOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<
|
|
FoldPadToTensorOp<tosa::Conv2DOp, ConvPadFoldAdaptor<tosa::Conv2DOp>>>(
|
|
context);
|
|
}
|
|
|
|
void DepthwiseConv2DOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<FoldPadToTensorOp<tosa::DepthwiseConv2DOp,
|
|
ConvPadFoldAdaptor<tosa::DepthwiseConv2DOp>>>(
|
|
context);
|
|
}
|
|
|
|
struct AvgPool2dAdaptiveToAvgPool2d
|
|
: public OpRewritePattern<tosa::AvgPool2dAdaptiveOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::AvgPool2dAdaptiveOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
llvm::SmallVector<int64_t> kernel;
|
|
llvm::SmallVector<int64_t> stride;
|
|
llvm::SmallVector<int64_t> pad;
|
|
if (!tosa::getConstShapeValues(op.getKernel().getDefiningOp(), kernel) ||
|
|
!tosa::getConstShapeValues(op.getStride().getDefiningOp(), stride) ||
|
|
!tosa::getConstShapeValues(op.getPad().getDefiningOp(), pad))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "expected constant kernel, stride, and pad operands");
|
|
|
|
auto replacement = tosa::AvgPool2dOp::create(
|
|
rewriter, op.getLoc(), op.getType(), op.getInput(), op.getInputZp(),
|
|
op.getOutputZp(), rewriter.getDenseI64ArrayAttr(kernel),
|
|
rewriter.getDenseI64ArrayAttr(stride),
|
|
rewriter.getDenseI64ArrayAttr(pad), op.getAccTypeAttr());
|
|
rewriter.replaceOp(op, replacement.getOutput());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void AvgPool2dAdaptiveOp::getCanonicalizationPatterns(
|
|
RewritePatternSet &results, MLIRContext *context) {
|
|
results.add<AvgPool2dAdaptiveToAvgPool2d>(context);
|
|
}
|
|
|
|
struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
Value input = op.getInput();
|
|
Value output = op.getOutput();
|
|
ShapedType inputType = llvm::cast<ShapedType>(input.getType());
|
|
ShapedType outputType = llvm::cast<ShapedType>(output.getType());
|
|
|
|
if (!inputType.hasStaticShape() || !outputType.hasStaticShape()) {
|
|
return failure();
|
|
}
|
|
|
|
// If the output and input shapes are 1x1, then this is a no op.
|
|
ArrayRef<int64_t> outputShape = outputType.getShape();
|
|
if (outputShape[1] != 1 || outputShape[2] != 1) {
|
|
return failure();
|
|
}
|
|
|
|
ArrayRef<int64_t> inputShape = inputType.getShape();
|
|
if (inputShape[1] != 1 || inputShape[2] != 1) {
|
|
return failure();
|
|
}
|
|
|
|
rewriter.replaceOp(op, input);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void MaxPool2dOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<MaxPool2dIsNoOp,
|
|
FoldPadToTensorOp<tosa::MaxPool2dOp,
|
|
PoolPadFoldAdaptor<tosa::MaxPool2dOp>>>(
|
|
context);
|
|
}
|
|
|
|
struct MaxPool2dAdaptiveToMaxPool2d
|
|
: public OpRewritePattern<tosa::MaxPool2dAdaptiveOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::MaxPool2dAdaptiveOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
llvm::SmallVector<int64_t> kernel;
|
|
llvm::SmallVector<int64_t> stride;
|
|
llvm::SmallVector<int64_t> pad;
|
|
if (!tosa::getConstShapeValues(op.getKernel().getDefiningOp(), kernel) ||
|
|
!tosa::getConstShapeValues(op.getStride().getDefiningOp(), stride) ||
|
|
!tosa::getConstShapeValues(op.getPad().getDefiningOp(), pad))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "expected constant kernel, stride, and pad operands");
|
|
|
|
auto replacement = tosa::MaxPool2dOp::create(
|
|
rewriter, op.getLoc(), op.getType(), op.getInput(),
|
|
rewriter.getDenseI64ArrayAttr(kernel),
|
|
rewriter.getDenseI64ArrayAttr(stride),
|
|
rewriter.getDenseI64ArrayAttr(pad), op.getNanModeAttr());
|
|
rewriter.replaceOp(op, replacement.getOutput());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void MaxPool2dAdaptiveOp::getCanonicalizationPatterns(
|
|
RewritePatternSet &results, MLIRContext *context) {
|
|
results.add<MaxPool2dAdaptiveToMaxPool2d>(context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Data Layout / Memory Reinterpretation.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
struct ConcatOptimization : public OpRewritePattern<tosa::ConcatOp> {
|
|
using OpRewritePattern<tosa::ConcatOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::ConcatOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (op.getInput1().size() != 1)
|
|
return failure();
|
|
if (op.getInput1().front().getType() != op.getType()) {
|
|
rewriter
|
|
.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(),
|
|
op.getInput1().front())
|
|
.getResult();
|
|
return success();
|
|
}
|
|
|
|
rewriter.replaceOp(op, op.getInput1().front());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct ConsecutiveConcatOptimization : public OpRewritePattern<tosa::ConcatOp> {
|
|
using OpRewritePattern<tosa::ConcatOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::ConcatOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
// Rewrite consecutive concats on the same axis into a single op.
|
|
// Keep track of the operands so we are able to construct a new concat
|
|
// later. Conservatively assume that we double the number of operands when
|
|
// canonicalizing
|
|
SmallVector<Value, 8> concatOperands;
|
|
concatOperands.reserve(2 * op.getNumOperands());
|
|
|
|
int32_t maxNumOperands = 0;
|
|
if (auto targetEnvAttr = tosa::lookupTargetEnv(op))
|
|
maxNumOperands =
|
|
getTosaLevelFromEnum(targetEnvAttr.getLevel()).MAX_TENSOR_LIST_SIZE;
|
|
|
|
// Find all operands that are foldable concats
|
|
bool foundRewritableConcat = false;
|
|
for (Value operand : op.getOperands()) {
|
|
concatOperands.emplace_back(operand);
|
|
|
|
auto producer = operand.getDefiningOp<tosa::ConcatOp>();
|
|
if (!producer)
|
|
continue;
|
|
|
|
// Not rewritable if axes are not the same
|
|
if (op.getAxis() != producer.getAxis())
|
|
continue;
|
|
|
|
// Replace the original operand with all incoming operands
|
|
foundRewritableConcat = true;
|
|
concatOperands.pop_back();
|
|
llvm::append_range(concatOperands, producer->getOperands());
|
|
}
|
|
|
|
if (!foundRewritableConcat)
|
|
return rewriter.notifyMatchFailure(op,
|
|
"No rewritable concat operand found.");
|
|
|
|
if (maxNumOperands > 0 &&
|
|
concatOperands.size() > static_cast<size_t>(maxNumOperands))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Rewriting would exceed the maximum number of operands for the "
|
|
"target environment level.");
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::ConcatOp>(
|
|
op, op.getType(), concatOperands, op.getAxisAttr());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void ConcatOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<ConcatOptimization, ConsecutiveConcatOptimization>(context);
|
|
}
|
|
|
|
LogicalResult SelectOp::canonicalize(SelectOp op, PatternRewriter &rewriter) {
|
|
auto notOp = op.getInput1().getDefiningOp<tosa::LogicalNotOp>();
|
|
if (!notOp)
|
|
return failure();
|
|
rewriter.modifyOpInPlace(op, [&]() {
|
|
op.getOperation()->setOperands(
|
|
{notOp.getInput1(), op.getOnFalse(), op.getOnTrue()});
|
|
});
|
|
return success();
|
|
}
|
|
|
|
struct ConsolidateTransposeOptimization
|
|
: public OpRewritePattern<tosa::TransposeOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::TransposeOp transposeOp,
|
|
PatternRewriter &rewriter) const override {
|
|
// Input is also TransposeOp - transpose(transpose(A)).
|
|
auto innerTranspose =
|
|
transposeOp.getInput1().getDefiningOp<tosa::TransposeOp>();
|
|
if (!innerTranspose)
|
|
return rewriter.notifyMatchFailure(transposeOp,
|
|
"input must be transpose operation");
|
|
|
|
const llvm::ArrayRef<int32_t> transposePerms = transposeOp.getPerms();
|
|
const llvm::ArrayRef<int32_t> innerTransposePerms =
|
|
innerTranspose.getPerms();
|
|
|
|
if (transposePerms.size() != innerTransposePerms.size())
|
|
return rewriter.notifyMatchFailure(
|
|
transposeOp,
|
|
"transpose and inner transpose perms sizes must be equal");
|
|
if (transposePerms.empty())
|
|
return rewriter.notifyMatchFailure(
|
|
transposeOp, "transpose perms sizes must be positive");
|
|
|
|
// Consolidate transposes into one transpose.
|
|
SmallVector<int32_t> perms(transposePerms.size());
|
|
for (int i = 0, s = transposePerms.size(); i < s; ++i)
|
|
perms[i] = innerTransposePerms[transposePerms[i]];
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::TransposeOp>(
|
|
transposeOp, transposeOp.getResult().getType(),
|
|
innerTranspose.getInput1(), rewriter.getDenseI32ArrayAttr(perms));
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// Determines the case when tosa.transpose is a tosa.reshape operation.
|
|
struct TransposeIsReshape : public OpRewritePattern<tosa::TransposeOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::TransposeOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
if (op.getInput1().getDefiningOp<tosa::TransposeOp>())
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Src is from transpose, can compose transposes");
|
|
|
|
Value result = op.getResult();
|
|
for (Operation *subop : result.getUsers()) {
|
|
if (isa_and_nonnull<tosa::TransposeOp>(subop))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Dest is used by transpose, can compose transposes");
|
|
}
|
|
|
|
auto input = op.getInput1();
|
|
auto inputTy = llvm::cast<ShapedType>(input.getType());
|
|
if (!inputTy.hasRank())
|
|
return rewriter.notifyMatchFailure(op, "Unranked input.");
|
|
|
|
int64_t numDynDims = 0;
|
|
for (int i = 0; i < inputTy.getRank(); ++i)
|
|
if (inputTy.isDynamicDim(i))
|
|
numDynDims++;
|
|
|
|
if (numDynDims > 1)
|
|
return rewriter.notifyMatchFailure(op, "Has more than one dynamic dim.");
|
|
|
|
const llvm::ArrayRef<int32_t> permValues = op.getPerms();
|
|
|
|
SmallVector<int64_t> nonZeroPerms;
|
|
nonZeroPerms.reserve(permValues.size());
|
|
for (auto idx : permValues) {
|
|
auto sz = inputTy.getDimSize(idx);
|
|
if (sz != 1)
|
|
nonZeroPerms.push_back(idx);
|
|
}
|
|
|
|
for (int i = 1, s = nonZeroPerms.size(); i < s; ++i)
|
|
if (nonZeroPerms[i - 1] > nonZeroPerms[i])
|
|
return rewriter.notifyMatchFailure(op,
|
|
"Transpose changes memory layout.");
|
|
|
|
SmallVector<int64_t> newShape;
|
|
newShape.reserve(inputTy.getRank());
|
|
for (int i = 0, s = inputTy.getRank(); i < s; ++i)
|
|
newShape.push_back(inputTy.getDimSize(permValues[i]));
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
|
|
op, op.getType(), op.getInput1(),
|
|
getTosaConstShape(rewriter, op.getLoc(), newShape));
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<ConsolidateTransposeOptimization, TransposeIsReshape>(context);
|
|
}
|
|
|
|
struct ClampIsNoOp : public OpRewritePattern<tosa::ClampOp> {
|
|
using OpRewritePattern::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::ClampOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
Value input = op.getInput();
|
|
auto inputType = llvm::cast<ShapedType>(op.getInput().getType());
|
|
auto inputElementType = inputType.getElementType();
|
|
|
|
if (isa<FloatType>(inputElementType)) {
|
|
// Unlike integer types, floating point types can represent infinity.
|
|
const auto minClamp =
|
|
llvm::cast<mlir::FloatAttr>(op.getMinValAttr()).getValue();
|
|
const auto maxClamp =
|
|
llvm::cast<mlir::FloatAttr>(op.getMaxValAttr()).getValue();
|
|
const bool isMin = minClamp.isNegInfinity();
|
|
const bool isMax = maxClamp.isInfinity();
|
|
|
|
if (isMin && isMax) {
|
|
rewriter.replaceOp(op, input);
|
|
return success();
|
|
}
|
|
return failure();
|
|
}
|
|
|
|
// i1 types are boolean in TOSA
|
|
const bool isBoolean = inputElementType.isInteger(1);
|
|
if (inputElementType.isUnsignedInteger() || isBoolean) {
|
|
const int64_t minClamp = llvm::cast<mlir::IntegerAttr>(op.getMinValAttr())
|
|
.getValue()
|
|
.getZExtValue();
|
|
const int64_t maxClamp = llvm::cast<mlir::IntegerAttr>(op.getMaxValAttr())
|
|
.getValue()
|
|
.getZExtValue();
|
|
|
|
const unsigned bitWidth = inputElementType.getIntOrFloatBitWidth();
|
|
const int64_t intMin = APInt::getMinValue(bitWidth).getZExtValue();
|
|
const int64_t intMax = APInt::getMaxValue(bitWidth).getZExtValue();
|
|
|
|
if (minClamp <= intMin && maxClamp >= intMax) {
|
|
rewriter.replaceOp(op, input);
|
|
return success();
|
|
}
|
|
return failure();
|
|
}
|
|
|
|
if (llvm::isa<IntegerType>(inputElementType)) {
|
|
const int64_t minClamp =
|
|
llvm::cast<mlir::IntegerAttr>(op.getMinValAttr()).getInt();
|
|
const int64_t maxClamp =
|
|
llvm::cast<mlir::IntegerAttr>(op.getMaxValAttr()).getInt();
|
|
|
|
const unsigned bitWidth = inputElementType.getIntOrFloatBitWidth();
|
|
const int64_t intMin = APInt::getSignedMinValue(bitWidth).getSExtValue();
|
|
const int64_t intMax = APInt::getSignedMaxValue(bitWidth).getSExtValue();
|
|
|
|
if (minClamp <= intMin && maxClamp >= intMax) {
|
|
rewriter.replaceOp(op, input);
|
|
return success();
|
|
}
|
|
return failure();
|
|
}
|
|
|
|
return failure();
|
|
}
|
|
};
|
|
|
|
// Attempts the following transformation:
|
|
//
|
|
// For integers a, b, a', and b' such that [a, b] ∩ [a', b'] ≠ ∅ and input
|
|
// tensor X the following identity holds:
|
|
//
|
|
// CLAMP(CLAMP(X, a, b), a', b') = CLAMP(X, max(a, a'), min(b, b'))
|
|
//
|
|
// subject to the following valid NaN propagation semantics:
|
|
// --------------------------------------------
|
|
// | OUTER CLAMP | INNER CLAMP | RESULT MODE |
|
|
// |-------------|--------------|-------------|
|
|
// | PROPAGATE | PROPAGATE | PROPAGATE |
|
|
// | PROPAGATE | IGNORE | IGNORE |
|
|
// | IGNORE | PROPAGATE | INVALID |
|
|
// | IGNORE | IGNORE | IGNORE |
|
|
// |------------------------------------------|
|
|
|
|
struct ClampClampOptimization : public OpRewritePattern<tosa::ClampOp> {
|
|
using OpRewritePattern<tosa::ClampOp>::OpRewritePattern;
|
|
|
|
// Helper structure to describe the range of a clamp operation.
|
|
template <typename T>
|
|
struct ClampRange {
|
|
ClampRange(const T &start, const T &end) : start(start), end(end) {}
|
|
T start;
|
|
T end;
|
|
|
|
// Helper function to determine if two Clamp ranges intersect.
|
|
bool intersects(const ClampRange<T> &otherRange) {
|
|
return start < otherRange.end && otherRange.start < end;
|
|
}
|
|
};
|
|
|
|
LogicalResult matchAndRewrite(tosa::ClampOp op,
|
|
PatternRewriter &rewriter) const override {
|
|
Value input = op.getInput();
|
|
|
|
// Check the input to the CLAMP op is itself a CLAMP.
|
|
auto clampOp = dyn_cast_if_present<tosa::ClampOp>(input.getDefiningOp());
|
|
if (!clampOp)
|
|
return failure();
|
|
|
|
// Check we have a valid NaN propagation combination.
|
|
const auto opNanMode = op.getNanMode();
|
|
const auto clampNanMode = clampOp.getNanMode();
|
|
if (opNanMode == NanPropagationMode::IGNORE &&
|
|
clampNanMode == NanPropagationMode::PROPAGATE)
|
|
return failure();
|
|
|
|
auto maxValAttr = op.getMaxValAttr();
|
|
auto minValAttr = op.getMinValAttr();
|
|
auto clampOpMaxValAttr = clampOp.getMaxValAttr();
|
|
auto clampOpMinValAttr = clampOp.getMinValAttr();
|
|
|
|
auto inputEType = llvm::cast<ShapedType>(input.getType()).getElementType();
|
|
if (auto quantType =
|
|
llvm::dyn_cast<mlir::quant::UniformQuantizedType>(inputEType)) {
|
|
inputEType = getStorageElementTypeFromQuantized(quantType);
|
|
}
|
|
|
|
Attribute newMinValAttr, newMaxValAttr;
|
|
if (mlir::isa<FloatType>(inputEType)) {
|
|
auto floatMaxValAttr = cast<mlir::FloatAttr>(maxValAttr);
|
|
auto floatMinValAttr = cast<mlir::FloatAttr>(minValAttr);
|
|
auto clampOpFloatMaxValAttr = cast<mlir::FloatAttr>(clampOpMaxValAttr);
|
|
auto clampOpFloatMinValAttr = cast<mlir::FloatAttr>(clampOpMinValAttr);
|
|
|
|
// Check we have intersecting ranges.
|
|
const auto opMinFloat = floatMinValAttr.getValue();
|
|
const auto opMaxFloat = floatMaxValAttr.getValue();
|
|
const auto clampOpMinFloat = clampOpFloatMinValAttr.getValue();
|
|
const auto clampOpMaxFloat = clampOpFloatMaxValAttr.getValue();
|
|
ClampRange<APFloat> opRangeFloatRange(opMinFloat, opMaxFloat);
|
|
ClampRange<APFloat> clampRangeFloatRange(clampOpMinFloat,
|
|
clampOpMaxFloat);
|
|
if (!opRangeFloatRange.intersects(clampRangeFloatRange))
|
|
return failure();
|
|
|
|
// Run the transformation.
|
|
auto newMinVal = std::max(opMinFloat, clampOpMinFloat);
|
|
auto newMaxVal = std::min(opMaxFloat, clampOpMaxFloat);
|
|
newMinValAttr = rewriter.getFloatAttr(inputEType, newMinVal);
|
|
newMaxValAttr = rewriter.getFloatAttr(inputEType, newMaxVal);
|
|
} else {
|
|
assert(mlir::isa<IntegerType>(inputEType));
|
|
auto intMaxValAttr = cast<mlir::IntegerAttr>(maxValAttr);
|
|
auto intMinValAttr = cast<mlir::IntegerAttr>(minValAttr);
|
|
auto clampOpIntMaxValAttr = cast<mlir::IntegerAttr>(clampOpMaxValAttr);
|
|
auto clampOpIntMinValAttr = cast<mlir::IntegerAttr>(clampOpMinValAttr);
|
|
|
|
if (inputEType.isUnsignedInteger()) {
|
|
// Check we have intersecting ranges.
|
|
const auto opMinInt = intMinValAttr.getUInt();
|
|
const auto opMaxInt = intMaxValAttr.getUInt();
|
|
const auto clampOpMinInt = clampOpIntMinValAttr.getUInt();
|
|
const auto clampOpMaxInt = clampOpIntMaxValAttr.getUInt();
|
|
ClampRange<std::uint64_t> opRangeIntRange(opMinInt, opMaxInt);
|
|
ClampRange<std::uint64_t> clampRangeIntRange(clampOpMinInt,
|
|
clampOpMaxInt);
|
|
if (!opRangeIntRange.intersects(clampRangeIntRange))
|
|
return failure();
|
|
|
|
// Run the transformation.
|
|
auto newMinVal = std::max(opMinInt, clampOpMinInt);
|
|
auto newMaxVal = std::min(opMaxInt, clampOpMaxInt);
|
|
newMinValAttr = rewriter.getIntegerAttr(inputEType, newMinVal);
|
|
newMaxValAttr = rewriter.getIntegerAttr(inputEType, newMaxVal);
|
|
} else {
|
|
// Check we have intersecting ranges.
|
|
const auto opMinInt = intMinValAttr.getInt();
|
|
const auto opMaxInt = intMaxValAttr.getInt();
|
|
const auto clampOpMinInt = clampOpIntMinValAttr.getInt();
|
|
const auto clampOpMaxInt = clampOpIntMaxValAttr.getInt();
|
|
ClampRange<std::int64_t> opRangeIntRange(opMinInt, opMaxInt);
|
|
ClampRange<std::int64_t> clampRangeIntRange(clampOpMinInt,
|
|
clampOpMaxInt);
|
|
if (!opRangeIntRange.intersects(clampRangeIntRange))
|
|
return failure();
|
|
|
|
// Run the transformation.
|
|
auto newMinVal = std::max(opMinInt, clampOpMinInt);
|
|
auto newMaxVal = std::min(opMaxInt, clampOpMaxInt);
|
|
newMinValAttr = rewriter.getIntegerAttr(inputEType, newMinVal);
|
|
newMaxValAttr = rewriter.getIntegerAttr(inputEType, newMaxVal);
|
|
}
|
|
}
|
|
|
|
auto newMode = (opNanMode != clampNanMode)
|
|
? tosa::NanPropagationMode::IGNORE
|
|
: opNanMode;
|
|
|
|
auto newModeAttr =
|
|
NanPropagationModeAttr::get(rewriter.getContext(), newMode);
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::ClampOp>(
|
|
op, op.getType(), clampOp.getInput(), newMinValAttr, newMaxValAttr,
|
|
newModeAttr);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void ClampOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<ClampIsNoOp>(context);
|
|
results.add<ClampClampOptimization>(context);
|
|
}
|
|
|
|
struct ConcatSliceOptimization : public OpRewritePattern<tosa::SliceOp> {
|
|
using OpRewritePattern<tosa::SliceOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::SliceOp sliceOp,
|
|
PatternRewriter &rewriter) const override {
|
|
Value sliceInput = sliceOp.getInput1();
|
|
auto concatOp = sliceInput.getDefiningOp<tosa::ConcatOp>();
|
|
if (!concatOp)
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "slice input must be concat operation");
|
|
|
|
OperandRange inputs = concatOp.getInput1();
|
|
auto concatType = dyn_cast<RankedTensorType>(concatOp.getType());
|
|
if (!concatType || !concatType.hasStaticShape())
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "slice input must be a static ranked tensor");
|
|
int32_t axis = concatOp.getAxis();
|
|
|
|
DenseElementsAttr startElems;
|
|
DenseElementsAttr sizeElems;
|
|
|
|
if (!matchPattern(sliceOp.getStart(), m_Constant(&startElems)))
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "start of slice must be a static ranked shape");
|
|
|
|
if (!matchPattern(sliceOp.getSize(), m_Constant(&sizeElems)))
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "size of slice must be a static ranked shape");
|
|
|
|
llvm::SmallVector<int64_t> sliceStarts =
|
|
llvm::to_vector(startElems.getValues<int64_t>());
|
|
llvm::SmallVector<int64_t> sliceSizes =
|
|
llvm::to_vector(sizeElems.getValues<int64_t>());
|
|
|
|
// Validate slice on the concatenated axis. Slicing along this
|
|
// axis should span only one of the inputs to the concatenate
|
|
// operation.
|
|
std::optional<Value> replaceWithSlice;
|
|
for (auto input : inputs) {
|
|
auto inputType = dyn_cast<RankedTensorType>(input.getType());
|
|
if (!inputType || !inputType.hasStaticShape())
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "concat input must be a static ranked tensor");
|
|
|
|
if (sliceStarts[axis] >= 0 && (sliceStarts[axis] + sliceSizes[axis]) <=
|
|
inputType.getDimSize(axis)) {
|
|
auto start_op =
|
|
getTosaConstShape(rewriter, sliceOp.getLoc(), sliceStarts);
|
|
auto size_op =
|
|
getTosaConstShape(rewriter, sliceOp.getLoc(), sliceSizes);
|
|
replaceWithSlice =
|
|
tosa::SliceOp::create(rewriter, sliceOp.getLoc(), sliceOp.getType(),
|
|
input, start_op, size_op)
|
|
.getResult();
|
|
break;
|
|
}
|
|
sliceStarts[axis] -= inputType.getDimSize(axis);
|
|
}
|
|
|
|
if (!replaceWithSlice)
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "corresponding concat input not found for slice");
|
|
|
|
rewriter.replaceOp(sliceOp, replaceWithSlice.value());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
struct PadSliceOptimization : public OpRewritePattern<tosa::SliceOp> {
|
|
using OpRewritePattern<tosa::SliceOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::SliceOp sliceOp,
|
|
PatternRewriter &rewriter) const override {
|
|
Value sliceInput = sliceOp.getInput1();
|
|
|
|
// Check if producer is a PadOp
|
|
auto padOp = sliceInput.getDefiningOp<tosa::PadOp>();
|
|
if (!padOp)
|
|
return rewriter.notifyMatchFailure(sliceOp,
|
|
"slice input must be a pad operation");
|
|
|
|
// Check PadOp has a single consumer
|
|
if (!padOp->hasOneUse())
|
|
return rewriter.notifyMatchFailure(sliceOp,
|
|
"pad shall have a single consumer");
|
|
|
|
// Check input is statically ranked
|
|
auto inputTy = dyn_cast<RankedTensorType>(padOp.getInput1().getType());
|
|
auto padTy = dyn_cast<RankedTensorType>(padOp.getType());
|
|
if (!inputTy || !padTy || !inputTy.hasRank())
|
|
return rewriter.notifyMatchFailure(sliceOp,
|
|
"slice input must be a ranked tensor");
|
|
|
|
// Validate and extract tosa::PadOp padding
|
|
DenseIntElementsAttr paddingElems;
|
|
if (!matchPattern(padOp.getPadding(), m_Constant(&paddingElems))) {
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp,
|
|
"`padding` input specified on the tosa::PadOp must be constant.");
|
|
}
|
|
llvm::SmallVector<int64_t> padPaddings =
|
|
llvm::to_vector(paddingElems.getValues<int64_t>());
|
|
|
|
// Extract slice parameters
|
|
DenseElementsAttr startElems;
|
|
if (!matchPattern(sliceOp.getStart(), m_Constant(&startElems)))
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "start of slice must be a static ranked shape");
|
|
llvm::SmallVector<int64_t> sliceStarts =
|
|
llvm::to_vector(startElems.getValues<int64_t>());
|
|
|
|
DenseElementsAttr sizeElems;
|
|
if (!matchPattern(sliceOp.getSize(), m_Constant(&sizeElems)))
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "size of slice must be a static ranked shape");
|
|
llvm::SmallVector<int64_t> sliceSizes =
|
|
llvm::to_vector(sizeElems.getValues<int64_t>());
|
|
|
|
// Check if dynamic dimensions are sliced
|
|
const int64_t rank = inputTy.getRank();
|
|
if (llvm::any_of(llvm::seq<int64_t>(0, rank), [&](int64_t i) {
|
|
const bool isDimDynamic = inputTy.isDynamicDim(i);
|
|
const bool isDimSliced =
|
|
(sliceStarts[i] != 0) || (sliceSizes[i] != kInferableDimSize);
|
|
|
|
return isDimDynamic && isDimSliced;
|
|
})) {
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "axis that are sliced shall be statically known.");
|
|
}
|
|
|
|
// Update the parameters
|
|
llvm::SmallVector<int64_t> newSliceStarts(rank, 0);
|
|
llvm::SmallVector<int64_t> newPadPaddings(2 * rank, 0);
|
|
llvm::SmallVector<int64_t> newPadShape(rank, ShapedType::kDynamic);
|
|
bool updated = false;
|
|
|
|
for (int64_t i = 0; i < rank; ++i) {
|
|
const int64_t padLo = padPaddings[i * 2];
|
|
const int64_t padHi = padPaddings[i * 2 + 1];
|
|
const int64_t sliceStart = sliceStarts[i];
|
|
const int64_t sliceSize = sliceSizes[i];
|
|
const int64_t sliceEnd = sliceStart + sliceSize;
|
|
|
|
// If dimension is dynamic pass-through
|
|
if (inputTy.isDynamicDim(i)) {
|
|
newPadPaddings[i * 2] = padLo;
|
|
newPadPaddings[i * 2 + 1] = padHi;
|
|
newSliceStarts[i] = sliceStart;
|
|
continue;
|
|
}
|
|
|
|
// Handle static dimensions
|
|
const int64_t dimSize = inputTy.getShape()[i];
|
|
const int64_t dimTotal = padLo + dimSize + padHi;
|
|
|
|
// Check slice within bounds
|
|
if (sliceStart < 0 || sliceEnd > dimTotal)
|
|
return rewriter.notifyMatchFailure(sliceOp, "slice is out-of-bounds");
|
|
|
|
// Compute updated slice start parameter
|
|
const int64_t newSliceStart = std::max<int64_t>(sliceStart - padLo, 0);
|
|
newSliceStarts[i] = newSliceStart;
|
|
updated |= newSliceStart != sliceStart;
|
|
|
|
// Compute updated pad parameters
|
|
const int64_t newPadLo = std::max<int64_t>(padLo - sliceStart, 0);
|
|
const int64_t newPadHi =
|
|
std::max<int64_t>(sliceEnd - (padLo + dimSize), 0);
|
|
newPadPaddings[i * 2] = newPadLo;
|
|
newPadPaddings[i * 2 + 1] = newPadHi;
|
|
updated |= (newPadLo != padLo) || (newPadHi != padHi);
|
|
|
|
// Calculate new pad output shape
|
|
newPadShape[i] =
|
|
newPadPaddings[i * 2] + dimSize + newPadPaddings[i * 2 + 1];
|
|
}
|
|
|
|
// Check that we actually need to proceed with the rewrite
|
|
if (!updated)
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "terminate condition; nothing to rewrite");
|
|
|
|
// Create a PadOp with updated padding
|
|
auto newPaddingsOp =
|
|
getTosaConstShape(rewriter, sliceOp.getLoc(), newPadPaddings);
|
|
auto newPadTy =
|
|
RankedTensorType::get(newPadShape, inputTy.getElementType());
|
|
auto newPadOp = tosa::PadOp::create(rewriter, padOp.getLoc(), newPadTy,
|
|
padOp.getInput1(), newPaddingsOp,
|
|
padOp.getPadConst());
|
|
|
|
// Update SliceOp and point to new PadOp
|
|
auto newStartOp =
|
|
getTosaConstShape(rewriter, sliceOp.getLoc(), newSliceStarts);
|
|
rewriter.replaceOpWithNewOp<tosa::SliceOp>(sliceOp, sliceOp.getType(),
|
|
newPadOp.getResult(), newStartOp,
|
|
sliceOp.getSize());
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// Update size operand of tosa.slice if size has dynamic dims but corresponding
|
|
// output dim is static
|
|
struct SliceDynamicSizeCanonicalization
|
|
: public OpRewritePattern<tosa::SliceOp> {
|
|
using OpRewritePattern<tosa::SliceOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::SliceOp sliceOp,
|
|
PatternRewriter &rewriter) const override {
|
|
ShapedType resultType = cast<ShapedType>(sliceOp.getType());
|
|
if (!resultType.hasRank())
|
|
return rewriter.notifyMatchFailure(sliceOp, "output must be ranked");
|
|
|
|
ElementsAttr sizeElems;
|
|
if (!matchPattern(sliceOp.getSize(), m_Constant(&sizeElems))) {
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "size of slice must be a static ranked shape");
|
|
}
|
|
|
|
llvm::SmallVector<int64_t> sliceSizes =
|
|
llvm::to_vector(sizeElems.getValues<int64_t>());
|
|
|
|
bool replaceSliceSize{false};
|
|
// if size op has kInferableDimSize indicating dynamic shape but
|
|
// corresponding dim on the output is statically known, update size to match
|
|
// with known output dim shape
|
|
for (const auto &[index, size] : llvm::enumerate(sliceSizes)) {
|
|
if (size == kInferableDimSize && !resultType.isDynamicDim(index)) {
|
|
sliceSizes[index] = resultType.getDimSize(index);
|
|
replaceSliceSize = true;
|
|
}
|
|
}
|
|
|
|
if (!replaceSliceSize) {
|
|
return rewriter.notifyMatchFailure(
|
|
sliceOp, "no dimension of size of slice is dynamic that resolves "
|
|
"to static output shape");
|
|
}
|
|
|
|
auto size_op = getTosaConstShape(rewriter, sliceOp.getLoc(), sliceSizes);
|
|
auto newSliceOp =
|
|
tosa::SliceOp::create(rewriter, sliceOp.getLoc(), sliceOp.getType(),
|
|
sliceOp.getInput1(), sliceOp.getStart(), size_op);
|
|
|
|
rewriter.replaceOp(sliceOp, newSliceOp.getResult());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void SliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<ConcatSliceOptimization, PadSliceOptimization,
|
|
SliceDynamicSizeCanonicalization>(context);
|
|
}
|
|
|
|
struct NonNarrowingCastsOptimization : public OpRewritePattern<tosa::CastOp> {
|
|
using OpRewritePattern<tosa::CastOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::CastOp castOp,
|
|
PatternRewriter &rewriter) const override {
|
|
const Value castInput = castOp.getInput();
|
|
auto innerCastOp = castInput.getDefiningOp<tosa::CastOp>();
|
|
if (!innerCastOp)
|
|
return rewriter.notifyMatchFailure(castOp,
|
|
"input must be cast operation");
|
|
|
|
const Value innerCastInput = innerCastOp.getInput();
|
|
|
|
const ShapedType innerInputType =
|
|
llvm::cast<ShapedType>(innerCastInput.getType());
|
|
const ShapedType innerOutputType =
|
|
llvm::cast<ShapedType>(innerCastOp.getType());
|
|
const ShapedType outerOutputType = llvm::cast<ShapedType>(castOp.getType());
|
|
|
|
const Type innerInputElemType = innerInputType.getElementType();
|
|
const Type innerOutputElemType = innerOutputType.getElementType();
|
|
const Type outerOutputElemType = outerOutputType.getElementType();
|
|
|
|
const SmallVector<Type, 3> types = {innerInputElemType, innerOutputElemType,
|
|
outerOutputElemType};
|
|
|
|
if (llvm::any_of(types, [](const Type type) {
|
|
// Support a specific set of floating point types since we need to be
|
|
// careful in not introducing unsupported type combinations
|
|
return !(type.isInteger() ||
|
|
llvm::isa<Float8E4M3FNType, Float8E5M2Type, BFloat16Type,
|
|
Float16Type, Float32Type>(type));
|
|
}))
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "only integer and f32, f16, bf16, f8E4M3FN, f8E5M2 types are "
|
|
"supported");
|
|
|
|
if (llvm::isa<Float8E5M2Type>(innerInputElemType) &&
|
|
llvm::isa<Float8E4M3FNType>(outerOutputElemType)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing f8E5M2 -> f8E4M3FN casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
if (llvm::isa<Float8E4M3FNType>(innerInputElemType) &&
|
|
llvm::isa<Float8E5M2Type>(outerOutputElemType)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing f8E4M3FN -> f8E5M2 casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
if (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(innerInputElemType) &&
|
|
outerOutputElemType.isInteger()) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing fp8 -> integer casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
if (innerInputElemType.isInteger() &&
|
|
llvm::isa<Float8E5M2Type, Float8E4M3FNType>(outerOutputElemType)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing integer -> fp8 casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
if (llvm::isa<Float16Type>(innerInputElemType) &&
|
|
llvm::isa<BFloat16Type>(outerOutputElemType)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing fp16 -> bf16 casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
if (llvm::isa<BFloat16Type>(innerInputElemType) &&
|
|
llvm::isa<Float16Type>(outerOutputElemType)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing bf16 -> fp16 casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
const auto isIntegerOneOfWidth = [](Type type, size_t bitwidth1,
|
|
size_t bitwidth2) {
|
|
return type.isInteger(bitwidth1) || type.isInteger(bitwidth2);
|
|
};
|
|
|
|
if (isIntegerOneOfWidth(innerInputElemType, 8, 16) &&
|
|
outerOutputElemType.isInteger(64)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing i8/i16 -> i64 casts which are not "
|
|
"legal in TOSA");
|
|
}
|
|
|
|
if (isIntegerOneOfWidth(innerInputElemType, 1, 64) &&
|
|
!outerOutputElemType.isInteger()) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing bool/i64 to float casts which are not "
|
|
"supported in all versions of TOSA");
|
|
}
|
|
|
|
if (!innerInputElemType.isInteger() &&
|
|
isIntegerOneOfWidth(outerOutputElemType, 1, 64)) {
|
|
return rewriter.notifyMatchFailure(
|
|
castOp, "avoid introducing float to bool/i64 casts which are not "
|
|
"supported in all versions of TOSA");
|
|
}
|
|
|
|
// Check that the cast we're considering for removal is non-narrowing
|
|
if (isNarrowingCast(innerInputType, innerOutputType))
|
|
return rewriter.notifyMatchFailure(castOp,
|
|
"inner cast operation is narrowing");
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::CastOp>(castOp, outerOutputType,
|
|
innerCastInput);
|
|
|
|
return success();
|
|
}
|
|
|
|
bool supportsNaN(const llvm::fltSemantics &semantics) const {
|
|
return semantics.nonFiniteBehavior !=
|
|
llvm::fltNonfiniteBehavior::FiniteOnly;
|
|
}
|
|
|
|
bool supportsInf(const llvm::fltSemantics &semantics) const {
|
|
return semantics.nonFiniteBehavior == llvm::fltNonfiniteBehavior::IEEE754;
|
|
}
|
|
|
|
bool isNarrowingCast(const ShapedType inType,
|
|
const ShapedType outType) const {
|
|
|
|
if (inType.getElementType().isInteger() &&
|
|
outType.getElementType().isInteger()) {
|
|
|
|
const auto inTypeSignedness =
|
|
cast<IntegerType>(inType.getElementType()).getSignedness();
|
|
const auto outTypeSignedness =
|
|
cast<IntegerType>(outType.getElementType()).getSignedness();
|
|
|
|
return (inTypeSignedness != outTypeSignedness ||
|
|
inType.getElementTypeBitWidth() >
|
|
outType.getElementTypeBitWidth());
|
|
}
|
|
|
|
if (inType.getElementType().isFloat() &&
|
|
outType.getElementType().isFloat()) {
|
|
|
|
FloatType inElemTy = cast<FloatType>(inType.getElementType());
|
|
FloatType outElemTy = cast<FloatType>(outType.getElementType());
|
|
llvm::fltSemantics inTypeSemantics = inElemTy.getFloatSemantics();
|
|
llvm::fltSemantics outTypeSemantics = outElemTy.getFloatSemantics();
|
|
|
|
// If the list of supported types needs to be updated in the future, the
|
|
// check down below will need to be revised, for example to account for
|
|
// unsigned floating point types, or types that use negative zero as the
|
|
// representation for NaN.
|
|
[[maybe_unused]] const auto isSupported = [](Type elemType) {
|
|
return llvm::isa<Float8E4M3FNType, Float8E5M2Type, BFloat16Type,
|
|
Float16Type, Float32Type>(elemType);
|
|
};
|
|
|
|
assert(isSupported(inElemTy) &&
|
|
"unsupported input element type in isNarrowingCast");
|
|
assert(isSupported(outElemTy) &&
|
|
"unsupported output element type in isNarrowingCast");
|
|
|
|
return (
|
|
inTypeSemantics.maxExponent > outTypeSemantics.maxExponent ||
|
|
inTypeSemantics.minExponent < outTypeSemantics.minExponent ||
|
|
inTypeSemantics.precision > outTypeSemantics.precision ||
|
|
(supportsNaN(inTypeSemantics) && !supportsNaN(outTypeSemantics)) ||
|
|
(supportsInf(inTypeSemantics) && !supportsInf(outTypeSemantics)));
|
|
}
|
|
|
|
// While some cases of int -> float casts can be non-narrowing, consider
|
|
// them narrowing for the purposes of this optimization
|
|
return true;
|
|
}
|
|
};
|
|
|
|
void CastOp::getCanonicalizationPatterns(RewritePatternSet &results,
|
|
MLIRContext *context) {
|
|
results.add<NonNarrowingCastsOptimization>(context);
|
|
}
|
|
|
|
struct CancellingBlockScaledCastsOptimization
|
|
: public OpRewritePattern<tosa::CastToBlockScaledOp> {
|
|
using OpRewritePattern<tosa::CastToBlockScaledOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(tosa::CastToBlockScaledOp castToBlockScaledOp,
|
|
PatternRewriter &rewriter) const override {
|
|
const Value castToBlockScaledInput = castToBlockScaledOp.getInputData();
|
|
auto castFromBlockScaledOp =
|
|
castToBlockScaledInput.getDefiningOp<tosa::CastFromBlockScaledOp>();
|
|
if (!castFromBlockScaledOp)
|
|
return rewriter.notifyMatchFailure(
|
|
castToBlockScaledOp,
|
|
"input must be cast_from_block_scaled operation");
|
|
|
|
const Value innerData = castFromBlockScaledOp.getInputData();
|
|
const Value innerScale = castFromBlockScaledOp.getInputScale();
|
|
const auto innerDataTy = llvm::cast<ShapedType>(innerData.getType());
|
|
const auto innerScaleTy = llvm::cast<ShapedType>(innerScale.getType());
|
|
|
|
const Value outerData = castToBlockScaledOp.getOutputData();
|
|
const Value outerScale = castToBlockScaledOp.getOutputScale();
|
|
const auto outerDataTy = llvm::cast<ShapedType>(outerData.getType());
|
|
const auto outerScaleTy = llvm::cast<ShapedType>(outerScale.getType());
|
|
|
|
if (innerDataTy != outerDataTy || innerScaleTy != outerScaleTy) {
|
|
return rewriter.notifyMatchFailure(
|
|
castToBlockScaledOp,
|
|
"inputs types to cast_from_block_scaled operation must match output "
|
|
"types to cast_to_block_scaled");
|
|
}
|
|
|
|
if (castFromBlockScaledOp.getBlockSize() !=
|
|
castToBlockScaledOp.getBlockSize()) {
|
|
return rewriter.notifyMatchFailure(
|
|
castToBlockScaledOp, "block sizes for cast_from_block_scaled and "
|
|
"cast_to_block_scaled must match");
|
|
}
|
|
|
|
rewriter.replaceOp(castToBlockScaledOp, {innerData, innerScale});
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
void CastToBlockScaledOp::getCanonicalizationPatterns(
|
|
RewritePatternSet &results, MLIRContext *context) {
|
|
results.add<CancellingBlockScaledCastsOptimization>(context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Operator Folders.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
template <typename Folder>
|
|
static DenseElementsAttr
|
|
binaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs, ShapedType returnTy,
|
|
bool foldDenseValues = false) {
|
|
if (!lhs || !rhs)
|
|
return {};
|
|
|
|
if (!returnTy.hasRank() || !returnTy.hasStaticShape())
|
|
return {};
|
|
|
|
const auto lETy = llvm::cast<ShapedType>(lhs.getType()).getElementType();
|
|
const auto rETy = llvm::cast<ShapedType>(rhs.getType()).getElementType();
|
|
if (lETy != rETy)
|
|
return {};
|
|
|
|
if (lhs.isSplat() && rhs.isSplat()) {
|
|
if (isa<FloatType>(lETy)) {
|
|
const APFloat l = lhs.getSplatValue<APFloat>();
|
|
const APFloat r = rhs.getSplatValue<APFloat>();
|
|
const auto maybeResult = Folder::fold(l, r);
|
|
if (failed(maybeResult))
|
|
return {};
|
|
return DenseElementsAttr::get(returnTy, maybeResult.value());
|
|
}
|
|
|
|
if (const auto lIntTy = llvm::dyn_cast<IntegerType>(lETy)) {
|
|
const APInt l = lhs.getSplatValue<APInt>();
|
|
const APInt r = rhs.getSplatValue<APInt>();
|
|
const auto maybeResult = Folder::fold(l, r, lIntTy.isUnsigned());
|
|
if (failed(maybeResult))
|
|
return {};
|
|
return DenseElementsAttr::get(returnTy, maybeResult.value());
|
|
}
|
|
}
|
|
|
|
if (foldDenseValues) {
|
|
assert(lETy.isIntOrIndex() &&
|
|
"Only integer types are currently supported.");
|
|
SmallVector<APInt> resultValues;
|
|
for (auto [l, r] :
|
|
llvm::zip(lhs.getValues<APInt>(), rhs.getValues<APInt>())) {
|
|
const auto maybeResult = Folder::fold(l, r, false);
|
|
if (failed(maybeResult))
|
|
return {};
|
|
resultValues.push_back(maybeResult.value());
|
|
}
|
|
return DenseElementsAttr::get(returnTy, resultValues);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
template <typename Folder>
|
|
static DenseElementsAttr unaryFolder(DenseElementsAttr val, ShapedType returnTy,
|
|
bool foldDenseValues = false) {
|
|
if (!val)
|
|
return {};
|
|
|
|
if (!returnTy.hasRank() || !returnTy.hasStaticShape())
|
|
return {};
|
|
|
|
const auto vETy = llvm::cast<ShapedType>(val.getType()).getElementType();
|
|
|
|
if (val.isSplat()) {
|
|
if (const auto vIntTy = llvm::dyn_cast<IntegerType>(vETy)) {
|
|
const APInt v = val.getSplatValue<APInt>();
|
|
const auto maybeResult = Folder::fold(v, vIntTy.isUnsigned());
|
|
if (failed(maybeResult))
|
|
return {};
|
|
return DenseElementsAttr::get(returnTy, maybeResult.value());
|
|
}
|
|
}
|
|
|
|
if (foldDenseValues) {
|
|
mlir::Type elemTy = val.getElementType();
|
|
if (elemTy.isIntOrIndex()) {
|
|
SmallVector<APInt> resultValues;
|
|
for (auto const &v : val.getValues<APInt>()) {
|
|
const auto maybeResult = Folder::fold(v, false);
|
|
if (failed(maybeResult))
|
|
return {};
|
|
resultValues.push_back(maybeResult.value());
|
|
}
|
|
return DenseElementsAttr::get(returnTy, resultValues);
|
|
}
|
|
}
|
|
|
|
// Folding arbitrarily sized tensor operations is not supported
|
|
return {};
|
|
}
|
|
|
|
static FailureOr<int64_t> getSingleI64From1ElementTensor(Value v) {
|
|
DenseIntElementsAttr dense{};
|
|
if (!matchPattern(v, m_Constant(&dense)))
|
|
return failure();
|
|
|
|
assert(dense.isSplat());
|
|
APInt a = dense.getSplatValue<APInt>();
|
|
return a.getSExtValue();
|
|
}
|
|
|
|
struct AddFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
const bool isUnsigned) {
|
|
bool overflow;
|
|
const APInt result =
|
|
isUnsigned ? lhs.uadd_ov(rhs, overflow) : lhs.sadd_ov(rhs, overflow);
|
|
if (overflow)
|
|
return failure();
|
|
return result;
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return lhs + rhs;
|
|
}
|
|
};
|
|
|
|
struct SubFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
const bool isUnsigned) {
|
|
bool overflow;
|
|
const APInt result =
|
|
isUnsigned ? lhs.usub_ov(rhs, overflow) : lhs.ssub_ov(rhs, overflow);
|
|
if (overflow)
|
|
return failure();
|
|
return result;
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return lhs - rhs;
|
|
}
|
|
};
|
|
|
|
struct MulFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
const bool isUnsigned) {
|
|
|
|
const unsigned originalWidth = lhs.getBitWidth();
|
|
|
|
// Check same type
|
|
if (lhs.getBitWidth() != rhs.getBitWidth()) {
|
|
return failure();
|
|
}
|
|
|
|
// If either is `0`
|
|
if (lhs == 0 || rhs == 0)
|
|
return APInt::getZero(originalWidth);
|
|
|
|
bool overflow = false;
|
|
APInt const result =
|
|
isUnsigned ? lhs.umul_ov(rhs, overflow) : lhs.smul_ov(rhs, overflow);
|
|
|
|
if (overflow)
|
|
return failure();
|
|
|
|
return result.trunc(originalWidth);
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return lhs * rhs;
|
|
}
|
|
};
|
|
|
|
static bool signsDiffer(const APInt &a, const APInt &b) {
|
|
return a.isNegative() != b.isNegative();
|
|
}
|
|
|
|
template <bool Ceil>
|
|
struct DivFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
bool isUnsigned) {
|
|
if (lhs.getBitWidth() != rhs.getBitWidth())
|
|
return failure();
|
|
if (rhs.isZero())
|
|
return failure();
|
|
|
|
if (isUnsigned) {
|
|
APInt q{};
|
|
APInt r{};
|
|
APInt::udivrem(lhs, rhs, q, r);
|
|
if (!r.isZero() && Ceil) {
|
|
return q + 1;
|
|
}
|
|
return q;
|
|
}
|
|
|
|
// Signed: start from trunc-toward-zero, then adjust to ceil.
|
|
bool overflow{false};
|
|
APInt const q = lhs.sdiv_ov(rhs, overflow);
|
|
if (overflow)
|
|
return failure();
|
|
APInt const r = lhs.srem(rhs);
|
|
|
|
if (Ceil && !r.isZero() && !signsDiffer(lhs, rhs)) {
|
|
// Same sign => exact quotient is positive; trunc is below ceil =>
|
|
// increment q.
|
|
return q + 1;
|
|
}
|
|
return q;
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return lhs / rhs;
|
|
}
|
|
};
|
|
|
|
struct ModFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
bool isUnsigned) {
|
|
if (lhs.getBitWidth() != rhs.getBitWidth())
|
|
return failure();
|
|
if (lhs.isNegative() || (!rhs.isStrictlyPositive()))
|
|
return failure();
|
|
|
|
if (isUnsigned) {
|
|
return lhs.urem(rhs);
|
|
}
|
|
|
|
return lhs.srem(rhs);
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
auto t = lhs;
|
|
auto const r = t.mod(rhs);
|
|
if (llvm::APFloatBase::opStatus::opOK == r) {
|
|
return t;
|
|
}
|
|
return failure();
|
|
}
|
|
};
|
|
|
|
struct MaxFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
bool isUnsigned) {
|
|
if (lhs.getBitWidth() != rhs.getBitWidth())
|
|
return failure();
|
|
return lhs.getSExtValue() >= rhs.getSExtValue() ? lhs : rhs;
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return lhs >= rhs ? lhs : rhs;
|
|
}
|
|
};
|
|
|
|
struct MinFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
bool isUnsigned) {
|
|
if (lhs.getBitWidth() != rhs.getBitWidth())
|
|
return failure();
|
|
return lhs.getSExtValue() <= rhs.getSExtValue() ? lhs : rhs;
|
|
}
|
|
|
|
static FailureOr<APFloat> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return lhs <= rhs ? lhs : rhs;
|
|
}
|
|
};
|
|
|
|
struct Exp2FoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &value, bool isUnsigned) {
|
|
auto const numBits = value.getBitWidth();
|
|
if (isUnsigned) {
|
|
auto const zextv = value.getZExtValue();
|
|
if (zextv >= numBits)
|
|
return failure();
|
|
return APInt::getOneBitSet(numBits, zextv);
|
|
}
|
|
auto const sextv = value.getSExtValue();
|
|
if (sextv < 0 || sextv >= numBits || (value.isNegative()))
|
|
return failure();
|
|
return APInt::getOneBitSet(numBits, sextv);
|
|
}
|
|
};
|
|
|
|
struct Log2CeilFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &value, bool isUnsigned) {
|
|
if (!value.isStrictlyPositive())
|
|
return failure();
|
|
return APInt(/*numBits=*/value.getBitWidth(), value.ceilLogBase2());
|
|
}
|
|
};
|
|
|
|
struct Log2FloorFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &value, bool isUnsigned) {
|
|
if (!value.isStrictlyPositive())
|
|
return failure();
|
|
return APInt(/*numBits=*/value.getBitWidth(), value.logBase2());
|
|
}
|
|
};
|
|
|
|
struct GreaterFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
const bool isUnsigned) {
|
|
return isUnsigned ? APInt(1, lhs.ugt(rhs)) : APInt(1, lhs.sgt(rhs));
|
|
}
|
|
|
|
static FailureOr<APInt> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return APInt(1, lhs > rhs);
|
|
}
|
|
};
|
|
|
|
struct GreaterEqualFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
const bool isUnsigned) {
|
|
return isUnsigned ? APInt(1, lhs.uge(rhs)) : APInt(1, lhs.sge(rhs));
|
|
}
|
|
|
|
static FailureOr<APInt> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return APInt(1, lhs >= rhs);
|
|
}
|
|
};
|
|
|
|
struct EqualFoldAdaptor {
|
|
static FailureOr<APInt> fold(const APInt &lhs, const APInt &rhs,
|
|
const bool isUnsigned) {
|
|
return APInt(1, lhs == rhs);
|
|
}
|
|
|
|
static FailureOr<APInt> fold(const APFloat &lhs, const APFloat &rhs) {
|
|
return APInt(1, lhs == rhs);
|
|
}
|
|
};
|
|
|
|
static bool isSplatZero(Type elemType, DenseElementsAttr val) {
|
|
if (llvm::isa<FloatType>(elemType))
|
|
return val && val.isSplat() && val.getSplatValue<APFloat>().isZero();
|
|
if (llvm::isa<IntegerType>(elemType))
|
|
return val && val.isSplat() && val.getSplatValue<APInt>().isZero();
|
|
return false;
|
|
}
|
|
|
|
static bool isSplatOne(Type elemType, DenseElementsAttr val, int64_t shift) {
|
|
if (llvm::isa<FloatType>(elemType))
|
|
return val && val.isSplat() &&
|
|
val.getSplatValue<APFloat>().isExactlyValue(1.0);
|
|
if (llvm::isa<IntegerType>(elemType)) {
|
|
const int64_t shifted = 1LL << shift;
|
|
return val && val.isSplat() &&
|
|
val.getSplatValue<APInt>().getSExtValue() == shifted;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
|
|
auto lhsTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
|
|
auto rhsTy = llvm::dyn_cast<RankedTensorType>(getInput2().getType());
|
|
auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
if (!lhsTy || !rhsTy || !resultTy)
|
|
return {};
|
|
|
|
// Cannot create an ElementsAttr from non-int/float/index types
|
|
if (!lhsTy.getElementType().isIntOrIndexOrFloat() ||
|
|
!rhsTy.getElementType().isIntOrIndexOrFloat())
|
|
return {};
|
|
|
|
auto resultETy = resultTy.getElementType();
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
|
|
const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
|
|
lhsTy.getShape(), rhsTy.getShape());
|
|
if (isBroadcastable && lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
|
|
return getInput1();
|
|
if (isBroadcastable && rhsTy == resultTy && isSplatZero(resultETy, lhsAttr))
|
|
return getInput2();
|
|
|
|
if (!lhsAttr || !rhsAttr)
|
|
return {};
|
|
|
|
return binaryFolder<AddFoldAdaptor>(lhsAttr, rhsAttr, resultTy);
|
|
}
|
|
|
|
OpFoldResult ArgMaxOp::fold(FoldAdaptor adaptor) {
|
|
auto inputTy = llvm::dyn_cast<RankedTensorType>(getInput().getType());
|
|
auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
if (!inputTy || !outputTy || !inputTy.hasStaticShape() ||
|
|
!outputTy.hasStaticShape())
|
|
return {};
|
|
|
|
const Type outputElementTy = getElementTypeOrSelf(outputTy);
|
|
if (inputTy.getDimSize(getAxis()) == 1 && outputElementTy.isInteger()) {
|
|
const auto outputElemIntTy = cast<IntegerType>(outputElementTy);
|
|
const APInt zero = APInt::getZero(outputElemIntTy.getWidth());
|
|
return DenseElementsAttr::get(outputTy, zero);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult IntDivOp::fold(FoldAdaptor adaptor) {
|
|
auto lhsTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
|
|
auto rhsTy = llvm::dyn_cast<RankedTensorType>(getInput2().getType());
|
|
auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
if (!lhsTy || !rhsTy || !resultTy)
|
|
return {};
|
|
if (lhsTy.getElementType() != rhsTy.getElementType())
|
|
return {};
|
|
|
|
// IntDivOp inputs must be integer type, no need to check for quantized
|
|
// type
|
|
auto resultETy = resultTy.getElementType();
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
if (lhsAttr && lhsAttr.isSplat()) {
|
|
if (llvm::isa<IntegerType>(resultETy) && resultTy.hasStaticShape() &&
|
|
lhsAttr.getSplatValue<APInt>().isZero())
|
|
return lhsAttr.resizeSplat(resultTy);
|
|
}
|
|
|
|
if (rhsAttr && rhsAttr.isSplat()) {
|
|
const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
|
|
lhsTy.getShape(), rhsTy.getShape());
|
|
if (isBroadcastable && lhsTy == resultTy &&
|
|
llvm::isa<IntegerType>(resultETy) &&
|
|
rhsAttr.getSplatValue<APInt>().isOne())
|
|
return getInput1();
|
|
}
|
|
|
|
if (rhsAttr && lhsAttr && rhsAttr.isSplat() && lhsAttr.isSplat() &&
|
|
llvm::isa<IntegerType>(resultETy)) {
|
|
APInt l = lhsAttr.getSplatValue<APInt>();
|
|
APInt r = rhsAttr.getSplatValue<APInt>();
|
|
if (!r.isZero()) {
|
|
auto intTy = dyn_cast<mlir::IntegerType>(resultETy);
|
|
auto const result =
|
|
DivFoldAdaptor</*Ceil*/ false>::fold(l, r, intTy.isUnsigned());
|
|
if (failed(result))
|
|
return {};
|
|
return DenseElementsAttr::get(resultTy, result.value());
|
|
}
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
namespace {
|
|
// calculate lhs * rhs >> shift according to TOSA Spec
|
|
// return nullopt if result is not in range of int32_t when shift > 0
|
|
std::optional<APInt> mulInt(APInt lhs, APInt rhs, int32_t shift,
|
|
unsigned bitwidth) {
|
|
bool overflow = false;
|
|
APInt result = lhs.sext(64).smul_ov(rhs.sext(64), overflow);
|
|
|
|
if (overflow)
|
|
return std::nullopt;
|
|
|
|
if (shift > 0) {
|
|
auto round = APInt(64, 1) << (shift - 1);
|
|
result += round;
|
|
result.ashrInPlace(shift);
|
|
// REQUIRE(product >= minimum_s<i32_t>() && product <=
|
|
// maximum_s<i32_t>())
|
|
if (!(result.getSExtValue() >= INT32_MIN &&
|
|
result.getSExtValue() <= INT32_MAX)) {
|
|
// REQUIRE failed
|
|
return std::nullopt;
|
|
}
|
|
}
|
|
|
|
return result.trunc(bitwidth);
|
|
}
|
|
|
|
DenseElementsAttr mulBinaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs,
|
|
RankedTensorType ty, int32_t shift) {
|
|
if (rhs && lhs && rhs.isSplat() && lhs.isSplat()) {
|
|
if (llvm::isa<IntegerType>(ty.getElementType())) {
|
|
APInt l = lhs.getSplatValue<APInt>();
|
|
APInt r = rhs.getSplatValue<APInt>();
|
|
|
|
if (shift == 0) {
|
|
return DenseElementsAttr::get(ty, l * r);
|
|
}
|
|
|
|
auto bitwidth = ty.getElementType().getIntOrFloatBitWidth();
|
|
const std::optional<APInt> result = mulInt(l, r, shift, bitwidth);
|
|
if (!result)
|
|
return {};
|
|
return DenseElementsAttr::get(ty, result.value());
|
|
}
|
|
|
|
if (llvm::isa<FloatType>(ty.getElementType())) {
|
|
APFloat l = lhs.getSplatValue<APFloat>();
|
|
APFloat r = rhs.getSplatValue<APFloat>();
|
|
APFloat result = l * r;
|
|
return DenseElementsAttr::get(ty, result);
|
|
}
|
|
}
|
|
|
|
return {};
|
|
}
|
|
} // namespace
|
|
|
|
OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
|
|
auto lhs = getInput1();
|
|
auto rhs = getInput2();
|
|
auto lhsTy = llvm::dyn_cast<RankedTensorType>(lhs.getType());
|
|
auto rhsTy = llvm::dyn_cast<RankedTensorType>(rhs.getType());
|
|
auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
if (!lhsTy || !rhsTy || !resultTy)
|
|
return {};
|
|
|
|
auto resultETy = resultTy.getElementType();
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
|
|
// Result right shift on i32_t data type only. For simplification,
|
|
// synthesize a zero shift for other data type.
|
|
int32_t shift = 0;
|
|
if (resultETy.isInteger(32)) {
|
|
ElementsAttr shift_elem;
|
|
if (getShift().getImpl()) {
|
|
if (!matchPattern(getShift(), m_Constant(&shift_elem)))
|
|
// cannot be folded when the shift value is unknown.
|
|
return {};
|
|
shift = shift_elem.getValues<IntegerAttr>()[0].getInt();
|
|
}
|
|
}
|
|
|
|
if (rhsTy == resultTy && isSplatZero(resultETy, lhsAttr) &&
|
|
resultTy.hasStaticShape())
|
|
// constant values can only be resized if resulting type is static
|
|
return lhsAttr.resizeSplat(resultTy);
|
|
if (lhsTy == resultTy && isSplatZero(resultETy, rhsAttr) &&
|
|
resultTy.hasStaticShape())
|
|
return rhsAttr.resizeSplat(resultTy);
|
|
|
|
const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
|
|
lhsTy.getShape(), rhsTy.getShape());
|
|
if (isBroadcastable && rhsTy == resultTy &&
|
|
isSplatOne(resultETy, lhsAttr, shift))
|
|
return rhs;
|
|
if (isBroadcastable && lhsTy == resultTy &&
|
|
isSplatOne(resultETy, rhsAttr, shift))
|
|
return lhs;
|
|
|
|
return mulBinaryFolder(lhsAttr, rhsAttr, resultTy, shift);
|
|
}
|
|
|
|
OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
|
|
auto lhsTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
|
|
auto rhsTy = llvm::dyn_cast<RankedTensorType>(getInput2().getType());
|
|
auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
if (!lhsTy || !rhsTy || !resultTy)
|
|
return {};
|
|
|
|
// Cannot create an ElementsAttr from non-int/float/index types
|
|
if (!lhsTy.getElementType().isIntOrIndexOrFloat() ||
|
|
!rhsTy.getElementType().isIntOrIndexOrFloat())
|
|
return {};
|
|
|
|
auto resultETy = resultTy.getElementType();
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
|
|
const bool isBroadcastable = OpTrait::util::staticallyKnownBroadcastable(
|
|
lhsTy.getShape(), rhsTy.getShape());
|
|
if (isBroadcastable && lhsTy == resultTy && isSplatZero(resultETy, rhsAttr))
|
|
return getInput1();
|
|
|
|
if (!lhsAttr || !rhsAttr)
|
|
return {};
|
|
|
|
return binaryFolder<SubFoldAdaptor>(lhsAttr, rhsAttr, resultTy);
|
|
}
|
|
|
|
OpFoldResult GreaterOp::fold(FoldAdaptor adaptor) {
|
|
auto resultTy = llvm::cast<ShapedType>(getType());
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
|
|
if (!lhsAttr || !rhsAttr)
|
|
return {};
|
|
|
|
return binaryFolder<GreaterFoldAdaptor>(lhsAttr, rhsAttr, resultTy);
|
|
}
|
|
|
|
OpFoldResult GreaterEqualOp::fold(FoldAdaptor adaptor) {
|
|
auto resultTy = llvm::cast<ShapedType>(getType());
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
|
|
if (!lhsAttr || !rhsAttr)
|
|
return {};
|
|
|
|
return binaryFolder<GreaterEqualFoldAdaptor>(lhsAttr, rhsAttr, resultTy);
|
|
}
|
|
|
|
OpFoldResult EqualOp::fold(FoldAdaptor adaptor) {
|
|
auto resultTy = llvm::cast<ShapedType>(getType());
|
|
auto lhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
|
|
auto rhsAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput2());
|
|
Value lhs = getInput1();
|
|
Value rhs = getInput2();
|
|
auto lhsTy = llvm::cast<ShapedType>(lhs.getType());
|
|
|
|
// If we are comparing an integer value to itself it is always true. We
|
|
// can not do this with float due to float values.
|
|
if (llvm::isa<IntegerType>(lhsTy.getElementType()) && resultTy.hasRank() &&
|
|
resultTy.hasStaticShape() && lhs == rhs) {
|
|
return DenseElementsAttr::get(resultTy, true);
|
|
}
|
|
|
|
if (!lhsAttr || !rhsAttr)
|
|
return {};
|
|
|
|
return binaryFolder<EqualFoldAdaptor>(lhsAttr, rhsAttr, resultTy);
|
|
}
|
|
|
|
OpFoldResult CastOp::fold(FoldAdaptor adaptor) {
|
|
if (getInput().getType() == getType())
|
|
return getInput();
|
|
|
|
auto operand = llvm::dyn_cast_if_present<ElementsAttr>(adaptor.getInput());
|
|
if (!operand)
|
|
return {};
|
|
|
|
auto inTy = llvm::cast<ShapedType>(getInput().getType());
|
|
auto outTy = llvm::cast<ShapedType>(getType());
|
|
if (!outTy.hasRank() || !outTy.hasStaticShape())
|
|
return {};
|
|
auto inETy = inTy.getElementType();
|
|
auto outETy = outTy.getElementType();
|
|
|
|
if (operand.isSplat()) {
|
|
if (llvm::isa<FloatType>(inETy) && llvm::isa<FloatType>(outETy)) {
|
|
bool overflow;
|
|
auto splatVal = operand.getSplatValue<APFloat>();
|
|
auto &semantics = llvm::cast<FloatType>(outETy).getFloatSemantics();
|
|
splatVal.convert(semantics, llvm::RoundingMode::NearestTiesToEven,
|
|
&overflow);
|
|
return SplatElementsAttr::get(outTy, splatVal);
|
|
}
|
|
|
|
if (llvm::isa<IntegerType>(inETy) && llvm::isa<FloatType>(outETy)) {
|
|
auto unsign = llvm::cast<IntegerType>(inETy).isUnsignedInteger();
|
|
APFloat splatVal(llvm::cast<FloatType>(outETy).getFloatSemantics());
|
|
splatVal.convertFromAPInt(operand.getSplatValue<APInt>(), !unsign,
|
|
llvm::RoundingMode::NearestTiesToEven);
|
|
return SplatElementsAttr::get(outTy, splatVal);
|
|
}
|
|
|
|
if (llvm::isa<FloatType>(inETy) && llvm::isa<IntegerType>(outETy)) {
|
|
auto unsign = llvm::cast<IntegerType>(outETy).isUnsignedInteger();
|
|
auto intVal = APSInt(
|
|
llvm::cast<IntegerType>(outETy).getIntOrFloatBitWidth(), unsign);
|
|
auto floatVal = operand.getSplatValue<APFloat>();
|
|
bool exact;
|
|
floatVal.convertToInteger(intVal, llvm::RoundingMode::NearestTiesToEven,
|
|
&exact);
|
|
return SplatElementsAttr::get(outTy, intVal);
|
|
}
|
|
|
|
if (llvm::isa<IntegerType>(inETy) && llvm::isa<IntegerType>(outETy)) {
|
|
const auto inIntType = llvm::cast<IntegerType>(inETy);
|
|
auto unsignIn = inIntType.isUnsignedInteger();
|
|
bool trunc =
|
|
inETy.getIntOrFloatBitWidth() > outETy.getIntOrFloatBitWidth();
|
|
auto intVal = operand.getSplatValue<APInt>();
|
|
auto bitwidth = outETy.getIntOrFloatBitWidth();
|
|
|
|
// i1 types are boolean in TOSA
|
|
if (outETy.isInteger(1)) {
|
|
intVal = APInt(bitwidth, intVal.isZero() ? 0 : 1);
|
|
} else if (trunc) {
|
|
intVal = intVal.trunc(bitwidth);
|
|
} else if (unsignIn || inIntType.isInteger(1)) {
|
|
intVal = intVal.zext(bitwidth);
|
|
} else {
|
|
intVal = intVal.sext(bitwidth);
|
|
}
|
|
|
|
return SplatElementsAttr::get(outTy, intVal);
|
|
}
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult ConstOp::fold(FoldAdaptor adaptor) { return getValuesAttr(); }
|
|
|
|
OpFoldResult ConstShapeOp::fold(FoldAdaptor adaptor) { return getValuesAttr(); }
|
|
|
|
#define REDUCE_FOLDER(OP) \
|
|
OpFoldResult OP::fold(FoldAdaptor adaptor) { \
|
|
ShapedType inputTy = llvm::cast<ShapedType>(getInput().getType()); \
|
|
if (!inputTy.hasRank()) \
|
|
return {}; \
|
|
if (inputTy != getType()) \
|
|
return {}; \
|
|
if (inputTy.getRank() == 0 || inputTy.getDimSize(getAxis()) == 1) \
|
|
return getInput(); \
|
|
return {}; \
|
|
}
|
|
|
|
REDUCE_FOLDER(ReduceAllOp)
|
|
REDUCE_FOLDER(ReduceAnyOp)
|
|
REDUCE_FOLDER(ReduceMaxOp)
|
|
REDUCE_FOLDER(ReduceMinOp)
|
|
REDUCE_FOLDER(ReduceProductOp)
|
|
REDUCE_FOLDER(ReduceSumOp)
|
|
#undef REDUCE_FOLDER
|
|
|
|
OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
|
|
auto inputTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
|
|
auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
|
|
if (!inputTy || !outputTy)
|
|
return {};
|
|
|
|
// Fold when the input and output types are the same. This is only safe
|
|
// when there is at most 1 dynamic dimension. For 2 or more dynamic
|
|
// dimensions, there may still be a productive reshape.
|
|
if (inputTy == outputTy && inputTy.getNumDynamicDims() < 2)
|
|
return getInput1();
|
|
|
|
// reshape(reshape(x)) -> reshape(x)
|
|
if (auto reshapeOp = llvm::dyn_cast_if_present<tosa::ReshapeOp>(
|
|
getInput1().getDefiningOp())) {
|
|
getInput1Mutable().assign(reshapeOp.getInput1());
|
|
return getResult();
|
|
}
|
|
|
|
// Cannot create an ElementsAttr from non-int/float/index types
|
|
if (!inputTy.getElementType().isIntOrIndexOrFloat())
|
|
return {};
|
|
|
|
// reshape(const(x)) -> const(reshape-attr(x))
|
|
if (auto operand =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1())) {
|
|
// Constants must have static shape.
|
|
if (!outputTy.hasStaticShape())
|
|
return {};
|
|
|
|
// Okay to duplicate splat constants.
|
|
if (operand.isSplat())
|
|
return SplatElementsAttr::get(outputTy,
|
|
operand.getSplatValue<Attribute>());
|
|
|
|
// Don't duplicate other constants.
|
|
if (!getInput1().hasOneUse())
|
|
return {};
|
|
|
|
llvm::SmallVector<int64_t> shapeVec;
|
|
if (!tosa::getConstShapeValues(getShape().getDefiningOp(), shapeVec))
|
|
return {};
|
|
|
|
return operand.reshape(
|
|
llvm::cast<ShapedType>(operand.getType()).clone(shapeVec));
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
|
|
// If the pad is all zeros we can fold this operation away.
|
|
if (adaptor.getPadding() && getInput1().getType() == getType()) {
|
|
auto densePad = llvm::dyn_cast<DenseElementsAttr>(adaptor.getPadding());
|
|
if (densePad && densePad.isSplat() &&
|
|
densePad.getSplatValue<APInt>().isZero()) {
|
|
return getInput1();
|
|
}
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
// Fold away cases where a tosa.resize operation returns a copy
|
|
// of the input image.
|
|
OpFoldResult ResizeOp::fold(FoldAdaptor adaptor) {
|
|
auto scaleAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getScale());
|
|
auto offsetAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getOffset());
|
|
auto borderAttr =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getBorder());
|
|
if (!scaleAttr || !offsetAttr || !borderAttr) {
|
|
return {};
|
|
}
|
|
|
|
auto scale = tosa::convertFromIntAttr(scaleAttr, /* rank = */ 4);
|
|
auto offset = tosa::convertFromIntAttr(offsetAttr, /* rank = */ 2);
|
|
auto border = tosa::convertFromIntAttr(borderAttr, /* rank = */ 2);
|
|
if (scale.size() != 4 || offset.size() != 2 || border.size() != 2) {
|
|
return {};
|
|
}
|
|
|
|
// Check unit scaling.
|
|
if (scale[0] != scale[1] || scale[2] != scale[3]) {
|
|
return {};
|
|
}
|
|
|
|
// There should be no offset.
|
|
if (offset[0] != 0 || offset[1] != 0) {
|
|
return {};
|
|
}
|
|
|
|
// There should be no border.
|
|
if (border[0] != 0 || border[1] != 0) {
|
|
return {};
|
|
}
|
|
|
|
return foldToInputIfTypeMatches(getType(), getInput());
|
|
}
|
|
|
|
OpFoldResult ReverseOp::fold(FoldAdaptor adaptor) {
|
|
auto operand = getInput1();
|
|
auto operandTy = llvm::cast<ShapedType>(operand.getType());
|
|
auto axis = getAxis();
|
|
// If the dim-length is 1, or reversing axis is unit-dim, also a no-op.
|
|
const bool isSplatInput =
|
|
llvm::isa_and_nonnull<SplatElementsAttr>(adaptor.getInput1());
|
|
if (!operandTy.hasRank() ||
|
|
(!isSplatInput && operandTy.getDimSize(axis) != 1))
|
|
return {};
|
|
return foldToInputIfTypeMatches(getType(), operand);
|
|
}
|
|
|
|
OpFoldResult SliceOp::fold(FoldAdaptor adaptor) {
|
|
auto inputTy = llvm::dyn_cast<RankedTensorType>(getInput1().getType());
|
|
auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
|
|
|
|
if (!inputTy || !outputTy)
|
|
return {};
|
|
|
|
if (inputTy == outputTy && inputTy.hasStaticShape())
|
|
return getInput1();
|
|
|
|
// Check if this is a no-op slice (starts at 0 and size matches input)
|
|
|
|
DenseElementsAttr startElems;
|
|
if (!matchPattern(getStart(), m_Constant(&startElems)))
|
|
return {};
|
|
|
|
// Check if all start values are zero
|
|
bool startIsZeros =
|
|
llvm::all_of(startElems.getValues<APInt>(),
|
|
[](const APInt &val) { return val.isZero(); });
|
|
|
|
if (startIsZeros) {
|
|
|
|
// Check if size matches input shape
|
|
DenseElementsAttr sizeElems;
|
|
if (!matchPattern(getSize(), m_Constant(&sizeElems)))
|
|
return {};
|
|
|
|
auto inputShape = inputTy.getShape();
|
|
auto sizeValues = sizeElems.getValues<APInt>();
|
|
|
|
bool sizeMatchesInput = true;
|
|
for (const auto &[i, sizeVal] : llvm::enumerate(sizeValues)) {
|
|
int64_t size = sizeVal.getSExtValue();
|
|
|
|
if (inputTy.isDynamicDim(i)) {
|
|
// For dynamic dimensions, check for kInferableDimSize indicating full
|
|
// dimension is sliced
|
|
if (size != kInferableDimSize) {
|
|
sizeMatchesInput = false;
|
|
break;
|
|
}
|
|
} else {
|
|
// For static dimensions, check that size must match exactly or be
|
|
// kInferableDimSize indicating full dimension is sliced
|
|
if (size != kInferableDimSize && size != inputShape[i]) {
|
|
sizeMatchesInput = false;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (sizeMatchesInput)
|
|
return getInput1();
|
|
}
|
|
|
|
// The following checks require the input to be a constant
|
|
if (!adaptor.getInput1())
|
|
return {};
|
|
|
|
// Cannot create an ElementsAttr from non-int/float/index types
|
|
if (!inputTy.getElementType().isIntOrIndexOrFloat() ||
|
|
!outputTy.getElementType().isIntOrIndexOrFloat())
|
|
return {};
|
|
|
|
auto operand = llvm::cast<ElementsAttr>(adaptor.getInput1());
|
|
if (operand.isSplat() && outputTy.hasStaticShape()) {
|
|
return SplatElementsAttr::get(outputTy, operand.getSplatValue<Attribute>());
|
|
}
|
|
|
|
if (inputTy.hasStaticShape() && outputTy.hasStaticShape() &&
|
|
outputTy.getNumElements() == 1) {
|
|
llvm::SmallVector<uint64_t> indices =
|
|
llvm::to_vector(startElems.getValues<uint64_t>());
|
|
if (auto values = operand.tryGetValues<Attribute>())
|
|
return SplatElementsAttr::get(outputTy, (*values)[indices]);
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult tosa::SelectOp::fold(FoldAdaptor adaptor) {
|
|
const Value pred = getPred();
|
|
const Value onTrue = getOnTrue();
|
|
const Value onFalse = getOnFalse();
|
|
|
|
const auto predTy = llvm::dyn_cast<RankedTensorType>(pred.getType());
|
|
const auto onTrueTy = llvm::dyn_cast<RankedTensorType>(onTrue.getType());
|
|
const auto onFalseTy = llvm::dyn_cast<RankedTensorType>(onFalse.getType());
|
|
if (!predTy || !onTrueTy || !onFalseTy)
|
|
return {};
|
|
|
|
const Type resultTy = getType();
|
|
|
|
const ArrayRef<int64_t> predShape = predTy.getShape();
|
|
const ArrayRef<int64_t> onTrueShape = onTrueTy.getShape();
|
|
|
|
if (onTrue == onFalse && onTrueTy == resultTy &&
|
|
OpTrait::util::staticallyKnownBroadcastable(predShape, onTrueShape))
|
|
return onTrue;
|
|
|
|
auto predicate =
|
|
llvm::dyn_cast_if_present<DenseIntElementsAttr>(adaptor.getInput1());
|
|
if (!predicate)
|
|
return {};
|
|
if (!predicate.isSplat())
|
|
return {};
|
|
|
|
const bool predicateValue = predicate.getSplatValue<APInt>().getBoolValue();
|
|
|
|
SmallVector<SmallVector<int64_t>, 3> shapes;
|
|
shapes.emplace_back(predShape);
|
|
shapes.emplace_back(onTrueShape);
|
|
shapes.emplace_back(onFalseTy.getShape());
|
|
const bool isBroadcastable =
|
|
OpTrait::util::staticallyKnownBroadcastable(shapes);
|
|
|
|
if (predicateValue == true && onTrueTy == resultTy && isBroadcastable)
|
|
return onTrue;
|
|
if (predicateValue == false && onFalseTy == resultTy && isBroadcastable)
|
|
return onFalse;
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult TileOp::fold(FoldAdaptor adaptor) {
|
|
if (getInput1().getType() == getType()) {
|
|
if (auto multiples = llvm::dyn_cast_if_present<DenseElementsAttr>(
|
|
adaptor.getMultiples())) {
|
|
if (multiples.isSplat() &&
|
|
multiples.getSplatValue<APInt>().getSExtValue() == 1)
|
|
return getInput1();
|
|
if (auto int_array_attr =
|
|
llvm::dyn_cast<DenseIntElementsAttr>(multiples)) {
|
|
if (llvm::all_of(int_array_attr.getValues<APInt>(),
|
|
[](APInt v) { return v.getSExtValue() == 1; }))
|
|
return getInput1();
|
|
}
|
|
}
|
|
}
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
|
|
auto resultTy = llvm::cast<ShapedType>(getType());
|
|
|
|
// Transposing splat values just means reshaping.
|
|
if (auto input =
|
|
llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1())) {
|
|
if (input.isSplat() && resultTy.hasRank() && resultTy.hasStaticShape() &&
|
|
input.getType().getElementType() == resultTy.getElementType())
|
|
return input.reshape(resultTy);
|
|
}
|
|
|
|
// Transpose is not the identity transpose.
|
|
const llvm::ArrayRef<int32_t> perms = getPerms();
|
|
|
|
if (!llvm::equal(llvm::seq<int32_t>(0, perms.size()), perms))
|
|
return {};
|
|
|
|
return foldToInputIfTypeMatches(getType(), getInput1());
|
|
}
|
|
|
|
OpFoldResult tosa::NegateOp::fold(FoldAdaptor adaptor) {
|
|
// Element-wise negate(negate(x)) = x
|
|
// iff all zero points are constant 0
|
|
auto definingOp = getInput1().getDefiningOp<tosa::NegateOp>();
|
|
if (!definingOp) {
|
|
// defining op of input1 is not a negate, cannot fold
|
|
return {};
|
|
}
|
|
|
|
if (FailureOr<int64_t> maybeIZp = getInput1ZeroPoint();
|
|
failed(maybeIZp) || *maybeIZp != 0) {
|
|
// input1 zero point is not constant 0, cannot fold
|
|
return {};
|
|
}
|
|
if (FailureOr<int64_t> maybeOZp = getOutputZeroPoint();
|
|
failed(maybeOZp) || *maybeOZp != 0) {
|
|
// output zero point is not constant 0, cannot fold
|
|
return {};
|
|
}
|
|
if (FailureOr<int64_t> maybeIZp = definingOp.getInput1ZeroPoint();
|
|
failed(maybeIZp) || *maybeIZp != 0) {
|
|
// definingOp's input1 zero point is not constant 0, cannot fold
|
|
return {};
|
|
}
|
|
if (FailureOr<int64_t> maybeOZp = definingOp.getOutputZeroPoint();
|
|
failed(maybeOZp) || *maybeOZp != 0) {
|
|
// definingOp's output zero point is not constant 0, cannot fold
|
|
return {};
|
|
}
|
|
|
|
return foldToInputIfTypeMatches(getType(), definingOp.getInput1());
|
|
}
|
|
|
|
OpFoldResult tosa::AbsOp::fold(FoldAdaptor adaptor) {
|
|
auto input = getInput1();
|
|
// Element-wise abs(abs(x)) = abs(x)
|
|
if (input.getDefiningOp<tosa::AbsOp>())
|
|
return foldToInputIfTypeMatches(getType(), input);
|
|
|
|
return {};
|
|
}
|
|
|
|
OpFoldResult tosa::ReciprocalOp::fold(FoldAdaptor adaptor) {
|
|
auto input = adaptor.getInput1();
|
|
|
|
auto inputAttr = llvm::dyn_cast_if_present<DenseElementsAttr>(input);
|
|
// Fold splat inputs only.
|
|
if (!inputAttr || !inputAttr.isSplat())
|
|
return {};
|
|
|
|
auto shapeType = llvm::cast<ShapedType>(getType());
|
|
if (!shapeType.hasRank() || !shapeType.hasStaticShape())
|
|
return {};
|
|
if (auto floatType = llvm::dyn_cast<FloatType>(inputAttr.getElementType())) {
|
|
auto floatVal = inputAttr.getSplatValue<APFloat>();
|
|
return DenseElementsAttr::get(shapeType,
|
|
ReciprocalOp::calcOneElement(floatVal));
|
|
}
|
|
|
|
return {};
|
|
}
|
|
|
|
template <typename Op, typename OpFoldAdaptor>
|
|
OpFoldResult unaryShapeFold(Op *op) {
|
|
auto input1ConstShape =
|
|
dyn_cast<tosa::ConstShapeOp>(op->getInput().getDefiningOp());
|
|
if (!input1ConstShape)
|
|
return {};
|
|
|
|
const auto input1Attr = cast<DenseElementsAttr>(input1ConstShape.getValues());
|
|
|
|
return unaryFolder<OpFoldAdaptor>(input1Attr, input1Attr.getType(),
|
|
/*foldDenseValues=*/true);
|
|
}
|
|
|
|
template <typename Op, typename OpFoldAdaptor>
|
|
OpFoldResult binaryFold(Op *op) {
|
|
auto input1ConstShape =
|
|
dyn_cast<tosa::ConstShapeOp>(op->getInput1().getDefiningOp());
|
|
auto input2ConstShape =
|
|
dyn_cast<tosa::ConstShapeOp>(op->getInput2().getDefiningOp());
|
|
if (!input1ConstShape || !input2ConstShape)
|
|
return {};
|
|
|
|
const auto input1Attr = cast<DenseElementsAttr>(input1ConstShape.getValues());
|
|
const auto input2Attr = cast<DenseElementsAttr>(input2ConstShape.getValues());
|
|
|
|
return binaryFolder<OpFoldAdaptor>(input1Attr, input2Attr,
|
|
input1Attr.getType(),
|
|
/*foldDenseValues=*/true);
|
|
}
|
|
|
|
OpFoldResult tosa::DimOp::fold(FoldAdaptor adaptor) {
|
|
const auto inputTy = llvm::dyn_cast<ShapedType>(getInput1().getType());
|
|
if (!inputTy || !inputTy.hasRank())
|
|
return {};
|
|
const int32_t axis = getAxis();
|
|
const int64_t dimSize = inputTy.getDimSize(axis);
|
|
if (ShapedType::isDynamic(dimSize))
|
|
return {};
|
|
|
|
OpBuilder builder(getContext());
|
|
const auto resultAttrTy =
|
|
RankedTensorType::get(/*rank=*/1, builder.getIndexType());
|
|
return DenseElementsAttr::get(resultAttrTy, dimSize);
|
|
}
|
|
|
|
OpFoldResult concatShapeFold(tosa::ConcatShapeOp *op) {
|
|
auto const inputs = op->getInput();
|
|
|
|
if (inputs.empty())
|
|
return {};
|
|
|
|
SmallVector<APInt> concatDims;
|
|
concatDims.reserve(/*max elem*/ 64);
|
|
for (auto const &v : inputs) {
|
|
auto vConstShape = dyn_cast<tosa::ConstShapeOp>(v.getDefiningOp());
|
|
if (!vConstShape)
|
|
return {};
|
|
|
|
const auto vAttr = cast<DenseElementsAttr>(vConstShape.getValues());
|
|
assert(vAttr);
|
|
|
|
auto const vAttrVals = vAttr.getValues<APInt>();
|
|
for (auto const &v : vAttrVals) {
|
|
concatDims.push_back(v);
|
|
}
|
|
}
|
|
|
|
auto *ctx = op->getContext();
|
|
assert(ctx != nullptr && "ctx is nullptr");
|
|
auto const rankedTy = RankedTensorType::get(
|
|
{static_cast<int64_t>(concatDims.size())}, IndexType::get(ctx));
|
|
|
|
return DenseElementsAttr::get(rankedTy, concatDims);
|
|
}
|
|
|
|
OpFoldResult sliceShapeFold(tosa::SliceShapeOp *op) {
|
|
auto const input1 = op->getInput();
|
|
auto const input2 = op->getStart();
|
|
auto const input3 = op->getSize();
|
|
|
|
auto input1ConstShape = dyn_cast<tosa::ConstShapeOp>(input1.getDefiningOp());
|
|
|
|
if (!input1ConstShape)
|
|
return {};
|
|
|
|
auto const input1Attr = cast<DenseElementsAttr>(input1ConstShape.getValues());
|
|
if (!input1Attr)
|
|
return {};
|
|
|
|
auto const input1Vals = input1Attr.getValues<APInt>();
|
|
auto const totalInput1 = input1Vals.size();
|
|
|
|
auto const start = getSingleI64From1ElementTensor(input2);
|
|
auto const size = getSingleI64From1ElementTensor(input3);
|
|
|
|
if (failed(start) || failed(size))
|
|
return {};
|
|
|
|
auto const startV = static_cast<int32_t>(start.value());
|
|
auto const sizeV = static_cast<int32_t>(size.value());
|
|
|
|
if ((sizeV <= 0) || (startV < 0) ||
|
|
(static_cast<size_t>(startV + sizeV) > totalInput1))
|
|
return {};
|
|
|
|
SmallVector<APInt> sliceOfInput;
|
|
sliceOfInput.reserve(totalInput1);
|
|
|
|
for (auto i = startV; i < (startV + sizeV); i++) {
|
|
sliceOfInput.push_back(input1Vals[i]);
|
|
}
|
|
|
|
auto *ctx = op->getContext();
|
|
assert(ctx != nullptr && "ctx is nullptr");
|
|
|
|
auto const rankedTy = RankedTensorType::get(
|
|
{static_cast<int64_t>(sliceOfInput.size())}, IndexType::get(ctx));
|
|
|
|
return DenseElementsAttr::get(rankedTy, sliceOfInput);
|
|
}
|
|
|
|
OpFoldResult tosa::AddShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<AddShapeOp, AddFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::SubShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<SubShapeOp, SubFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::MulShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<MulShapeOp, MulFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::DivCeilShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<DivCeilShapeOp, DivFoldAdaptor</*Ceil*/ true>>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::DivFloorShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<DivFloorShapeOp, DivFoldAdaptor</*Ceil*/ false>>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::ModShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<ModShapeOp, ModFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::MaxShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<MaxShapeOp, MaxFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::MinShapeOp::fold(FoldAdaptor adaptor) {
|
|
return binaryFold<MinShapeOp, MinFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::Exp2ShapeOp::fold(FoldAdaptor adaptor) {
|
|
return unaryShapeFold<Exp2ShapeOp, Exp2FoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::Log2CeilShapeOp::fold(FoldAdaptor adaptor) {
|
|
return unaryShapeFold<Log2CeilShapeOp, Log2CeilFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::Log2FloorShapeOp::fold(FoldAdaptor adaptor) {
|
|
return unaryShapeFold<Log2FloorShapeOp, Log2FloorFoldAdaptor>(this);
|
|
}
|
|
|
|
OpFoldResult tosa::ConcatShapeOp::fold(FoldAdaptor adaptor) {
|
|
return concatShapeFold(this);
|
|
}
|
|
|
|
OpFoldResult tosa::SliceShapeOp::fold(FoldAdaptor adaptor) {
|
|
return sliceShapeFold(this);
|
|
}
|