// Copyright 2021 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
//===- TypePropagation.cpp -----------------------------------------------===//
//
// The dispatch regions passed to the backends legalizes the bitwidth of
// element types used for the input/output buffers. To avoid illegal load/stores
// within the dispatch, the type needs to be propagated to avoid having tensors
// of illegal bitwidths.
//
// This pass works in two phases.
//
// In the first phase, it uses the dialect conversion framework to propagate
// the types to ensure all tensors have legal bitwidths:
// - All ops are marked dynamically illegal if their operands/result uses
//   unsupported element type.
// - A generic pattern is added to legalize all such ops that triggers on every
//   operation.
//   - For operations with illegal result types, it creates a new
//     operations with legalized return types.
//   - This pattern uses the generic operation creation methods to be
//     op-agnostic.
// - For ops that need specific handling, patterns are added with higher
// benefit,
//   so that they trigger first during legalization.
//
// In the second phase, encoding attributes that don't need further
// materialization after the propagation are dropped from types, so they don't
// need to be handled in later stages of Codegen.
//
// This needs to happen in a separate, second phase due to the fix-point
// iterative nature of the dialect conversion. If the encoding attributes were
// removed in the first phase, multiple visits to the same operation might
// change the bitwidth due to missing encoding information
// (e.g., 'packed_storage').
//
// The second phase works similar to the first one, it also uses the dialect
// conversion framework, marking all operations dynamically illegal if operands
// or results have an encoding attribute that should be be dropped and using
// generic patterns to update operations.
//
//===---------------------------------------------------------------------===//

#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Dialect/Encoding/Utils/ElementPackingUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Transforms/DialectConversion.h"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_TYPEPROPAGATIONPASS
#include "iree/compiler/Codegen/Common/Passes.h.inc"

/// Insert instructions to convert from one element type to another.
static Value convertElementType(OpBuilder &b, Location loc, Type targetType,
                                Value source) {
  Type sourceType = source.getType();
  if (sourceType == targetType) {
    return source;
  }
  if (isa<IntegerType>(sourceType) && isa<IntegerType>(targetType)) {
    unsigned sourceBitWidth = sourceType.getIntOrFloatBitWidth();
    unsigned destBitWidth = targetType.getIntOrFloatBitWidth();
    if (sourceBitWidth > destBitWidth) {
      return arith::TruncIOp::create(b, loc, targetType, source);
    } else {
      return arith::ExtUIOp::create(b, loc, targetType, source);
    }
  }
  return nullptr;
}

/// Legalizes the given type. If the type is already legal, returns
/// std::nullopt.
static std::optional<Type> getLegalizedType(Type t) {
  if (auto shapedType = dyn_cast<RankedTensorType>(t)) {
    std::optional<Type> legalizedElementType =
        legalizeStorageElementType(shapedType);
    if (!legalizedElementType) {
      return std::nullopt;
    }
    return RankedTensorType::get(shapedType.getShape(),
                                 legalizedElementType.value(),
                                 shapedType.getEncoding());
  }
  return std::nullopt;
}

static Type dropEncodingWithoutMaterialization(Type t) {
  if (auto shapedType = dyn_cast<RankedTensorType>(t)) {
    // TODO(#21185): Remove special casing and replace with more generic logic.
    if (auto packedStorage =
            dyn_cast_if_present<IREE::Encoding::PackedStorageAttr>(
                shapedType.getEncoding())) {
      return shapedType.dropEncoding();
    }
  }
  return t;
}

namespace {

/// Materialize
Value materializeAsConvertElementType(OpBuilder &builder, Type type,
                                      ValueRange inputs, Location loc) {
  assert(inputs.size() == 1 && "expected exactly one input");
  return convertElementType(builder, loc, type, inputs[0]);
}

/// Type converter to use for type propagation.
struct TypePropagationTypeConverter : TypeConverter {
  TypePropagationTypeConverter() {
    addConversion([](Type t) {
      auto convertedType = getLegalizedType(t);
      if (!convertedType) {
        return t;
      }
      return convertedType.value();
    });
  }
};

/// Type converter used to drop encodings that need no materialization.
struct DropEncodingTypeConverter : TypeConverter {
  DropEncodingTypeConverter() {
    addConversion([](Type t) -> Type {
      if (auto dispatchTensor =
              dyn_cast<IREE::TensorExt::DispatchTensorType>(t)) {
        Type oldBoundTy = dispatchTensor.getBoundType();
        Type newBoundTy = dropEncodingWithoutMaterialization(oldBoundTy);
        if (newBoundTy == oldBoundTy) {
          return t;
        }
        return IREE::TensorExt::DispatchTensorType::get(
            dispatchTensor.getAccess(), newBoundTy);
      }
      return dropEncodingWithoutMaterialization(t);
    });
  }
};

/// Base class for patterns that handle individual operations.
template <typename T>
struct TypePropagationPattern : OpConversionPattern<T> {
  TypePropagationPattern(TypeConverter &typeConverter, MLIRContext *context)
      : OpConversionPattern<T>(typeConverter, context, 100) {}
};

/// Type conversion for arith.constant operands.
struct ConstantOpTypeConversion : TypePropagationPattern<arith::ConstantOp> {
  using TypePropagationPattern<arith::ConstantOp>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(arith::ConstantOp constantOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    auto attr = cast<ElementsAttr>(constantOp.getValue());
    auto attrType = dyn_cast<RankedTensorType>(attr.getType());
    if (!attrType) {
      return rewriter.notifyMatchFailure(
          constantOp, "expected attribute type to be shaped type");
    }
    std::optional<Type> legalizedElementType =
        legalizeStorageElementType(attrType);
    if (!legalizedElementType) {
      return rewriter.notifyMatchFailure(constantOp,
                                         "cannot legalize elementType");
    }
    if (!legalizedElementType->isIntOrFloat()) {
      return rewriter.notifyMatchFailure(
          constantOp, "expected legalized type to be integer or float type");
    }
    SmallVector<APInt> legalizedValues;
    unsigned numElements = attr.isSplat() ? 1 : attr.getNumElements();
    legalizedValues.reserve(numElements);
    unsigned bitWidth = legalizedElementType->getIntOrFloatBitWidth();
    for (auto value : attr.getValues<APInt>()) {
      legalizedValues.emplace_back(bitWidth, value.getZExtValue());
    }
    auto newAttrType = RankedTensorType::get(attrType.getShape(),
                                             legalizedElementType.value());
    auto newAttr = DenseElementsAttr::get(newAttrType, legalizedValues);
    rewriter.replaceOpWithNewOp<arith::ConstantOp>(constantOp, newAttrType,
                                                   newAttr);
    return success();
  }
};

/// Type conversion for Linalg named op. Interface patterns are not used
/// here cause the region of the operation cannot be cloned. Instead create
/// a new operation with the operands of the correct type.
template <typename OpTy>
struct NamedOpTypePropagation : TypePropagationPattern<OpTy> {
  using TypePropagationPattern<OpTy>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(OpTy namedOp, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    SmallVector<Type> resultTypes;
    resultTypes.reserve(namedOp->getNumResults());
    for (auto resultType : namedOp->getResultTypes()) {
      Type legalizedType = this->getTypeConverter()->convertType(resultType);
      resultTypes.push_back(legalizedType);
    }
    rewriter.replaceOpWithNewOp<OpTy>(namedOp, resultTypes, adaptor.getInputs(),
                                      adaptor.getOutputs(),
                                      linalg::getPrunedAttributeList(namedOp));
    return success();
  }
};

/// Propagates the type for `linalg.generic` operation.
/// - Convert operands whose type has changed.
/// - Convert corresponding basic block argument type and introduce element
/// conversion ops to get back the original type.
/// - Convert the result type if the `outs` operand has changed.
struct GenericOpTypePropagation : TypePropagationPattern<linalg::GenericOp> {
  using TypePropagationPattern<linalg::GenericOp>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(linalg::GenericOp genericOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    llvm::SmallSetVector<unsigned, 8> modifiedOperandIndex;
    SmallVector<Type> resultTypes;

    // 1. Check if any of the operands needs to be legalized.
    for (auto [index, operand] : llvm::enumerate(genericOp->getOpOperands())) {
      Type operandType = operand.get().getType();
      Type legalizedType = this->getTypeConverter()->convertType(operandType);
      if (operandType != legalizedType) {
        modifiedOperandIndex.insert(index);
      }
      // If the operand is an `outs` tensor, its type needs to be changed.
      if (genericOp.isDpsInit(&operand)) {
        resultTypes.push_back(legalizedType);
      }
    }

    // 2. If there are no operands modified, just return failure.
    assert(!modifiedOperandIndex.empty() &&
           "unexpected all types legal within conversion pattern");

    // 3. Create a clone of the operation without cloning its regions.
    auto linalgOp = cast<linalg::LinalgOp>(genericOp.getOperation());
    auto modifiedOp = cast<linalg::LinalgOp>(mlir::cloneWithoutRegions(
        rewriter, linalgOp, resultTypes, adaptor.getOperands()));

    if (genericOp->getNumRegions() != 1) {
      return genericOp.emitOpError("unhandled linalg op with numRegions != 1");
    }

    // 4. Inline the region from the original operation into the new
    // operation.
    rewriter.inlineRegionBefore(genericOp->getRegions().front(),
                                modifiedOp->getRegions().front(),
                                modifiedOp->getRegions().front().begin());
    Region &modifiedOpRegion = modifiedOp->getRegions().front();

    // 5. Convert the signature of the region to use the corresponding element
    // type.
    TypeConverter::SignatureConversion signatureConverter(
        modifiedOpRegion.getNumArguments());
    for (auto [index, arg] : llvm::enumerate(modifiedOpRegion.getArguments())) {
      Type argType = arg.getType();
      if (!modifiedOperandIndex.count(index)) {
        signatureConverter.addInputs(index, argType);
        continue;
      }
      auto tensorArgTy =
          dyn_cast<RankedTensorType>(modifiedOp->getOperand(index).getType());
      if (!tensorArgTy) {
        return genericOp->emitOpError()
               << "failed to legalize arg because operand is not a tensor";
      }
      Type legalizedArgType = legalizeStorageElementType(tensorArgTy);
      signatureConverter.addInputs(index, legalizedArgType);
    }
    rewriter.applySignatureConversion(&modifiedOpRegion.front(),
                                      signatureConverter, getTypeConverter());

    // 6. Introduce scalar conversion operations to convert back to the
    // original scalar type.
    {
      OpBuilder::InsertionGuard g(rewriter);
      Block *entryBlock = modifiedOp.getBlock();

      // 6b. If any of the operands modified were outputs, the yield values
      // need to be modified as well.
      Operation *yieldOp = entryBlock->getTerminator();
      rewriter.setInsertionPoint(yieldOp);
      bool modifyYield = false;
      SmallVector<Value> yieldOperands(yieldOp->operand_begin(),
                                       yieldOp->operand_end());
      for (auto modifiedOperandIndex : modifiedOperandIndex) {
        OpOperand *modifiedOpOperand =
            &modifiedOp->getOpOperand(modifiedOperandIndex);
        if (modifiedOp.isDpsInit(modifiedOpOperand)) {
          modifyYield = true;
          OpOperand *yieldOperand =
              modifiedOp.getMatchingYieldValue(modifiedOpOperand);
          auto tensorOpTy =
              dyn_cast<RankedTensorType>(modifiedOpOperand->get().getType());
          Type legalizedType = legalizeStorageElementType(tensorOpTy);
          yieldOperands[yieldOperand->getOperandNumber()] = convertElementType(
              rewriter, yieldOp->getLoc(), legalizedType, yieldOperand->get());
        }
      }
      if (modifyYield) {
        rewriter.replaceOpWithNewOp<linalg::YieldOp>(yieldOp, yieldOperands);
      }
    }

    rewriter.replaceOp(genericOp, modifiedOp->getResults());
    return success();
  }
};

/// Legalizes `linalg.fill` operation.
struct LinalgFillTypePropagation : TypePropagationPattern<linalg::FillOp> {
  using TypePropagationPattern<linalg::FillOp>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(linalg::FillOp fillOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    Value value = adaptor.getInputs().front();
    auto tensorResultTy =
        dyn_cast<RankedTensorType>(fillOp.getResultTypes().front());
    if (!tensorResultTy) {
      return fillOp->emitOpError() << "failed to legalize type for fill value";
    }
    Type legalizedElementType = legalizeStorageElementType(tensorResultTy);
    if (!legalizedElementType) {
      return fillOp.emitOpError("failed to get legalized type for value");
    }
    Value legalizedValue = convertElementType(rewriter, fillOp->getLoc(),
                                              legalizedElementType, value);
    rewriter.replaceOpWithNewOp<linalg::FillOp>(
        fillOp, ValueRange{legalizedValue}, ValueRange{adaptor.getOutputs()});
    return success();
  }
};

/// Pattern to legalize `tensor.extract` operations.
struct TensorExtractTypePropagation
    : TypePropagationPattern<tensor::ExtractOp> {
  using TypePropagationPattern<tensor::ExtractOp>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(tensor::ExtractOp extractOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    Location loc = extractOp.getLoc();
    Value newExtract = tensor::ExtractOp::create(
        rewriter, loc, adaptor.getTensor(), adaptor.getIndices());
    Value replacement = convertElementType(
        rewriter, loc, extractOp.getResult().getType(), newExtract);
    rewriter.replaceOp(extractOp, replacement);
    return success();
  }
};

/// Pattern to legalize `iree_linalg_ext.scatter` operations.
struct IREELinalgExtScatterTypePropagation
    : TypePropagationPattern<IREE::LinalgExt::ScatterOp> {
  using TypePropagationPattern<
      IREE::LinalgExt::ScatterOp>::TypePropagationPattern;
  LogicalResult
  matchAndRewrite(IREE::LinalgExt::ScatterOp scatterOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    Type inputType = scatterOp.getUpdates().getType();
    Type legalizedInputType = this->getTypeConverter()->convertType(inputType);
    Type maskType =
        scatterOp.getMask() ? scatterOp.getMask().getType() : Type();
    Type legalizedMaskType =
        maskType ? this->getTypeConverter()->convertType(maskType) : Type();
    Type resultType = scatterOp.getOriginal().getType();
    Type legalizedResultType =
        this->getTypeConverter()->convertType(resultType);

    if (inputType == legalizedInputType && maskType == legalizedMaskType &&
        resultType == legalizedResultType) {
      return scatterOp.emitOpError(
          "unexpected all types legal within conversion pattern");
    }

    // Create a clone of the operation without cloning its regions.
    auto modifiedOp =
        cast<IREE::LinalgExt::ScatterOp>(mlir::cloneWithoutRegions(
            rewriter, scatterOp, {legalizedResultType}, adaptor.getOperands()));

    // Inline the region from the original operation into the new operation.
    rewriter.inlineRegionBefore(scatterOp->getRegions().front(),
                                modifiedOp->getRegions().front(),
                                modifiedOp->getRegions().front().begin());
    Region &modifiedOpRegion = modifiedOp->getRegions().front();

    // Convert the signature of the region to use the corresponding element
    // type.
    TypeConverter::SignatureConversion signatureConverter(
        modifiedOpRegion.getNumArguments());
    auto updateTensorTy =
        dyn_cast<RankedTensorType>(modifiedOp.getUpdateType());
    if (!updateTensorTy) {
      return scatterOp->emitOpError() << "failed to legalize type for argument";
    }
    Type legalizedArgType = legalizeStorageElementType(updateTensorTy);
    signatureConverter.addInputs(0, legalizedArgType);
    signatureConverter.addInputs(1, legalizedArgType);
    rewriter.applySignatureConversion(&modifiedOpRegion.front(),
                                      signatureConverter, getTypeConverter());

    {
      // Introduce scalar conversion operations to convert back to the original
      // scalar type.
      OpBuilder::InsertionGuard g(rewriter);
      Block *entryBlock = &modifiedOp->getRegion(0).getBlocks().front();

      // If the output is of an illegal type, the yield value needs to be
      // modified
      auto yieldOp = entryBlock->getTerminator();

      rewriter.setInsertionPoint(yieldOp);
      OpOperand *modifiedOpOperand = &yieldOp->getOpOperand(0);

      auto yieldOperand =
          convertElementType(rewriter, yieldOp->getLoc(), legalizedArgType,
                             modifiedOpOperand->get());

      rewriter.replaceOpWithNewOp<IREE::LinalgExt::YieldOp>(yieldOp,
                                                            yieldOperand);
    }
    rewriter.replaceOp(scatterOp, modifiedOp->getResults());
    return success();
  }
};

/// Pattern to legalize `iree_linalg_ext.sort` operations.
struct IREELinalgExtSortTypePropagation
    : TypePropagationPattern<IREE::LinalgExt::SortOp> {
  using TypePropagationPattern<IREE::LinalgExt::SortOp>::TypePropagationPattern;
  LogicalResult
  matchAndRewrite(IREE::LinalgExt::SortOp sortOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    SmallVector<Type> legalizedResultTypes;
    for (Type resultType : sortOp->getResultTypes()) {
      Type legalizedType = this->getTypeConverter()->convertType(resultType);
      legalizedResultTypes.push_back(legalizedType);
    }

    // Create a clone of the operation without cloning its regions.
    auto modifiedOp = cast<IREE::LinalgExt::SortOp>(mlir::cloneWithoutRegions(
        rewriter, sortOp, {legalizedResultTypes}, adaptor.getOperands()));

    // Inline the region from the original operation into the new operation.
    rewriter.inlineRegionBefore(sortOp->getRegions().front(),
                                modifiedOp->getRegions().front(),
                                modifiedOp->getRegions().front().begin());
    Region &modifiedOpRegion = modifiedOp->getRegions().front();

    // Convert the signature of the region to use the corresponding element
    // type.
    TypeConverter::SignatureConversion signatureConverter(
        modifiedOpRegion.getNumArguments());
    for (auto [index, arg] : llvm::enumerate(modifiedOpRegion.getArguments())) {
      size_t operandIndex = index / 2;
      auto operandTensorType =
          dyn_cast<RankedTensorType>(modifiedOp.getOperandType(operandIndex));
      Type legalizedArgType = legalizeStorageElementType(operandTensorType);
      if (!legalizedArgType) {
        return sortOp.emitOpError("failed to get legalized type for argument");
      }
      signatureConverter.addInputs(index, legalizedArgType);
    }
    rewriter.applySignatureConversion(&modifiedOpRegion.front(),
                                      signatureConverter, getTypeConverter());
    rewriter.replaceOp(sortOp, modifiedOp->getResults());
    return success();
  }
};

/// Simple rewrite pattern that just forwards the source as the result if the
/// result type is not legal (but source type is)
template <typename OpTy>
struct ForwardSourceType : TypePropagationPattern<OpTy> {
  using TypePropagationPattern<OpTy>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    if (op->getNumResults() != 1 || adaptor.getOperands().size() != 1) {
      return rewriter.notifyMatchFailure(
          op, "unhandled op with multiple operands/results");
    }
    Value input = adaptor.getOperands()[0];
    rewriter.replaceOp(op, input);
    return success();
  }
};

/// Rewrite pattern to replace the element type (if it is not legal) with the
/// legal element type.
struct LegalizeResultElementType : ConversionPattern {
  LegalizeResultElementType(TypeConverter &typeConverter, MLIRContext *context)
      : ConversionPattern(typeConverter, MatchAnyOpTypeTag(), /*benefit=*/1,
                          context) {}

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> convertedOperands,
                  ConversionPatternRewriter &rewriter) const final {
    Location loc = op->getLoc();
    SmallVector<Type> resultTypes;
    for (Type resultType : op->getResultTypes()) {
      Type legalizedType = this->typeConverter->convertType(resultType);
      assert(legalizedType && "Failed to drop encoding from type");
      resultTypes.push_back(legalizedType);
    }
    OperationState state(loc, op->getName(), convertedOperands, resultTypes,
                         op->getAttrs());
    for (unsigned i = 0, e = op->getNumRegions(); i != e; ++i) {
      state.addRegion();
    }
    for (auto successor : op->getSuccessors()) {
      state.addSuccessors(successor);
    }
    Operation *newOp = rewriter.create(state);

    // Move all the regions from the old op to the new op and legalize its
    // signature.
    for (const auto &[index, region] : llvm::enumerate(op->getRegions())) {
      Region &newOpRegion = newOp->getRegion(index);
      rewriter.inlineRegionBefore(region, newOpRegion, newOpRegion.begin());
      TypeConverter::SignatureConversion signatureConverter(
          newOpRegion.getNumArguments());
      bool doSignatureConversion = false;
      for (const auto &[argIndex, arg] :
           llvm::enumerate(newOpRegion.getArguments())) {
        Type argType = arg.getType();
        Type legalizedType = this->typeConverter->convertType(argType);
        signatureConverter.addInputs(argIndex, legalizedType);
        doSignatureConversion |= argType != legalizedType;
      }
      if (doSignatureConversion) {
        rewriter.applySignatureConversion(&newOpRegion.front(),
                                          signatureConverter);
      }
    }
    rewriter.replaceOp(op, newOp->getResults());
    return success();
  }
};

// Rewrite pattern for converting the signature of all basic blocks in the
// top-level operation.
template <typename OpTy>
struct LegalizeBasicBlocks : TypePropagationPattern<OpTy> {
  using TypePropagationPattern<OpTy>::TypePropagationPattern;

  LogicalResult
  matchAndRewrite(OpTy funcOp, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const final {
    FailureOr<Block *> newEntry =
        rewriter.convertRegionTypes(&funcOp.getBody(), *this->typeConverter);
    if (failed(newEntry)) {
      return rewriter.notifyMatchFailure(funcOp,
                                         "failed to convert region types");
    }
    rewriter.modifyOpInPlace(funcOp, []() {});
    return success();
  }
};

static void configureConversionTarget(ConversionTarget &target,
                                      TypeConverter &typeConverter) {
  target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp funcOp) {
    Region &body = funcOp.getBody();
    for (Block &block : body) {
      for (auto arg : block.getArguments()) {
        Type argType = arg.getType();
        Type convertedArgType = typeConverter.convertType(argType);
        if (convertedArgType != argType) {
          return false;
        }
      }
    }
    return true;
  });
  target.markUnknownOpDynamicallyLegal([&](Operation *op) {
    for (auto operand : op->getOperands()) {
      Type operandType = operand.getType();
      if (operandType != typeConverter.convertType(operandType)) {
        return false;
      }
    }
    for (auto result : op->getResults()) {
      Type resultType = result.getType();
      if (resultType != typeConverter.convertType(resultType)) {
        return false;
      }
    }
    return true;
  });
}

struct TypePropagationPass final
    : impl::TypePropagationPassBase<TypePropagationPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<arith::ArithDialect>();
  }
  void runOnOperation() override {
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);

    // First phase, legalize element type bitwidth for all tensors by
    // propagating types.
    TypePropagationTypeConverter typeConverter;
    typeConverter.addSourceMaterialization(materializeAsConvertElementType);
    typeConverter.addTargetMaterialization(materializeAsConvertElementType);

    patterns.insert<
        ConstantOpTypeConversion, ForwardSourceType<arith::ExtUIOp>,
        ForwardSourceType<arith::TruncIOp>, GenericOpTypePropagation,
        LinalgFillTypePropagation, LegalizeBasicBlocks<func::FuncOp>,
        LegalizeResultElementType,
        NamedOpTypePropagation<linalg::BatchMatmulOp>,
        NamedOpTypePropagation<linalg::MatmulOp>,
        NamedOpTypePropagation<linalg::MatvecOp>,
        NamedOpTypePropagation<linalg::DotOp>, TensorExtractTypePropagation,
        IREELinalgExtScatterTypePropagation, IREELinalgExtSortTypePropagation>(
        typeConverter, context);

    ConversionTarget target(*context);
    configureConversionTarget(target, typeConverter);

    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns)))) {
      signalPassFailure();
    }

    // Second phase, drop encoding attributes that are no longer needed.
    DropEncodingTypeConverter dropConverter;
    RewritePatternSet dropPatterns(context);
    // We need less specific patterns here because we simply drop attributes
    // from result and operand types.
    dropPatterns
        .insert<LegalizeBasicBlocks<func::FuncOp>, LegalizeResultElementType,
                NamedOpTypePropagation<linalg::BatchMatmulOp>,
                NamedOpTypePropagation<linalg::MatmulOp>,
                NamedOpTypePropagation<linalg::MatvecOp>,
                NamedOpTypePropagation<linalg::DotOp>>(dropConverter, context);
    ConversionTarget dropTarget(*context);
    configureConversionTarget(dropTarget, dropConverter);
    if (failed(applyPartialConversion(getOperation(), dropTarget,
                                      std::move(dropPatterns)))) {
      signalPassFailure();
    }
  }
};
} // namespace
} // namespace mlir::iree_compiler
