// Copyright 2022 The IREE Authors
//
// Licensed 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

#include "iree/compiler/ExternalInterfaces/UtilExternalModels.h"

#include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"

#include <numeric>

namespace mlir::iree_compiler {

namespace {

//===----------------------------------------------------------------------===//
// InferIntDivisibilityOpInterface
//===----------------------------------------------------------------------===//

static IREE::Util::ConstantIntDivisibility
getDivisibilityOfOperand(Value v,
                         IREE::Util::IntegerDivisibility divisibility) {
  if (!divisibility.isUninitialized()) {
    return divisibility.getValue();
  }
  APInt intVal;
  if (matchPattern(v, m_ConstantInt(&intVal))) {
    uint64_t udiv = intVal.getZExtValue();
    uint64_t sdiv = std::abs(intVal.getSExtValue());
    return IREE::Util::ConstantIntDivisibility(udiv, sdiv);
  }
  return IREE::Util::ConstantIntDivisibility(1, 1);
}

/// Visits affine expressions and recursively calculates the divisibilities of
/// each subexpression. The final divisibilities of the expression and its
/// subexpressions will be stored in the map for which a reference is provided
/// to the AffineExprDivisibilityFinder (i.e., `divisibilityMap`).
class AffineExprDivisibilityFinder
    : public AffineExprVisitor<AffineExprDivisibilityFinder,
                               IREE::Util::ConstantIntDivisibility> {
public:
  using ExprDivisibilityMap =
      llvm::DenseMap<AffineExpr, IREE::Util::ConstantIntDivisibility>;
  AffineExprDivisibilityFinder(ExprDivisibilityMap &divisibilityMap)
      : divisibilityMap(divisibilityMap) {}

  IREE::Util::ConstantIntDivisibility
  visitConstantExpr(AffineConstantExpr expr) {
    // Constant expressions are trivial, since they are always static.
    uint64_t constValue = std::abs(expr.getValue());
    return IREE::Util::ConstantIntDivisibility(constValue, constValue);
  }

  IREE::Util::ConstantIntDivisibility visitDimExpr(AffineDimExpr expr) {
    // Dim expressions cannot be analyzed further, so return the divisibility
    // in `divisibilityMap` if it has been populated by the caller, or fallback
    // to the minimum divisibility.
    if (divisibilityMap.contains(expr)) {
      return divisibilityMap[expr];
    }
    return IREE::Util::IntegerDivisibility::getMinDivisibility().getValue();
  }

  IREE::Util::ConstantIntDivisibility visitSymbolExpr(AffineSymbolExpr expr) {
    // Symbol expressions cannot be analyzed further, so return the divisibility
    // in `divisibilityMap` if it has been populated by the caller, or fallback
    // to the minimum divisibility.
    if (divisibilityMap.contains(expr)) {
      return divisibilityMap[expr];
    }
    return IREE::Util::IntegerDivisibility::getMinDivisibility().getValue();
  }

  /// Infer the divisibility of an addition or subtraction expression by
  /// recursively visiting the LHS and RHS, and then unioning the results.
  IREE::Util::ConstantIntDivisibility visitAddExpr(AffineBinaryOpExpr expr) {
    if (divisibilityMap.contains(expr)) {
      return divisibilityMap[expr];
    }
    // The divisibility of an addition is the GCD of its constituents'
    // divisibilities.
    IREE::Util::ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
    IREE::Util::ConstantIntDivisibility rhsDiv = visit(expr.getRHS());
    return lhsDiv.getUnion(rhsDiv);
  }

  /// Infer the divisibility of a multiplication expression by recursively
  /// visiting the LHS and RHS, and then multiplying the results.
  IREE::Util::ConstantIntDivisibility visitMulExpr(AffineBinaryOpExpr expr) {
    if (divisibilityMap.contains(expr)) {
      return divisibilityMap[expr];
    }
    // The divisibility of a multiplication is the product of its constituents'
    // divisibilities.
    IREE::Util::ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
    IREE::Util::ConstantIntDivisibility rhsDiv = visit(expr.getRHS());
    return IREE::Util::ConstantIntDivisibility(lhsDiv.udiv() * rhsDiv.udiv(),
                                               lhsDiv.sdiv() * rhsDiv.sdiv());
  }

  IREE::Util::ConstantIntDivisibility
  visitFloorDivExpr(AffineBinaryOpExpr expr) {
    return visitDivExpr(expr);
  }

  IREE::Util::ConstantIntDivisibility
  visitCeilDivExpr(AffineBinaryOpExpr expr) {
    return visitDivExpr(expr);
  }

  /// Infer the divisibility of a mod expression. If the RHS is a constant,
  /// the result divisibility is gcd(lhs_divisibility, rhs_constant), since
  /// (d * k) mod c is always divisible by gcd(d, c). Furthermore, if the
  /// LHS divisibility is itself divisible by the constant (i.e., d % c == 0),
  /// then (d * k) mod c is always zero, represented as divisibility 0.
  IREE::Util::ConstantIntDivisibility visitModExpr(AffineBinaryOpExpr expr) {
    if (divisibilityMap.contains(expr)) {
      return divisibilityMap[expr];
    }
    auto constRhs = dyn_cast<AffineConstantExpr>(expr.getRHS());
    if (!constRhs || constRhs.getValue() == 0) {
      return IREE::Util::ConstantIntDivisibility(1, 1);
    }
    auto constValue = static_cast<uint64_t>(std::abs(constRhs.getValue()));
    IREE::Util::ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
    // If the LHS is always a multiple of constValue, x mod constValue is
    // always zero. Divisibility 0 is the lattice top ("divides everything").
    uint64_t modUDiv = (lhsDiv.udiv() % constValue == 0)
                           ? 0
                           : std::gcd(lhsDiv.udiv(), constValue);
    uint64_t modSDiv = (lhsDiv.sdiv() % constValue == 0)
                           ? 0
                           : std::gcd(lhsDiv.sdiv(), constValue);
    return IREE::Util::ConstantIntDivisibility(modUDiv, modSDiv);
  }

private:
  IREE::Util::ConstantIntDivisibility
  visitInvalidExpr(AffineBinaryOpExpr expr) {
    return IREE::Util::IntegerDivisibility::getMinDivisibility().getValue();
  }

  /// Helper shared by ceildiv and floordiv implementations. Returns the minimum
  /// divisibility as a fallback if the divisor is not a constant, because the
  /// divisibility cannot be inferred in this case. If the divisor is a
  /// constant, then this function recursively visits the dividend, and returns
  /// the quotient of the dividend's divisibility with the divisor.
  IREE::Util::ConstantIntDivisibility visitDivExpr(AffineBinaryOpExpr expr) {
    if (divisibilityMap.contains(expr)) {
      return divisibilityMap[expr];
    }
    auto constRhs = dyn_cast<AffineConstantExpr>(expr.getRHS());
    // Division by zero is undefined, so return the minimum divisibility.
    if (!constRhs || constRhs.getValue() == 0) {
      return IREE::Util::ConstantIntDivisibility(1, 1);
    }
    auto constValue = static_cast<uint64_t>(std::abs(constRhs.getValue()));
    IREE::Util::ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
    uint64_t divUDiv =
        lhsDiv.udiv() % constValue == 0 ? lhsDiv.udiv() / constValue : 1;
    uint64_t divSDiv =
        lhsDiv.sdiv() % constValue == 0 ? lhsDiv.sdiv() / constValue : 1;
    return IREE::Util::ConstantIntDivisibility(divUDiv, divSDiv);
  }

  ExprDivisibilityMap &divisibilityMap;
};

/// Returns the divisibilities of each AffineMap result based on the
/// divisibilities of its dims and symbols. The `dimAndSymbolDivisibilities`
/// should contain the divisibilities of the dims, followed by the
/// divisibilities of the symbols in ascending order by their positions.
static SmallVector<IREE::Util::ConstantIntDivisibility> getResultDivisibilities(
    AffineMap map,
    ArrayRef<IREE::Util::ConstantIntDivisibility> dimAndSymbolDivisibilities) {
  // Seed the AffineExprDivisibilityFinder with the dimAndSymbolDivisibilities.
  llvm::DenseMap<AffineExpr, IREE::Util::ConstantIntDivisibility>
      exprDivisibilityMap;
  SmallVector<AffineExpr> inputExprs;
  inputExprs.append(llvm::map_to_vector(
      llvm::seq<int64_t>(map.getNumDims()),
      [&](int64_t dim) { return getAffineDimExpr(dim, map.getContext()); }));
  inputExprs.append(llvm::map_to_vector(
      llvm::seq<int64_t>(map.getNumSymbols()),
      [&](int64_t sym) { return getAffineSymbolExpr(sym, map.getContext()); }));
  for (auto [expr, divisibility] :
       llvm::zip_equal(inputExprs, dimAndSymbolDivisibilities)) {
    exprDivisibilityMap[expr] = divisibility;
  }
  AffineExprDivisibilityFinder divisibilityFinder(exprDivisibilityMap);

  // Walk each result expression and compute their divisibilities.
  SmallVector<IREE::Util::ConstantIntDivisibility> resultDivisibilities;
  for (AffineExpr resultExpr : map.getResults()) {
    resultDivisibilities.push_back(divisibilityFinder.visit(resultExpr));
  }
  return resultDivisibilities;
}

struct AffineApplyInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          AffineApplyInferIntDivisibilityOpInterface, affine::AffineApplyOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto affineApplyOp = cast<affine::AffineApplyOp>(op);
    SmallVector<IREE::Util::ConstantIntDivisibility> operandDivisibilities;
    for (auto [operand, divisibility] :
         llvm::zip(affineApplyOp.getOperands(), argDivs)) {
      operandDivisibilities.push_back(
          getDivisibilityOfOperand(operand, divisibility));
    }

    SmallVector<IREE::Util::ConstantIntDivisibility> resultDivisibilities =
        getResultDivisibilities(affineApplyOp.getMap(), operandDivisibilities);
    for (auto [result, divisibility] :
         llvm::zip_equal(affineApplyOp->getResults(), resultDivisibilities)) {
      setResultDivs(result, divisibility);
    }
  }
};

/// Infer the result divisibility of an affine.min or affine.max operation
/// based on its operand divisibilities. The result divisibility is the GCD
/// of the divisibilities of each of the affine map results, because the result
/// of the affine.min/max op could be any of these results.
template <typename MinOrMaxTy>
static void inferAffineMinOrMaxResultDivisibility(
    MinOrMaxTy minOrMaxOp, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
    IREE::Util::SetIntDivisibilityFn setResultDivs) {
  static_assert(
      llvm::is_one_of<MinOrMaxTy, affine::AffineMinOp,
                      affine::AffineMaxOp>::value,
      "MinOrMaxTy must be affine::AffineMinOp or affine::AffineMaxOp");
  SmallVector<IREE::Util::ConstantIntDivisibility> operandDivisibilities;
  for (auto [operand, divisibility] :
       llvm::zip(minOrMaxOp.getOperands(), argDivs)) {
    operandDivisibilities.push_back(
        getDivisibilityOfOperand(operand, divisibility));
  }

  SmallVector<IREE::Util::ConstantIntDivisibility> resultDivisibilities =
      getResultDivisibilities(minOrMaxOp.getMap(), operandDivisibilities);

  IREE::Util::ConstantIntDivisibility resultDivisibility =
      resultDivisibilities.pop_back_val();
  for (auto divisibility : resultDivisibilities) {
    resultDivisibility = resultDivisibility.getUnion(divisibility);
  }
  setResultDivs(minOrMaxOp.getResult(), resultDivisibility);
}

struct AffineMinInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          AffineMinInferIntDivisibilityOpInterface, affine::AffineMinOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto affineMinOp = cast<affine::AffineMinOp>(op);
    inferAffineMinOrMaxResultDivisibility(affineMinOp, argDivs, setResultDivs);
  }
};

struct AffineMaxInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          AffineMaxInferIntDivisibilityOpInterface, affine::AffineMaxOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto affineMaxOp = cast<affine::AffineMaxOp>(op);
    inferAffineMinOrMaxResultDivisibility(affineMaxOp, argDivs, setResultDivs);
  }
};

struct AffineDelinearizeIndexInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          AffineDelinearizeIndexInferIntDivisibilityOpInterface,
          affine::AffineDelinearizeIndexOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto delinearizeOp = cast<affine::AffineDelinearizeIndexOp>(op);
    MLIRContext *ctx = op->getContext();

    // Operands are: [linear_index, dynamic_basis_values...]
    IREE::Util::ConstantIntDivisibility linearDiv =
        getDivisibilityOfOperand(delinearizeOp.getLinearIndex(), argDivs[0]);

    ArrayRef<int64_t> staticBasis = delinearizeOp.getStaticBasis();
    int64_t numResults = delinearizeOp.getNumResults();

    // Build affine expressions for each result.
    // Dim 0 = linear index, symbols = dynamic basis values.
    AffineExpr linearExpr = getAffineDimExpr(0, ctx);

    // Collect operand divisibilities: [linear_index_div, dynamic_basis_divs...]
    SmallVector<IREE::Util::ConstantIntDivisibility> operandDivs;
    operandDivs.push_back(linearDiv);

    // Map static/dynamic basis values to affine expressions.
    int64_t dynIdx = 0;
    SmallVector<AffineExpr> basisExprs;
    for (int64_t i = 0, e = static_cast<int64_t>(staticBasis.size()); i < e;
         ++i) {
      if (ShapedType::isDynamic(staticBasis[i])) {
        basisExprs.push_back(getAffineSymbolExpr(dynIdx, ctx));
        operandDivs.push_back(getDivisibilityOfOperand(
            delinearizeOp.getDynamicBasis()[dynIdx], argDivs[1 + dynIdx]));
        dynIdx++;
      } else {
        basisExprs.push_back(getAffineConstantExpr(staticBasis[i], ctx));
      }
    }

    // The computation basis skips the outer bound if present.
    bool hasOuter = delinearizeOp.hasOuterBound();
    int64_t basisStart = hasOuter ? 1 : 0;

    // Each result[i] can be expressed as an affine expression of the linear
    // index using the effective basis (after dropping outer bound if present).
    // Effective basis B[k] = basisExprs[basisStart + k], for k = 0..N-2.
    // Stride s[i] = product of B[i..N-2] = product of
    //               basisExprs[basisStart+i .. end].
    //
    // result[0]   = x floordiv s[0]
    // result[i>0] = (x floordiv s[i]) mod B[i-1]
    // For i=N-1, s[N-1]=1, so result[N-1] = x mod B[N-2].

    AffineExpr stride = getAffineConstantExpr(1, ctx);
    for (int64_t i = numResults - 1; i >= 0; --i) {
      AffineExpr resultExpr;
      if (i == 0) {
        resultExpr = linearExpr.floorDiv(stride);
      } else {
        resultExpr =
            (linearExpr.floorDiv(stride)) % basisExprs[basisStart + i - 1];
      }

      AffineMap resultMap = AffineMap::get(1, dynIdx, resultExpr, ctx);
      SmallVector<IREE::Util::ConstantIntDivisibility> divs =
          getResultDivisibilities(resultMap, operandDivs);
      setResultDivs(delinearizeOp.getResult(i), divs[0]);

      if (i > 0) {
        stride = basisExprs[basisStart + i - 1] * stride;
      }
    }
  }
};

/// Helper for binary arith ops whose result divisibility is the GCD (union) of
/// their operands' divisibilities. This covers add, sub, min, and max.
template <typename OpTy>
struct ArithBinaryGCDInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          ArithBinaryGCDInferIntDivisibilityOpInterface<OpTy>, OpTy> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto binOp = cast<OpTy>(op);
    auto lhsDiv = getDivisibilityOfOperand(binOp.getLhs(), argDivs[0]);
    auto rhsDiv = getDivisibilityOfOperand(binOp.getRhs(), argDivs[1]);
    setResultDivs(binOp.getResult(), lhsDiv.getUnion(rhsDiv));
  }
};

/// For arith.select, the result divisibility is the GCD of the true and false
/// operands' divisibilities. The condition (operand 0) is i1 and irrelevant.
struct ArithSelectInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          ArithSelectInferIntDivisibilityOpInterface, arith::SelectOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto selectOp = cast<arith::SelectOp>(op);
    // argDivs[0] is the condition (i1), argDivs[1] is true, argDivs[2] is
    // false.
    auto trueDiv =
        getDivisibilityOfOperand(selectOp.getTrueValue(), argDivs[1]);
    auto falseDiv =
        getDivisibilityOfOperand(selectOp.getFalseValue(), argDivs[2]);
    setResultDivs(selectOp.getResult(), trueDiv.getUnion(falseDiv));
  }
};

struct ArithConstantInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          ArithConstantInferIntDivisibilityOpInterface, arith::ConstantOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto constOp = cast<arith::ConstantOp>(op);
    auto constAttr = dyn_cast_if_present<IntegerAttr>(constOp.getValue());
    if (constAttr) {
      const APInt &value = constAttr.getValue();
      uint64_t udiv = value.getZExtValue();
      uint64_t sdiv = std::abs(value.getSExtValue());
      setResultDivs(constOp.getResult(),
                    IREE::Util::ConstantIntDivisibility(udiv, sdiv));
    }
  }
};

struct ArithMulIInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          ArithMulIInferIntDivisibilityOpInterface, arith::MulIOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto mulOp = cast<arith::MulIOp>(op);

    auto lhsDivisibility = getDivisibilityOfOperand(mulOp.getLhs(), argDivs[0]);
    auto rhsDivisibility = getDivisibilityOfOperand(mulOp.getRhs(), argDivs[1]);

    uint64_t mulUDiv = lhsDivisibility.udiv() * rhsDivisibility.udiv();
    uint64_t mulSDiv = lhsDivisibility.sdiv() * rhsDivisibility.sdiv();

    setResultDivs(mulOp.getResult(),
                  IREE::Util::ConstantIntDivisibility(mulUDiv, mulSDiv));
  }
};

struct ArithDivUIInferIntDivisibilityOpInterface
    : IREE::Util::InferIntDivisibilityOpInterface::ExternalModel<
          ArithDivUIInferIntDivisibilityOpInterface, arith::DivUIOp> {

  void inferResultDivisibility(
      Operation *op, ArrayRef<IREE::Util::IntegerDivisibility> argDivs,
      IREE::Util::SetIntDivisibilityFn setResultDivs) const {
    auto divOp = cast<arith::DivUIOp>(op);

    APInt intVal;
    if (!matchPattern(divOp.getRhs(), m_ConstantInt(&intVal))) {
      return;
    }

    auto lhsDivisibility = getDivisibilityOfOperand(divOp.getLhs(), argDivs[0]);

    uint64_t divUDiv = lhsDivisibility.udiv() % intVal.getZExtValue() == 0
                           ? lhsDivisibility.udiv() / intVal.getZExtValue()
                           : 1;
    uint64_t divSDiv =
        lhsDivisibility.sdiv() % std::abs(intVal.getSExtValue()) == 0
            ? lhsDivisibility.sdiv() / std::abs(intVal.getSExtValue())
            : 1;

    setResultDivs(divOp, IREE::Util::ConstantIntDivisibility(divUDiv, divSDiv));
  }
};

//===----------------------------------------------------------------------===//
// ValueBoundsOpInterface
//===----------------------------------------------------------------------===//

/// For some reason, this interface has to be done as an external model.
struct UtilAssumeIntValueBoundsOpInterface
    : ValueBoundsOpInterface::ExternalModel<UtilAssumeIntValueBoundsOpInterface,
                                            IREE::Util::AssumeIntOp> {
  void populateBoundsForIndexValue(Operation *op, Value value,
                                   ValueBoundsConstraintSet &cstr) const {
    auto assumeOp = cast<IREE::Util::AssumeIntOp>(op);
    auto result = cast<OpResult>(value);
    assert(result.getOwner() == op && "value is a result of this index op");
    auto [min, max] =
        assumeOp.getUnionedUnsignedRange(result.getResultNumber());

    std::optional<int64_t> udiv =
        assumeOp.getUnionedUnsignedDivisor(result.getResultNumber());

    if (min) {
      cstr.bound(result) >= *min;
    }
    if (max) {
      cstr.bound(result) <= *max;
    }
    if (udiv) {
      // To represent the divisibility guarantee, emit a bound clamping the
      // value to the udiv value. i.e.
      //
      // v == floordiv(v, udiv) * udiv
      //
      // Mod/divide folders can cleanup such terms with the appropriate bounds
      // query.
      AffineExpr expr =
          cstr.getExpr(assumeOp.getOperand(result.getResultNumber()));
      AffineExpr udivCst =
          getAffineConstantExpr(udiv.value(), op->getContext());
      AffineExpr clampExpr = expr.floorDiv(udivCst) * udivCst;
      cstr.bound(result) == clampExpr;
    }
  }
};

//===----------------------------------------------------------------------===//
// GlobalOpInterface
//===----------------------------------------------------------------------===//

struct GlobalOpInterfaceExternalModel
    : IREE::Util::GlobalOpInterface::ExternalModel<
          GlobalOpInterfaceExternalModel, ml_program::GlobalOp> {
  Attribute getGlobalInitialValue(Operation *op) const {
    return cast<ml_program::GlobalOp>(op).getValueAttr();
  }
  void setGlobalInitialValue(Operation *op, Attribute value) const {
    if (value) {
      cast<ml_program::GlobalOp>(op).setValueAttr(value);
    } else {
      cast<ml_program::GlobalOp>(op).removeValueAttr();
    }
  }

  IREE::Util::InliningPolicyAttrInterface
  getGlobalInliningPolicy(Operation *op) const {
    if (op->hasAttr("noinline")) {
      return IREE::Util::InlineNeverAttr::get(op->getContext());
    }
    return {};
  }
  void
  setGlobalInliningPolicy(Operation *op,
                          IREE::Util::InliningPolicyAttrInterface value) const {
    if (isa_and_nonnull<IREE::Util::InlineNeverAttr>(value)) {
      op->setAttr("noinline", UnitAttr::get(op->getContext()));
    } else {
      op->removeAttr("noinline");
    }
  }

  IREE::Util::GlobalLoadOpInterface createLoadOp(Operation *op, Location loc,
                                                 OpBuilder &builder) const {
    auto globalOp = cast<ml_program::GlobalOp>(op);
    if (globalOp.getIsMutable()) {
      return cast<IREE::Util::GlobalLoadOpInterface>(
          ml_program::GlobalLoadOp::create(builder, loc, globalOp.getType(),
                                           FlatSymbolRefAttr::get(globalOp))
              .getOperation());
    } else {
      return cast<IREE::Util::GlobalLoadOpInterface>(
          ml_program::GlobalLoadConstOp::create(
              builder, loc, globalOp.getType(),
              FlatSymbolRefAttr::get(globalOp))
              .getOperation());
    }
  }

  IREE::Util::GlobalStoreOpInterface createStoreOp(Operation *op, Location loc,
                                                   Value value,
                                                   OpBuilder &builder) const {
    auto globalOp = cast<ml_program::GlobalOp>(op);
    return cast<IREE::Util::GlobalStoreOpInterface>(
        ml_program::GlobalStoreOp ::create(
            builder, loc, FlatSymbolRefAttr::get(globalOp), value)
            .getOperation());
  }
};

//===----------------------------------------------------------------------===//
// NumericCastOpInterface
//===----------------------------------------------------------------------===//

// Since all details of the interface are provided via default implementations,
// we can just have one templated external model to apply per op, vs one
// explicit model per op.
struct GenericNumericCastExternalModel {
  template <typename OpTy>
  struct ExternalModel
      : IREE::Util::NumericCastOpInterface::ExternalModel<ExternalModel<OpTy>,
                                                          OpTy> {};

  template <typename OpTy>
  static void add(MLIRContext *context) {
    OpTy::template attachInterface<ExternalModel<OpTy>>(*context);
  }

  template <typename OpTy1, typename OpTy2, typename... More>
  static void add(MLIRContext *context) {
    add<OpTy1>(context);
    add<OpTy2, More...>(context);
  }
};

//===----------------------------------------------------------------------===//
// TiedOpInterface
//===----------------------------------------------------------------------===//

struct InsertSliceOpTiedOpInterface
    : IREE::Util::TiedOpInterface::ExternalModel<InsertSliceOpTiedOpInterface,
                                                 tensor::InsertSliceOp> {
  Value getTiedResult(Operation *op, unsigned resultIndex) const {
    auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
    return IREE::Util::TiedOpInterface::findTiedBaseValue(
        insertSliceOp.getDest());
  }

  ::std::optional<unsigned>
  getTiedResultOperandIndex(Operation *op, unsigned resultIndex) const {
    return {1}; // dest
  }

  SmallVector<int64_t> getTiedResultOperandIndices(Operation *op) const {
    return {1}; // dest
  }
};

template <typename OpTy>
struct LinalgOpTiedOpInterface
    : IREE::Util::TiedOpInterface::ExternalModel<LinalgOpTiedOpInterface<OpTy>,
                                                 OpTy> {
  Value getTiedResult(Operation *op, unsigned resultIndex) const {
    auto linalgOp = cast<OpTy>(op);
    return IREE::Util::TiedOpInterface::findTiedBaseValue(
        linalgOp.getDpsInits()[resultIndex]);
  }

  ::std::optional<unsigned>
  getTiedResultOperandIndex(Operation *op, unsigned resultIndex) const {
    auto linalgOp = cast<OpTy>(op);
    return {linalgOp.getDpsInitsMutable()[resultIndex].getOperandNumber()};
  }

  SmallVector<int64_t> getTiedResultOperandIndices(Operation *op) const {
    SmallVector<int64_t> result;
    for (unsigned i = 0; i < op->getNumResults(); ++i) {
      result.push_back(*getTiedResultOperandIndex(op, i));
    }
    return result;
  }
};

/// Helper structure that iterates over all LinalgOps in `OpTys` and registers
/// the `TiedOpInterface` with each of them.
template <typename... Ops>
struct LinalgOpTiedOpInterfaceHelper {
  static void registerOpInterface(MLIRContext *context) {
    (void)std::initializer_list<int>{
        0,
        (Ops::template attachInterface<LinalgOpTiedOpInterface<Ops>>(*context),
         0)...};
  }
};

//===----------------------------------------------------------------------===//
// HoistableOpInterface
//===----------------------------------------------------------------------===//

template <typename OpTy>
struct UnhoistableOpInterface : IREE::Util::HoistableOpInterface::ExternalModel<
                                    UnhoistableOpInterface<OpTy>, OpTy> {
  bool isHoistableOp(Operation *) const { return false; }
  bool isHoistableLeafOp(Operation *) const { return false; }
};

template <typename OpTy>
struct HoistableNonLeafOpInterface
    : IREE::Util::HoistableOpInterface::ExternalModel<
          HoistableNonLeafOpInterface<OpTy>, OpTy> {
  bool isHoistableLeafOp(Operation *) const { return false; }
};

// The default interface is always hoistable. This acts as an override
// for other default hoistability checks as the interface is checked
// first.
template <typename OpTy>
struct AlwaysHoistableOpInterface
    : IREE::Util::HoistableOpInterface::ExternalModel<
          AlwaysHoistableOpInterface<OpTy>, OpTy> {};

template <typename OpTy>
struct HoistableLinalgOpInterface
    : IREE::Util::HoistableOpInterface::ExternalModel<
          HoistableLinalgOpInterface<OpTy>, OpTy> {
  bool isHoistableOp(Operation *) const { return true; }

  // Determines if a linalg op is a hoistable leaf, based on heuristics.
  bool isHoistableLeafOp(Operation *op) const {
    // Don't hoist bit extend ops because fusing them with their
    // consumers prevents materializing the high bit-width tensor and they
    // preform very little real computation.
    if (IREE::LinalgExt::isBitExtendOp(op)) {
      return false;
    }

    // Hoist all non-generic linalg ops except for fill ops which should be
    // fused with their consumers.
    auto genericOp = dyn_cast<linalg::GenericOp>(op);
    if (!genericOp) {
      return !isa<linalg::FillOp>(op);
    }

    // Don't hoist ops with no tensor inputs. They are likely to be fill-like
    // or sequences (from `linalg.index`) which can be fused with their
    // consumers.
    if (IREE::LinalgExt::hasOnlyScalarInputs(genericOp)) {
      return false;
    }

    // Don't hoist broadcast-like ops because fusing them makes the new
    // op cheaper.
    if (linalg::isaBroadcastOpInterface(genericOp).has_value()) {
      return false;
    }

    // Hoist all other ops.
    return true;
  }
  bool isAtomicallyHoistableOp(Operation *) const { return true; }
  bool isOperandHoistable(Operation *, OpOperand *) const { return true; }
};

/// Helper structures that iterates over all Op types in `OpTys` and registers
/// the associated Hoistable___OpInterface.
template <typename... Ops>
struct UnhoistableOpInterfaceHelper {
  static void registerOpInterface(MLIRContext *context) {
    (Ops::template attachInterface<UnhoistableOpInterface<Ops>>(*context), ...);
  }
};

template <typename... Ops>
struct HoistableNonLeafOpInterfaceHelper {
  static void registerOpInterface(MLIRContext *context) {
    (Ops::template attachInterface<HoistableNonLeafOpInterface<Ops>>(*context),
     ...);
  }
};

template <typename... Ops>
struct AlwaysHoistableOpInterfaceHelper {
  static void registerOpInterface(MLIRContext *context) {
    (Ops::template attachInterface<AlwaysHoistableOpInterface<Ops>>(*context),
     ...);
  }
};

template <typename... Ops>
struct HoistableLinalgOpInterfaceHelper {
  static void registerOpInterface(MLIRContext *context) {
    (Ops::template attachInterface<HoistableLinalgOpInterface<Ops>>(*context),
     ...);
  }
};

//===----------------------------------------------------------------------===//
// MutableRegionBranchOpInterface
//===----------------------------------------------------------------------===//

// External model for scf.for operation.
struct SCFForOpMutableRegionBranchOpInterface
    : IREE::Util::MutableRegionBranchOpInterface::ExternalModel<
          SCFForOpMutableRegionBranchOpInterface, scf::ForOp> {
  Operation *rebuildWithExpandedTypes(
      Operation *op,
      llvm::function_ref<void(Type, SmallVectorImpl<Type> &)> expandTypeFn,
      llvm::function_ref<void(Value, SmallVectorImpl<Value> &, OpBuilder &)>
          expandOperandFn,
      std::optional<
          llvm::function_ref<Value(Location, Value, ValueRange, OpBuilder &)>>
          wrapExpandedBlockArgFn,
      llvm::function_ref<void(Region &, bool /*canModifyEntryBlock*/)>
          expandRegionFn,
      OpBuilder &builder) const {
    auto forOp = cast<scf::ForOp>(op);
    Location loc = forOp.getLoc();

    // Expand iter_args operands (lb, ub, step are never expanded).
    SmallVector<Value> newIterArgs;
    for (Value iterArg : forOp.getInitArgs()) {
      expandOperandFn(iterArg, newIterArgs, builder);
    }

    // Expand result types.
    SmallVector<Type> newResultTypes;
    for (Type resultType : forOp.getResultTypes()) {
      expandTypeFn(resultType, newResultTypes);
    }

    // Create new for loop with expanded signature.
    auto newForOp =
        scf::ForOp::create(builder, loc, forOp.getLowerBound(),
                           forOp.getUpperBound(), forOp.getStep(), newIterArgs);

    // Move the body from old to new.
    newForOp.getBodyRegion().takeBody(forOp.getBodyRegion());

    // The new op now has the old body but with unexpanded block arguments.
    // We need to update the block arguments to match the expanded iter args.
    Block &body = *newForOp.getBody();

    // Insert new block arguments for expanded types.
    // Start from the back to preserve indices.
    // Note we skip the induction var at 0.
    for (int i = body.getNumArguments() - 1; i >= 1; --i) {
      auto arg = body.getArgument(i);
      SmallVector<Type> expandedTypes;
      expandTypeFn(arg.getType(), expandedTypes);
      if (expandedTypes.size() > 1) {
        // This type expands to multiple values.
        // Insert new arguments after the current one.
        SmallVector<Value> expandedArgs;
        expandedArgs.push_back(arg);
        for (unsigned j = 1; j < expandedTypes.size(); ++j) {
          expandedArgs.push_back(
              body.insertArgument(i + j, expandedTypes[j], arg.getLoc()));
        }

        // If wrapper callback provided, create replacement value and update
        // uses.
        if (wrapExpandedBlockArgFn) {
          OpBuilder blockBuilder(&body, body.begin());
          Value replacement = (*wrapExpandedBlockArgFn)(
              arg.getLoc(), arg, expandedArgs, blockBuilder);
          arg.replaceAllUsesExcept(replacement, replacement.getDefiningOp());
        }
      }
    }

    // Recursively expand operations in the body.
    expandRegionFn(newForOp.getBodyRegion(), /*canModifyEntryBlock=*/false);

    // Update the yield to match expanded result types.
    if (body.mightHaveTerminator()) {
      auto *terminator = body.getTerminator();
      if (auto yieldOp = dyn_cast<scf::YieldOp>(terminator)) {
        SmallVector<Value> newYieldOperands;
        builder.setInsertionPoint(yieldOp);
        for (Value operand : yieldOp.getOperands()) {
          expandOperandFn(operand, newYieldOperands, builder);
        }
        scf::YieldOp::create(builder, yieldOp.getLoc(), newYieldOperands);
        yieldOp.erase();
      }
    }

    return newForOp;
  }

  llvm::BitVector getPreservedBlockArguments(Operation *op,
                                             unsigned regionIndex) const {
    auto forOp = cast<scf::ForOp>(op);
    llvm::BitVector preserved(forOp.getBody()->getNumArguments());
    preserved.set(0); // Preserve induction variable.
    return preserved;
  }

  OperandRange getRegionEntryOperands(Operation *op,
                                      unsigned regionIndex) const {
    auto forOp = cast<scf::ForOp>(op);
    return forOp.getInitArgs();
  }

  ResultRange getRegionExitResults(Operation *op, unsigned regionIndex) const {
    return op->getResults();
  }

  OperandRange getExpandableTerminatorOperands(Operation *op,
                                               Operation *terminator,
                                               unsigned regionIndex) const {
    // For scf.yield, all operands are expandable.
    return terminator->getOperands();
  }
};

// External model for scf.if operation.
struct SCFIfOpMutableRegionBranchOpInterface
    : IREE::Util::MutableRegionBranchOpInterface::ExternalModel<
          SCFIfOpMutableRegionBranchOpInterface, scf::IfOp> {
  Operation *rebuildWithExpandedTypes(
      Operation *op,
      llvm::function_ref<void(Type, SmallVectorImpl<Type> &)> expandTypeFn,
      llvm::function_ref<void(Value, SmallVectorImpl<Value> &, OpBuilder &)>
          expandOperandFn,
      std::optional<
          llvm::function_ref<Value(Location, Value, ValueRange, OpBuilder &)>>
          wrapExpandedBlockArgFn,
      llvm::function_ref<void(Region &, bool /*canModifyEntryBlock*/)>
          expandRegionFn,
      OpBuilder &builder) const {
    // Note: scf.if has no block arguments, so wrapExpandedBlockArgFn is unused.
    (void)wrapExpandedBlockArgFn;
    auto ifOp = cast<scf::IfOp>(op);
    Location loc = ifOp.getLoc();

    // Expand result types.
    SmallVector<Type> newResultTypes;
    for (Type resultType : ifOp.getResultTypes()) {
      expandTypeFn(resultType, newResultTypes);
    }

    // Create new if op with expanded result types.
    auto newIfOp =
        scf::IfOp::create(builder, loc, newResultTypes, ifOp.getCondition(),
                          /*withElseRegion=*/!ifOp.getElseRegion().empty());

    // Copy regions from old to new op.
    // Use takeBody to move the region contents.
    newIfOp.getThenRegion().takeBody(ifOp.getThenRegion());
    if (!ifOp.getElseRegion().empty()) {
      newIfOp.getElseRegion().takeBody(ifOp.getElseRegion());
    }

    // Now recursively expand operations in the regions.
    // This will handle the expansion of operations inside.
    expandRegionFn(newIfOp.getThenRegion(), /*canModifyEntryBlock=*/false);
    if (!newIfOp.getElseRegion().empty()) {
      expandRegionFn(newIfOp.getElseRegion(), /*canModifyEntryBlock=*/false);
    }

    // Now we need to update the yield operations to match the expanded result
    // types. The yield ops need to yield expanded values for resource results.
    auto updateYield = [&](Block &block) {
      if (!block.mightHaveTerminator()) {
        return;
      }
      auto *terminator = block.getTerminator();
      if (!terminator) {
        return;
      }
      auto yieldOp = dyn_cast<scf::YieldOp>(terminator);
      if (!yieldOp) {
        return;
      }

      SmallVector<Value> newYieldOperands;
      builder.setInsertionPoint(yieldOp);
      for (Value operand : yieldOp.getOperands()) {
        expandOperandFn(operand, newYieldOperands, builder);
      }
      scf::YieldOp::create(builder, yieldOp.getLoc(), newYieldOperands);
      yieldOp.erase();
    };

    updateYield(newIfOp.getThenRegion().front());
    if (!newIfOp.getElseRegion().empty()) {
      updateYield(newIfOp.getElseRegion().front());
    }

    return newIfOp;
  }

  llvm::BitVector getPreservedBlockArguments(Operation *op,
                                             unsigned regionIndex) const {
    // scf.if has no block arguments.
    return llvm::BitVector(0);
  }

  OperandRange getRegionEntryOperands(Operation *op,
                                      unsigned regionIndex) const {
    // No operands map to block args.
    return OperandRange(op->operand_end(), op->operand_end());
  }

  ResultRange getRegionExitResults(Operation *op, unsigned regionIndex) const {
    return op->getResults();
  }

  OperandRange getExpandableTerminatorOperands(Operation *op,
                                               Operation *terminator,
                                               unsigned regionIndex) const {
    // For scf.yield, all operands are expandable.
    return terminator->getOperands();
  }
};

// External model for scf.while operation.
struct SCFWhileOpMutableRegionBranchOpInterface
    : IREE::Util::MutableRegionBranchOpInterface::ExternalModel<
          SCFWhileOpMutableRegionBranchOpInterface, scf::WhileOp> {
  Operation *rebuildWithExpandedTypes(
      Operation *op,
      llvm::function_ref<void(Type, SmallVectorImpl<Type> &)> expandTypeFn,
      llvm::function_ref<void(Value, SmallVectorImpl<Value> &, OpBuilder &)>
          expandOperandFn,
      std::optional<
          llvm::function_ref<Value(Location, Value, ValueRange, OpBuilder &)>>
          wrapExpandedBlockArgFn,
      llvm::function_ref<void(Region &, bool /*canModifyEntryBlock*/)>
          expandRegionFn,
      OpBuilder &builder) const {
    auto whileOp = cast<scf::WhileOp>(op);
    Location loc = whileOp.getLoc();

    // Expand initial operands.
    SmallVector<Value> newInits;
    for (Value init : whileOp.getInits()) {
      expandOperandFn(init, newInits, builder);
    }

    // Expand result types.
    SmallVector<Type> newResultTypes;
    for (Type resultType : whileOp.getResultTypes()) {
      expandTypeFn(resultType, newResultTypes);
    }

    // scf.while builder needs the before and after region builders.
    // We'll create it without them and then manually populate.
    auto newWhileOp =
        scf::WhileOp::create(builder, loc, newResultTypes, newInits,
                             /*beforeBuilder=*/nullptr,
                             /*afterBuilder=*/nullptr);

    // Move the regions from old to new.
    newWhileOp.getBefore().takeBody(whileOp.getBefore());
    newWhileOp.getAfter().takeBody(whileOp.getAfter());

    // Helper lambda to expand block arguments and call wrapper.
    auto expandBlockArgs = [&](Block &block) {
      for (int i = block.getNumArguments() - 1; i >= 0; --i) {
        auto arg = block.getArgument(i);
        SmallVector<Type> expandedTypes;
        expandTypeFn(arg.getType(), expandedTypes);
        if (expandedTypes.size() > 1) {
          // This type expands to multiple values.
          // Insert new arguments after the current one.
          SmallVector<Value> expandedArgs;
          expandedArgs.push_back(arg);
          for (unsigned j = 1; j < expandedTypes.size(); ++j) {
            expandedArgs.push_back(
                block.insertArgument(i + j, expandedTypes[j], arg.getLoc()));
          }

          // If wrapper callback provided, create replacement value and update
          // uses.
          if (wrapExpandedBlockArgFn) {
            OpBuilder blockBuilder(&block, block.begin());
            Value replacement = (*wrapExpandedBlockArgFn)(
                arg.getLoc(), arg, expandedArgs, blockBuilder);
            arg.replaceAllUsesExcept(replacement, replacement.getDefiningOp());
          }
        }
      }
    };

    // Expand block arguments in both regions.
    Block &beforeBlock = newWhileOp.getBefore().front();
    expandBlockArgs(beforeBlock);

    Block &afterBlock = newWhileOp.getAfter().front();
    expandBlockArgs(afterBlock);

    // Recursively expand operations in both regions.
    expandRegionFn(newWhileOp.getBefore(), /*canModifyEntryBlock=*/false);
    expandRegionFn(newWhileOp.getAfter(), /*canModifyEntryBlock=*/false);

    // Rebuild the condition with expanded operands.
    if (beforeBlock.mightHaveTerminator()) {
      if (auto condOp =
              dyn_cast<scf::ConditionOp>(beforeBlock.getTerminator())) {
        SmallVector<Value> newCondArgs;
        builder.setInsertionPoint(condOp);
        for (Value operand : condOp.getArgs()) {
          expandOperandFn(operand, newCondArgs, builder);
        }
        scf::ConditionOp::create(builder, condOp.getLoc(),
                                 condOp.getCondition(), newCondArgs);
        condOp.erase();
      }
    }

    // Rebuild the yield in the after region with expanded operands.
    if (afterBlock.mightHaveTerminator()) {
      if (auto yieldOp = dyn_cast<scf::YieldOp>(afterBlock.getTerminator())) {
        SmallVector<Value> newYieldOperands;
        builder.setInsertionPoint(yieldOp);
        for (Value operand : yieldOp.getOperands()) {
          expandOperandFn(operand, newYieldOperands, builder);
        }
        scf::YieldOp::create(builder, yieldOp.getLoc(), newYieldOperands);
        yieldOp.erase();
      }
    }

    return newWhileOp;
  }

  llvm::BitVector getPreservedBlockArguments(Operation *op,
                                             unsigned regionIndex) const {
    // All block arguments in scf.while can be expanded.
    auto whileOp = cast<scf::WhileOp>(op);
    Region &region =
        regionIndex == 0 ? whileOp.getBefore() : whileOp.getAfter();
    return llvm::BitVector(region.front().getNumArguments());
  }

  OperandRange getRegionEntryOperands(Operation *op,
                                      unsigned regionIndex) const {
    auto whileOp = cast<scf::WhileOp>(op);
    if (regionIndex == 0) {
      // Before region gets the init operands.
      return whileOp.getInits();
    }
    // After region receives operands from the condition region's scf.condition
    // op, which are not direct operands of the scf.while itself. Return empty
    // range.
    return OperandRange(op->operand_end(), op->operand_end());
  }

  ResultRange getRegionExitResults(Operation *op, unsigned regionIndex) const {
    if (regionIndex == 1) {
      // After region produces the results.
      return op->getResults();
    }
    // Before region doesn't directly produce results.
    return ResultRange(op->result_end(), op->result_end());
  }

  OperandRange getExpandableTerminatorOperands(Operation *op,
                                               Operation *terminator,
                                               unsigned regionIndex) const {
    if (auto condOp = dyn_cast<scf::ConditionOp>(terminator)) {
      // For condition, expand the args but not the condition boolean.
      return condOp.getArgs();
    }
    // For scf.yield, all operands are expandable.
    return terminator->getOperands();
  }
};

// External model for scf.index_switch operation.
struct SCFIndexSwitchOpMutableRegionBranchOpInterface
    : IREE::Util::MutableRegionBranchOpInterface::ExternalModel<
          SCFIndexSwitchOpMutableRegionBranchOpInterface, scf::IndexSwitchOp> {
  Operation *rebuildWithExpandedTypes(
      Operation *op,
      llvm::function_ref<void(Type, SmallVectorImpl<Type> &)> expandTypeFn,
      llvm::function_ref<void(Value, SmallVectorImpl<Value> &, OpBuilder &)>
          expandOperandFn,
      std::optional<
          llvm::function_ref<Value(Location, Value, ValueRange, OpBuilder &)>>
          wrapExpandedBlockArgFn,
      llvm::function_ref<void(Region &, bool /*canModifyEntryBlock*/)>
          expandRegionFn,
      OpBuilder &builder) const {
    // Note: scf.index_switch has no block arguments, so wrapExpandedBlockArgFn
    // is unused.
    (void)wrapExpandedBlockArgFn;
    auto switchOp = cast<scf::IndexSwitchOp>(op);
    Location loc = switchOp.getLoc();

    // Expand result types.
    SmallVector<Type> newResultTypes;
    for (Type resultType : switchOp.getResultTypes()) {
      expandTypeFn(resultType, newResultTypes);
    }

    // Create new index_switch op with expanded result types.
    auto newSwitchOp = scf::IndexSwitchOp::create(
        builder, loc, newResultTypes, switchOp.getArg(), switchOp.getCases(),
        switchOp.getCases().size());

    // Clone each case region.
    for (unsigned i = 0; i < switchOp.getCaseRegions().size(); ++i) {
      Region &oldRegion = switchOp.getCaseRegions()[i];
      Region &newRegion = newSwitchOp.getCaseRegions()[i];
      if (oldRegion.empty()) {
        continue;
      }
      builder.createBlock(&newRegion);
      builder.setInsertionPointToStart(&newRegion.front());

      IRMapping mapping;
      for (Operation &op : oldRegion.front()) {
        if (auto yieldOp = dyn_cast<scf::YieldOp>(op)) {
          SmallVector<Value> newYieldOperands;
          for (Value operand : yieldOp.getOperands()) {
            if (mapping.contains(operand)) {
              expandOperandFn(mapping.lookup(operand), newYieldOperands,
                              builder);
            } else {
              expandOperandFn(operand, newYieldOperands, builder);
            }
          }
          scf::YieldOp::create(builder, yieldOp.getLoc(), newYieldOperands);
        } else {
          builder.clone(op, mapping);
        }
      }
    }

    // Clone default region.
    Region &oldDefault = switchOp.getDefaultRegion();
    Region &newDefault = newSwitchOp.getDefaultRegion();
    if (!oldDefault.empty()) {
      builder.createBlock(&newDefault);
      builder.setInsertionPointToStart(&newDefault.front());

      IRMapping mapping;
      for (Operation &op : oldDefault.front()) {
        if (auto yieldOp = dyn_cast<scf::YieldOp>(op)) {
          SmallVector<Value> newYieldOperands;
          for (Value operand : yieldOp.getOperands()) {
            if (mapping.contains(operand)) {
              expandOperandFn(mapping.lookup(operand), newYieldOperands,
                              builder);
            } else {
              expandOperandFn(operand, newYieldOperands, builder);
            }
          }
          scf::YieldOp::create(builder, yieldOp.getLoc(), newYieldOperands);
        } else {
          builder.clone(op, mapping);
        }
      }
    }

    // Recursively expand operations in all regions.
    for (Region &region : newSwitchOp.getCaseRegions()) {
      if (!region.empty()) {
        expandRegionFn(region, /*canModifyEntryBlock=*/false);
      }
    }
    if (!newSwitchOp.getDefaultRegion().empty()) {
      expandRegionFn(newSwitchOp.getDefaultRegion(),
                     /*canModifyEntryBlock=*/false);
    }

    return newSwitchOp;
  }

  llvm::BitVector getPreservedBlockArguments(Operation *op,
                                             unsigned regionIndex) const {
    // scf.index_switch has no block arguments.
    return llvm::BitVector(0);
  }

  OperandRange getRegionEntryOperands(Operation *op,
                                      unsigned regionIndex) const {
    // No operands map to block args.
    return OperandRange(op->operand_end(), op->operand_end());
  }

  ResultRange getRegionExitResults(Operation *op, unsigned regionIndex) const {
    return op->getResults();
  }

  OperandRange getExpandableTerminatorOperands(Operation *op,
                                               Operation *terminator,
                                               unsigned regionIndex) const {
    // For scf.yield, all operands are expandable.
    return terminator->getOperands();
  }
};

// Hoistable interface for region-containing control flow operations.
// Control flow is hoistable if control operands are constant and
// nested operations are hoistable (checked via atomic hoisting).
template <typename OpTy>
struct RegionControlFlowHoistableOpInterface
    : IREE::Util::HoistableOpInterface::ExternalModel<
          RegionControlFlowHoistableOpInterface<OpTy>, OpTy> {
  bool isHoistableOp(Operation *op) const {
    // Control flow is hoistable if all nested operations are hoistable.
    for (Region &region : op->getRegions()) {
      WalkResult result = region.walk([](Operation *nestedOp) {
        // Check if nested op is hoistable.
        bool isHoistable = false;
        if (auto hoistable =
                dyn_cast<IREE::Util::HoistableOpInterface>(nestedOp)) {
          isHoistable = hoistable.isHoistableOp();
        } else {
          // Ops without interface must be memory-effect-free to be hoistable.
          isHoistable = mlir::isMemoryEffectFree(nestedOp);
        }
        if (!isHoistable) {
          return WalkResult::interrupt();
        }
        // Don't descend into IsolatedFromAbove ops - treat them atomically.
        return nestedOp->hasTrait<OpTrait::IsIsolatedFromAbove>()
                   ? WalkResult::skip()
                   : WalkResult::advance();
      });
      if (result.wasInterrupted()) {
        return false;
      }
    }
    return true;
  }

  bool isHoistableLeafOp(Operation *) const { return false; }
  bool isAtomicallyHoistableOp(Operation *) const { return true; }
  bool isOperandHoistable(Operation *, OpOperand *) const { return true; }
};

template <typename... Ops>
struct RegionControlFlowHoistableOpInterfaceHelper {
  static void registerOpInterface(MLIRContext *context) {
    (Ops::template attachInterface<RegionControlFlowHoistableOpInterface<Ops>>(
         *context),
     ...);
  }
};

} // namespace

void registerUtilExternalModels(DialectRegistry &registry) {
  // Must ensure that any dependent dialects are registered.
  registry.insert<affine::AffineDialect>();
  registry.insert<arith::ArithDialect>();
  registry.insert<linalg::LinalgDialect>();
  registry.insert<ml_program::MLProgramDialect>();
  registry.insert<scf::SCFDialect>();
  registry.insert<tensor::TensorDialect>();

  registry.addExtension(
      +[](MLIRContext *context, ml_program::MLProgramDialect *dialect) {
        ml_program::GlobalOp::attachInterface<GlobalOpInterfaceExternalModel>(
            *context);
      });

  registry.addExtension(+[](MLIRContext *context,
                            arith::ArithDialect *dialect) {
    GenericNumericCastExternalModel::add<
        arith::BitcastOp, arith::ExtFOp, arith::ExtUIOp, arith::ExtSIOp,
        arith::FPToSIOp, arith::FPToUIOp, arith::IndexCastOp, arith::TruncFOp,
        arith::TruncIOp, arith::SIToFPOp, arith::UIToFPOp>(context);
    arith::ConstantOp::attachInterface<
        ArithConstantInferIntDivisibilityOpInterface>(*context);
    arith::MulIOp::attachInterface<ArithMulIInferIntDivisibilityOpInterface>(
        *context);
    arith::DivUIOp::attachInterface<ArithDivUIInferIntDivisibilityOpInterface>(
        *context);
    arith::AddIOp::attachInterface<
        ArithBinaryGCDInferIntDivisibilityOpInterface<arith::AddIOp>>(*context);
    arith::SubIOp::attachInterface<
        ArithBinaryGCDInferIntDivisibilityOpInterface<arith::SubIOp>>(*context);
    arith::MinUIOp::attachInterface<
        ArithBinaryGCDInferIntDivisibilityOpInterface<arith::MinUIOp>>(
        *context);
    arith::MaxUIOp::attachInterface<
        ArithBinaryGCDInferIntDivisibilityOpInterface<arith::MaxUIOp>>(
        *context);
    arith::MinSIOp::attachInterface<
        ArithBinaryGCDInferIntDivisibilityOpInterface<arith::MinSIOp>>(
        *context);
    arith::MaxSIOp::attachInterface<
        ArithBinaryGCDInferIntDivisibilityOpInterface<arith::MaxSIOp>>(
        *context);
    arith::SelectOp::attachInterface<
        ArithSelectInferIntDivisibilityOpInterface>(*context);
  });

  registry.addExtension(
      +[](MLIRContext *context, affine::AffineDialect *dialect) {
        affine::AffineApplyOp::attachInterface<
            AffineApplyInferIntDivisibilityOpInterface>(*context);
        affine::AffineMinOp::attachInterface<
            AffineMinInferIntDivisibilityOpInterface>(*context);
        affine::AffineMaxOp::attachInterface<
            AffineMaxInferIntDivisibilityOpInterface>(*context);
        affine::AffineDelinearizeIndexOp::attachInterface<
            AffineDelinearizeIndexInferIntDivisibilityOpInterface>(*context);
      });

  registry.addExtension(
      +[](MLIRContext *context, tensor::TensorDialect *dialect) {
        tensor::InsertSliceOp::attachInterface<InsertSliceOpTiedOpInterface>(
            *context);
      });

  registry.addExtension(
      +[](MLIRContext *context, linalg::LinalgDialect *dialect) {
        // Register all Linalg structured ops. `LinalgOp` is an interface and it
        // is not possible to attach an external interface to an existing
        // interface. Therefore, attach the `TiedOpInterface` to all ops
        // one-by-one.
        LinalgOpTiedOpInterfaceHelper<
#define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
            >::registerOpInterface(context);
      });

  registry.addExtension(+[](MLIRContext *context,
                            IREE::LinalgExt::IREELinalgExtDialect *dialect) {
    LinalgOpTiedOpInterfaceHelper<
#define GET_OP_LIST
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp.inc"
        >::registerOpInterface(context);
  });

  // Hoistable Op Interface registration.

  // Register hoistable op interfaces for Encoding ops.
  registry.addExtension(
      +[](MLIRContext *context, IREE::Encoding::IREEEncodingDialect *dialect) {
        UnhoistableOpInterfaceHelper<
            IREE::Encoding::SetEncodingOp>::registerOpInterface(context);
      });

  // Register hoistable op interfaces for Flow ops.
  registry.addExtension(
      +[](MLIRContext *context, IREE::Flow::FlowDialect *dialect) {
        UnhoistableOpInterfaceHelper<
            IREE::Flow::DispatchWorkgroupCountOp>::registerOpInterface(context);

        AlwaysHoistableOpInterfaceHelper<
            IREE::Flow::TensorEncodeOp>::registerOpInterface(context);
      });

  // Register hoistable op interfaces for linalg ops.
  // We have a specific allow-list for Linalg ops because we want to consider
  // new additions carefully.
  registry.addExtension(
      +[](MLIRContext *context, linalg::LinalgDialect *dialect) {
        // Structured op implementations and a handful of pure ops are included.
        // Notably: IndexOp is not included because it establishes a hidden
        // dependency to the iterator and is non-const.

        // Register all LinalgOps ops. `LinalgOp` is an interface and it is
        // not possible to attach an external interface to an existing
        // interface. Therefore, attach the `HoistableLinalgOpInterface` to all
        // ops one-by-one.
        HoistableLinalgOpInterfaceHelper<
#define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
            >::registerOpInterface(context);
        UnhoistableOpInterfaceHelper<
#define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
            >::registerOpInterface(context);

        AlwaysHoistableOpInterfaceHelper<
            linalg::PackOp, linalg::UnPackOp>::registerOpInterface(context);
      });
  // Register hoistable op interfaces for tensor ops.
  registry.addExtension(
      +[](MLIRContext *context, tensor::TensorDialect *dialect) {
        // Never hoist empty and other pure metadata ops as a leaf. It's fine to
        // hoist them as a part of a larger constant tree that does actual work.
        HoistableNonLeafOpInterfaceHelper<
            tensor::EmptyOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp,
            tensor::ExtractSliceOp>::registerOpInterface(context);
        // Cases of trivial pack/unpack should be handled as canonicalizations
        // before we get here, thus we're safe to always hoist.
        AlwaysHoistableOpInterfaceHelper<tensor::PadOp>::registerOpInterface(
            context);
      });
  registry.addExtension(
      +[](MLIRContext *context, IREE::Util::UtilDialect *dialect) {
        IREE::Util::AssumeIntOp::attachInterface<
            UtilAssumeIntValueBoundsOpInterface>(*context);
      });

  // Register MutableRegionBranchOpInterface for SCF ops.
  // Register hoistable op interfaces for SCF control flow ops.
  registry.addExtension(+[](MLIRContext *context, scf::SCFDialect *dialect) {
    scf::ForOp::attachInterface<SCFForOpMutableRegionBranchOpInterface>(
        *context);
    scf::IfOp::attachInterface<SCFIfOpMutableRegionBranchOpInterface>(*context);
    scf::WhileOp::attachInterface<SCFWhileOpMutableRegionBranchOpInterface>(
        *context);
    scf::IndexSwitchOp::attachInterface<
        SCFIndexSwitchOpMutableRegionBranchOpInterface>(*context);
    RegionControlFlowHoistableOpInterfaceHelper<
        scf::ForOp, scf::IfOp, scf::WhileOp>::registerOpInterface(context);
  });
}

} // namespace mlir::iree_compiler
