//===- 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);
}