From 04536c4984206eb20d93355145f9f4e9a0f0d266 Mon Sep 17 00:00:00 2001 From: Pradeep Kumar Date: Thu, 30 Jul 2026 18:45:58 +0530 Subject: [PATCH] [MLIR][APFloat] Add Type definitions for UE5M3 type (#212712) This commit adds UE5M3 type definitions in MLIR based on the Float8E5M3FNU APFloat type --- mlir/include/mlir-c/BuiltinTypes.h | 12 +++++++++++ mlir/include/mlir/Bindings/Python/IRTypes.h | 13 ++++++++++++ mlir/include/mlir/IR/Builders.h | 1 + mlir/include/mlir/IR/BuiltinTypes.td | 20 +++++++++++++++++++ mlir/include/mlir/IR/CommonTypeConstraints.td | 2 ++ mlir/lib/AsmParser/TokenKinds.def | 1 + mlir/lib/AsmParser/TypeParser.cpp | 4 ++++ mlir/lib/Bindings/Python/IRTypes.cpp | 11 ++++++++++ mlir/lib/CAPI/IR/BuiltinTypes.cpp | 16 +++++++++++++++ .../Conversion/LLVMCommon/TypeConverter.cpp | 2 +- mlir/lib/IR/AsmPrinter.cpp | 1 + mlir/lib/IR/Builders.cpp | 4 ++++ mlir/lib/IR/BuiltinTypes.cpp | 1 + mlir/python/mlir/extras/types.py | 2 ++ mlir/test/IR/attribute.mlir | 4 ++++ mlir/test/Target/LLVMIR/llvmir.mlir | 3 +++ mlir/test/python/ir/builtin_types.py | 9 +++++++++ mlir/utils/lldb-scripts/mlirDataFormatters.py | 1 + mlir/utils/tree-sitter-mlir/grammar.js | 2 +- 19 files changed, 107 insertions(+), 2 deletions(-) diff --git a/mlir/include/mlir-c/BuiltinTypes.h b/mlir/include/mlir-c/BuiltinTypes.h index f6c30f375cb1..0a56f126a6c1 100644 --- a/mlir/include/mlir-c/BuiltinTypes.h +++ b/mlir/include/mlir-c/BuiltinTypes.h @@ -215,6 +215,18 @@ MLIR_CAPI_EXPORTED MlirType mlirFloat8E8M0FNUTypeGet(MlirContext ctx); MLIR_CAPI_EXPORTED MlirStringRef mlirFloat8E8M0FNUTypeGetName(void); +/// Returns the typeID of a Float8E5M3FNU type. +MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E5M3FNUTypeGetTypeID(void); + +/// Checks whether the given type is an f8E5M3FNU type. +MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M3FNU(MlirType type); + +/// Creates an f8E5M3FNU type in the given context. The type is owned by the +/// context. +MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M3FNUTypeGet(MlirContext ctx); + +MLIR_CAPI_EXPORTED MlirStringRef mlirFloat8E5M3FNUTypeGetName(void); + /// Returns the typeID of an BFloat16 type. MLIR_CAPI_EXPORTED MlirTypeID mlirBFloat16TypeGetTypeID(void); diff --git a/mlir/include/mlir/Bindings/Python/IRTypes.h b/mlir/include/mlir/Bindings/Python/IRTypes.h index c84ed456de30..2cd18099d55c 100644 --- a/mlir/include/mlir/Bindings/Python/IRTypes.h +++ b/mlir/include/mlir/Bindings/Python/IRTypes.h @@ -211,6 +211,19 @@ public: static void bindDerived(ClassTy &c); }; +/// Floating Point Type subclass - Float8E5M3FNUType. +class MLIR_PYTHON_API_EXPORTED PyFloat8E5M3FNUType + : public PyConcreteType { +public: + static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M3FNU; + static constexpr GetTypeIDFunctionTy getTypeIdFunction = + mlirFloat8E5M3FNUTypeGetTypeID; + static constexpr const char *pyClassName = "Float8E5M3FNUType"; + using PyConcreteType::PyConcreteType; + + static void bindDerived(ClassTy &c); +}; + /// Floating Point Type subclass - BF16Type. class MLIR_PYTHON_API_EXPORTED PyBF16Type : public PyConcreteType { diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h index a6cb1456544b..ab5974770da0 100644 --- a/mlir/include/mlir/IR/Builders.h +++ b/mlir/include/mlir/IR/Builders.h @@ -62,6 +62,7 @@ public: // Types. FloatType getF8E8M0Type(); + FloatType getF8E5M3FNUType(); FloatType getF8E4M3FNType(); FloatType getF8E5M2Type(); FloatType getBF16Type(); diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td index 40ccaefa6de3..d5d53dc8b953 100644 --- a/mlir/include/mlir/IR/BuiltinTypes.td +++ b/mlir/include/mlir/IR/BuiltinTypes.td @@ -433,6 +433,26 @@ def Builtin_Float8E8M0FNU : Builtin_FloatType<"Float8E8M0FNU", "f8E8M0FNU"> { }]; } +//===----------------------------------------------------------------------===// +// Float8E5M3FNUType +//===----------------------------------------------------------------------===// + +def Builtin_Float8E5M3FNU : Builtin_FloatType<"Float8E5M3FNU", "f8E5M3FNU"> { + let summary = "8-bit unsigned floating point with 5-bit exponent, 3-bit mantissa"; + let description = [{ + An 8-bit floating point type with 0 sign bit, 5 bits exponent and 3 bits + mantissa. This is not a standard type as defined by IEEE-754, but it + follows similar conventions with the following characteristics: + + * bit encoding: S0E5M3 + * exponent bias: 15 + * infinities: Not supported + * NaNs: Supported with all bits set to 1 + * Zero: Supported + * denormals when exponent is 0 + }]; +} + //===----------------------------------------------------------------------===// // BFloat16Type //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td index f83eea488eb2..af57542fde84 100644 --- a/mlir/include/mlir/IR/CommonTypeConstraints.td +++ b/mlir/include/mlir/IR/CommonTypeConstraints.td @@ -377,6 +377,8 @@ def F6E3M2FN : Type($_self)">, "f6E BuildableType<"$_builder.getType()">; def F8E8M0FNU : Type($_self)">, "f8E8M0FNU type">, BuildableType<"$_builder.getType()">; +def F8E5M3FNU : Type($_self)">, "f8E5M3FNU type">, + BuildableType<"$_builder.getType()">; def AnyComplex : Type($_self)">, "complex-type", "::mlir::ComplexType">; diff --git a/mlir/lib/AsmParser/TokenKinds.def b/mlir/lib/AsmParser/TokenKinds.def index f5e5c25832a3..e9a883bdee92 100644 --- a/mlir/lib/AsmParser/TokenKinds.def +++ b/mlir/lib/AsmParser/TokenKinds.def @@ -106,6 +106,7 @@ TOK_KEYWORD(f4E2M1FN) TOK_KEYWORD(f6E2M3FN) TOK_KEYWORD(f6E3M2FN) TOK_KEYWORD(f8E8M0FNU) +TOK_KEYWORD(f8E5M3FNU) TOK_KEYWORD(f128) TOK_KEYWORD(false) TOK_KEYWORD(floordiv) diff --git a/mlir/lib/AsmParser/TypeParser.cpp b/mlir/lib/AsmParser/TypeParser.cpp index 2cdec14d65fa..0617e7331718 100644 --- a/mlir/lib/AsmParser/TypeParser.cpp +++ b/mlir/lib/AsmParser/TypeParser.cpp @@ -49,6 +49,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) { case Token::kw_f8E4M3B11FNUZ: case Token::kw_f8E3M4: case Token::kw_f8E8M0FNU: + case Token::kw_f8E5M3FNU: case Token::kw_bf16: case Token::kw_f16: case Token::kw_tf32: @@ -340,6 +341,9 @@ Type Parser::parseNonFunctionType() { case Token::kw_f8E8M0FNU: consumeToken(Token::kw_f8E8M0FNU); return builder.getType(); + case Token::kw_f8E5M3FNU: + consumeToken(Token::kw_f8E5M3FNU); + return builder.getType(); case Token::kw_bf16: consumeToken(Token::kw_bf16); return builder.getType(); diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp index 75fd55c90c2b..28796b733670 100644 --- a/mlir/lib/Bindings/Python/IRTypes.cpp +++ b/mlir/lib/Bindings/Python/IRTypes.cpp @@ -239,6 +239,16 @@ void PyFloat8E8M0FNUType::bindDerived(ClassTy &c) { nb::arg("context") = nb::none(), "Create a float8_e8m0fnu type."); } +void PyFloat8E5M3FNUType::bindDerived(ClassTy &c) { + c.def_static( + "get", + [](DefaultingPyMlirContext context) { + MlirType t = mlirFloat8E5M3FNUTypeGet(context->get()); + return PyFloat8E5M3FNUType(context->getRef(), t); + }, + nb::arg("context") = nb::none(), "Create a float8_e5m3fnu type."); +} + void PyBF16Type::bindDerived(ClassTy &c) { c.def_static( "get", @@ -936,6 +946,7 @@ void populateIRTypes(nb::module_ &m) { PyFloat8E5M2FNUZType::bind(m); PyFloat8E3M4Type::bind(m); PyFloat8E8M0FNUType::bind(m); + PyFloat8E5M3FNUType::bind(m); PyBF16Type::bind(m); PyF16Type::bind(m); PyTF32Type::bind(m); diff --git a/mlir/lib/CAPI/IR/BuiltinTypes.cpp b/mlir/lib/CAPI/IR/BuiltinTypes.cpp index 6464fef4653e..073fe112d614 100644 --- a/mlir/lib/CAPI/IR/BuiltinTypes.cpp +++ b/mlir/lib/CAPI/IR/BuiltinTypes.cpp @@ -265,6 +265,22 @@ MlirStringRef mlirFloat8E8M0FNUTypeGetName(void) { return wrap(Float8E8M0FNUType::name); } +MlirTypeID mlirFloat8E5M3FNUTypeGetTypeID() { + return wrap(Float8E5M3FNUType::getTypeID()); +} + +bool mlirTypeIsAFloat8E5M3FNU(MlirType type) { + return llvm::isa(unwrap(type)); +} + +MlirType mlirFloat8E5M3FNUTypeGet(MlirContext ctx) { + return wrap(Float8E5M3FNUType::get(unwrap(ctx))); +} + +MlirStringRef mlirFloat8E5M3FNUTypeGetName(void) { + return wrap(Float8E5M3FNUType::name); +} + MlirTypeID mlirBFloat16TypeGetTypeID() { return wrap(BFloat16Type::getTypeID()); } diff --git a/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp b/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp index a60ecc97aaee..b267623d9b80 100644 --- a/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp +++ b/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp @@ -301,7 +301,7 @@ Type LLVMTypeConverter::convertFloatType(FloatType type) const { if (isa(type)) + Float8E8M0FNUType, Float8E5M3FNUType>(type)) return IntegerType::get(&getContext(), type.getWidth()); // Other floating-point types: A custom type conversion rule must be diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp index ca5c2d2a88ee..b95ab00bd5fd 100644 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -2805,6 +2805,7 @@ void AsmPrinter::Impl::printTypeImpl(Type type) { .Case([&](Type) { os << "f8E4M3B11FNUZ"; }) .Case([&](Type) { os << "f8E3M4"; }) .Case([&](Type) { os << "f8E8M0FNU"; }) + .Case([&](Type) { os << "f8E5M3FNU"; }) .Case([&](Type) { os << "bf16"; }) .Case([&](Type) { os << "f16"; }) .Case([&](Type) { os << "tf32"; }) diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp index a08101b20e12..98fdfb5dfd04 100644 --- a/mlir/lib/IR/Builders.cpp +++ b/mlir/lib/IR/Builders.cpp @@ -34,6 +34,10 @@ Location Builder::getFusedLoc(ArrayRef locs, Attribute metadata) { FloatType Builder::getF8E8M0Type() { return Float8E8M0FNUType::get(context); } +FloatType Builder::getF8E5M3FNUType() { + return Float8E5M3FNUType::get(context); +} + FloatType Builder::getF8E4M3FNType() { return Float8E4M3FNType::get(context); } FloatType Builder::getF8E5M2Type() { return Float8E5M2Type::get(context); } diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp index 786c30851a07..1c3021940ef0 100644 --- a/mlir/lib/IR/BuiltinTypes.cpp +++ b/mlir/lib/IR/BuiltinTypes.cpp @@ -192,6 +192,7 @@ FLOAT_TYPE_SEMANTICS(Float8E4M3FNUZType, Float8E4M3FNUZ) FLOAT_TYPE_SEMANTICS(Float8E4M3B11FNUZType, Float8E4M3B11FNUZ) FLOAT_TYPE_SEMANTICS(Float8E3M4Type, Float8E3M4) FLOAT_TYPE_SEMANTICS(Float8E8M0FNUType, Float8E8M0FNU) +FLOAT_TYPE_SEMANTICS(Float8E5M3FNUType, Float8E5M3FNU) FLOAT_TYPE_SEMANTICS(BFloat16Type, BFloat) FLOAT_TYPE_SEMANTICS(Float16Type, IEEEhalf) FLOAT_TYPE_SEMANTICS(FloatTF32Type, FloatTF32) diff --git a/mlir/python/mlir/extras/types.py b/mlir/python/mlir/extras/types.py index b875d639e9d4..394dd9208871 100644 --- a/mlir/python/mlir/extras/types.py +++ b/mlir/python/mlir/extras/types.py @@ -21,6 +21,7 @@ from ..ir import ( Float8E4M3Type, Float8E5M2Type, Float8E8M0FNUType, + Float8E5M3FNUType, FloatTF32Type, FunctionType, IndexType, @@ -84,6 +85,7 @@ f4E2M1FN = lambda: Float4E2M1FNType.get() f6E2M3FN = lambda: Float6E2M3FNType.get() f6E3M2FN = lambda: Float6E3M2FNType.get() f8E8M0FNU = lambda: Float8E8M0FNUType.get() +f8E5M3FNU = lambda: Float8E5M3FNUType.get() none = lambda: NoneType.get() diff --git a/mlir/test/IR/attribute.mlir b/mlir/test/IR/attribute.mlir index edb7357e4e04..b6b57e2d587c 100644 --- a/mlir/test/IR/attribute.mlir +++ b/mlir/test/IR/attribute.mlir @@ -80,6 +80,10 @@ func.func @float_attrs_pass() { // CHECK: float_attr = 2.000000e+00 : f8E8M0FNU float_attr = 2. : f8E8M0FNU } : () -> () + "test.float_attrs"() { + // CHECK: float_attr = 2.000000e+00 : f8E5M3FNU + float_attr = 2. : f8E5M3FNU + } : () -> () "test.float_attrs"() { // CHECK: float_attr = 2.000000e+00 : f16 float_attr = 2. : f16 diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir index b9db5aa3e318..5ffbb3eef7bb 100644 --- a/mlir/test/Target/LLVMIR/llvmir.mlir +++ b/mlir/test/Target/LLVMIR/llvmir.mlir @@ -99,6 +99,9 @@ llvm.mlir.global internal @f8E4M3B11FNUZ_global_as_i8(1.5 : f8E4M3B11FNUZ) : i8 // CHECK: @f8E8M0FNU_global_as_i8 = internal global i8 127 llvm.mlir.global internal @f8E8M0FNU_global_as_i8(1.0 : f8E8M0FNU) : i8 +// CHECK: @f8E5M3FNU_global_as_i8 = internal global i8 120 +llvm.mlir.global internal @f8E5M3FNU_global_as_i8(1.0 : f8E5M3FNU) : i8 + // CHECK: @bf16_global_as_i16 = internal global i16 16320 llvm.mlir.global internal @bf16_global_as_i16(1.5 : bf16) : i16 diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py index 3fa93f9d0463..acfcc6eaa9bc 100644 --- a/mlir/test/python/ir/builtin_types.py +++ b/mlir/test/python/ir/builtin_types.py @@ -135,6 +135,8 @@ def testFloatTypeSubclasses(): # CHECK: True print(isinstance(Type.parse("f8E8M0FNU", ctx), FloatType)) # CHECK: True + print(isinstance(Type.parse("f8E5M3FNU", ctx), FloatType)) + # CHECK: True print(isinstance(Type.parse("f16", ctx), FloatType)) # CHECK: True print(isinstance(Type.parse("bf16", ctx), FloatType)) @@ -277,6 +279,8 @@ def testFloatType(): print("float:", Float8E4M3B11FNUZType.get()) # CHECK: float: f8E8M0FNU print("float:", Float8E8M0FNUType.get()) + # CHECK: float: f8E5M3FNU + print("float:", Float8E5M3FNUType.get()) # CHECK: float: bf16 print("float:", BF16Type.get()) # CHECK: float: f16 @@ -676,6 +680,7 @@ def testTypeIDs(): (Float8E4M3B11FNUZType, Float8E4M3B11FNUZType.get()), (Float8E5M2FNUZType, Float8E5M2FNUZType.get()), (Float8E8M0FNUType, Float8E8M0FNUType.get()), + (Float8E5M3FNUType, Float8E5M3FNUType.get()), (BF16Type, BF16Type.get()), (F16Type, F16Type.get()), (F32Type, F32Type.get()), @@ -706,6 +711,7 @@ def testTypeIDs(): # CHECK: Float8E4M3B11FNUZType(f8E4M3B11FNUZ) # CHECK: Float8E5M2FNUZType(f8E5M2FNUZ) # CHECK: Float8E8M0FNUType(f8E8M0FNU) + # CHECK: Float8E5M3FNUType(f8E5M3FNU) # CHECK: BF16Type(bf16) # CHECK: F16Type(f16) # CHECK: F32Type(f32) @@ -858,6 +864,9 @@ def testConcreteTypesRoundTrip(): # CHECK: Float8E8M0FNUType # CHECK: Float8E8M0FNUType(f8E8M0FNU) print_downcasted(Float8E8M0FNUType.get()) + # CHECK: Float8E5M3FNUType + # CHECK: Float8E5M3FNUType(f8E5M3FNU) + print_downcasted(Float8E5M3FNUType.get()) # CHECK: BF16Type # CHECK: BF16Type(bf16) print_downcasted(BF16Type.get()) diff --git a/mlir/utils/lldb-scripts/mlirDataFormatters.py b/mlir/utils/lldb-scripts/mlirDataFormatters.py index 38e8278eefbb..dec3a599c75e 100644 --- a/mlir/utils/lldb-scripts/mlirDataFormatters.py +++ b/mlir/utils/lldb-scripts/mlirDataFormatters.py @@ -61,6 +61,7 @@ builtin_attr_type_mnemonics = { "mlir::Float8E4M3B11FNUZType": '"f8E4M3B11FNUZ"', "mlir::Float8E3M4Type": '"f8E3M4"', "mlir::Float8E8M0FNUType": '"f8E8M0FNU"', + "mlir::Float8E5M3FNUType": '"f8E5M3FNU"', "mlir::BFloat16Type": '"bf16"', "mlir::Float16Type": '"f16"', "mlir::FloatTF32Type": '"tf32"', diff --git a/mlir/utils/tree-sitter-mlir/grammar.js b/mlir/utils/tree-sitter-mlir/grammar.js index 2dadd46c4760..8a8b8e8037c3 100644 --- a/mlir/utils/tree-sitter-mlir/grammar.js +++ b/mlir/utils/tree-sitter-mlir/grammar.js @@ -232,7 +232,7 @@ const common = { float_type : $ => token( choice('f16', 'f32', 'f64', 'f80', 'f128', 'bf16', 'f8E3M4', 'f8E4M3FN', 'f8E4M3', 'f8E5M2', 'f4E2M1FN', 'f6E2M3FN', 'f6E3M2FN', - 'f8E8M0FNU')), + 'f8E8M0FNU', 'f8E5M3FNU')), index_type : $ => token('index'), none_type : $ => token('none'), complex_type : $ => seq(token('complex'), '<', $._prim_type, '>'),