#include <sstream>

#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Support/LLVM.h"
#include "triton/Dialect/Triton/IR/Dialect.h"
#include "triton/Dialect/Triton/IR/Types.h"
#include "triton/Dialect/Triton/IR/Utility.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/MathExtras.h"

namespace mlir {
namespace triton {

void LoadOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  effects.emplace_back(MemoryEffects::Read::get(), &getPtrMutable(),
                       GlobalMemory::get());
  if (getIsVolatile())
    effects.emplace_back(MemoryEffects::Write::get());
}

} // namespace triton
} // namespace mlir

#define GET_OP_CLASSES
#include "triton/Dialect/Triton/IR/Ops.cpp.inc"

// enum attribute definitions
#include "triton/Dialect/Triton/IR/OpsEnums.cpp.inc"

#include "TritonCanonicalize.inc"

namespace mlir {
namespace triton {

//-- LoadOp --
void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr,
                   CacheModifier cache, EvictionPolicy evict, bool isVolatile) {
  LoadOp::build(builder, state, ptr, /*mask=*/{}, /*other=*/{}, cache, evict,
                isVolatile);
}

void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr,
                   Value mask, CacheModifier cache, EvictionPolicy evict,
                   bool isVolatile) {
  LoadOp::build(builder, state, ptr, mask, /*other=*/{}, cache, evict,
                isVolatile);
}

Value LoadOp::getPredicateOperand() { return getMask(); }

void LoadOp::setPredicateOperand(Value pred) { getMaskMutable().assign(pred); }

Type LoadOp::getPredicateOperandTypeLike() { return getPtr().getType(); }

// load(ptr, splat(1), ...)        -> load(ptr, ...)
// load(ptr, splat(0), other, ...) -> other
struct CanonicalizeMaskedLoadPattern : public OpRewritePattern<LoadOp> {
  CanonicalizeMaskedLoadPattern(MLIRContext *context)
      : OpRewritePattern<LoadOp>(context, 1) {}

  LogicalResult matchAndRewrite(LoadOp loadOp,
                                PatternRewriter &rewriter) const override {
    auto mask = loadOp.getMask();
    if (!mask)
      return failure();

    auto constantMask = mask.getDefiningOp<arith::ConstantOp>();
    if (!constantMask)
      return failure();

    auto splatMask = mlir::dyn_cast<SplatElementsAttr>(constantMask.getValue());
    if (!splatMask)
      return failure();

    if (splatMask.getSplatValue<IntegerAttr>().getValue() == true) {
      // mask = splat(1)
      rewriter.replaceOpWithNewOp<LoadOp>(
          loadOp, loadOp.getType(), loadOp.getPtr(), Value(), Value(),
          loadOp.getCache(), loadOp.getEvict(), loadOp.getIsVolatile());
    } else {
      // mask = splat(0)

      // If there's no "other", the value is "undef".  Perhaps we want to
      // optimize it in the future.x
      auto otherVal = loadOp.getOther();
      if (!otherVal)
        return failure();
      rewriter.replaceOp(loadOp, otherVal);
    }
    return success();
  }
};

void LoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                         MLIRContext *context) {
  results.add<CanonicalizeMaskedLoadPattern>(context);
}

//-- StoreOp --
void StoreOp::build(OpBuilder &builder, OperationState &state, Value ptr,
                    Value value, CacheModifier cache, EvictionPolicy evict) {
  return StoreOp::build(builder, state, ptr, value, /*mask=*/{}, cache, evict);
}

Value StoreOp::getPredicateOperand() { return getMask(); }

void StoreOp::setPredicateOperand(Value pred) { getMaskMutable().assign(pred); }

Type StoreOp::getPredicateOperandTypeLike() { return getPtr().getType(); }

// store(ptr, value, splat(1), ...) -> store(ptr, value, ...)
// store(ptr, value, splat(0), ...) -> [none]
struct CanonicalizeMaskedStorePattern : public OpRewritePattern<StoreOp> {
  CanonicalizeMaskedStorePattern(MLIRContext *context)
      : OpRewritePattern<StoreOp>(context, 1) {}

  LogicalResult matchAndRewrite(StoreOp storeOp,
                                PatternRewriter &rewriter) const override {
    auto mask = storeOp.getMask();
    if (!mask)
      return failure();

    auto constantMask = mask.getDefiningOp<arith::ConstantOp>();
    if (!constantMask)
      return failure();

    auto splatMask = mlir::dyn_cast<SplatElementsAttr>(constantMask.getValue());
    if (!splatMask)
      return failure();

    if (splatMask.getSplatValue<IntegerAttr>().getValue() == true) {
      // mask = splat(1)
      rewriter.replaceOpWithNewOp<StoreOp>(
          storeOp, storeOp.getPtr(), storeOp.getValue(), storeOp.getCache(),
          storeOp.getEvict());
    } else {
      // mask = splat(0)
      rewriter.eraseOp(storeOp);
    }
    return success();
  }
};

Value AtomicRMWOp::getPredicateOperand() { return getMask(); }

void AtomicRMWOp::setPredicateOperand(Value pred) {
  getMaskMutable().assign(pred);
}

Type AtomicRMWOp::getPredicateOperandTypeLike() { return getPtr().getType(); }

LogicalResult AtomicPollOp::verify() {
  if (getSem() != MemSemantic::ACQUIRE && getSem() != MemSemantic::RELAXED)
    return emitOpError("only supports acquire and relaxed semantics");

  unsigned bitWidth = getExpected().getType().getIntOrFloatBitWidth();
  if (bitWidth != 16 && bitWidth != 32 && bitWidth != 64)
    return emitOpError(
        "only supports integer elements with width {16, 32, 64}");
  return success();
}

void StoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
  results.add<CanonicalizeMaskedStorePattern>(context);
}

//-- TransOp --
OpFoldResult TransOp::fold(FoldAdaptor adaptor) {
  // transpose(x, order=[0, 1, ...]) -> x
  if (isIota(getOrder())) {
    // If the source and result types are the same, we can return the source
    // If their layout is different (even if structurally equivalent), we need
    // to insert a convert_layout in between as otherwise ::fold complains
    // We do this in CanonicalizeConvertFromTranspose
    if (getSrc().getType() == getType()) {
      return getSrc();
    }
  }

  // transpose(transpose(x)) -> transpose(x)
  if (auto innerTrans = getSrc().getDefiningOp<TransOp>()) {
    setOrder(applyPermutation(innerTrans.getOrder(), getOrder()));
    setOperand(innerTrans.getSrc());
    return getResult();
  }

  // Eliminate splat constant transpose ops.
  if (auto attr =
          llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSrc()))
    return attr.reshape(getType());

  return {};
}

LogicalResult TransOp::verify() {
  auto order = getOrder();
  auto srcTy = cast<RankedTensorType>(getSrc().getType());
  if (order.size() != srcTy.getShape().size()) {
    return emitError("order must have the same size as the source tensor");
  }
  if (!isPermutationOfIota(order)) {
    return emitError("order must be a permutation of 0..n-1");
  }
  SmallVector<int64_t> retShape = applyPermutation(srcTy.getShape(), order);
  if (retShape != getType().getShape()) {
    return emitError(
        "result shape must match the permutation of the source shape");
  }
  return success();
}

LogicalResult
TransOp::inferReturnTypes(MLIRContext *context, std::optional<Location> loc,
                          TransOp::Adaptor adaptor,
                          SmallVectorImpl<Type> &inferredReturnTypes) {

  // type is the same as the input
  auto argTy = cast<RankedTensorType>(adaptor.getSrc().getType());
  auto shape = argTy.getShape();
  auto order = adaptor.getOrder();
  SmallVector<int64_t> retShape = applyPermutation(shape, order);

  auto retEltTy = argTy.getElementType();
  Attribute argEncoding = argTy.getEncoding();
  Attribute retEncoding;
  if (argEncoding) {
    Dialect &dialect = argEncoding.getDialect();
    auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
    if (failed(inferLayoutInterface->inferTransOpEncoding(
            argEncoding, shape, order, retEncoding, loc))) {
      return failure();
    }
  }
  inferredReturnTypes.push_back(
      RankedTensorType::get(retShape, retEltTy, retEncoding));
  return success();
}

//-- DotOp --
LogicalResult
DotOp::inferReturnTypes(MLIRContext *context, std::optional<Location> location,
                        ValueRange operands, DictionaryAttr attributes,
                        PropertyRef properties, RegionRange regions,
                        SmallVectorImpl<Type> &inferredReturnTypes) {
  // type is the same as the accumulator
  auto accTy = cast<RankedTensorType>(operands[2].getType());
  inferredReturnTypes.push_back(accTy);

  // verify encodings
  auto aEnc = cast<RankedTensorType>(operands[0].getType()).getEncoding();
  auto bEnc = cast<RankedTensorType>(operands[1].getType()).getEncoding();
  auto retEnc = accTy.getEncoding();
  if (aEnc) {
    assert(bEnc && retEnc);
    Dialect &dialect = retEnc.getDialect();
    auto interface = cast<DialectInferLayoutInterface>(&dialect);
    if (interface->inferDotOpEncoding(aEnc, 0, retEnc, location).failed())
      return failure();
    if (interface->inferDotOpEncoding(bEnc, 1, retEnc, location).failed())
      return failure();
  }
  return success();
}

LogicalResult DotOp::verify() {
  auto aTy = getA().getType();
  auto bTy = getB().getType();
  if (aTy.getElementType().getIntOrFloatBitWidth() !=
      bTy.getElementType().getIntOrFloatBitWidth())
    return emitError(
        "element types of operands A and B must have same bit width");
  auto aEncoding = aTy.getEncoding();
  auto bEncoding = bTy.getEncoding();
  if (!aEncoding && !bEncoding)
    return success();
  // Verify that the encodings are valid.
  if (!aEncoding || !bEncoding)
    return emitError("mismatching encoding between A and B operands");
  auto accTy = getC().getType();
  auto retEnc = accTy.getEncoding();
  if (!retEnc)
    return emitError("miss encoding of C operand");
  Dialect &dialect = retEnc.getDialect();
  auto interface = cast<DialectInferLayoutInterface>(&dialect);
  return interface->verifyDotOpEncodingCompatibility(getOperation(), aEncoding,
                                                     bEncoding);
}

//-- DotScaledOp --
bool DotScaledOp::verifyDims() {
  auto aShape = this->getA().getType().getShape();
  auto bShape = this->getB().getType().getShape();

  auto aKdim = aShape[aShape.size() - 1];
  auto bKdim = bShape[aShape.size() - 2];
  if (this->getAElemType() == ScaleDotElemType::E2M1) {
    if (this->getLhsKPack())
      aKdim *= 2;
  }
  if (this->getBElemType() == ScaleDotElemType::E2M1) {
    if (this->getRhsKPack())
      bKdim *= 2;
  }

  return aKdim == bKdim;
}

bool DotScaledOp::verifyOutputDims() {
  auto cShape = this->getC().getType().getShape();
  auto oMdim = cShape[cShape.size() - 2];
  auto oNdim = cShape[cShape.size() - 1];
  auto aShape = this->getA().getType().getShape();
  auto bShape = this->getB().getType().getShape();
  auto adim = aShape[aShape.size() - 2];
  auto bdim = bShape[bShape.size() - 1];
  if (this->getAElemType() == ScaleDotElemType::E2M1) {
    if (!this->getLhsKPack())
      adim *= 2;
  }
  if (this->getBElemType() == ScaleDotElemType::E2M1) {
    if (!this->getRhsKPack())
      bdim *= 2;
  }
  if (adim != oMdim || bdim != oNdim)
    return false;
  return true;
}

LogicalResult DotScaledOp::verify() {
  auto aShape = this->getA().getType().getShape();
  int64_t rank = aShape.size();
  if (rank < 2)
    return this->emitError("operands must be at least 2D");

  auto k = aShape[rank - 1];
  if (this->getAElemType() == ScaleDotElemType::E2M1) {
    if (this->getLhsKPack())
      k *= 2;
  }
  auto cShape = this->getC().getType().getShape();
  int64_t mDim = cShape[cShape.size() - 2];
  int64_t nDim = cShape[cShape.size() - 1];
  int32_t scaleFactor;
  std::string scaleErr;
  if (failed(deduceScaleFactor(
          getA(), getAScale(), getAElemType(), getLhsKPack(), getB(),
          getBScale(), getBElemType(), getRhsKPack(), scaleFactor, scaleErr)))
    return this->emitError(scaleErr);

  if (getAScale()) {
    auto aScaleShape = getAScale().getType().getShape();
    if (aScaleShape[rank - 2] != mDim)
      return this->emitError(
          "scales M dimension must match the operand M dimension");
    if (aScaleShape[rank - 1] != k / scaleFactor)
      return this->emitError("scales K dimension must match the operand K "
                             "divided by the scale factor");
  }
  if (getBScale()) {
    auto bScaleShape = getBScale().getType().getShape();
    if (bScaleShape[rank - 2] != nDim)
      return this->emitError(
          "scales N dimension must match the operand N dimension");
    if (bScaleShape[rank - 1] != k / scaleFactor)
      return this->emitError("scales K dimension must match the operand K "
                             "divided by the scale factor");
  }
  return success();
}

LogicalResult deduceScaleFactor(ArrayRef<int64_t> lhsShape,
                                std::optional<ArrayRef<int64_t>> lhsScaleShape,
                                ScaleDotElemType lhsFormat, bool lhsKPack,
                                ArrayRef<int64_t> rhsShape,
                                std::optional<ArrayRef<int64_t>> rhsScaleShape,
                                ScaleDotElemType rhsFormat, bool rhsKPack,
                                int32_t &scaleFactor, std::string &errMsg) {
  auto deduceByShape = [&errMsg](ArrayRef<int64_t> operandShape,
                                 std::optional<ArrayRef<int64_t>> scaleShape,
                                 int opIdx, ScaleDotElemType format,
                                 bool kPack) -> int32_t {
    if (!scaleShape)
      return 0;
    if (llvm::product_of(*scaleShape) == 1)
      return 0;

    int64_t unpackFactor = (format == ScaleDotElemType::E2M1 && kPack) ? 2 : 1;
    int64_t kdim = operandShape[opIdx == 0 ? operandShape.size() - 1
                                           : operandShape.size() - 2] *
                   unpackFactor;
    int32_t scaleFactor = kdim / (*scaleShape)[scaleShape->size() - 1];
    if (scaleFactor != 16 && scaleFactor != 32) {
      std::ostringstream oss;
      oss << "scale factor must be 16 or 32. Got " << scaleFactor;
      errMsg = oss.str();
      return 0;
    }
    return scaleFactor;
  };

  errMsg.clear();
  int32_t scaleFactorA =
      deduceByShape(lhsShape, lhsScaleShape, 0, lhsFormat, lhsKPack);
  if (!errMsg.empty())
    return failure();
  int32_t scaleFactorB =
      deduceByShape(rhsShape, rhsScaleShape, 1, rhsFormat, rhsKPack);
  if (!errMsg.empty())
    return failure();

  if (scaleFactorA == 0 && scaleFactorB == 0) {
    scaleFactor = 32;
    return success();
  }
  if (scaleFactorA != 0 && scaleFactorB != 0) {
    if (scaleFactorA != scaleFactorB) {
      std::ostringstream oss;
      oss << "Operands must have the same scale factor; (lhs: " << scaleFactorA
          << " vs rhs: " << scaleFactorB << ")";
      errMsg = oss.str();
      return failure();
    }
    scaleFactor = scaleFactorA;
    return success();
  }
  scaleFactor = scaleFactorA != 0 ? scaleFactorA : scaleFactorB;
  return success();
}

LogicalResult DotScaledOp::deduceScaleFactor(
    Value lhs, Value lhsScale, ScaleDotElemType lhsFormat, bool lhsKPack,
    Value rhs, Value rhsScale, ScaleDotElemType rhsFormat, bool rhsKPack,
    int32_t &scaleFactor, std::string &errMsg) {
  auto getScaleShape = [](Value scale) -> std::optional<ArrayRef<int64_t>> {
    if (!scale) {
      return std::nullopt;
    } else {
      return cast<RankedTensorType>(scale.getType()).getShape();
    }
  };

  auto lhsShape = cast<RankedTensorType>(lhs.getType()).getShape();
  auto rhsShape = cast<RankedTensorType>(rhs.getType()).getShape();

  return triton::deduceScaleFactor(lhsShape, getScaleShape(lhsScale), lhsFormat,
                                   lhsKPack, rhsShape, getScaleShape(rhsScale),
                                   rhsFormat, rhsKPack, scaleFactor, errMsg);
}

int32_t DotScaledOp::deduceScaleFactor() {
  int32_t scaleFactor;
  std::string errMsg;
  if (failed(deduceScaleFactor(
          getA(), getAScale(), getAElemType(), getLhsKPack(), getB(),
          getBScale(), getBElemType(), getRhsKPack(), scaleFactor, errMsg)))
    llvm::report_fatal_error(errMsg.c_str());
  return scaleFactor;
}

//-- MakeRangeOp --
OpFoldResult MakeRangeOp::fold(FoldAdaptor adaptor) {
  // make_range(start, start + 1) -> constant(start)
  if (adaptor.getStart() + 1 == adaptor.getEnd()) {
    auto shapedType = cast<ShapedType>(getType());
    return SplatElementsAttr::get(shapedType, adaptor.getStartAttr());
  }
  return {};
}

LogicalResult MakeRangeOp::verify() {
  int64_t start = getStartAttr().getInt();
  int64_t end = getEndAttr().getInt();
  if (start >= end) {
    return this->emitOpError() << "start must be less than end";
  }
  auto ty = getType();
  if (ty.getShape().size() != 1) {
    return this->emitOpError() << "return type must be a 1D tensor";
  }
  if (end - start != ty.getShape()[0]) {
    return this->emitOpError()
           << "number of elements in returned tensor, " << ty.getShape()[0]
           << ", must match size of range [" << start << ", " << end
           << "), which has " << end - start << " elements";
  }
  if (!ty.getElementType().isInteger(32)) {
    return this->emitOpError() << "returned tensor must have i32 elements";
  }
  return success();
}

//-- ReduceOp --
static LogicalResult
inferReduceReturnShape(std::optional<Location> loc, RankedTensorType argTy,
                       Type retEltTy, int axis,
                       SmallVectorImpl<Type> &inferredReturnTypes) {
  auto retShape = argTy.getShape().vec();
  retShape.erase(retShape.begin() + axis);
  if (retShape.empty()) {
    // 0d-tensor -> scalar
    inferredReturnTypes.push_back(retEltTy);
  } else {
    // nd-tensor where n >= 1
    // infer encoding
    Attribute argEncoding = argTy.getEncoding();
    Attribute retEncoding;
    if (argEncoding) {
      Dialect &dialect = argEncoding.getDialect();
      auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
      if (failed(inferLayoutInterface->inferReduceOpEncoding(
              argEncoding, axis, retEncoding, loc))) {
        return failure();
      }
    }
    // create type
    inferredReturnTypes.push_back(
        RankedTensorType::get(retShape, retEltTy, retEncoding));
  }
  return success();
}

LogicalResult
ReduceOp::inferReturnTypes(MLIRContext *context, std::optional<Location> loc,
                           ValueRange operands, DictionaryAttr attributes,
                           PropertyRef properties, RegionRange regions,
                           SmallVectorImpl<Type> &inferredReturnTypes) {
  Properties *prop = properties.as<Properties *>();
  int axis = prop->axis.getInt();
  for (auto arg : operands) {
    auto argTy = cast<RankedTensorType>(arg.getType());
    auto retEltTy = argTy.getElementType();
    if (failed(inferReduceReturnShape(loc, argTy, retEltTy, axis,
                                      inferredReturnTypes))) {
      return failure();
    }
  }
  return success();
}

// Helpers for Reductions and Scans
namespace {

template <class Op> LogicalResult verifyReduceScan(Op &op) {
  if (op.getOperands().empty()) {
    return op.emitOpError() << "must have at least 1 operand";
  }
  if (op.getNumOperands() != op.getNumResults()) {
    return op.emitOpError() << "must have the same number of inputs as outputs";
  }
  auto axis = op.getAxis();
  auto firstRank = 0;
  for (auto tensorTy : op.getInputTypes()) {
    int64_t rank = tensorTy.getRank();
    if (axis < 0 || axis >= rank)
      return op.emitOpError() << "axis out of bounds for operand rank " << rank;
    if (firstRank == 0)
      firstRank = rank;
    else if (rank != firstRank)
      return op.emitOpError()
             << "all operands must have the same rank, but got ranks "
             << firstRank << " and " << rank;
  }
  for (auto [opElemTy, resTy] :
       llvm::zip(op.getElementTypes(), op.getResultTypes())) {
    if (opElemTy != getElementTypeOrSelf(resTy)) {
      return op.emitOpError() << "operand types and result types must agree";
    }
  }
  return success();
}

template <class ReturnOp, class Op> LogicalResult verifyRegionsImpl(Op &op) {
  auto argElementTypes = op.getElementTypes();
  const auto &operands = op.getOperands();
  const auto numArgs = 2 * operands.size();
  auto &block = *op.getBody();
  if (block.getNumArguments() != numArgs) {
    return op.emitOpError() << "nested block must take " << numArgs
                            << " arguments, but given block with "
                            << block.getNumArguments() << " arguments";
  }
  const auto &blockArgTypes = block.getArgumentTypes();
  for (unsigned i = 0; i < numArgs; ++i) {
    const auto &blockArgTy = blockArgTypes[i];
    const auto &argElemTy = argElementTypes[i % operands.size()];
    if (blockArgTy != argElemTy) {
      return op.emitOpError()
             << "type mismatch on combine operation. Expected argument " << i
             << " to have type " << argElemTy << " but got " << blockArgTy;
    }
  }

  auto terminator = dyn_cast<ReturnOp>(block.getTerminator());
  if (!terminator) {
    return op.emitOpError()
           << "combine operation must be terminated "
           << "with a ReduceReturnOp but got " << block.getTerminator();
  }
  const auto &combineResults = terminator->getOperands();
  if (combineResults.size() != operands.size()) {
    return op.emitOpError()
           << "expected combine operation to return " << operands.size()
           << " values but got " << combineResults.size();
  }
  for (unsigned i = 0; i < combineResults.size(); ++i) {
    const auto &resultTy = combineResults[i].getType();
    const auto &argElemTy = argElementTypes[i];
    if (resultTy != argElemTy) {
      return op.emitOpError()
             << "type mismatch on combine operation. Expected argument " << i
             << " to have type " << argElemTy << " but got " << resultTy;
    }
  }
  return success();
}

llvm::SmallVector<RankedTensorType>
getInputTypesImpl(const Operation::operand_range &operands) {
  llvm::SmallVector<RankedTensorType> srcTys;
  srcTys.reserve(operands.size());
  for (const auto &ty : operands.getTypes()) {
    srcTys.push_back(cast<RankedTensorType>(ty));
  }
  return srcTys;
}

template <typename ValueRange>
llvm::SmallVector<Type> getElementTypesImpl(const ValueRange &operands) {
  llvm::SmallVector<Type> srcElemTys;
  srcElemTys.reserve(operands.size());
  for (const auto &op : operands) {
    srcElemTys.push_back(cast<RankedTensorType>(op.getType()).getElementType());
  }
  return srcElemTys;
}

} // namespace

LogicalResult ReduceOp::verify() { return verifyReduceScan(*this); }

LogicalResult ReduceOp::verifyRegions() {
  return verifyRegionsImpl<ReduceReturnOp>(*this);
}

llvm::SmallVector<RankedTensorType> ReduceOp::getInputTypes() {
  return getInputTypesImpl(this->getOperands());
}

llvm::SmallVector<Type> ReduceOp::getElementTypes() {
  return getElementTypesImpl(this->getOperands());
}

::mlir::Operation *ReduceOp::getSingleCombiner() {
  if (getNumOperands() != 1 || getNumResults() != 1)
    return nullptr;
  Block *block = &(*getCombineOp().begin());
  Operation *yield = block->getTerminator();
  Operation *reduceOp = yield->getOperand(0).getDefiningOp();
  if (!reduceOp || reduceOp->getNumOperands() != 2 ||
      reduceOp->getNumResults() != 1)
    return nullptr;
  Value arg0 = block->getArgument(0), arg1 = block->getArgument(1);
  Value lhs = reduceOp->getOperand(0), rhs = reduceOp->getOperand(1);
  // For commutative combiners, the reversed argument-operand mapping is
  // equivalent.
  bool reversedMapping = (lhs == arg1 && rhs == arg0) &&
                         reduceOp->hasTrait<OpTrait::IsCommutative>();
  if (!(lhs == arg0 && rhs == arg1) && !reversedMapping)
    return nullptr;

  return reduceOp;
}

//-- ScanOp --
void ScanOp::build(OpBuilder &builder, OperationState &state,
                   ValueRange operands, int axis, bool reverse) {
  SmallVector<Type> inferredReturnTypes;
  for (auto arg : operands)
    inferredReturnTypes.push_back(arg.getType());
  ScanOp::build(builder, state, inferredReturnTypes, operands, axis, reverse);
}

LogicalResult
ScanOp::inferReturnTypes(MLIRContext *context, std::optional<Location> location,
                         ValueRange operands, DictionaryAttr attributes,
                         PropertyRef properties, RegionRange regions,
                         SmallVectorImpl<Type> &inferredReturnTypes) {
  for (auto arg : operands)
    inferredReturnTypes.push_back(arg.getType());
  return success();
}

LogicalResult ScanOp::verify() { return verifyReduceScan(*this); }

LogicalResult ScanOp::verifyRegions() {
  return verifyRegionsImpl<ScanReturnOp>(*this);
}

llvm::SmallVector<RankedTensorType> ScanOp::getInputTypes() {
  return getInputTypesImpl(this->getOperands());
}

llvm::SmallVector<Type> ScanOp::getElementTypes() {
  return getElementTypesImpl(this->getOperands());
}

//-- MapElementwiseOp
LogicalResult MapElementwiseOp::verify() {
  if (getOperands().empty()) {
    return emitOpError() << "MapElementwiseOp must have at least 1 operand";
  }
  if (!llvm::isPowerOf2_32(getPack())) {
    return emitOpError() << "Pack must be a power of 2";
  }
  return success();
}

template <typename T>
static SmallVector<T> repeatInterleave(const SmallVectorImpl<T> &vs,
                                       int nRepeat) {
  SmallVector<T> result;
  result.reserve(vs.size() * nRepeat);
  for (auto v : vs)
    for (int i = 0; i < nRepeat; ++i)
      result.push_back(v);
  return result;
}

LogicalResult MapElementwiseOp::verifyRegions() {
  // Verify signature
  auto *firstBlock = &getRegion().getBlocks().front();
  if (firstBlock->getNumArguments() != getNumOperands() * getPack()) {
    return emitOpError() << "region has wrong number of arguments";
  }

  auto expectedArgTypes =
      repeatInterleave(getElementTypesImpl(getOperands()), getPack());
  if (firstBlock->getArgumentTypes() != expectedArgTypes) {
    return emitError() << "argument types did not match";
  }
  auto expectedReturnTypes =
      repeatInterleave(getElementTypesImpl(getResults()), getPack());
  auto walkRes = getRegion().walk([&](Operation *op) -> WalkResult {
    auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
    // Ban stores as we won't get the redundant masking correct by treating it
    // as a scalar.
    if (memEffects && memEffects.hasEffect<MemoryEffects::Write>()) {
      return op->emitOpError()
             << "Stores are not supported inside map_elementwise";
    }
    if (isa<MapElementwiseReturnOp>(op) &&
        op->getOperandTypes() != expectedReturnTypes) {
      return op->emitError()
             << "region return does not match map_elementwise result";
    }
    return WalkResult::advance();
  });
  return success(!walkRes.wasInterrupted());
}

//-- SplatOp --
OpFoldResult SplatOp::fold(FoldAdaptor adaptor) {
  auto value = adaptor.getSrc();
  if (!value)
    return {};
  if (!isa<FloatAttr, IntegerAttr>(value))
    return {};
  auto shapedType = cast<ShapedType>(getType());
  auto ret = SplatElementsAttr::get(shapedType, ArrayRef<Attribute>(value));
  return ret;
}

//-- UnsplatOp --
LogicalResult UnsplatOp::verify() {
  auto srcShape = getSrc().getType().getShape();
  if (product(srcShape) != 1) {
    return emitError("source tensor must have exactly one element");
  }
  return success();
}

LogicalResult UnsplatOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  auto dstTy = cast<RankedTensorType>(operands[0].getType()).getElementType();
  inferredReturnTypes.push_back(dstTy);
  return success();
}

//-- ExpandDimsOp --
LogicalResult ExpandDimsOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> loc, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  // infer shape
  auto arg = operands[0];
  auto argTy = cast<RankedTensorType>(arg.getType());
  auto retShape = argTy.getShape().vec();
  Properties *prop = properties.as<Properties *>();
  int axis = prop->axis.getInt();
  retShape.insert(retShape.begin() + axis, 1);
  // infer encoding
  Attribute argEncoding = argTy.getEncoding();
  Attribute retEncoding;
  if (argEncoding) {
    Dialect &dialect = argEncoding.getDialect();
    auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
    if (failed(inferLayoutInterface->inferExpandDimsOpEncoding(
            argEncoding, axis, retEncoding, loc)))
      return emitOptionalError(loc, "failed to infer layout for ExpandDimsOp");
  }
  // create type
  auto argEltTy = argTy.getElementType();
  inferredReturnTypes.push_back(
      RankedTensorType::get(retShape, argEltTy, retEncoding));
  return success();
}

LogicalResult ExpandDimsOp::canonicalize(ExpandDimsOp op,
                                         PatternRewriter &rewriter) {
  auto definingOp = op.getSrc().getDefiningOp();
  if (!definingOp) {
    return failure();
  }
  // expand_dims(splat) -> splat
  if (auto splat = dyn_cast<SplatOp>(definingOp)) {
    rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), splat.getSrc());
    return success();
  }
  // expand_dims(broadcast(x)) -> broadcast(expand_dims(x))
  //
  // On its own this doesn't do much, but consider
  //    broadcast(expand_dims(broadcast))
  // -> broadcast(broadcast(expand_dims))
  // -> broadcast(expand_dims)
  if (auto broadcast = dyn_cast<BroadcastOp>(definingOp)) {
    auto src = broadcast.getSrc();
    auto srcTy = src.getType();
    SmallVector<int64_t> newExpandShape(srcTy.getShape());
    newExpandShape.insert(newExpandShape.begin() + op.getAxis(), 1);

    // Infer the encoding of the new expand op, if encodings are present.
    Attribute newExpandEnc;
    if (auto srcEnc = srcTy.getEncoding()) {
      Dialect &dialect = srcEnc.getDialect();
      auto inferLayoutInterface = cast<DialectInferLayoutInterface>(&dialect);
      if (failed(inferLayoutInterface->inferExpandDimsOpEncoding(
              srcEnc, op.getAxis(), newExpandEnc, op.getLoc()))) {
        return emitOptionalError(op.getLoc(),
                                 "failed to infer layout for ExpandDimsOp");
      }
    }

    auto newExpandTy = RankedTensorType::get(
        newExpandShape, srcTy.getElementType(), newExpandEnc);
    auto newExpand = ExpandDimsOp::create(rewriter, op.getLoc(), newExpandTy,
                                          src, op.getAxis());
    auto newBroadcast = BroadcastOp::create(
        rewriter, broadcast.getLoc(), op.getType(), newExpand.getResult());
    rewriter.replaceOp(op, {newBroadcast.getResult()});
    return success();
  }

  return failure();
}

template <typename ViewLikeOp>
static OpFoldResult foldViewLikeOp(ViewLikeOp op, Attribute value) {
  if (!value)
    return {};

  auto shapedType = cast<ShapedType>(op.getType());
  if (auto denseElemsAttr = dyn_cast<DenseElementsAttr>(value)) {
    if (denseElemsAttr.isSplat()) {
      return denseElemsAttr.resizeSplat(shapedType);
    } else {
      return denseElemsAttr.reshape(shapedType);
    }
  }
  return {};
}

OpFoldResult ExpandDimsOp::fold(FoldAdaptor adaptor) {
  return foldViewLikeOp(*this, adaptor.getSrc());
}

//-- CatOp --
LogicalResult CatOp::verify() {
  RankedTensorType lhsTy = getLhs().getType();
  RankedTensorType resultTy = getType();

  int64_t operandElements = lhsTy.getNumElements() * 2;
  if (resultTy.getNumElements() != operandElements) {
    return emitOpError("result element count must equal the sum of the "
                       "operand element counts, expected ")
           << operandElements << " but got " << resultTy.getNumElements();
  }

  Attribute operandEnc = lhsTy.getEncoding();
  Attribute resultEnc = resultTy.getEncoding();
  if (!!operandEnc != !!resultEnc) {
    return emitOpError("requires that either (a) operands and result all have "
                       "encodings, or (b) none do.");
  }
  if (!resultEnc)
    return success();

  auto interface = cast<DialectInferLayoutInterface>(&resultEnc.getDialect());
  return interface->verifyCatOpEncodingCompatibility(getOperation());
}

//-- ReshapeOp --

std::optional<unsigned> ReshapeOp::getExpandDimsAxis() {
  auto srcTy = getSrc().getType();
  auto dstTy = getType();
  if (srcTy.getRank() + 1 != dstTy.getRank())
    return std::nullopt;
  auto srcShape = srcTy.getShape();
  auto dstShape = dstTy.getShape();
  for (unsigned axis = 0; axis < dstShape.size(); ++axis) {
    if (dstShape[axis] == 1 &&
        srcShape.take_front(axis) == dstShape.take_front(axis) &&
        srcShape.drop_front(axis) == dstShape.drop_front(axis + 1))
      return axis;
  }
  return std::nullopt;
}

void ReshapeOp::build(OpBuilder &builder, OperationState &state,
                      ArrayRef<int64_t> shape, Value src, bool allowReorder) {
  auto srcTy = cast<RankedTensorType>(src.getType());
  auto srcEnc = srcTy.getEncoding();
  Attribute dstEnc;
  if (srcEnc) {
    auto result =
        cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
            ->inferReshapeOpEncoding(srcTy.getShape(), srcEnc, shape, dstEnc,
                                     allowReorder, state.location);
    assert(succeeded(result));
  }
  auto dstTy = RankedTensorType::get(shape, srcTy.getElementType(), dstEnc);
  build(builder, state, dstTy, src, allowReorder);
}

LogicalResult ReshapeOp::canonicalize(ReshapeOp op, PatternRewriter &rewriter) {
  if (op.getEfficientLayout())
    return failure();

  auto definingOp = op.getSrc().getDefiningOp();
  if (!definingOp) {
    return failure();
  }

  // reshape(expand_dims(x)) -> x
  if (auto expandDims = dyn_cast<ExpandDimsOp>(definingOp)) {
    if (!op.getAllowReorder() &&
        op.getType() == expandDims.getSrc().getType()) {
      rewriter.replaceOp(op, expandDims.getSrc());
      return success();
    }
  }

  // reshape(reshape) -> reshape
  if (auto parentReshape = dyn_cast<ReshapeOp>(definingOp)) {
    // Allow reorder if either reshape allowed it
    const bool allowReorder =
        (op.getAllowReorder() || parentReshape.getAllowReorder());
    rewriter.replaceOpWithNewOp<ReshapeOp>(op, op.getType(),
                                           parentReshape.getSrc(), allowReorder,
                                           op.getEfficientLayout());
    return success();
  }

  // reshape(splat) -> splat
  if (auto splat = dyn_cast<SplatOp>(definingOp)) {
    rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), splat.getSrc());
    return success();
  }

  return failure();
}

OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
  if (getType() == getSrc().getType() && !getAllowReorder()) {
    // no-op
    return getSrc();
  }

  return foldViewLikeOp(*this, adaptor.getSrc());
}

LogicalResult ReshapeOp::verify() {
  auto dstTy = getType();
  auto srcTy = getSrc().getType();
  if (getType().getNumElements() != srcTy.getNumElements()) {
    return emitError(
        "number of src and dst elements of reshape must be the same");
  }

  Attribute srcEnc = srcTy.getEncoding();
  Attribute dstEnc = dstTy.getEncoding();
  if (!!srcEnc != !!dstEnc) {
    return emitError("Op requires that either (a) src and dst both have "
                     "encodings, or (b) neither does.");
  }

  if (!srcEnc) {
    return success();
  }

  // Check that we can infer the dst encoding from the src encoding and that the
  // inferred dst encoding is the same as the given dst encoding. We pass the
  // current dst encoding as a hint so that allowReorder reshapes are guaranteed
  // to produce the current encoding iff it is valid.
  Attribute inferredDstEnc = dstEnc;
  auto layoutInterface =
      cast<DialectInferLayoutInterface>(&srcEnc.getDialect());
  auto result = layoutInterface->inferReshapeOpEncoding(
      srcTy.getShape(), srcEnc, dstTy.getShape(), inferredDstEnc,
      getAllowReorder(), getLoc());
  if (failed(result))
    return failure();
  return layoutInterface->verifyLayoutsAreEqual(
      dstTy.getShape(), inferredDstEnc, dstEnc, getLoc());
}

//-- FpToFpOp --

// Fold FpToFpOp when the input operand is a constant zero.
OpFoldResult FpToFpOp::fold(FoldAdaptor adaptor) {
  auto srcVal = getSrc();
  auto dstTy = getType();
  // Fold trivial cast
  if (srcVal.getType() == dstTy) {
    return srcVal;
  }

  auto resElemType = cast<FloatType>(getElementTypeOrSelf(getType()));
  const llvm::fltSemantics &semantic = resElemType.getFloatSemantics();

  if (matchPattern(srcVal, m_PosZeroFloat())) {
    llvm::APFloat posZero =
        llvm::APFloat::getZero(semantic, /*negative=*/false);
    if (auto tensorTy = dyn_cast<RankedTensorType>(dstTy))
      return DenseElementsAttr::get(tensorTy, posZero);
    return Builder(getContext()).getFloatAttr(resElemType, posZero);
  }

  if (matchPattern(srcVal, m_NegZeroFloat())) {
    llvm::APFloat negZero = llvm::APFloat::getZero(semantic, /*negative=*/true);
    if (auto tensorTy = dyn_cast<RankedTensorType>(dstTy))
      return DenseElementsAttr::get(tensorTy, negZero);
    return Builder(getContext()).getFloatAttr(resElemType, negZero);
  }

  return {};
}

LogicalResult FpToFpOp::verify() {
  auto dstType = getType();
  auto srcType = getSrc().getType();
  if (auto dstTensorType = dyn_cast<RankedTensorType>(dstType))
    dstType = dstTensorType.getElementType();
  if (auto srcTensorType = dyn_cast<RankedTensorType>(srcType))
    srcType = srcTensorType.getElementType();
  if ((dstType.getIntOrFloatBitWidth() < srcType.getIntOrFloatBitWidth()) &&
      (!getRounding().has_value())) {
    return emitError("Rounding mode is required for FP downcast");
  }
  return success();
}

//-- BitcastOp --
OpFoldResult BitcastOp::fold(FoldAdaptor adaptor) {
  if (getSrc().getType() == getType())
    return getSrc();
  return {};
}

LogicalResult BitcastOp::verify() {
  // Bitcast only allows conversion between types with the same bit width.
  Type dstType = getType();
  Type srcType = getSrc().getType();
  // Strip tensor shapes; SameOperandsAndResultShape guarantees shapes match.
  if (auto dstTensorType = dyn_cast<RankedTensorType>(dstType))
    dstType = dstTensorType.getElementType();
  if (auto srcTensorType = dyn_cast<RankedTensorType>(srcType))
    srcType = srcTensorType.getElementType();
  bool dstIsPtr = isa<triton::PointerType>(dstType);
  bool srcIsPtr = isa<triton::PointerType>(srcType);
  if (dstIsPtr || srcIsPtr) {
    // Bitcast supports pointer-to-pointer conversions but not
    // pointer-to-scalar.
    if (dstIsPtr && srcIsPtr) {
      if (triton::getAddressSpace(dstType) != triton::getAddressSpace(srcType))
        return emitError(
            "Cannot bitcast pointer between different address spaces");
      return success();
    }
    return emitError("Cannot bitcast pointer to non-pointer type");
  }
  unsigned dstBits = dstType.getIntOrFloatBitWidth();
  unsigned srcBits = srcType.getIntOrFloatBitWidth();
  if (dstBits != srcBits) {
    return emitError("Cannot bitcast data-type of size ")
           << srcBits << " to data-type of size " << dstBits;
  }
  return success();
}

//-- BroadcastOp --
void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
  results.add<BroadcastSplatPattern, BroadcastBroadcastPattern>(context);
}

OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
  if (getType() == getSrc().getType()) {
    // no-op
    return getSrc();
  }

  auto value = adaptor.getSrc();
  if (!value)
    return {};

  if (auto denseElemsAttr = dyn_cast<SplatElementsAttr>(value)) {
    auto shapedType = cast<ShapedType>(getType());
    return denseElemsAttr.resizeSplat(shapedType);
  }
  return {};
}

LogicalResult BroadcastOp::verify() {
  auto src = getSrc();
  auto srcTensorType = cast<RankedTensorType>(src.getType());
  auto srcShape = srcTensorType.getShape();
  auto result = getResult();
  auto resultTensorType = cast<RankedTensorType>(result.getType());
  auto resultShape = resultTensorType.getShape();
  if (srcShape.size() != resultShape.size()) {
    return emitError("rank of source must be same as rank of result");
  }
  for (size_t i = 0; i < srcShape.size(); i++) {
    if (srcShape[i] != 1 && srcShape[i] != resultShape[i]) {
      return emitError("Different dimensions at index ")
             << i << " between source and result.  "
             << "Broadcast requires the source dimension to be 1.";
    }
  }
  return success();
}

//-- AddPtrOp --
OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) {
  // addptr(ptr, 0) -> ptr
  if (matchPattern(adaptor.getOffset(), m_Zero())) {
    return getPtr();
  }
  return {};
}

//-- MakeTensorDescOp --
void MakeTensorDescOp::build(OpBuilder &builder, OperationState &state,
                             Value base, ValueRange shape, ValueRange strides,
                             ArrayRef<int32_t> blockShape, bool isSignedInteger,
                             triton::PaddingOption padding) {
  auto ptrTy = dyn_cast<triton::PointerType>(base.getType());
  if (!ptrTy) {
    llvm::report_fatal_error("Expected pointer type");
  }
  auto elemTy = ptrTy.getPointeeType();
  SmallVector<int64_t> blockShape64(blockShape);
  auto descTy = TensorDescType::get(blockShape64, elemTy, isSignedInteger);
  auto paddingAttr = PaddingOptionAttr::get(builder.getContext(), padding);
  return build(builder, state, descTy, base, shape, strides, paddingAttr);
}

//-- IntToPtrOp --
// Pattern 1: int_to_ptr(ptr_to_int(ptr)) -> ptr
// Eliminates round-trip pointer conversions
struct CanonicalizeIntToPtrOfPtrToInt : public OpRewritePattern<IntToPtrOp> {
  CanonicalizeIntToPtrOfPtrToInt(MLIRContext *context)
      : OpRewritePattern<IntToPtrOp>(context, 1) {}

  LogicalResult matchAndRewrite(IntToPtrOp intToPtrOp,
                                PatternRewriter &rewriter) const override {
    // Match: int_to_ptr(ptr_to_int(ptr))
    auto ptrToIntOp = intToPtrOp.getSrc().getDefiningOp<PtrToIntOp>();
    if (!ptrToIntOp)
      return failure();
    rewriter.replaceOpWithNewOp<BitcastOp>(intToPtrOp, intToPtrOp.getType(),
                                           ptrToIntOp.getSrc());
    return success();
  }
};

// Pattern 2: int_to_ptr(addi(val, constant_offset)) -> addptr(int_to_ptr(val),
// element_offset). Only when offset is constant and divisible by element size
struct CanonicalizeIntToPtrWithAdd : public OpRewritePattern<IntToPtrOp> {
  CanonicalizeIntToPtrWithAdd(MLIRContext *context)
      : OpRewritePattern<IntToPtrOp>(context, 1) {}

  LogicalResult matchAndRewrite(IntToPtrOp intToPtrOp,
                                PatternRewriter &rewriter) const override {
    // Match: int_to_ptr(addi(val, constant_offset))
    auto addOp = intToPtrOp.getSrc().getDefiningOp<arith::AddIOp>();
    if (!addOp)
      return failure();

    Value intValue = addOp.getLhs();
    Value offsetValue = addOp.getRhs();

    // Get the element size from the result pointer type
    auto resultType = intToPtrOp.getType();
    auto ptrType = cast<PointerType>(getElementTypeOrSelf(resultType));
    int64_t elemSizeBits = triton::getPointeeBitWidth(ptrType);
    int64_t elemSizeBytes = std::max<int64_t>(1, elemSizeBits / 8);

    // Check if offset is a constant (either directly or via splat)
    // Only apply canonicalization for constant offsets
    std::optional<int64_t> constantByteOffset;
    if (auto constOp = offsetValue.getDefiningOp<arith::ConstantOp>()) {
      if (auto intAttr = dyn_cast<IntegerAttr>(constOp.getValue())) {
        constantByteOffset = intAttr.getValue().getSExtValue();
      } else if (auto splatAttr =
                     dyn_cast<SplatElementsAttr>(constOp.getValue())) {
        constantByteOffset =
            splatAttr.getSplatValue<IntegerAttr>().getValue().getSExtValue();
      }
    }

    if (!constantByteOffset.has_value())
      return failure(); // Only handle constant offsets

    // Check if the byte offset is divisible by element size
    if (constantByteOffset.value() % elemSizeBytes != 0)
      return failure();

    // Compute element offset at compile time
    int64_t elementOffset = constantByteOffset.value() / elemSizeBytes;

    // Create int_to_ptr(val) for the base
    auto loc = intToPtrOp.getLoc();
    Value basePtr = IntToPtrOp::create(rewriter, loc, resultType, intValue);

    // Create the element offset constant
    Value elementOffsetValue;

    // Get the integer type from the offset value to match its type
    Type offsetElemType;
    if (auto tensorType = dyn_cast<RankedTensorType>(offsetValue.getType())) {
      offsetElemType = tensorType.getElementType();
    } else {
      offsetElemType = offsetValue.getType();
    }

    if (auto tensorType = dyn_cast<RankedTensorType>(resultType)) {
      // Create a splat constant for tensor types, matching the offset's type
      auto offsetAttr = rewriter.getIntegerAttr(offsetElemType, elementOffset);
      auto splatType = RankedTensorType::get(
          tensorType.getShape(), offsetElemType, tensorType.getEncoding());
      auto splatAttr = SplatElementsAttr::get(splatType, offsetAttr);
      elementOffsetValue = arith::ConstantOp::create(rewriter, loc, splatAttr);
    } else {
      // Scalar case
      elementOffsetValue = arith::ConstantOp::create(
          rewriter, loc,
          rewriter.getIntegerAttr(offsetElemType, elementOffset));
    }

    // Replace with addptr(int_to_ptr(val), element_offset)
    rewriter.replaceOpWithNewOp<AddPtrOp>(intToPtrOp, resultType, basePtr,
                                          elementOffsetValue);
    return success();
  }
};

void IntToPtrOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
  results.add<CanonicalizeIntToPtrOfPtrToInt, CanonicalizeIntToPtrWithAdd>(
      context);
}

// The following ops, including `call`, `func`, and `return` are copied and
// modified from
// https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Func/IR/FuncOps.cpp
// We could revert it back once MLIR has a better inliner interface.
//-- FuncOp --
void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,
                   FunctionType type, ArrayRef<NamedAttribute> attrs,
                   ArrayRef<DictionaryAttr> argAttrs) {
  state.addAttribute(SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(name));
  state.addAttribute(getFunctionTypeAttrName(state.name), TypeAttr::get(type));
  state.attributes.append(attrs.begin(), attrs.end());
  state.addRegion();

  if (argAttrs.empty())
    return;
  assert(type.getNumInputs() == argAttrs.size());
  call_interface_impl::addArgAndResultAttrs(
      builder, state, argAttrs, /*resultAttrs=*/{},
      getArgAttrsAttrName(state.name), getResAttrsAttrName(state.name));
}

ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) {
  auto buildFuncType =
      [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
         function_interface_impl::VariadicFlag,
         std::string &) { return builder.getFunctionType(argTypes, results); };

  return function_interface_impl::parseFunctionOp(
      parser, result, /*allowVariadic=*/false,
      getFunctionTypeAttrName(result.name), buildFuncType,
      getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
}

void FuncOp::print(OpAsmPrinter &printer) {
  function_interface_impl::printFunctionOp(
      printer, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),
      getArgAttrsAttrName(), getResAttrsAttrName());
}

// -- CallOp --
LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  // Check that the callee attribute was specified.
  auto fnAttr = (*this).getProperties().callee;
  if (!fnAttr)
    return emitOpError("requires a 'callee' symbol reference attribute");
  FuncOp fn = symbolTable.lookupNearestSymbolFrom<FuncOp>(*this, fnAttr);
  if (!fn)
    return emitOpError() << "'" << fnAttr.getValue()
                         << "' does not reference a valid function";

  // Verify that the operand and result types match the callee.
  auto fnType = fn.getFunctionType();
  if (fnType.getNumInputs() != getNumOperands())
    return emitOpError("incorrect number of operands for callee");

  for (unsigned i = 0, e = fnType.getNumInputs(); i != e; ++i)
    if (getOperand(i).getType() != fnType.getInput(i))
      return emitOpError("operand type mismatch: expected operand type ")
             << fnType.getInput(i) << ", but provided "
             << getOperand(i).getType() << " for operand number " << i;

  if (fnType.getNumResults() != getNumResults())
    return emitOpError("incorrect number of results for callee");

  for (unsigned i = 0, e = fnType.getNumResults(); i != e; ++i)
    if (getResult(i).getType() != fnType.getResult(i)) {
      auto diag = emitOpError("result type mismatch at index ") << i;
      diag.attachNote() << "      op result types: " << getResultTypes();
      diag.attachNote() << "function result types: " << fnType.getResults();
      return diag;
    }

  return success();
}

// -- ReturnOp --
LogicalResult ReturnOp::verify() {
  auto function = cast<FuncOp>((*this)->getParentOp());

  // The operand number and types must match the function signature.
  const auto &results = function.getFunctionType().getResults();
  if (getNumOperands() != results.size())
    return emitOpError("has ")
           << getNumOperands() << " operands, but enclosing function (@"
           << function.getName() << ") returns " << results.size();

  for (unsigned i = 0, e = results.size(); i != e; ++i)
    if (getOperand(i).getType() != results[i])
      return emitError() << "type of return operand " << i << " ("
                         << getOperand(i).getType()
                         << ") doesn't match function result type ("
                         << results[i] << ")"
                         << " in function @" << function.getName();

  return success();
}

// -- JoinOp --

void JoinOp::build(OpBuilder &builder, OperationState &state, Value lhs,
                   Value rhs) {
  auto lhsTy = cast<RankedTensorType>(lhs.getType());
  SmallVector<int64_t> retShape(lhsTy.getShape());
  retShape.push_back(2);

  Attribute srcEnc = lhsTy.getEncoding();
  Attribute retEnc;
  if (srcEnc) {
    if (failed(cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
                   ->inferDefaultJoinOpEncoding(
                       srcEnc, retEnc, lhsTy.getShape(), state.location))) {
      llvm_unreachable("failed to infer join encoding");
    }
  }
  auto retTy = RankedTensorType::get(retShape, lhsTy.getElementType(), retEnc);
  JoinOp::build(builder, state, retTy, lhs, rhs);
}

LogicalResult JoinOp::verify() {
  RankedTensorType srcTy = getLhs().getType();
  SmallVector<int64_t> retShape(srcTy.getShape());
  retShape.push_back(2);

  RankedTensorType retTy = getType();
  if (SmallVector<int64_t>(retTy.getShape()) != retShape) {
    return emitOpError("result shape must be (")
           << retShape << "), but got " << retTy.getShape();
  }
  if (retTy.getElementType() != srcTy.getElementType()) {
    return emitOpError("result element type must match the input element type");
  }
  Attribute retEnc = retTy.getEncoding();
  if (!retEnc) {
    if (srcTy.getEncoding()) {
      return emitOpError("result encoding must be specified");
    }
    return success();
  }
  // There are multiple correct destination layout for a given source layout but
  // there is only one correct source layout (modulo broadcasting) for a given
  // destination layout. So we verify that the source layout match the
  // destination layout.
  Attribute srcEnc;
  Location location = getLoc();
  if (cast<DialectInferLayoutInterface>(&retEnc.getDialect())
          ->inferSplitOpEncoding(retEnc, srcEnc, retShape, location)
          .failed()) {
    return failure();
  }

  if (cast<triton::DialectInferLayoutInterface>(&srcEnc.getDialect())
          ->verifyLayoutsAreEqual(srcTy.getShape(), srcEnc, srcTy.getEncoding(),
                                  {}, /*ignoreRegBroadcast=*/true)
          .failed()) {
    return emitOpError("incompatible join layout");
  }
  return success();
}

// -- SplitOp --
bool SplitOp::isCompatibleReturnTypes(TypeRange lhs, TypeRange rhs) {
  for (auto [lhs, rhs] : llvm::zip_equal(lhs, rhs)) {
    auto lhsTy = cast<RankedTensorType>(lhs);
    auto rhsTy = cast<RankedTensorType>(rhs);
    if (lhsTy.getShape() != rhsTy.getShape() ||
        lhsTy.getElementType() != rhsTy.getElementType())
      return false;

    auto lhsEnc = lhsTy.getEncoding();
    auto rhsEnc = rhsTy.getEncoding();
    if (!lhsEnc || !rhsEnc) {
      if (lhsEnc || rhsEnc)
        return false;
      continue;
    }

    if (failed(cast<DialectInferLayoutInterface>(&lhsEnc.getDialect())
                   ->verifyLayoutsAreEqual(lhsTy.getShape(), lhsEnc, rhsEnc, {},
                                           /*ignoreRegBroadcast=*/true)))
      return false;
  }
  return true;
}

LogicalResult SplitOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location,
    SplitOp::Adaptor adaptor, SmallVectorImpl<Type> &inferredReturnTypes) {
  auto srcTy = cast<RankedTensorType>(adaptor.getSrc().getType());
  auto srcShape = srcTy.getShape();

  if (srcShape.empty() || srcShape.back() != 2) {
    return emitOptionalError(location,
                             "last dimension of input tensor must be 2");
  }
  ArrayRef<int64_t> retShape(srcShape.begin(), srcShape.end() - 1);

  Attribute srcEnc = srcTy.getEncoding();
  Attribute retEnc;
  if (srcEnc) {
    if (cast<DialectInferLayoutInterface>(&srcEnc.getDialect())
            ->inferSplitOpEncoding(srcEnc, retEnc, srcTy.getShape(), location)
            .failed()) {
      return failure();
    }
  }
  auto retTy = RankedTensorType::get(retShape, srcTy.getElementType(), retEnc);
  inferredReturnTypes.push_back(retTy);
  inferredReturnTypes.push_back(retTy);
  return success();
}

// -- ElementwiseInlineAsmOp --
void ElementwiseInlineAsmOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  if (getPure())
    return;
  effects.emplace_back(MemoryEffects::Write::get());
  effects.emplace_back(MemoryEffects::Read::get());
}

Speculation::Speculatability ElementwiseInlineAsmOp::getSpeculatability() {
  if (getPure())
    return Speculation::Speculatable;
  return Speculation::NotSpeculatable;
}

LogicalResult ElementwiseInlineAsmOp::verify() {
  if (getNumOperands() >= 1) {
    auto tensorType = dyn_cast<RankedTensorType>(getOperand(0).getType());
    size_t numInputElems = tensorType ? tensorType.getNumElements() : 0;
    if (numInputElems % this->getPackedElement() != 0) {
      return emitError("number of input elements ")
             << numInputElems
             << " must be a multiple of the op's packed_element attribute, "
             << getPackedElement();
    }
  }
  return success();
}

// -- ExternElementwiseOp --
void ExternElementwiseOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  if (getPure())
    return;
  effects.emplace_back(MemoryEffects::Write::get());
  effects.emplace_back(MemoryEffects::Read::get());
}

Speculation::Speculatability ExternElementwiseOp::getSpeculatability() {
  if (getPure())
    return Speculation::Speculatable;
  return Speculation::NotSpeculatable;
}

// -- GatherOp --
LogicalResult GatherOp::verify() {
  RankedTensorType indicesTy = getIndices().getType();
  RankedTensorType srcTy = getSrc().getType();
  RankedTensorType resTy = getResult().getType();

  if (indicesTy.getShape() != resTy.getShape()) {
    return emitOpError("indices and output shapes must match");
  }
  if (indicesTy.getEncoding() != resTy.getEncoding()) {
    return emitOpError("indices and output encodings must match");
  }
  if (srcTy.getElementType() != resTy.getElementType()) {
    return emitOpError("input and output element types must match");
  }
  if (srcTy.getRank() != indicesTy.getRank()) {
    return emitOpError("input and indices ranks must match");
  }
  if (getAxis() >= srcTy.getRank()) {
    return emitOpError("gather dimension must be less than the input rank");
  }
  for (uint32_t dim = 0; dim < indicesTy.getRank(); ++dim) {
    if (dim == getAxis())
      continue;
    if (indicesTy.getShape()[dim] != srcTy.getShape()[dim]) {
      return emitOpError("indices dimension ")
             << dim << " must match the corresponding input dimension";
    }
  }

  return success();
}

LogicalResult GatherOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  GatherOpAdaptor adaptor(operands, attributes, properties, regions);
  auto indicesType = cast<RankedTensorType>(adaptor.getIndices().getType());
  auto srcType = cast<RankedTensorType>(adaptor.getSrc().getType());

  // Shape and encoding of the indices with the element type of the src.
  inferredReturnTypes.push_back(indicesType.clone(srcType.getElementType()));
  return success();
}

// -- DescriptorGatherOp
LogicalResult verifyGatherScatterResultType(Operation *op,
                                            ShapedType resultType,
                                            ShapedType indicesType) {
  if (indicesType.getRank() != 1)
    return op->emitOpError("x offsets must be a 1D tensor, but got ")
           << indicesType;
  if (resultType.getRank() != 2)
    return op->emitOpError("result must be a 2D tensor, but got ")
           << resultType;

  // The swizzling of TMA accesses matches that of the MMAv3 shared memory
  // layouts. However, these have minimum size requirements.
  // TODO: We can support smaller gather sizes by padding the `local_alloc` this
  // lowers to to the nearest minimum tile size.
  if (unsigned rows = resultType.getShape()[0]; rows < 8) {
    return op->emitOpError("must have at least 8 rows, but got ") << rows;
  }

  Type dtype = resultType.getElementType();
  if (dtype.getIntOrFloatBitWidth() > 32)
    return op->emitOpError("TMA dtype cannot be greater than 32 bits");

  unsigned minCols = 32 / dtype.getIntOrFloatBitWidth() * 8;
  if (unsigned cols = resultType.getShape()[1]; cols < minCols) {
    return op->emitOpError("must have at least ")
           << minCols << " columns for " << dtype << ", but got " << cols;
  }

  if (resultType.getShape()[0] != indicesType.getShape()[0]) {
    return op->emitOpError("result tensor must have as many rows as indices (")
           << indicesType.getShape()[0] << "), but got " << resultType;
  }

  return success();
}

LogicalResult verifyGatherScatterOp(Operation *op, ShapedType blockType,
                                    ShapedType resultType,
                                    ShapedType indicesType) {
  // Gather from `!tt.tensordesc<1xMxdtype>`.
  if (blockType.getRank() != 2) {
    return op->emitOpError("descriptor block must be a 2D tensor, but got ")
           << blockType;
  }
  if (blockType.getShape()[0] != 1) {
    return op->emitOpError("descriptor block must have exactly 1 row, but got ")
           << blockType;
  }

  // With x offsets `tensor<Nxinttype>` into `tensor<NxMxdtype>`.
  if (failed(verifyGatherScatterResultType(op, resultType, indicesType)))
    return failure();

  if (resultType.getShape()[1] != blockType.getShape()[1]) {
    return op->emitOpError("result tensor number of columns must match block (")
           << blockType.getShape()[1] << "), but got " << resultType;
  }
  if (resultType.getElementType() != blockType.getElementType()) {
    return op->emitOpError("result tensor element type must match block (")
           << blockType.getElementType() << "), but got " << resultType;
  }

  return success();
}

LogicalResult DescriptorGatherOp::verify() {
  return verifyGatherScatterOp(*this,
                               getDesc().getType().getSignlessBlockType(),
                               getResult().getType(), getXOffsets().getType());
}

// -- DescriptorScatterOp --
LogicalResult DescriptorScatterOp::verify() {
  return verifyGatherScatterOp(*this,
                               getDesc().getType().getSignlessBlockType(),
                               getSrc().getType(), getXOffsets().getType());
}

// -- DescriptorLoadOp --
LogicalResult verifyDescriptorLoadStoreOp(Operation *op,
                                          TensorDescInterface desc,
                                          ShapedType tensor) {
  RankedTensorType block = desc.getSignlessBlockType();
  if (block.getElementType() != tensor.getElementType()) {
    return op->emitOpError("descriptor block and tensor element types must "
                           "match, but got descriptor element type ")
           << block.getElementType() << " and tensor element type "
           << tensor.getElementType();
  }

  ArrayRef<int64_t> blockShape = block.getShape();
  ArrayRef<int64_t> tensorShape = tensor.getShape();
  unsigned blockNumels = product(blockShape);
  unsigned tensorNumels = product(tensorShape);
  if (blockNumels != tensorNumels) {
    return op->emitOpError("descriptor block and tensor must have the same "
                           "number of elements, but got descriptor block "
                           "with ")
           << blockNumels << " elements tensor with " << tensorNumels
           << " elements";
  }
  return success();
}

LogicalResult DescriptorLoadOp::verify() {
  return verifyDescriptorLoadStoreOp(*this, getDesc().getType(), getType());
}

// -- DescriptorStoreOp --
LogicalResult DescriptorStoreOp::verify() {
  return verifyDescriptorLoadStoreOp(*this, getDesc().getType(),
                                     getSrc().getType());
}

// -- DescriptorReduceOp --
LogicalResult DescriptorReduceOp::verify() {
  return verifyDescriptorLoadStoreOp(*this, getDesc().getType(),
                                     getSrc().getType());
}

} // namespace triton
} // namespace mlir
