Skip to content

Commit 73d83f2

Browse files
sergey-kozubsergey-kozub
authored
[MLIR] Add f6E2M3FN type (#107999)
This PR adds `f6E2M3FN` type to mlir. `f6E2M3FN` type is proposed in [OpenCompute MX Specification](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf). It defines a 6-bit floating point number with bit layout S1E2M3. Unlike IEEE-754 types, there are no infinity or NaN values. ```c f6E2M3FN - Exponent bias: 1 - Maximum stored exponent value: 3 (binary 11) - Maximum unbiased exponent value: 3 - 1 = 2 - Minimum stored exponent value: 1 (binary 01) - Minimum unbiased exponent value: 1 − 1 = 0 - Has Positive and Negative zero - Doesn't have infinity - Doesn't have NaNs Additional details: - Zeros (+/-): S.00.000 - Max normal number: S.11.111 = ±2^(2) x (1 + 0.875) = ±7.5 - Min normal number: S.01.000 = ±2^(0) = ±1.0 - Max subnormal number: S.00.111 = ±2^(0) x 0.875 = ±0.875 - Min subnormal number: S.00.001 = ±2^(0) x 0.125 = ±0.125 ``` Related PRs: - [PR-94735](#94735) [APFloat] Add APFloat support for FP6 data types - [PR-105573](#105573) [MLIR] Add f6E3M2FN type - was used as a template for this PR
1 parent 09e3a36 commit 73d83f2

File tree

24 files changed

+134
-8
lines changed

24 files changed

+134
-8
lines changed

‎mlir/include/mlir-c/BuiltinTypes.h

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,16 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat(MlirType type);
7979
/// Returns the bitwidth of a floating-point type.
8080
MLIR_CAPI_EXPORTED unsigned mlirFloatTypeGetWidth(MlirType type);
8181

82+
/// Returns the typeID of an Float6E2M3FN type.
83+
MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E2M3FNTypeGetTypeID(void);
84+
85+
/// Checks whether the given type is an f6E2M3FN type.
86+
MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat6E2M3FN(MlirType type);
87+
88+
/// Creates an f6E2M3FN type in the given context. The type is owned by the
89+
/// context.
90+
MLIR_CAPI_EXPORTED MlirType mlirFloat6E2M3FNTypeGet(MlirContext ctx);
91+
8292
/// Returns the typeID of an Float6E3M2FN type.
8393
MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E3M2FNTypeGetTypeID(void);
8494

‎mlir/include/mlir/IR/Builders.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ class Builder {
6060
Attribute metadata = Attribute());
6161

6262
// Types.
63+
FloatType getFloat6E2M3FNType();
6364
FloatType getFloat6E3M2FNType();
6465
FloatType getFloat8E5M2Type();
6566
FloatType getFloat8E4M3Type();

‎mlir/include/mlir/IR/BuiltinTypes.h

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ class FloatType : public Type {
6767
static FloatType getFloat8E4M3FNUZ(MLIRContext *ctx);
6868
static FloatType getFloat8E4M3B11FNUZ(MLIRContext *ctx);
6969
static FloatType getFloat8E3M4(MLIRContext *ctx);
70+
static FloatType getFloat6E2M3FN(MLIRContext *ctx);
7071
static FloatType getFloat6E3M2FN(MLIRContext *ctx);
7172

7273
/// Methods for support type inquiry through isa, cast, and dyn_cast.
@@ -414,11 +415,15 @@ inline bool BaseMemRefType::isValidElementType(Type type) {
414415
}
415416

416417
inline bool FloatType::classof(Type type) {
417-
return llvm::isa<Float6E3M2FNType, Float8E5M2Type, Float8E4M3Type,
418-
Float8E4M3FNType, Float8E5M2FNUZType, Float8E4M3FNUZType,
419-
Float8E4M3B11FNUZType, Float8E3M4Type, BFloat16Type,
420-
Float16Type, FloatTF32Type, Float32Type, Float64Type,
421-
Float80Type, Float128Type>(type);
418+
return llvm::isa<Float6E2M3FNType, Float6E3M2FNType, Float8E5M2Type,
419+
Float8E4M3Type, Float8E4M3FNType, Float8E5M2FNUZType,
420+
Float8E4M3FNUZType, Float8E4M3B11FNUZType, Float8E3M4Type,
421+
BFloat16Type, Float16Type, FloatTF32Type, Float32Type,
422+
Float64Type, Float80Type, Float128Type>(type);
423+
}
424+
425+
inline FloatType FloatType::getFloat6E2M3FN(MLIRContext *ctx) {
426+
return Float6E2M3FNType::get(ctx);
422427
}
423428

424429
inline FloatType FloatType::getFloat6E3M2FN(MLIRContext *ctx) {

‎mlir/include/mlir/IR/BuiltinTypes.td

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,11 +233,32 @@ def Builtin_Float8E3M4 : Builtin_FloatType<"Float8E3M4", "f8E3M4"> {
233233
}];
234234
}
235235

236+
//===----------------------------------------------------------------------===//
237+
// Float6E2M3FNType
238+
239+
def Builtin_Float6E2M3FN : Builtin_FloatType<"Float6E2M3FN", "f6E2M3FN"> {
240+
let summary = "6-bit floating point with 2-bit exponent and 3-bit mantissa";
241+
let description = [{
242+
An 6-bit floating point type with 1 sign bit, 2 bits exponent and 3 bits
243+
mantissa. This is not a standard type as defined by IEEE-754, but it
244+
follows similar conventions with the following characteristics:
245+
246+
* bit encoding: S1E2M3
247+
* exponent bias: 1
248+
* infinities: Not supported
249+
* NaNs: Not supported
250+
* denormals when exponent is 0
251+
252+
Open Compute Project (OCP) microscaling formats (MX) specification:
253+
https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
254+
}];
255+
}
256+
236257
//===----------------------------------------------------------------------===//
237258
// Float6E3M2FNType
238259

239260
def Builtin_Float6E3M2FN : Builtin_FloatType<"Float6E3M2FN", "f6E3M2FN"> {
240-
let summary = "6-bit floating point with 3 bits exponent and 2 bit mantissa";
261+
let summary = "6-bit floating point with 3-bit exponent and 2-bit mantissa";
241262
let description = [{
242263
An 6-bit floating point type with 1 sign bit, 3 bits exponent and 2 bits
243264
mantissa. This is not a standard type as defined by IEEE-754, but it

‎mlir/include/mlir/IR/CommonTypeConstraints.td

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,8 @@ def F8E5M2FNUZ : Type<CPred<"$_self.isFloat8E5M2FNUZ()">, "f8E5M2FNUZ type">,
347347
BuildableType<"$_builder.getFloat8E5M2FNUZType()">;
348348
def F8E3M4 : Type<CPred<"$_self.isFloat8E3M4()">, "f8E3M4 type">,
349349
BuildableType<"$_builder.getFloat8E3M4Type()">;
350+
def F6E2M3FN : Type<CPred<"$_self.isFloat6E2M3FN()">, "f6E2M3FN type">,
351+
BuildableType<"$_builder.getFloat6E2M3FNType()">;
350352
def F6E3M2FN : Type<CPred<"$_self.isFloat6E3M2FN()">, "f6E3M2FN type">,
351353
BuildableType<"$_builder.getFloat6E3M2FNType()">;
352354

‎mlir/include/mlir/IR/Types.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ class Type {
125125
// Convenience predicates. This is only for floating point types,
126126
// derived types should use isa/dyn_cast.
127127
bool isIndex() const;
128+
bool isFloat6E2M3FN() const;
128129
bool isFloat6E3M2FN() const;
129130
bool isFloat8E5M2() const;
130131
bool isFloat8E4M3() const;

‎mlir/lib/AsmParser/TokenKinds.def

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ TOK_KEYWORD(f8E5M2FNUZ)
101101
TOK_KEYWORD(f8E4M3FNUZ)
102102
TOK_KEYWORD(f8E4M3B11FNUZ)
103103
TOK_KEYWORD(f8E3M4)
104+
TOK_KEYWORD(f6E2M3FN)
104105
TOK_KEYWORD(f6E3M2FN)
105106
TOK_KEYWORD(f128)
106107
TOK_KEYWORD(false)

‎mlir/lib/AsmParser/TypeParser.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) {
3939
case Token::kw_tuple:
4040
case Token::kw_vector:
4141
case Token::inttype:
42+
case Token::kw_f6E2M3FN:
4243
case Token::kw_f6E3M2FN:
4344
case Token::kw_f8E5M2:
4445
case Token::kw_f8E4M3:
@@ -304,6 +305,9 @@ Type Parser::parseNonFunctionType() {
304305
}
305306

306307
// float-type
308+
case Token::kw_f6E2M3FN:
309+
consumeToken(Token::kw_f6E2M3FN);
310+
return builder.getFloat6E2M3FNType();
307311
case Token::kw_f6E3M2FN:
308312
consumeToken(Token::kw_f6E3M2FN);
309313
return builder.getFloat6E3M2FNType();

‎mlir/lib/Bindings/Python/IRTypes.cpp

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,27 @@ class PyFloatType : public PyConcreteType<PyFloatType> {
124124
}
125125
};
126126

127+
/// Floating Point Type subclass - Float6E2M3FNType.
128+
class PyFloat6E2M3FNType
129+
: public PyConcreteType<PyFloat6E2M3FNType, PyFloatType> {
130+
public:
131+
static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat6E2M3FN;
132+
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
133+
mlirFloat6E2M3FNTypeGetTypeID;
134+
static constexpr const char *pyClassName = "Float6E2M3FNType";
135+
using PyConcreteType::PyConcreteType;
136+
137+
static void bindDerived(ClassTy &c) {
138+
c.def_static(
139+
"get",
140+
[](DefaultingPyMlirContext context) {
141+
MlirType t = mlirFloat6E2M3FNTypeGet(context->get());
142+
return PyFloat6E2M3FNType(context->getRef(), t);
143+
},
144+
py::arg("context") = py::none(), "Create a float6_e2m3fn type.");
145+
}
146+
};
147+
127148
/// Floating Point Type subclass - Float6E3M2FNType.
128149
class PyFloat6E3M2FNType
129150
: public PyConcreteType<PyFloat6E3M2FNType, PyFloatType> {
@@ -901,6 +922,7 @@ void mlir::python::populateIRTypes(py::module &m) {
901922
PyIntegerType::bind(m);
902923
PyFloatType::bind(m);
903924
PyIndexType::bind(m);
925+
PyFloat6E2M3FNType::bind(m);
904926
PyFloat6E3M2FNType::bind(m);
905927
PyFloat8E4M3FNType::bind(m);
906928
PyFloat8E5M2Type::bind(m);

‎mlir/lib/CAPI/IR/BuiltinTypes.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,18 @@ unsigned mlirFloatTypeGetWidth(MlirType type) {
8585
return llvm::cast<FloatType>(unwrap(type)).getWidth();
8686
}
8787

88+
MlirTypeID mlirFloat6E2M3FNTypeGetTypeID() {
89+
return wrap(Float6E2M3FNType::getTypeID());
90+
}
91+
92+
bool mlirTypeIsAFloat6E2M3FN(MlirType type) {
93+
return unwrap(type).isFloat6E2M3FN();
94+
}
95+
96+
MlirType mlirFloat6E2M3FNTypeGet(MlirContext ctx) {
97+
return wrap(FloatType::getFloat6E2M3FN(unwrap(ctx)));
98+
}
99+
88100
MlirTypeID mlirFloat6E3M2FNTypeGetTypeID() {
89101
return wrap(Float6E3M2FNType::getTypeID());
90102
}

0 commit comments

Comments
 (0)