// Copyright 2025 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
//
//===-------------- ConvertUnsupportedFloatArith.cpp ---------------------===//
//
//   Emulate arith and vector floating point operations that use float types
//   which are unsupported on a target by inserting extf/truncf pairs around all
//   such operations in order to produce arithmetic that can be performed while
//   preserving the original rounding behavior.
//
//===---------------------------------------------------------------------===//

#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/DebugLog.h"
#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-convert-unsupported-float-arith"

namespace mlir::iree_compiler {

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

namespace {

/// Detector for T::get(MLIRContext*) used by llvm::is_detected.
template <typename T>
using hasContextGet = decltype(T::get(std::declval<MLIRContext *>()));

/// Helpers to append types to a vector if they are small float types (fp4/fp8).
template <typename T>
static void maybeAppendSmallFloatType(MLIRContext *ctx,
                                      SmallVectorImpl<Type> &types) {
  if constexpr (llvm::is_detected<hasContextGet, T>::value) {
    Type t = T::get(ctx);
    if (isa<FloatType>(t)) {
      unsigned bitWidth = t.getIntOrFloatBitWidth();
      if (bitWidth == 4 || bitWidth == 8) {
        types.push_back(t);
      }
    }
  }
}

template <typename... Ts>
static void appendSmallFloatTypes(MLIRContext *ctx,
                                  SmallVectorImpl<Type> &types) {
  (maybeAppendSmallFloatType<Ts>(ctx, types), ...);
}

//===----------------------------------------------------------------------===//
// Helper for float emulation patterns
//===----------------------------------------------------------------------===//
//
// This follows the same approach as IREE's
// runtime/src/iree/base/internal/math.h for floating-point conversions. The
// compiler emits equivalent logic using MLIR arith ops instead of C control
// flow.
//
// IEEE 754 floating-point format: [sign | exponent | mantissa]
//   - Sign: 1 bit (0 = positive, 1 = negative)
//   - Exponent: biased (stored = actual + bias)
//   - Mantissa: fractional bits with implicit leading 1 for normal values
//
// Special values:
//   - Zero: exp=0, mantissa=0 (signed zero if format supports it)
//   - Denormal: exp=0, mantissa!=0, value = mantissa * 2^(1-bias-mantissa_bits)
//   - Inf: exp=max, mantissa=0 (IEEE types only)
//   - NaN: exp=max, mantissa!=0 (IEEE), or sign bit only (FNUZ)

// F32 format constants (IEEE 754 binary32).
constexpr int kF32MantBits = 23;
constexpr int kF32Bias = 127;

/// Extracted components from an f32 value stored as i32 bits.
struct F32Fields {
  Value sign;      // Sign bit (0 or 1) shifted to bit 0.
  Value biasedExp; // Biased exponent (8 bits).
  Value mantissa;  // Mantissa (23 bits, no implicit leading 1).
};

/// Helper class for emulating small float (e.g., fp4, fp8) conversions to/from
/// f32 using integer bit manipulation. Handles both scalar and vector types.
///
/// Takes a small float type (scalar or vector), queries its semantics via
/// APFloat, and provides methods that return Value constants for the format
/// parameters.
class FloatEmulationHelper {
public:
  /// Constructor for use with small float types (fp4, fp8).
  /// The smallFloatBitWidth parameter determines the packed integer type.
  FloatEmulationHelper(RewriterBase &rewriter, Location loc, Type type,
                       unsigned smallFloatBitWidth)
      : rewriter(rewriter), loc(loc), vecType(dyn_cast<VectorType>(type)),
        sem(cast<FloatType>(getElementTypeOrSelf(type)).getFloatSemantics()) {
    // Setup scalar and vector types for i32, small int (i4/i8), f32.
    Type i32Scalar = rewriter.getI32Type();
    Type smallIntScalar = rewriter.getIntegerType(smallFloatBitWidth);
    Type f32Scalar = rewriter.getF32Type();
    i32Type =
        vecType ? VectorType::get(vecType.getShape(), i32Scalar) : i32Scalar;
    smallIntType = vecType ? VectorType::get(vecType.getShape(), smallIntScalar)
                           : smallIntScalar;
    f32Type =
        vecType ? VectorType::get(vecType.getShape(), f32Scalar) : f32Scalar;

    // Derive format parameters from APFloat semantics.
    smallMantBits = llvm::APFloat::semanticsPrecision(sem) - 1;
    int totalBits = llvm::APFloat::semanticsSizeInBits(sem);
    smallExpBits = totalBits - 1 - smallMantBits;
    smallBias = 1 - llvm::APFloat::semanticsMinExponent(sem);

    // Query format capabilities.
    smallHasInf = llvm::APFloat::semanticsHasInf(sem);
    smallHasNan = llvm::APFloat::semanticsHasNaN(sem);

    // Check for FNUZ types where NaN is encoded as sign bit only
    // (e.g., 0x80 for fp8, 0x8 for fp4). These types have no negative zero
    // and no infinity.
    if (llvm::APFloat::semanticsHasZero(sem)) {
      llvm::APFloat negZero = llvm::APFloat::getZero(sem, /*Negative=*/true);
      smallHasNegZero = negZero.isZero() && negZero.isNegative();
    } else {
      smallHasNegZero = false;
    }
    smallNanIsNegZero = !smallHasNegZero && !smallHasInf && smallHasNan;
  }

  //===--------------------------------------------------------------------===//
  // Generic constant creation
  //===--------------------------------------------------------------------===//

  /// Creates an i32 constant, splatted if working with vectors.
  Value createI32Const(int64_t value) {
    auto attr = rewriter.getIntegerAttr(rewriter.getI32Type(), value);
    if (vecType) {
      auto splatAttr = SplatElementsAttr::get(cast<ShapedType>(i32Type), attr);
      return rewriter.createOrFold<arith::ConstantOp>(loc, i32Type, splatAttr);
    }
    return rewriter.createOrFold<arith::ConstantOp>(loc, i32Type, attr);
  }

  /// Creates an f32 constant, splatted if working with vectors.
  Value createF32Const(float value) {
    auto attr = rewriter.getF32FloatAttr(value);
    if (vecType) {
      auto splatAttr = SplatElementsAttr::get(cast<ShapedType>(f32Type), attr);
      return rewriter.createOrFold<arith::ConstantOp>(loc, f32Type, splatAttr);
    }
    return rewriter.createOrFold<arith::ConstantOp>(loc, f32Type, attr);
  }

  //===--------------------------------------------------------------------===//
  // F32 format constants (as Value)
  //===--------------------------------------------------------------------===//

  Value getF32MantBitsConst() { return createI32Const(kF32MantBits); }
  Value getF32BiasConst() { return createI32Const(kF32Bias); }
  Value getF32MantMaskConst() {
    return createI32Const((1 << kF32MantBits) - 1);
  }
  Value getF32MaxExpConst() { return createI32Const(0xFF); }
  Value getF32ImplicitBitConst() { return createI32Const(1 << kF32MantBits); }
  Value getF32NaNConst() { return createI32Const(0x7FC00000); }
  Value getF32InfConst() { return createI32Const(0x7F800000); }

  //===--------------------------------------------------------------------===//
  // Small float format constants (as Value)
  //===--------------------------------------------------------------------===//

  Value getSmallMantBitsConst() { return createI32Const(smallMantBits); }
  Value getSmallBiasConst() { return createI32Const(smallBias); }
  Value getSmallSignShiftConst() {
    return createI32Const(smallExpBits + smallMantBits);
  }
  Value getSmallMantMaskConst() {
    return createI32Const((1u << smallMantBits) - 1);
  }
  Value getSmallExpMaskConst() {
    return createI32Const(((1u << smallExpBits) - 1) << smallMantBits);
  }
  Value getSmallSignMaskConst() {
    return createI32Const(1u << (smallExpBits + smallMantBits));
  }
  Value getSmallMaxExpConst() {
    return createI32Const((1 << smallExpBits) - 1);
  }
  Value getSmallMaxNormalExpConst() {
    int maxExp = (1 << smallExpBits) - 1;
    return createI32Const(smallHasInf ? maxExp - 1 : maxExp);
  }

  /// Returns the mantissa shift between f32 and small float.
  Value getMantShiftConst() {
    return createI32Const(kF32MantBits - smallMantBits);
  }

  /// Returns the bias difference (f32_bias - small_bias).
  Value getBiasDiffConst() { return createI32Const(kF32Bias - smallBias); }

  /// Returns the subnormal exponent constant (1 - bias).
  Value getSubnormalExpConst() { return createI32Const(1 - smallBias); }

  /// Returns the NaN encoding for the small float type.
  Value getNaNEncodingConst() {
    if (smallNanIsNegZero) {
      // FNUZ types: NaN = sign bit only (e.g., 0x80 for fp8, 0x8 for fp4).
      return getSmallSignMaskConst();
    }
    // IEEE and FN types: NaN = all exp bits + some mantissa bits.
    unsigned expMask = ((1u << smallExpBits) - 1) << smallMantBits;
    unsigned mantMask = (1u << smallMantBits) - 1;
    return createI32Const(expMask | mantMask);
  }

  /// Returns the Inf encoding for the small float type (0 if no Inf support).
  Value getInfEncodingConst() {
    if (!smallHasInf) {
      return createI32Const(0);
    }
    return getSmallExpMaskConst();
  }

  /// Returns the max finite value encoding for the small float type.
  Value getMaxFiniteConst() {
    unsigned maxExpCode = (1u << smallExpBits) - 1;
    unsigned mantMask = (1u << smallMantBits) - 1;
    unsigned expMask = ((1u << smallExpBits) - 1) << smallMantBits;

    if (smallHasInf) {
      return createI32Const(((maxExpCode - 1) << smallMantBits) | mantMask);
    }
    if (smallHasNan && !smallNanIsNegZero) {
      // FN types: max exp is valid but max mantissa is NaN.
      return createI32Const(expMask | (mantMask - 1));
    }
    // FNUZ or no-NaN types: all bit patterns except NaN are valid.
    return createI32Const(expMask | mantMask);
  }

  /// Returns the denormal scale factor for extf: 2^(1 - bias - mantBits).
  Value getDenormalScaleConst() {
    float scale = std::ldexp(1.0f, 1 - smallBias - smallMantBits);
    return createF32Const(scale);
  }

  //===--------------------------------------------------------------------===//
  // Format capability queries
  //===--------------------------------------------------------------------===//

  bool hasInf() const { return smallHasInf; }
  bool hasNan() const { return smallHasNan; }
  bool hasNegZero() const { return smallHasNegZero; }
  bool isNanEncodedAsNegZero() const { return smallNanIsNegZero; }

  //===--------------------------------------------------------------------===//
  // F32 field extraction
  //===--------------------------------------------------------------------===//

  /// Extracts sign, exponent, and mantissa from an f32 value (as i32 bits).
  F32Fields extractF32Fields(Value i32Val) {
    Value cMantBits = getF32MantBitsConst();
    Value cExpMask = getF32MaxExpConst();
    Value cMantMask = getF32MantMaskConst();
    Value cSignShift = createI32Const(31);

    F32Fields fields;
    fields.sign = arith::ShRUIOp::create(rewriter, loc, i32Val, cSignShift);
    fields.biasedExp = arith::AndIOp::create(
        rewriter, loc, arith::ShRUIOp::create(rewriter, loc, i32Val, cMantBits),
        cExpMask);
    fields.mantissa = arith::AndIOp::create(rewriter, loc, i32Val, cMantMask);

    return fields;
  }

  //===--------------------------------------------------------------------===//
  // Rounding support
  //===--------------------------------------------------------------------===//

  /// Adds a bias to input for round-to-nearest-even before right-shifting.
  /// This matches runtime/src/iree/base/internal/math.h's bias_to_nearest_even:
  ///   even_bit = 1 << shift_amount
  ///   odd_bit = even_bit >> 1
  ///   bias = (input & even_bit) ? odd_bit : (odd_bit - 1)
  ///   return input + bias
  ///
  /// The caller should right-shift the result by shift_amount after this call.
  Value biasForRoundToNearestEven(Value input, Value shiftAmount) {
    Value c0 = createI32Const(0);
    Value c1 = createI32Const(1);
    Value evenBit = arith::ShLIOp::create(rewriter, loc, c1, shiftAmount);
    Value oddBit = arith::ShRUIOp::create(rewriter, loc, evenBit, c1);
    Value oddBitMinus1 = arith::SubIOp::create(rewriter, loc, oddBit, c1);
    Value hasEvenBit = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::ne,
        arith::AndIOp::create(rewriter, loc, input, evenBit), c0);
    Value bias = arith::SelectOp::create(rewriter, loc, hasEvenBit, oddBit,
                                         oddBitMinus1);
    return arith::AddIOp::create(rewriter, loc, input, bias);
  }

  //===--------------------------------------------------------------------===//
  // Type accessors
  //===--------------------------------------------------------------------===//

  Type getI32Type() const { return i32Type; }
  Type getSmallIntType() const { return smallIntType; }
  Type getF32Type() const { return f32Type; }

private:
  RewriterBase &rewriter;
  Location loc;
  VectorType vecType;
  const llvm::fltSemantics &sem;

  // Derived types for scalar/vector operations.
  Type i32Type;
  Type smallIntType;
  Type f32Type;

  // Small float format parameters (derived from semantics).
  int smallExpBits;
  int smallMantBits;
  int smallBias;
  bool smallHasInf;
  bool smallHasNan;
  bool smallHasNegZero;
  bool smallNanIsNegZero;
};

//===----------------------------------------------------------------------===//
// TruncF to small float emulation pattern
//===----------------------------------------------------------------------===//

/// Emulates arith.truncf from f32 to small floats (fp4, fp8) using integer bit
/// manipulation. This implementation follows IREE's
/// runtime/src/iree/base/internal/math.h
/// (specifically iree_math_truncate_f32_to_bits_rounding_to_nearest_even).
///
/// Features:
///   - Round-to-nearest-even (IEEE 754 default rounding mode).
///   - Proper denormal/subnormal generation for underflow cases.
///   - Correct handling of all format variants (IEEE, FN, FNUZ).
///
/// Supported format categories:
///
/// 1. FNUZ types (f8E5M2FNUZ, f8E4M3FNUZ): No Inf, no negative zero.
///    - NaN is encoded as 0x80 (sign=1, exp=0, mantissa=0).
///    - Overflow produces NaN; zero is always positive.
///
/// 2. IEEE types (f8E5M2): Has Inf and negative zero.
///    - Standard IEEE-like encoding with Inf at max exponent.
///
/// 3. FN types (f8E4M3FN, f4E2M1FN): No Inf, but may have negative zero.
///    - Max exponent values are valid finite numbers (except NaN encoding).
///    - f4E2M1FN has no NaN and no negative zero.
struct TruncFToSmallFloat final : OpRewritePattern<arith::TruncFOp> {
  TruncFToSmallFloat(MLIRContext *ctx, ArrayRef<Type> sourceTypes)
      : OpRewritePattern(ctx), sourceTypes(sourceTypes) {}

  LogicalResult matchAndRewrite(arith::TruncFOp op,
                                PatternRewriter &rewriter) const override {
    Type resultType = op.getResult().getType();
    Type resultElemType = getElementTypeOrSelf(resultType);
    Type inputElemType = getElementTypeOrSelf(op.getIn().getType());

    unsigned resultBitWidth = resultElemType.getIntOrFloatBitWidth();
    if (!isa<Float32Type>(inputElemType) ||
        (resultBitWidth != 4 && resultBitWidth != 8)) {
      return failure();
    }

    // Only match types that are in our source types list.
    if (!llvm::is_contained(sourceTypes, resultElemType)) {
      return failure();
    }

    Location loc = op.getLoc();
    FloatEmulationHelper helper(rewriter, loc, resultType, resultBitWidth);

    // Get constants from helper.
    Value c0 = helper.createI32Const(0);
    Value c1 = helper.createI32Const(1);
    Value cF32MantBits = helper.getF32MantBitsConst();
    Value cF32MantMask = helper.getF32MantMaskConst();
    Value cF32MaxExp = helper.getF32MaxExpConst();
    Value cF32Bias = helper.getF32BiasConst();
    Value cDstBias = helper.getSmallBiasConst();
    Value cDstMantBits = helper.getSmallMantBitsConst();
    Value cDstExpShift = cDstMantBits;
    Value cDstSignShift = helper.getSmallSignShiftConst();
    Value cDstMantMask = helper.getSmallMantMaskConst();
    Value cDstMaxNormalExp = helper.getSmallMaxNormalExpConst();
    Value cMantShift = helper.getMantShiftConst();
    Value cNaN = helper.getNaNEncodingConst();
    Value cDstSignMask = helper.getSmallSignMaskConst();

    // Bitcast f32 to i32 and extract fields.
    Value i32Val = arith::BitcastOp::create(rewriter, loc, helper.getI32Type(),
                                            op.getIn());
    F32Fields f32Fields = helper.extractF32Fields(i32Val);

    // Compute destination sign.
    Value dstSign =
        arith::ShLIOp::create(rewriter, loc, f32Fields.sign, cDstSignShift);

    // Check for NaN/Inf in source.
    Value srcIsNanOrInf =
        arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq,
                              f32Fields.biasedExp, cF32MaxExp);
    Value srcMantIsZero = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::eq, f32Fields.mantissa, c0);
    Value srcIsInf =
        arith::AndIOp::create(rewriter, loc, srcIsNanOrInf, srcMantIsZero);
    // Derive srcIsNan from srcIsInf to avoid redundant comparison.
    // srcIsNan = srcIsNanOrInf XOR srcIsInf
    Value srcIsNan =
        arith::XOrIOp::create(rewriter, loc, srcIsNanOrInf, srcIsInf);

    // Check for zero or subnormal in source (exp == 0).
    Value srcExpIsZero = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::eq, f32Fields.biasedExp, c0);

    // Compute arithmetic exponent (unbiased).
    Value arithmeticExp =
        arith::SubIOp::create(rewriter, loc, f32Fields.biasedExp, cF32Bias);

    // Check overflow: biased_dst_exp > max_normal_exp.
    Value biasedDstExp =
        arith::AddIOp::create(rewriter, loc, arithmeticExp, cDstBias);
    Value isOverflow =
        arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::sgt,
                              biasedDstExp, cDstMaxNormalExp);

    // Check underflow: biased_dst_exp <= 0 means subnormal or zero.
    Value isUnderflow = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::sle, biasedDstExp, c0);

    // Check if rounding caused mantissa overflow (carry into exponent).
    Value biasedMant =
        helper.biasForRoundToNearestEven(f32Fields.mantissa, cMantShift);
    Value mantOverflowed = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::ugt, biasedMant, cF32MantMask);
    biasedDstExp = arith::SelectOp::create(
        rewriter, loc, mantOverflowed,
        arith::AddIOp::create(rewriter, loc, biasedDstExp, c1), biasedDstExp);
    biasedMant =
        arith::SelectOp::create(rewriter, loc, mantOverflowed, c0, biasedMant);

    // Re-check overflow after rounding increment.
    isOverflow = arith::OrIOp::create(
        rewriter, loc, isOverflow,
        arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::sgt,
                              biasedDstExp, cDstMaxNormalExp));

    // Shift mantissa to destination width.
    Value dstMant =
        arith::ShRUIOp::create(rewriter, loc, biasedMant, cMantShift);
    dstMant = arith::AndIOp::create(rewriter, loc, dstMant, cDstMantMask);

    // Pack normal result.
    Value dstExp =
        arith::ShLIOp::create(rewriter, loc, biasedDstExp, cDstExpShift);
    Value normalResult = arith::OrIOp::create(
        rewriter, loc, arith::OrIOp::create(rewriter, loc, dstSign, dstExp),
        dstMant);

    // Underflow case: generate subnormal or zero.
    // shift_amount = f32_mant_bits - dst_mant_bits - arithmetic_exp + (1 -
    // dst_bias)
    Value dstSubnormalExp = helper.getSubnormalExpConst();
    Value shiftAmount = arith::SubIOp::create(
        rewriter, loc,
        arith::SubIOp::create(rewriter, loc, cF32MantBits, cDstMantBits),
        arith::SubIOp::create(rewriter, loc, arithmeticExp, dstSubnormalExp));

    // Add implicit leading 1 to f32 mantissa for the shift.
    Value cImplicitBit = helper.getF32ImplicitBitConst();
    Value effectiveMant =
        arith::OrIOp::create(rewriter, loc, f32Fields.mantissa, cImplicitBit);

    // Compute round-to-nearest-even for subnormal.
    Value subnormalMantRounded =
        helper.biasForRoundToNearestEven(effectiveMant, shiftAmount);

    // Shift to get subnormal mantissa.
    Value subnormalMant = arith::ShRUIOp::create(
        rewriter, loc, subnormalMantRounded, shiftAmount);
    subnormalMant =
        arith::AndIOp::create(rewriter, loc, subnormalMant, cDstMantMask);

    // Subnormal result has exp=0.
    Value subnormalResult =
        arith::OrIOp::create(rewriter, loc, dstSign, subnormalMant);

    // Check if shift is too large (complete underflow to zero).
    Value shiftTooLarge = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::sgt, shiftAmount,
        helper.createI32Const(kF32MantBits + 1));
    // Zero (signed if supported).
    Value zeroValue = helper.isNanEncodedAsNegZero() ? c0 : dstSign;
    subnormalResult = arith::SelectOp::create(rewriter, loc, shiftTooLarge,
                                              zeroValue, subnormalResult);

    // Select cascade for final result.
    //
    // Unlike runtime code which uses early returns, SSA form requires computing
    // all paths and selecting between them. The ORDER of selects matters:
    // later selects override earlier ones. We order from lowest to highest
    // priority so the final select (NaN) always wins.
    //
    // Priority (lowest to highest):
    //   1. Normal/subnormal computation (base case)
    //   2. Source zero/subnormal -> zero
    //   3. Negative zero correction (FNUZ only, must be before NaN handling.)
    //   4. Overflow -> Inf or NaN
    //   5. Source Inf -> Inf or NaN
    //   6. Source NaN -> NaN (highest priority, always wins)
    Value result = arith::SelectOp::create(rewriter, loc, isUnderflow,
                                           subnormalResult, normalResult);

    // F32 subnormals (exp=0) become zero in small floats (much smaller than
    // min).
    Value zeroResult = helper.isNanEncodedAsNegZero() ? c0 : dstSign;
    result = arith::SelectOp::create(rewriter, loc, srcExpIsZero, zeroResult,
                                     result);

    // FNUZ: Negative zero (sign bit only, e.g., 0x80 for fp8, 0x8 for fp4)
    // must become positive zero. CRITICAL: This must happen BEFORE
    // NaN/Inf/overflow handling. For FNUZ types, the sign-bit-only pattern
    // is the NaN encoding, not negative zero. If we did this after, we would
    // incorrectly convert NaN to zero.
    if (helper.isNanEncodedAsNegZero()) {
      Value resultIsNegZero = arith::CmpIOp::create(
          rewriter, loc, arith::CmpIPredicate::eq, result, cDstSignMask);
      result =
          arith::SelectOp::create(rewriter, loc, resultIsNegZero, c0, result);
    }

    // Overflow and source Inf both map to the same result:
    // Inf (IEEE) or NaN (FN/FNUZ) or saturate to max finite (no Inf/NaN).
    Value infOrOverflowResult;
    if (helper.hasInf()) {
      infOrOverflowResult = arith::OrIOp::create(rewriter, loc, dstSign,
                                                 helper.getInfEncodingConst());
    } else if (helper.hasNan()) {
      infOrOverflowResult = cNaN;
    } else {
      // No Inf, no NaN: saturate to max finite.
      infOrOverflowResult = arith::OrIOp::create(rewriter, loc, dstSign,
                                                 helper.getMaxFiniteConst());
    }
    result = arith::SelectOp::create(rewriter, loc, isOverflow,
                                     infOrOverflowResult, result);
    result = arith::SelectOp::create(rewriter, loc, srcIsInf,
                                     infOrOverflowResult, result);

    // Handle source NaN last so it takes precedence.
    Value nanResult = helper.hasNan() ? cNaN : c0;
    result =
        arith::SelectOp::create(rewriter, loc, srcIsNan, nanResult, result);

    // Truncate to small int type and bitcast to small float.
    result = arith::TruncIOp::create(rewriter, loc, helper.getSmallIntType(),
                                     result);
    result = arith::BitcastOp::create(rewriter, loc, resultType, result);
    rewriter.replaceOp(op, result);
    return success();
  }

private:
  SmallVector<Type> sourceTypes;
};

//===----------------------------------------------------------------------===//
// ExtF from small float emulation pattern
//===----------------------------------------------------------------------===//

/// Emulates arith.extf from small floats (fp4, fp8) to f32 using integer bit
/// manipulation. This implementation follows IREE's
/// runtime/src/iree/base/internal/math.h
/// (specifically iree_math_make_f32_from_bits).
///
/// For normal values: adjust exponent bias and shift mantissa.
///
/// For denormals (exp=0, mantissa!=0):
///   value = mantissa * 2^(1 - bias - mantissa_bits)
/// We implement this using uitofp + mulf with a precomputed scale factor,
/// which is simpler and more general than enumerating all possible values.
struct ExtFFromSmallFloat final : OpRewritePattern<arith::ExtFOp> {
  ExtFFromSmallFloat(MLIRContext *ctx, ArrayRef<Type> sourceTypes)
      : OpRewritePattern(ctx), sourceTypes(sourceTypes) {}

  LogicalResult matchAndRewrite(arith::ExtFOp op,
                                PatternRewriter &rewriter) const override {
    Type resultType = op.getResult().getType();
    Type inputType = op.getIn().getType();
    Type resultElemType = getElementTypeOrSelf(resultType);
    Type inputElemType = getElementTypeOrSelf(inputType);

    unsigned inputBitWidth = inputElemType.getIntOrFloatBitWidth();
    if ((inputBitWidth != 4 && inputBitWidth != 8) ||
        !isa<Float32Type>(resultElemType)) {
      return failure();
    }

    // Only match types that are in our source types list.
    if (!llvm::is_contained(sourceTypes, inputElemType)) {
      return failure();
    }

    Location loc = op.getLoc();
    FloatEmulationHelper helper(rewriter, loc, inputType, inputBitWidth);

    // Get constants from helper.
    Value c0 = helper.createI32Const(0);
    Value cBiasDiff = helper.getBiasDiffConst();
    Value cMantShift = helper.getMantShiftConst();
    Value cSrcMaxExp = helper.getSmallMaxExpConst();
    Value cF32NaN = helper.getF32NaNConst();
    Value cF32Inf = helper.getF32InfConst();
    Value cF32MantBits = helper.getF32MantBitsConst();

    // Bitcast small float to small int, extend to i32, and extract fields.
    Value smallIntVal = arith::BitcastOp::create(
        rewriter, loc, helper.getSmallIntType(), op.getIn());
    Value i32Val =
        arith::ExtUIOp::create(rewriter, loc, helper.getI32Type(), smallIntVal);

    // Extract fields from small float.
    Value cSrcMantBits = helper.getSmallMantBitsConst();
    Value cSrcExpMask = helper.getSmallMaxExpConst();
    Value cSrcMantMask = helper.getSmallMantMaskConst();
    Value cSrcSignShift = helper.getSmallSignShiftConst();

    Value sign = arith::ShRUIOp::create(rewriter, loc, i32Val, cSrcSignShift);
    Value biasedExpSrc = arith::AndIOp::create(
        rewriter, loc,
        arith::ShRUIOp::create(rewriter, loc, i32Val, cSrcMantBits),
        cSrcExpMask);
    Value mantSrc = arith::AndIOp::create(rewriter, loc, i32Val, cSrcMantMask);

    // Compute f32 sign bit.
    Value f32Sign =
        arith::ShLIOp::create(rewriter, loc, sign, helper.createI32Const(31));

    // Precompute mantissa comparisons (used by multiple checks below).
    Value mantIsZero = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::eq, mantSrc, c0);
    Value mantIsNonZero = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::ne, mantSrc, c0);

    // Detect special values based on format type.
    Value isNaN;
    Value isInf; // Only used for IEEE types (hasInf=true).
    if (!helper.hasNan()) {
      // Types without NaN (e.g., f4E2M1FN): isNaN is always false.
      // We use a constant false comparison that will be optimized away.
      isNaN = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ne, c0,
                                    c0);
    } else if (helper.isNanEncodedAsNegZero()) {
      // FNUZ: NaN = sign bit only (sign=1, exp=0, mantissa=0), no Inf.
      Value cFNUZNaN = helper.getNaNEncodingConst();
      isNaN = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq,
                                    i32Val, cFNUZNaN);
    } else if (helper.hasInf()) {
      // IEEE types: NaN = exp==max && mantissa!=0, Inf = exp==max &&
      // mantissa==0.
      Value expIsMax = arith::CmpIOp::create(
          rewriter, loc, arith::CmpIPredicate::eq, biasedExpSrc, cSrcMaxExp);
      isNaN = arith::AndIOp::create(rewriter, loc, expIsMax, mantIsNonZero);
      isInf = arith::AndIOp::create(rewriter, loc, expIsMax, mantIsZero);
    } else {
      // FN types with NaN: NaN only at specific encoding, no Inf.
      Value cNaNEncoding = helper.getNaNEncodingConst();
      isNaN = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::eq,
                                    i32Val, cNaNEncoding);
    }

    // Check for zero (exp=0, mantissa=0).
    Value expIsZero = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::eq, biasedExpSrc, c0);
    Value isZero = arith::AndIOp::create(rewriter, loc, expIsZero, mantIsZero);

    // Check for denormal (exp=0, mantissa!=0).
    Value isDenormal =
        arith::AndIOp::create(rewriter, loc, expIsZero, mantIsNonZero);

    // Normal value conversion:
    // f32_exp = src_exp + bias_diff, f32_mant = src_mant << mant_shift
    Value normalF32Exp =
        arith::AddIOp::create(rewriter, loc, biasedExpSrc, cBiasDiff);
    Value normalF32Mant =
        arith::ShLIOp::create(rewriter, loc, mantSrc, cMantShift);
    Value normalResult = arith::OrIOp::create(
        rewriter, loc, f32Sign,
        arith::OrIOp::create(
            rewriter, loc,
            arith::ShLIOp::create(rewriter, loc, normalF32Exp, cF32MantBits),
            normalF32Mant));

    // For denormals: value = mantissa * 2^(1 - bias - mant_bits)
    // Use mulf with a precomputed scale factor instead of enumerating values.
    Value mantF32 =
        arith::UIToFPOp::create(rewriter, loc, helper.getF32Type(), mantSrc);
    Value scale = helper.getDenormalScaleConst();
    Value denormalF32 = arith::MulFOp::create(rewriter, loc, mantF32, scale);
    // Apply sign: bitcast to i32, OR with sign bit.
    Value denormalI32 = arith::BitcastOp::create(
        rewriter, loc, helper.getI32Type(), denormalF32);
    Value denormalResult =
        arith::OrIOp::create(rewriter, loc, f32Sign, denormalI32);

    // Select cascade for final result (same ordering logic as TruncF).
    // Later selects override earlier ones. NaN must be last (highest priority).
    Value result = arith::SelectOp::create(rewriter, loc, isDenormal,
                                           denormalResult, normalResult);

    // Zero: use signed zero if format supports negative zero, else +0.
    Value zeroResult = helper.hasNegZero() ? f32Sign : c0;
    result = arith::SelectOp::create(rewriter, loc, isZero, zeroResult, result);

    // Inf (IEEE types only): preserve sign.
    if (helper.hasInf()) {
      Value signedInf = arith::OrIOp::create(rewriter, loc, f32Sign, cF32Inf);
      result = arith::SelectOp::create(rewriter, loc, isInf, signedInf, result);
    }

    // NaN: always canonical quiet NaN (sign bit ignored).
    // Must be last to take precedence over zero for FNUZ (where sign bit is
    // NaN).
    result = arith::SelectOp::create(rewriter, loc, isNaN, cF32NaN, result);

    // Bitcast to f32.
    result = arith::BitcastOp::create(rewriter, loc, resultType, result);
    rewriter.replaceOp(op, result);

    return success();
  }

private:
  SmallVector<Type> sourceTypes;
};

struct ConvertUnsupportedFloatArithPass final
    : impl::ConvertUnsupportedFloatArithPassBase<
          ConvertUnsupportedFloatArithPass> {
  void runOnOperation() override;
  using Base::Base;
};

/// Returns the types that need extf/truncf emulation (bit manipulation) for
/// the given GPU target. Types with hardware conversion instructions are
/// excluded since ArithToAMDGPU patterns in ConvertToROCDL handle those.
///
/// Note: ALL small float types need arithmetic emulation (wrapping addf/mulf
/// with extf/truncf) because no GPU has native fp4/fp8 arithmetic instructions.
/// This function only determines which types need SOFTWARE conversion.
static SmallVector<Type>
getTypesNeedingConversionEmulationForGPU(MLIRContext *context,
                                         IREE::GPU::TargetAttr gpuAttr) {
  SmallVector<Type> types;
  appendSmallFloatTypes<
#define GET_TYPEDEF_LIST
#include "mlir/IR/BuiltinTypes.cpp.inc"
      >(context, types);

  // Remove types that have hardware conversion support on this chip.
  StringRef chipset = gpuAttr.getArch();
  FailureOr<amdgpu::Chipset> maybeChipset = amdgpu::Chipset::parse(chipset);
  if (failed(maybeChipset)) {
    LDBG() << "Invalid chip name";
    return types;
  }
  constexpr amdgpu::Chipset kGfx942{9, 4, 2};
  constexpr amdgpu::Chipset kGfx950{9, 5, 0};
  if (*maybeChipset >= kGfx942 && *maybeChipset < kGfx950) {
    // gfx942 has hardware conversion for FNUZ types.
    llvm::erase(types, Float8E4M3FNUZType::get(context));
    llvm::erase(types, Float8E5M2FNUZType::get(context));
  }
  if (amdgpu::hasOcpFp8(*maybeChipset)) {
    // gfx950+ and gfx12+ have hardware conversion for OCP types.
    llvm::erase(types, Float8E4M3FNType::get(context));
    llvm::erase(types, Float8E5M2Type::get(context));
    llvm::erase(types, Float4E2M1FNType::get(context));
  }
  return types;
}

} // namespace

void ConvertUnsupportedFloatArithPass::runOnOperation() {
  MLIRContext *context = &getContext();
  FunctionOpInterface funcOp = getOperation();
  Type targetType = Float32Type::get(context);

  // All small float types (fp4, fp8) need arithmetic emulation unless the
  // hardware has native arithmetic instructions for these types. Operations
  // like addf/mulf usually have to be wrapped with extf/truncf to compute in
  // f32.
  SmallVector<Type> allSmallFloatTypes;
  appendSmallFloatTypes<
#define GET_TYPEDEF_LIST
#include "mlir/IR/BuiltinTypes.cpp.inc"
      >(context, allSmallFloatTypes);

  // Apply the standard float emulation patterns. This inserts extf/truncf pairs
  // around unsupported float operations.
  {
    TypeConverter converter;
    arith::populateEmulateUnsupportedFloatsConversions(
        converter, allSmallFloatTypes, targetType);
    RewritePatternSet patterns(context);
    arith::populateEmulateUnsupportedFloatsPatterns(patterns, converter);
    ConversionTarget target(*context);
    arith::populateEmulateUnsupportedFloatsLegality(target, converter);

    // Mark scaling ops as legal - they have their own expansion patterns in
    // arith::populateExpandScalingExtTruncPatterns that run in later passes.
    // We don't want to insert extf/truncf pairs around them.
    target.addLegalOp<arith::ScalingExtFOp, arith::ScalingTruncFOp>();

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

  // Determine which types need software conversion emulation. By default, all
  // the small float types need the emulation. For GPU targets, some types may
  // have hardware conversion support and can be skipped.
  SmallVector<Type> typesNeedingConversionEmulation = allSmallFloatTypes;
  if (IREE::GPU::TargetAttr gpuAttr = getGPUTargetAttr(funcOp)) {
    typesNeedingConversionEmulation =
        getTypesNeedingConversionEmulationForGPU(context, gpuAttr);
  }

  // Emulate extf/truncf to/from small float types using integer bit
  // manipulation. Only for types without hardware conversion support.
  // This is gated by the enableExtTruncEmulation flag.
  if (enableExtTruncEmulation && !typesNeedingConversionEmulation.empty()) {
    RewritePatternSet emulationPatterns(context);
    emulationPatterns.add<TruncFToSmallFloat, ExtFFromSmallFloat>(
        context, typesNeedingConversionEmulation);
    walkAndApplyPatterns(funcOp, std::move(emulationPatterns));
  }
}

} // namespace mlir::iree_compiler
