namespace mlir::quant::stablehlo {

bool IsLargeFloatType(Type type) {
  type = getElementTypeOrSelf(type);
  return isa<FloatType>(type) && type.getIntOrFloatBitWidth() > 16;
}

Type ToBfloat16Type(Type type) {
  if (auto shaped = mlir::dyn_cast<ShapedType>(type)) {
    const Type elem = shaped.getElementType();
    if (IsLargeFloatType(elem)) {
      return shaped.clone(BFloat16Type::get(type.getContext()));
    }

namespace mlir::quant::stablehlo {
namespace {

class BFloat16TypeConverter : public TypeConverter {
 public:
  BFloat16TypeConverter() {
    addConversion([](const Type type) -> Type {
      return IsLargeFloatType(type) ? ToBfloat16Type(type) : type;
    });
  }
};

// This helper function makes legality check easier. Both convert ops in the
// patterns below are considered legal:
//  - `BitcastConvertOp` (i32 -> f32) + `ConvertOp` (f32 -> bf16)