// 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 <memory>
#include <utility>

#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "iree/compiler/InputConversion/Common/Passes.h"
#include "iree/compiler/Utils/ConversionUtils.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/Transforms/StructuralTypeConversions.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir::iree_compiler::InputConversion {

#define GEN_PASS_DEF_DEMOTEF32TOF16PASS
#define GEN_PASS_DEF_DEMOTEF64TOF32PASS
#define GEN_PASS_DEF_DEMOTEI64TOI32PASS
#define GEN_PASS_DEF_PROMOTEBF16TOF32PASS
#define GEN_PASS_DEF_PROMOTEF16TOF32PASS
#include "iree/compiler/InputConversion/Common/Passes.h.inc"

namespace {

Value convertRankedFloat(OpBuilder &builder, Type type, ValueRange inputs,
                         Location loc) {
  Type eTy = getElementTypeOrSelf(type);
  Type inputETy = getElementTypeOrSelf(inputs[0].getType());
  if (!isa<FloatType>(getElementTypeOrSelf(type))) {
    return nullptr;
  }

  if (inputETy.getIntOrFloatBitWidth() > eTy.getIntOrFloatBitWidth()) {
    return arith::TruncFOp::create(builder, loc, type, inputs[0]);
  }

  if (inputETy.getIntOrFloatBitWidth() < eTy.getIntOrFloatBitWidth()) {
    return arith::ExtFOp::create(builder, loc, type, inputs[0]);
  }

  return nullptr;
};

Value convertRankedInteger(OpBuilder &builder, Type type, ValueRange inputs,
                           Location loc) {
  Type eTy = getElementTypeOrSelf(type);
  Type inputETy = getElementTypeOrSelf(inputs[0].getType());
  if (!isa<FloatType>(getElementTypeOrSelf(type))) {
    return nullptr;
  }
  bool isUnsigned = eTy.isUnsignedInteger();

  int64_t inBitwidth = inputETy.getIntOrFloatBitWidth();
  int64_t outBitwidth = eTy.getIntOrFloatBitWidth();

  if (inBitwidth > outBitwidth) {
    return arith::TruncIOp::create(builder, loc, type, inputs[0]);
  }

  if (inBitwidth < outBitwidth && isUnsigned) {
    return arith::ExtUIOp::create(builder, loc, type, inputs[0]);
  }

  if (inBitwidth < outBitwidth && !isUnsigned) {
    return arith::ExtSIOp::create(builder, loc, type, inputs[0]);
  }

  return nullptr;
};

// Converts from |SourceType| to |TargetType|.
template <typename SourceType, typename TargetType>
struct PrimitiveTypeConverter : TypeConverter {
  explicit PrimitiveTypeConverter() {
    addConversion([](Type type) { return type; });
    addConversion([&](SourceType type) -> Type {
      if (!isSourceType(type)) {
        return type;
      }
      return getTargetType(type);
    });
    addConversion([&](ComplexType type) {
      return ComplexType::get(convertType(type.getElementType()));
    });
    addConversion([&](RankedTensorType type) {
      return RankedTensorType::get(type.getShape(),
                                   convertType(type.getElementType()),
                                   type.getEncoding());
    });
    addConversion([&](VectorType type) {
      return VectorType::get(type.getShape(),
                             convertType(type.getElementType()));
    });
    addConversion([&](IREE::Util::PtrType ptrType) {
      return IREE::Util::PtrType::get(convertType(ptrType.getTargetType()));
    });
  }

  virtual ~PrimitiveTypeConverter() = default;

  // Returns true if |type| matches the expected source type.
  // Subclasses can override to restrict their conversion to specific subtypes.
  virtual bool isSourceType(SourceType type) { return true; }

  // Returns the newly converted type of |type|.
  // Subclasses can override to pass additional type parameters.
  virtual Type getTargetType(SourceType type) = 0;
};

template <typename SourceType, typename TargetType>
struct FloatTypeConverter : PrimitiveTypeConverter<SourceType, TargetType> {
  explicit FloatTypeConverter() {
    this->addSourceMaterialization(convertRankedFloat);
    this->addTargetMaterialization(convertRankedFloat);
  }
};

template <typename SourceType, typename TargetType>
struct IntegerTypeConverter : PrimitiveTypeConverter<SourceType, TargetType> {
  explicit IntegerTypeConverter() {
    this->addSourceMaterialization(convertRankedInteger);
    this->addTargetMaterialization(convertRankedInteger);
  }
};

// Tries to completely convert a generic Operation.
// This will process attributes, result types, and nested regions.
struct GenericTypeConversionPattern : ConversionPattern {
  GenericTypeConversionPattern(MLIRContext *context,
                               TypeConverter &typeConverter)
      : ConversionPattern(typeConverter, MatchAnyOpTypeTag(), 0, context) {}
  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    // Convert attributes only if this is a constant-like op.
    // This is because some ops use typed attributes for structural information
    // - like linalg ops using i64 for dimension indices - and if we converted
    // them all the ops would become invalid. This may still be too broad,
    // though, if some constant ops include attributes with both the type we
    // want to convert and structural information in the same type.
    llvm::SmallVector<NamedAttribute> newAttrs;
    if (op->hasTrait<OpTrait::ConstantLike>()) {
      for (auto attr : op->getAttrs()) {
        auto newAttr = convertAttribute(op->getLoc(), attr.getValue(),
                                        *getTypeConverter());
        newAttrs.push_back(NamedAttribute(attr.getName(), newAttr));
      }
    } else {
      newAttrs.append(op->getAttrs().begin(), op->getAttrs().end());
    }

    llvm::SmallVector<Type> newResults;
    (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);

    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
                         newResults, newAttrs, op->getSuccessors());

    for (Region &r : op->getRegions()) {
      Region *newRegion = state.addRegion();
      rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
      (void)getTypeConverter()->convertSignatureArgs(
          newRegion->getArgumentTypes(), result);
      rewriter.applySignatureConversion(&newRegion->front(), result);
    }

    Operation *newOp = rewriter.create(state);
    rewriter.replaceOp(op, newOp->getResults());
    return success();
  }
};

struct GlobalOpConversionPattern
    : OpInterfaceConversionPattern<IREE::Util::GlobalOpInterface> {
  GlobalOpConversionPattern(MLIRContext *context, TypeConverter &typeConverter)
      : OpInterfaceConversionPattern(typeConverter, context) {}
  LogicalResult
  matchAndRewrite(IREE::Util::GlobalOpInterface op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    llvm::SmallVector<NamedAttribute> newAttrs;
    for (auto attr : op->getAttrs()) {
      auto newAttr =
          convertAttribute(op->getLoc(), attr.getValue(), *getTypeConverter());
      newAttrs.push_back(NamedAttribute(attr.getName(), newAttr));
    }
    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
                         {}, newAttrs, op->getSuccessors());
    Operation *newOp = rewriter.create(state);
    rewriter.replaceOp(op, newOp->getResults());
    return success();
  }
};

template <typename OpTy, typename TypeTy,
          typename OperandToResultWidthLegalityRelation>
struct ConvertTypeSensitiveArithCastOp : OpConversionPattern<OpTy> {
  using OpConversionPattern<OpTy>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto resultType =
        this->getTypeConverter()->convertType(op.getResult().getType());
    auto operandType =
        this->getTypeConverter()->convertType(op.getOperand().getType());

    auto resultEType = cast<TypeTy>(getElementTypeOrSelf(resultType));
    auto operandEType = cast<TypeTy>(getElementTypeOrSelf(operandType));
    // If post-conversion, the types would be equal, then the op becomes a
    // no-op. Note that the op does not itself allow such a configuration, so we
    // have to catch this before creating the new op.
    if (resultEType == operandEType) {
      rewriter.replaceOp(op, adaptor.getOperands()[0]);
      return success();
    }
    // If after conversion the op becomes invalid, but not same-type (which we
    // can fold above), then bail out.
    // TODO: In some cases, we can repair the situation here, but for integer
    // truncation, we don't know whether we should invert with signed or
    // unsigned extension.
    if (!OperandToResultWidthLegalityRelation()(operandEType.getWidth(),
                                                resultEType.getWidth())) {
      return rewriter.notifyMatchFailure(op, "invalid width combination");
    }
    rewriter.replaceOpWithNewOp<OpTy>(op, resultType, op.getOperand());
    return success();
  }
};

template <typename Base, typename Converter>
struct ConvertTypesPass : Base {
  using Base::Base;
  void runOnOperation() override {
    MLIRContext *context = &this->getContext();

    // Scan the module to detect external functions with types that would be
    // converted. This pass cannot be used with them.
    auto moduleOp = this->getOperation();
    SmallVector<std::pair<mlir::FunctionOpInterface, FunctionType>>
        exportedFuncOps;
    for (auto funcOp : moduleOp.template getOps<mlir::FunctionOpInterface>()) {
      const auto funcType = dyn_cast<FunctionType>(funcOp.getFunctionType());
      if (!funcType) {
        funcOp.emitError()
            << "cannot determine function type of function; do not use the "
               "pass or manually convert the function prior to running it";
        return this->signalPassFailure();
      }
      if (funcOp.isExternal() && !typeConverter.isSignatureLegal(funcType)) {
        funcOp.emitError()
            << "external functions with types that are being demoted are not "
               "allowed; do not use the pass or manually convert the function "
               "signature as required prior to running it";
        return this->signalPassFailure();
      }
      if (funcOp.isPublic()) {
        exportedFuncOps.push_back({funcOp, funcType});
      }
    }

    RewritePatternSet patterns(context);
    patterns.insert<GenericTypeConversionPattern>(context, typeConverter);
    patterns.insert<GlobalOpConversionPattern>(context, typeConverter);
    patterns.insert<ConvertTypeSensitiveArithCastOp<arith::TruncFOp, FloatType,
                                                    std::greater<unsigned>>>(
        typeConverter, context);
    patterns.insert<ConvertTypeSensitiveArithCastOp<arith::ExtFOp, FloatType,
                                                    std::less<unsigned>>>(
        typeConverter, context);
    patterns.insert<ConvertTypeSensitiveArithCastOp<
        arith::TruncIOp, IntegerType, std::less<unsigned>>>(typeConverter,
                                                            context);
    patterns.insert<ConvertTypeSensitiveArithCastOp<arith::ExtUIOp, IntegerType,
                                                    std::less<unsigned>>>(
        typeConverter, context);
    patterns.insert<ConvertTypeSensitiveArithCastOp<arith::ExtSIOp, IntegerType,
                                                    std::less<unsigned>>>(
        typeConverter, context);
    ConversionTarget target(*context);

    populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
        patterns, typeConverter);
    populateFunctionOpInterfaceTypeConversionPattern<IREE::Util::InitializerOp>(
        patterns, typeConverter);
    populateFunctionOpInterfaceTypeConversionPattern<IREE::Util::FuncOp>(
        patterns, typeConverter);
    cf::populateCFStructuralTypeConversionsAndLegality(typeConverter, patterns,
                                                       target);

    // Operations are legal if they don't contain any illegal type.
    target.markUnknownOpDynamicallyLegal([&](Operation *op) {
      if (auto globalOp = dyn_cast<IREE::Util::GlobalOpInterface>(op)) {
        return typeConverter.isLegal(globalOp.getGlobalType());
      } else if (auto funcOp = dyn_cast<mlir::FunctionOpInterface>(op)) {
        for (Type type : funcOp.getArgumentTypes()) {
          if (!typeConverter.isLegal(type)) {
            return false;
          }
        }
        for (Type type : funcOp.getResultTypes()) {
          if (!typeConverter.isLegal(type)) {
            return false;
          }
        }
      }
      for (Type type : op->getResultTypes()) {
        if (!typeConverter.isLegal(type)) {
          return false;
        }
      }
      for (Type type : op->getOperandTypes()) {
        if (!typeConverter.isLegal(type)) {
          return false;
        }
      }
      return true;
    });

    // Note that this will fail if we can't convert any types.
    if (failed(applyFullConversion(this->getOperation(), target,
                                   std::move(patterns)))) {
      return this->signalPassFailure();
    }

    // Warn any public functions changed as part of the conversion.
    bool hasWarned = false;
    for (auto [funcOp, oldType] : exportedFuncOps) {
      const auto newType = cast<FunctionType>(funcOp.getFunctionType());
      if (newType != oldType) {
        if (!hasWarned) {
          hasWarned = true;
          llvm::errs()
              << "\n"
              << "WARNING: ConvertTypesPass (--iree-input-demote-*-to-*) "
                 "changed public function signatures; callers at runtime must "
                 "match the new expected I/O types:\n";
        }
        llvm::errs() << "\n"
                     << "  Old signature:\n"
                     << "    @" << funcOp.getName() << oldType << "\n"
                     << "  New signature:\n"
                     << "    @" << funcOp.getName() << newType << "\n";
      }
    }
    if (hasWarned) {
      llvm::errs() << "\n";
    }
  }

  Converter typeConverter;
};
} // namespace

namespace {
struct DemoteI64ToI32Converter
    : PrimitiveTypeConverter<IntegerType, IntegerType> {
  bool isSourceType(IntegerType type) override { return type.isInteger(64); }
  Type getTargetType(IntegerType type) override {
    return IntegerType::get(type.getContext(), 32, type.getSignedness());
  }
};
class DemoteI64ToI32Pass final
    : public ConvertTypesPass<impl::DemoteI64ToI32PassBase<DemoteI64ToI32Pass>,
                              DemoteI64ToI32Converter> {};
} // namespace

namespace {
struct DemoteF32ToF16Converter
    : PrimitiveTypeConverter<Float32Type, Float16Type> {
  Type getTargetType(Float32Type type) override {
    return Float16Type::get(type.getContext());
  }
};
class DemoteF32ToF16Pass final
    : public ConvertTypesPass<impl::DemoteF32ToF16PassBase<DemoteF32ToF16Pass>,
                              DemoteF32ToF16Converter> {};
} // namespace

namespace {
struct PromoteF16ToF32Converter
    : PrimitiveTypeConverter<Float16Type, Float32Type> {
  Type getTargetType(Float16Type type) override {
    return Float32Type::get(type.getContext());
  }
};
class PromoteF16ToF32Pass final
    : public ConvertTypesPass<
          impl::PromoteF16ToF32PassBase<PromoteF16ToF32Pass>,
          PromoteF16ToF32Converter> {};
} // namespace

namespace {
struct PromoteBF16ToF32Converter
    : FloatTypeConverter<BFloat16Type, Float32Type> {
  Type getTargetType(BFloat16Type type) override {
    return Float32Type::get(type.getContext());
  }
};
class PromoteBF16ToF32Pass final
    : public ConvertTypesPass<
          impl::PromoteBF16ToF32PassBase<PromoteBF16ToF32Pass>,
          PromoteBF16ToF32Converter> {};
} // namespace

namespace {
struct DemoteF64ToF32Converter
    : PrimitiveTypeConverter<Float64Type, Float32Type> {
  Type getTargetType(Float64Type type) override {
    return Float32Type::get(type.getContext());
  }
};
class DemoteF64ToF32Pass final
    : public ConvertTypesPass<impl::DemoteF64ToF32PassBase<DemoteF64ToF32Pass>,
                              DemoteF64ToF32Converter> {};
} // namespace

} // namespace mlir::iree_compiler::InputConversion
