Add Float8E5M2FNUZ and Float8E4M3FNUZ types to MLIR

[jakeh-gc]

Summary

Float8E5M2FNUZ and Float8E4M3FNUZ have been added to APFloat in D141863.
This change adds these types as MLIR builtin types alongside Float8E5M2 and Float8E4M3FN (added in D133823 and D138075).

[krzysz00]

I don't see any reason not to land this, so go ahead.

mlir/include/mlir-c/BuiltinTypes.h

 
 /// Checks whether the given type is an f8E4M3FN type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3FN(MlirType type);
 
 /// Creates an f8E4M3FN type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3FNTypeGet(MlirContext ctx);
 
+/// Checks whether the given type is an f8E5M2FNUZ type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M2FNUZ(MlirType type);
+
+/// Creates an f8E5M2FNUZ type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M2FNUZTypeGet(MlirContext ctx);
+
+/// Checks whether the given type is an f8E4M3FNUZ type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3FNUZ(MlirType type);
+
+/// Creates an f8E4M3FNUZ type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3FNUZTypeGet(MlirContext ctx);
+
 /// Checks whether the given type is a bf16 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsABF16(MlirType type);
 
 /// Creates a bf16 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirBF16TypeGet(MlirContext ctx);
 
 /// Checks whether the given type is an f16 type.

mlir/include/mlir/IR/Builders.h

public:
   // Locations.
   Location getUnknownLoc();
   Location getFusedLoc(ArrayRef<Location> locs,
                        Attribute metadata = Attribute());
 
   // Types.
   FloatType getFloat8E5M2Type();
   FloatType getFloat8E4M3FNType();
+  FloatType getFloat8E5M2FNUZType();
+  FloatType getFloat8E4M3FNUZType();
   FloatType getBF16Type();
   FloatType getF16Type();
   FloatType getF32Type();
   FloatType getF64Type();
   FloatType getF80Type();
   FloatType getF128Type();
 
   IndexType getIndexType();

mlir/include/mlir/IR/BuiltinTypes.h

public:
   static FloatType getBF16(MLIRContext *ctx);
   static FloatType getF16(MLIRContext *ctx);
   static FloatType getF32(MLIRContext *ctx);
   static FloatType getF64(MLIRContext *ctx);
   static FloatType getF80(MLIRContext *ctx);
   static FloatType getF128(MLIRContext *ctx);
   static FloatType getFloat8E5M2(MLIRContext *ctx);
   static FloatType getFloat8E4M3FN(MLIRContext *ctx);
+  static FloatType getFloat8E5M2FNUZ(MLIRContext *ctx);
+  static FloatType getFloat8E4M3FNUZ(MLIRContext *ctx);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Type type);
 
   /// Return the bitwidth of this float type.
   unsigned getWidth();
 
   /// Return the width of the mantissa of this type.

 
 inline bool BaseMemRefType::isValidElementType(Type type) {
   return type.isIntOrIndexOrFloat() ||
          type.isa<ComplexType, MemRefType, VectorType, UnrankedMemRefType>() ||
          type.isa<MemRefElementTypeInterface>();
 }
 
 inline bool FloatType::classof(Type type) {
-  return type.isa<Float8E5M2Type, Float8E4M3FNType, BFloat16Type, Float16Type,
-                  Float32Type, Float64Type, Float80Type, Float128Type>();
+  return type.isa<Float8E5M2Type, Float8E4M3FNType, Float8E5M2FNUZType,
+                  Float8E4M3FNUZType, BFloat16Type, Float16Type, Float32Type,
+                  Float64Type, Float80Type, Float128Type>();
 }
 
 inline FloatType FloatType::getFloat8E5M2(MLIRContext *ctx) {
   return Float8E5M2Type::get(ctx);
 }
 
 inline FloatType FloatType::getFloat8E4M3FN(MLIRContext *ctx) {
   return Float8E4M3FNType::get(ctx);
 }
 
+inline FloatType FloatType::getFloat8E5M2FNUZ(MLIRContext *ctx) {
+  return Float8E5M2FNUZType::get(ctx);
+}
+
+inline FloatType FloatType::getFloat8E4M3FNUZ(MLIRContext *ctx) {
+  return Float8E4M3FNUZType::get(ctx);
+}
+
 inline FloatType FloatType::getBF16(MLIRContext *ctx) {
   return BFloat16Type::get(ctx);
 }
 
 inline FloatType FloatType::getF16(MLIRContext *ctx) {
   return Float16Type::get(ctx);
 }
 

mlir/include/mlir/IR/BuiltinTypes.td

characteristics:
       * NaNs: supported with exponent bits and mantissa bits set to all 1s
       * denormals when exponent is 0
 
     Described in: https://arxiv.org/abs/2209.05433
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// Float8E5M2FNUZType
+
+def Builtin_Float8E5M2FNUZ : Builtin_FloatType<"Float8E5M2FNUZ"> {
+  let summary = "8-bit floating point with 2 bit mantissa";
+  let description = [{
+    An 8-bit floating point type with 1 sign bit, 5 bits exponent and 2 bits
+    mantissa. This is not a standard type as defined by IEEE-754, but it follows
+    similar conventions, with the exception that there are no infinity values,
+    no negative zero, and only one NaN representation. This type has the
+    following characteristics:
+
+      * bit encoding: S1E5M2
+      * exponent bias: 16
+      * infinities: Not supported
+      * NaNs: Supported with sign bit set to 1, exponent bits and mantissa bits set to all 0s
+      * denormals when exponent is 0
+
+    Described in: https://arxiv.org/abs/2206.02915
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Float8E4M3FNUZType
+
+def Builtin_Float8E4M3FNUZ : Builtin_FloatType<"Float8E4M3FNUZ"> {
+  let summary = "8-bit floating point with 3 bit mantissa";
+  let description = [{
+    An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits
+    mantissa. This is not a standard type as defined by IEEE-754, but it follows
+    similar conventions, with the exception that there are no infinity values,
+    no negative zero, and only one NaN representation. This type has the
+    following characteristics:
+
+      * bit encoding: S1E4M3
+      * exponent bias: 8
+      * infinities: Not supported
+      * NaNs: Supported with sign bit set to 1, exponent bits and mantissa bits set to all 0s
+      * denormals when exponent is 0
+
+    Described in: https://arxiv.org/abs/2209.05433
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // BFloat16Type
 
 def Builtin_BFloat16 : Builtin_FloatType<"BFloat16"> {
   let summary = "bfloat16 floating-point type";
 }
 
 //===----------------------------------------------------------------------===//
 // Float16Type

mlir/include/mlir/IR/OpBase.td

 def F128 : F<128>;
 
 def BF16 : Type<CPred<"$_self.isBF16()">, "bfloat16 type">,
            BuildableType<"$_builder.getBF16Type()">;
 def F8E4M3FN : Type<CPred<"$_self.isFloat8E4M3FN()">, "f8E4M3FN type">,
                BuildableType<"$_builder.getFloat8E4M3FNType()">;
 def F8E5M2 : Type<CPred<"$_self.isFloat8E5M2()">, "f8E5M2 type">,
              BuildableType<"$_builder.getFloat8E5M2Type()">;
+def F8E4M3FNUZ : Type<CPred<"$_self.isFloat8E4M3FNUZ()">, "f8E4M3FNUZ type">,
+                 BuildableType<"$_builder.getFloat8E4M3FNUZType()">;
+def F8E5M2FNUZ : Type<CPred<"$_self.isFloat8E5M2FNUZ()">, "f8E5M2FNUZ type">,
+                 BuildableType<"$_builder.getFloat8E5M2FNUZType()">;
 
 def AnyComplex : Type<CPred<"$_self.isa<::mlir::ComplexType>()">,
                       "complex-type", "::mlir::ComplexType">;
 
 class Complex<Type type>
     : ConfinedType<AnyComplex, [
           SubstLeaves<"$_self",
                       "$_self.cast<::mlir::ComplexType>().getElementType()",

mlir/include/mlir/IR/Types.h

public:
   /// Get the dialect this type is registered to.
   Dialect &getDialect() const { return impl->getAbstractType().getDialect(); }
 
   // Convenience predicates.  This is only for floating point types,
   // derived types should use isa/dyn_cast.
   bool isIndex() const;
   bool isFloat8E5M2() const;
   bool isFloat8E4M3FN() const;
+  bool isFloat8E5M2FNUZ() const;
+  bool isFloat8E4M3FNUZ() const;
   bool isBF16() const;
   bool isF16() const;
   bool isF32() const;
   bool isF64() const;
   bool isF80() const;
   bool isF128() const;
 
   /// Return true if this is an integer type with the specified width.

mlir/lib/AsmParser/TokenKinds.def

 TOK_KEYWORD(dense)
 TOK_KEYWORD(dense_resource)
 TOK_KEYWORD(f16)
 TOK_KEYWORD(f32)
 TOK_KEYWORD(f64)
 TOK_KEYWORD(f80)
 TOK_KEYWORD(f8E5M2)
 TOK_KEYWORD(f8E4M3FN)
+TOK_KEYWORD(f8E5M2FNUZ)
+TOK_KEYWORD(f8E4M3FNUZ)
 TOK_KEYWORD(f128)
 TOK_KEYWORD(false)
 TOK_KEYWORD(floordiv)
 TOK_KEYWORD(for)
 TOK_KEYWORD(func)
 TOK_KEYWORD(index)
 TOK_KEYWORD(loc)
 TOK_KEYWORD(max)

mlir/lib/AsmParser/TypeParser.cpp

OptionalParseResult Parser::parseOptionalType(Type &type) {
   case Token::kw_memref:
   case Token::kw_tensor:
   case Token::kw_complex:
   case Token::kw_tuple:
   case Token::kw_vector:
   case Token::inttype:
   case Token::kw_f8E5M2:
   case Token::kw_f8E4M3FN:
+  case Token::kw_f8E5M2FNUZ:
+  case Token::kw_f8E4M3FNUZ:
   case Token::kw_bf16:
   case Token::kw_f16:
   case Token::kw_f32:
   case Token::kw_f64:
   case Token::kw_f80:
   case Token::kw_f128:
   case Token::kw_index:
   case Token::kw_none:
Type Parser::parseNonFunctionType() {
 
   // float-type
   case Token::kw_f8E5M2:
     consumeToken(Token::kw_f8E5M2);
     return builder.getFloat8E5M2Type();
   case Token::kw_f8E4M3FN:
     consumeToken(Token::kw_f8E4M3FN);
     return builder.getFloat8E4M3FNType();
+  case Token::kw_f8E5M2FNUZ:
+    consumeToken(Token::kw_f8E5M2FNUZ);
+    return builder.getFloat8E5M2FNUZType();
+  case Token::kw_f8E4M3FNUZ:
+    consumeToken(Token::kw_f8E4M3FNUZ);
+    return builder.getFloat8E4M3FNUZType();
   case Token::kw_bf16:
     consumeToken(Token::kw_bf16);
     return builder.getBF16Type();
   case Token::kw_f16:
     consumeToken(Token::kw_f16);
     return builder.getF16Type();
   case Token::kw_f32:
     consumeToken(Token::kw_f32);

mlir/lib/Bindings/Python/IRTypes.cpp

c.def_static(
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E5M2TypeGet(context->get());
           return PyFloat8E5M2Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e5m2 type.");
   }
 };
 
+/// Floating Point Type subclass - Float8E4M3FNUZ.
+class PyFloat8E4M3FNUZType : public PyConcreteType<PyFloat8E4M3FNUZType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3FNUZ;
+  static constexpr const char *pyClassName = "Float8E4M3FNUZType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](DefaultingPyMlirContext context) {
+          MlirType t = mlirFloat8E4M3FNUZTypeGet(context->get());
+          return PyFloat8E4M3FNUZType(context->getRef(), t);
+        },
+        py::arg("context") = py::none(), "Create a float8_e4m3fnuz type.");
+  }
+};
+
+/// Floating Point Type subclass - Float8E5M2FNUZ.
+class PyFloat8E5M2FNUZType : public PyConcreteType<PyFloat8E5M2FNUZType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2FNUZ;
+  static constexpr const char *pyClassName = "Float8E5M2FNUZType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get",
+        [](DefaultingPyMlirContext context) {
+          MlirType t = mlirFloat8E5M2FNUZTypeGet(context->get());
+          return PyFloat8E5M2FNUZType(context->getRef(), t);
+        },
+        py::arg("context") = py::none(), "Create a float8_e5m2fnuz type.");
+  }
+};
+
 /// Floating Point Type subclass - BF16Type.
 class PyBF16Type : public PyConcreteType<PyBF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsABF16;
   static constexpr const char *pyClassName = "BF16Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {

 
 } // namespace
 
 void mlir::python::populateIRTypes(py::module &m) {
   PyIntegerType::bind(m);
   PyIndexType::bind(m);
   PyFloat8E4M3FNType::bind(m);
   PyFloat8E5M2Type::bind(m);
+  PyFloat8E4M3FNUZType::bind(m);
+  PyFloat8E5M2FNUZType::bind(m);
   PyBF16Type::bind(m);
   PyF16Type::bind(m);
   PyF32Type::bind(m);
   PyF64Type::bind(m);
   PyNoneType::bind(m);
   PyComplexType::bind(m);
   PyShapedType::bind(m);
   PyVectorType::bind(m);
   PyRankedTensorType::bind(m);
   PyUnrankedTensorType::bind(m);
   PyMemRefType::bind(m);
   PyUnrankedMemRefType::bind(m);
   PyTupleType::bind(m);
   PyFunctionType::bind(m);
   PyOpaqueType::bind(m);
 }

mlir/lib/CAPI/IR/BuiltinTypes.cpp

 bool mlirTypeIsAFloat8E4M3FN(MlirType type) {
   return unwrap(type).isFloat8E4M3FN();
 }
 
 MlirType mlirFloat8E4M3FNTypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E4M3FN(unwrap(ctx)));
 }
 
+bool mlirTypeIsAFloat8E5M2FNUZ(MlirType type) {
+  return unwrap(type).isFloat8E5M2FNUZ();
+}
+
+MlirType mlirFloat8E5M2FNUZTypeGet(MlirContext ctx) {
+  return wrap(FloatType::getFloat8E5M2FNUZ(unwrap(ctx)));
+}
+
+bool mlirTypeIsAFloat8E4M3FNUZ(MlirType type) {
+  return unwrap(type).isFloat8E4M3FNUZ();
+}
+
+MlirType mlirFloat8E4M3FNUZTypeGet(MlirContext ctx) {
+  return wrap(FloatType::getFloat8E4M3FNUZ(unwrap(ctx)));
+}
+
 bool mlirTypeIsABF16(MlirType type) { return unwrap(type).isBF16(); }
 
 MlirType mlirBF16TypeGet(MlirContext ctx) {
   return wrap(FloatType::getBF16(unwrap(ctx)));
 }
 
 bool mlirTypeIsAF16(MlirType type) { return unwrap(type).isF16(); }
 

mlir/lib/IR/AsmPrinter.cpp

void AsmPrinter::Impl::printTypeImpl(Type type) {
   TypeSwitch<Type>(type)
       .Case<OpaqueType>([&](OpaqueType opaqueTy) {
         printDialectSymbol(os, "!", opaqueTy.getDialectNamespace(),
                            opaqueTy.getTypeData());
       })
       .Case<IndexType>([&](Type) { os << "index"; })
       .Case<Float8E5M2Type>([&](Type) { os << "f8E5M2"; })
       .Case<Float8E4M3FNType>([&](Type) { os << "f8E4M3FN"; })
+      .Case<Float8E5M2FNUZType>([&](Type) { os << "f8E5M2FNUZ"; })
+      .Case<Float8E4M3FNUZType>([&](Type) { os << "f8E4M3FNUZ"; })
       .Case<BFloat16Type>([&](Type) { os << "bf16"; })
       .Case<Float16Type>([&](Type) { os << "f16"; })
       .Case<Float32Type>([&](Type) { os << "f32"; })
       .Case<Float64Type>([&](Type) { os << "f64"; })
       .Case<Float80Type>([&](Type) { os << "f80"; })
       .Case<Float128Type>([&](Type) { os << "f128"; })
       .Case<IntegerType>([&](IntegerType integerTy) {
         if (integerTy.isSigned())

mlir/lib/IR/Builders.cpp

 FloatType Builder::getFloat8E5M2Type() {
   return FloatType::getFloat8E5M2(context);
 }
 
 FloatType Builder::getFloat8E4M3FNType() {
   return FloatType::getFloat8E4M3FN(context);
 }
 
+FloatType Builder::getFloat8E5M2FNUZType() {
+  return FloatType::getFloat8E5M2FNUZ(context);
+}
+
+FloatType Builder::getFloat8E4M3FNUZType() {
+  return FloatType::getFloat8E4M3FNUZ(context);
+}
+
 FloatType Builder::getBF16Type() { return FloatType::getBF16(context); }
 
 FloatType Builder::getF16Type() { return FloatType::getF16(context); }
 
 FloatType Builder::getF32Type() { return FloatType::getF32(context); }
 
 FloatType Builder::getF64Type() { return FloatType::getF64(context); }
 

mlir/lib/IR/BuiltinTypes.cpp

IntegerType IntegerType::scaleElementBitwidth(unsigned scale) {
   return IntegerType::get(getContext(), scale * getWidth(), getSignedness());
 }
 
 //===----------------------------------------------------------------------===//
 // Float Type
 //===----------------------------------------------------------------------===//
 
 unsigned FloatType::getWidth() {
-  if (isa<Float8E5M2Type, Float8E4M3FNType>())
+  if (isa<Float8E5M2Type, Float8E4M3FNType, Float8E5M2FNUZType,
+          Float8E4M3FNUZType>())
     return 8;
   if (isa<Float16Type, BFloat16Type>())
     return 16;
   if (isa<Float32Type>())
     return 32;
   if (isa<Float64Type>())
     return 64;
   if (isa<Float80Type>())
     return 80;
   if (isa<Float128Type>())
     return 128;
   llvm_unreachable("unexpected float type");
 }
 
 /// Returns the floating semantics for the given type.
 const llvm::fltSemantics &FloatType::getFloatSemantics() {
   if (isa<Float8E5M2Type>())
     return APFloat::Float8E5M2();
   if (isa<Float8E4M3FNType>())
     return APFloat::Float8E4M3FN();
+  if (isa<Float8E5M2FNUZType>())
+    return APFloat::Float8E5M2FNUZ();
+  if (isa<Float8E4M3FNUZType>())
+    return APFloat::Float8E4M3FNUZ();
   if (isa<BFloat16Type>())
     return APFloat::BFloat();
   if (isa<Float16Type>())
     return APFloat::IEEEhalf();
   if (isa<Float32Type>())
     return APFloat::IEEEsingle();
   if (isa<Float64Type>())
     return APFloat::IEEEdouble();

mlir/lib/IR/MLIRContext.cpp

#endif
   //===--------------------------------------------------------------------===//
 
   DenseMap<TypeID, AbstractType *> registeredTypes;
   StorageUniquer typeUniquer;
 
   /// Cached Type Instances.
   Float8E5M2Type f8E5M2Ty;
   Float8E4M3FNType f8E4M3FNTy;
+  Float8E5M2FNUZType f8E5M2FNUZTy;
+  Float8E4M3FNUZType f8E4M3FNUZTy;
   BFloat16Type bf16Ty;
   Float16Type f16Ty;
   Float32Type f32Ty;
   Float64Type f64Ty;
   Float80Type f80Ty;
   Float128Type f128Ty;
   IndexType indexTy;
   IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty;
MLIRContext::MLIRContext(const DialectRegistry &registry, Threading setting)
 
   // Initialize several common attributes and types to avoid the need to lock
   // the context when accessing them.
 
   //// Types.
   /// Floating-point Types.
   impl->f8E5M2Ty = TypeUniquer::get<Float8E5M2Type>(this);
   impl->f8E4M3FNTy = TypeUniquer::get<Float8E4M3FNType>(this);
+  impl->f8E5M2FNUZTy = TypeUniquer::get<Float8E5M2FNUZType>(this);
+  impl->f8E4M3FNUZTy = TypeUniquer::get<Float8E4M3FNUZType>(this);
   impl->bf16Ty = TypeUniquer::get<BFloat16Type>(this);
   impl->f16Ty = TypeUniquer::get<Float16Type>(this);
   impl->f32Ty = TypeUniquer::get<Float32Type>(this);
   impl->f64Ty = TypeUniquer::get<Float64Type>(this);
   impl->f80Ty = TypeUniquer::get<Float80Type>(this);
   impl->f128Ty = TypeUniquer::get<Float128Type>(this);
   /// Index Type.
   impl->indexTy = TypeUniquer::get<IndexType>(this);

 StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; }
 
 Float8E5M2Type Float8E5M2Type::get(MLIRContext *context) {
   return context->getImpl().f8E5M2Ty;
 }
 Float8E4M3FNType Float8E4M3FNType::get(MLIRContext *context) {
   return context->getImpl().f8E4M3FNTy;
 }
+Float8E5M2FNUZType Float8E5M2FNUZType::get(MLIRContext *context) {
+  return context->getImpl().f8E5M2FNUZTy;
+}
+Float8E4M3FNUZType Float8E4M3FNUZType::get(MLIRContext *context) {
+  return context->getImpl().f8E4M3FNUZTy;
+}
 BFloat16Type BFloat16Type::get(MLIRContext *context) {
   return context->getImpl().bf16Ty;
 }
 Float16Type Float16Type::get(MLIRContext *context) {
   return context->getImpl().f16Ty;
 }
 Float32Type Float32Type::get(MLIRContext *context) {
   return context->getImpl().f32Ty;

mlir/lib/IR/Types.cpp

 //===----------------------------------------------------------------------===//
 // Type
 //===----------------------------------------------------------------------===//
 
 MLIRContext *Type::getContext() const { return getDialect().getContext(); }
 
 bool Type::isFloat8E5M2() const { return isa<Float8E5M2Type>(); }
 bool Type::isFloat8E4M3FN() const { return isa<Float8E4M3FNType>(); }
+bool Type::isFloat8E5M2FNUZ() const { return isa<Float8E5M2FNUZType>(); }
+bool Type::isFloat8E4M3FNUZ() const { return isa<Float8E4M3FNUZType>(); }
 bool Type::isBF16() const { return isa<BFloat16Type>(); }
 bool Type::isF16() const { return isa<Float16Type>(); }
 bool Type::isF32() const { return isa<Float32Type>(); }
 bool Type::isF64() const { return isa<Float64Type>(); }
 bool Type::isF80() const { return isa<Float80Type>(); }
 bool Type::isF128() const { return isa<Float128Type>(); }
 
 bool Type::isIndex() const { return isa<IndexType>(); }

mlir/python/mlir/_mlir_libs/_mlir/ir.pyi

__all__ = [
     "DialectDescriptor",
     "Dialects",
     "Diagnostic",
     "DiagnosticHandler",
     "DiagnosticSeverity",
     "DictAttr",
     "Float8E4M3FNType",
     "Float8E5M2Type",
+    "Float8E4M3FNUZType",
+    "Float8E5M2FNUZType",
     "F16Type",
     "F32Type",
     "F64Type",
     "FlatSymbolRefAttr",
     "FloatAttr",
     "FunctionType",
     "IndexType",
     "InferTypeOpInterface",

 
 class Float8E5M2Type(Type):
     def __init__(self, cast_from_type: Type) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> Float8E5M2Type: ...
     @staticmethod
     def isinstance(arg: Any) -> bool: ...
 
+class Float8E4M3FNUZType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    @staticmethod
+    def get(*args, **kwargs) -> Float8E4M3FNUZType: ...
+    @staticmethod
+    def isinstance(arg: Any) -> bool: ...
+
+class Float8E5M2FNUZType(Type):
+    def __init__(self, cast_from_type: Type) -> None: ...
+    @staticmethod
+    def get(*args, **kwargs) -> Float8E5M2FNUZType: ...
+    @staticmethod
+    def isinstance(arg: Any) -> bool: ...
+
 # TODO: Auto-generated. Audit and fix.
 class F16Type(Type):
     def __init__(self, cast_from_type: Type) -> None: ...
     @staticmethod
     def get(*args, **kwargs) -> F16Type: ...
     @staticmethod
     def isinstance(arg: Any) -> bool: ...
 

mlir/test/IR/attribute.mlir

"test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : f8E5M2
     float_attr = 2. : f8E5M2
   } : () -> ()
   "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : f8E4M3FN
     float_attr = 2. : f8E4M3FN
   } : () -> ()
   "test.float_attrs"() {
+    // CHECK: float_attr = 2.000000e+00 : f8E5M2FNUZ
+    float_attr = 2. : f8E5M2FNUZ
+  } : () -> ()
+  "test.float_attrs"() {
+    // CHECK: float_attr = 2.000000e+00 : f8E4M3FNUZ
+    float_attr = 2. : f8E4M3FNUZ
+  } : () -> ()
+  "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : f16
     float_attr = 2. : f16
   } : () -> ()
   "test.float_attrs"() {
     // CHECK: float_attr = 2.000000e+00 : bf16
     float_attr = 2. : bf16
   } : () -> ()
   "test.float_attrs"() {

mlir/test/python/ir/builtin_types.py

 # CHECK-LABEL: TEST: testFloatType
 @run
 def testFloatType():
   with Context():
     # CHECK: float: f8E4M3FN
     print("float:", Float8E4M3FNType.get())
     # CHECK: float: f8E5M2
     print("float:", Float8E5M2Type.get())
+    # CHECK: float: f8E5M2FNUZ
+    print("float:", Float8E5M2FNUZType.get())
+    # CHECK: float: f8E4M3FNUZ
+    print("float:", Float8E4M3FNUZType.get())
     # CHECK: float: bf16
     print("float:", BF16Type.get())
     # CHECK: float: f16
     print("float:", F16Type.get())
     # CHECK: float: f32
     print("float:", F32Type.get())
     # CHECK: float: f64
     print("float:", F64Type.get())

mlir/utils/lldb-scripts/mlirDataFormatters.py

builtin_attr_type_mnemonics = {
     "mlir::StringAttr": '""...""',
     "mlir::StridedLayout": '"strided_layout"',
     "mlir::UnitAttr": '"unit"',
     "mlir::CallSiteLoc": '"loc(callsite(...))"',
     "mlir::FusedLoc": '"loc(fused<...>[...])"',
     "mlir::UnknownLoc": '"loc(unknown)"',
     "mlir::Float8E5M2Type": '"f8E5M2"',
     "mlir::Float8E4M3FNType": '"f8E4M3FN"',
+    "mlir::Float8E5M2FNUZType": '"f8E5M2FNUZ"',
+    "mlir::Float8E4M3FNUZType": '"f8E4M3FNUZ"',
     "mlir::BFloat16Type": '"bf16"',
     "mlir::Float16Type": '"f16"',
     "mlir::Float32Type": '"f32"',
     "mlir::Float64Type": '"f64"',
     "mlir::Float80Type": '"f80"',
     "mlir::Float128Type": '"f128"',
     "mlir::FunctionType": '"(...) -> (...)"',
     "mlir::IndexType": '"index"',