/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file dtype_conversion.h
 * \brief Header file of data type conversion routines.
 */
#ifndef TVM_TIR_TRANSFORM_DTYPE_CONVERSION_H_
#define TVM_TIR_TRANSFORM_DTYPE_CONVERSION_H_

#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>

namespace tvm {

namespace tirx {

/*!
 * \brief Rounding mode: https://en.wikipedia.org/wiki/Rounding
 */
enum class RoundingMode {
  // Round half to nearest even
  kHalfToEven = 0U,
  // Round down
  kDown = 1U,
  // Round up
  kUp = 2U,
  // Round towards zero
  kTowardsZero = 3U,
};

/*!
 * \brief Floating point internal representation.
 */
class FloatConfig {
 public:
  /*!
   * \brief Style of infinite number representation.
   */
  enum class InftyStyle {
    // Exponent all ones, mantissa all zeros
    kIEEE = 0U,
    // No representation of infinity
    kNone = 1U
  };
  /*!
   * \brief Style of NaN (not-a-number) representation.
   */
  enum class NaNStyle {
    // Exponent all ones, mantissa non zeros
    // - quiet NaN : 1XXXXX...
    // - signaling NaN : 0XXXXX...
    kIEEE = 0U,
    // No representation of infinity
    kNone = 1U,
    // Both exponent bits and mantissa bits are all ones.
    kAllOnes = 2U,
  };
  // The number of exponent bits.
  int exponent;
  // The number of mantissa (also know as fraction in IEEE format) bits.
  int mantissa;
  // The exponent bias in IEEE format.
  int bias;
  // The representation of infinity.
  InftyStyle infty_style;
  // The representation of NaN (Not a Number).
  NaNStyle nan_style;

  FloatConfig(int exponent, int mantissa, int bias, InftyStyle infty_style, NaNStyle nan_style)
      : exponent(exponent),
        mantissa(mantissa),
        bias(bias),
        infty_style(infty_style),
        nan_style(nan_style) {}

  inline int bits() const { return mantissa + exponent + 1; }

  /*!
   * \brief Create float config from data type.
   * \param dtype The data type, must be a floating point.
   * \return The FloatConfig class containing internal floating point representation.
   */
  static FloatConfig FromDataType(PrimType dtype) {
    DLDataTypeCode code = dtype.code();
    TVM_FFI_ICHECK(dtype.MatchesCode(
        DLDataTypeCode::kDLFloat, DLDataTypeCode::kDLBfloat, DLDataTypeCode::kDLFloat8_e3m4,
        DLDataTypeCode::kDLFloat8_e4m3, DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
        DLDataTypeCode::kDLFloat8_e4m3fn, DLDataTypeCode::kDLFloat8_e4m3fnuz,
        DLDataTypeCode::kDLFloat8_e5m2, DLDataTypeCode::kDLFloat8_e5m2fnuz,
        DLDataTypeCode::kDLFloat8_e8m0fnu, DLDataTypeCode::kDLFloat6_e2m3fn,
        DLDataTypeCode::kDLFloat6_e3m2fn, DLDataTypeCode::kDLFloat4_e2m1fn))
        << "FloatConfig is only applicable to floating point data types, got " << dtype
        << " instead.";
    if (code == DLDataTypeCode::kDLFloat) {
      // IEEE 754 binary formats
      // Reference: https://en.wikipedia.org/wiki/Floating-point_arithmetic
      switch (dtype.bits()) {
        case 16:
          return FloatConfig(5, 10, 15, InftyStyle::kIEEE, NaNStyle::kIEEE);
        case 32:
          return FloatConfig(8, 23, 127, InftyStyle::kIEEE, NaNStyle::kIEEE);
        default:
          // float64
          return FloatConfig(11, 52, 1023, InftyStyle::kIEEE, NaNStyle::kIEEE);
      }
    } else if (dtype.MatchesCode(DLDataTypeCode::kDLBfloat)) {
      // bfloat16,
      return FloatConfig(8, 7, 127, InftyStyle::kIEEE, NaNStyle::kIEEE);
    } else if (dtype.MatchesCode(DLDataTypeCode::kDLFloat8_e3m4, DLDataTypeCode::kDLFloat8_e4m3,
                                 DLDataTypeCode::kDLFloat8_e4m3b11fnuz,
                                 DLDataTypeCode::kDLFloat8_e4m3fn,
                                 DLDataTypeCode::kDLFloat8_e4m3fnuz, DLDataTypeCode::kDLFloat8_e5m2,
                                 DLDataTypeCode::kDLFloat8_e5m2fnuz,
                                 DLDataTypeCode::kDLFloat8_e8m0fnu)) {  // float8
      // NVIDIA/Arm/Intel's FP8 formats for Deep Learning
      // Reference: https://arxiv.org/abs/2209.05433
      switch (code) {
        case DLDataTypeCode::kDLFloat8_e3m4:
          // E3M4 format, not consistent with IEEE-754
          return FloatConfig(3, 4, 3, InftyStyle::kNone, NaNStyle::kAllOnes);
        case DLDataTypeCode::kDLFloat8_e4m3:
          // E4M3 format, not consistent with IEEE-754
          return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kAllOnes);
        case DLDataTypeCode::kDLFloat8_e4m3b11fnuz:
          // E4M3 variant with b11 encoding, not consistent with IEEE-754
          return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kAllOnes);
        case DLDataTypeCode::kDLFloat8_e4m3fn:
          // E4M3 format, not consistent with IEEE-754
          return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kAllOnes);
        case DLDataTypeCode::kDLFloat8_e4m3fnuz:
          // UE4M3 format, not consistent with IEEE-754
          return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kAllOnes);
        case DLDataTypeCode::kDLFloat8_e5m2:
          // UE5M2 format, consistent with IEEE-754
          return FloatConfig(5, 2, 15, InftyStyle::kIEEE, NaNStyle::kIEEE);
        case DLDataTypeCode::kDLFloat8_e5m2fnuz:
          // UE5M2 format, not consistent with IEEE-754
          return FloatConfig(5, 2, 15, InftyStyle::kNone, NaNStyle::kAllOnes);
        case DLDataTypeCode::kDLFloat8_e8m0fnu:
          // UE8M0 format, not consistent with IEEE-754
          return FloatConfig(8, 0, 127, InftyStyle::kNone, NaNStyle::kAllOnes);
        default:
          TVM_FFI_THROW(InternalError) << "Unknown float8 variant: " << dtype;
      }
    } else if (dtype.MatchesCode(DLDataTypeCode::kDLFloat6_e2m3fn,
                                 DLDataTypeCode::kDLFloat6_e3m2fn)) {  // float6
      switch (code) {
        case DLDataTypeCode::kDLFloat6_e2m3fn:
          // E2M3 format, not consistent with IEEE-754
          return FloatConfig(2, 3, 1, InftyStyle::kNone, NaNStyle::kNone);
        case DLDataTypeCode::kDLFloat6_e3m2fn:
          // E3M2 format, not consistent with IEEE-754
          return FloatConfig(3, 2, 3, InftyStyle::kNone, NaNStyle::kNone);
        default:
          TVM_FFI_THROW(InternalError) << "Unknown float6 variant: " << dtype;
      }
    } else {
      // float4
      // E2M1 format, not consistent with IEEE-754
      return FloatConfig(2, 1, 1, InftyStyle::kNone, NaNStyle::kNone);
    }
  }
};

/*!
 * \brief Reinterpret value as unsigned integer with equal number of bits.
 * \param value The value to interpret.
 * \return The reinterpreted uint value.
 */
PrimExpr ReinterpretAsUInt(PrimExpr value);

/*!
 * \brief Get the unsigned integer data type used as storage when the specified dtype is not
 *   supported natively.
 * \param dtype The data type.
 * \return The uint data type, the number of bits is
 *   the same as input dtype.
 */
PrimType GetStorageUIntDType(PrimType dtype);

/*!
 * \brief Conversion routine from value stored in one floating point data type to another floating
 *   point data type.
 * \param src_value The floating point value to be converted.
 * \param tgt_dtype The target floating point data type.
 * \param round_mode The rounding mode to use, defaults to kHalfToEven.
 * \return The converted value in target floating point data type.
 * \note Used when there is no native data type conversion implementation.
 */
PrimExpr DTypeConversion(PrimExpr src_value, PrimType tgt_dtype,
                         RoundingMode round_mode = RoundingMode::kHalfToEven);

}  // namespace tirx
}  // namespace tvm
#endif  // TVM_TIR_TRANSFORM_DTYPE_CONVERSION_H_
