Skip to content

Commit e0dbf2b

Browse files
apivovarovapivovarov
committed
Add f8E4M3 IEEE 754 type to mlir
1 parent 38a8850 commit e0dbf2b

File tree

24 files changed

+136
-10
lines changed

24 files changed

+136
-10
lines changed

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,16 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M2(MlirType type);
8989
/// context.
9090
MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M2TypeGet(MlirContext ctx);
9191

92+
/// Returns the typeID of an Float8E4M3 type.
93+
MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E4M3TypeGetTypeID(void);
94+
95+
/// Checks whether the given type is an f8E4M3 type.
96+
MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3(MlirType type);
97+
98+
/// Creates an f8E4M3 type in the given context. The type is owned by the
99+
/// context.
100+
MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3TypeGet(MlirContext ctx);
101+
92102
/// Returns the typeID of an Float8E4M3FN type.
93103
MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E4M3FNTypeGetTypeID(void);
94104

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ class Builder {
6161

6262
// Types.
6363
FloatType getFloat8E5M2Type();
64+
FloatType getFloat8E4M3Type();
6465
FloatType getFloat8E4M3FNType();
6566
FloatType getFloat8E5M2FNUZType();
6667
FloatType getFloat8E4M3FNUZType();

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ class FloatType : public Type {
5656
static FloatType getF80(MLIRContext *ctx);
5757
static FloatType getF128(MLIRContext *ctx);
5858
static FloatType getFloat8E5M2(MLIRContext *ctx);
59+
static FloatType getFloat8E4M3(MLIRContext *ctx);
5960
static FloatType getFloat8E4M3FN(MLIRContext *ctx);
6061
static FloatType getFloat8E5M2FNUZ(MLIRContext *ctx);
6162
static FloatType getFloat8E4M3FNUZ(MLIRContext *ctx);
@@ -405,16 +406,20 @@ inline bool BaseMemRefType::isValidElementType(Type type) {
405406
}
406407

407408
inline bool FloatType::classof(Type type) {
408-
return llvm::isa<Float8E5M2Type, Float8E4M3FNType, Float8E5M2FNUZType,
409-
Float8E4M3FNUZType, Float8E4M3B11FNUZType, BFloat16Type,
410-
Float16Type, FloatTF32Type, Float32Type, Float64Type,
411-
Float80Type, Float128Type>(type);
409+
return llvm::isa<
410+
Float8E5M2Type, Float8E4M3Type, Float8E4M3FNType, Float8E5M2FNUZType,
411+
Float8E4M3FNUZType, Float8E4M3B11FNUZType, BFloat16Type, Float16Type,
412+
FloatTF32Type, Float32Type, Float64Type, Float80Type, Float128Type>(type);
412413
}
413414

414415
inline FloatType FloatType::getFloat8E5M2(MLIRContext *ctx) {
415416
return Float8E5M2Type::get(ctx);
416417
}
417418

419+
inline FloatType FloatType::getFloat8E4M3(MLIRContext *ctx) {
420+
return Float8E4M3Type::get(ctx);
421+
}
422+
418423
inline FloatType FloatType::getFloat8E4M3FN(MLIRContext *ctx) {
419424
return Float8E4M3FNType::get(ctx);
420425
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,25 @@ def Builtin_Float8E5M2 : Builtin_FloatType<"Float8E5M2", "f8E5M2"> {
9797
}];
9898
}
9999

100+
//===----------------------------------------------------------------------===//
101+
// Float8E4M3Type
102+
103+
def Builtin_Float8E4M3 : Builtin_FloatType<"Float8E4M3", "f8E4M3"> {
104+
let summary = "8-bit floating point with 3 bit mantissa";
105+
let description = [{
106+
An 8-bit floating point type with 1 sign bit, 4 bits exponent and 3 bits
107+
mantissa. This is not a standard type as defined by IEEE-754, but it
108+
follows similar conventions with the following characteristics:
109+
110+
* bit encoding: S1E4M3
111+
* exponent bias: 7
112+
* infinities: supported with exponent set to all 1s and mantissa 0s
113+
* NaNs: supported with exponent bits set to all 1s and mantissa of
114+
(001, 010, 011, 100, 101, 110, 111)
115+
* denormals when exponent is 0
116+
}];
117+
}
118+
100119
//===----------------------------------------------------------------------===//
101120
// Float8E4M3FNType
102121

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,8 @@ def F8E4M3FN : Type<CPred<"$_self.isFloat8E4M3FN()">, "f8E4M3FN type">,
331331
BuildableType<"$_builder.getFloat8E4M3FNType()">;
332332
def F8E5M2 : Type<CPred<"$_self.isFloat8E5M2()">, "f8E5M2 type">,
333333
BuildableType<"$_builder.getFloat8E5M2Type()">;
334+
def F8E4M3 : Type<CPred<"$_self.isFloat8E4M3()">, "f8E4M3 type">,
335+
BuildableType<"$_builder.getFloat8E4M3Type()">;
334336
def F8E4M3FNUZ : Type<CPred<"$_self.isFloat8E4M3FNUZ()">, "f8E4M3FNUZ type">,
335337
BuildableType<"$_builder.getFloat8E4M3FNUZType()">;
336338
def F8E4M3B11FNUZ : Type<CPred<"$_self.isFloat8E4M3B11FNUZ()">, "f8E4M3B11FNUZ type">,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ class Type {
126126
// derived types should use isa/dyn_cast.
127127
bool isIndex() const;
128128
bool isFloat8E5M2() const;
129+
bool isFloat8E4M3() const;
129130
bool isFloat8E4M3FN() const;
130131
bool isFloat8E5M2FNUZ() const;
131132
bool isFloat8E4M3FNUZ() const;

‎mlir/lib/AsmParser/TokenKinds.def

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ TOK_KEYWORD(f32)
9595
TOK_KEYWORD(f64)
9696
TOK_KEYWORD(f80)
9797
TOK_KEYWORD(f8E5M2)
98+
TOK_KEYWORD(f8E4M3)
9899
TOK_KEYWORD(f8E4M3FN)
99100
TOK_KEYWORD(f8E5M2FNUZ)
100101
TOK_KEYWORD(f8E4M3FNUZ)

‎mlir/lib/AsmParser/TypeParser.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) {
4141
case Token::kw_vector:
4242
case Token::inttype:
4343
case Token::kw_f8E5M2:
44+
case Token::kw_f8E4M3:
4445
case Token::kw_f8E4M3FN:
4546
case Token::kw_f8E5M2FNUZ:
4647
case Token::kw_f8E4M3FNUZ:
@@ -305,6 +306,9 @@ Type Parser::parseNonFunctionType() {
305306
case Token::kw_f8E5M2:
306307
consumeToken(Token::kw_f8E5M2);
307308
return builder.getFloat8E5M2Type();
309+
case Token::kw_f8E4M3:
310+
consumeToken(Token::kw_f8E4M3);
311+
return builder.getFloat8E4M3Type();
308312
case Token::kw_f8E4M3FN:
309313
consumeToken(Token::kw_f8E4M3FN);
310314
return builder.getFloat8E4M3FNType();

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

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ class PyFloat8E4M3FNType
143143
}
144144
};
145145

146-
/// Floating Point Type subclass - Float8M5E2Type.
146+
/// Floating Point Type subclass - Float8E5M2Type.
147147
class PyFloat8E5M2Type : public PyConcreteType<PyFloat8E5M2Type, PyFloatType> {
148148
public:
149149
static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2;
@@ -163,6 +163,26 @@ class PyFloat8E5M2Type : public PyConcreteType<PyFloat8E5M2Type, PyFloatType> {
163163
}
164164
};
165165

166+
/// Floating Point Type subclass - Float8E4M3Type.
167+
class PyFloat8E4M3Type : public PyConcreteType<PyFloat8E4M3Type, PyFloatType> {
168+
public:
169+
static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3;
170+
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
171+
mlirFloat8E4M3TypeGetTypeID;
172+
static constexpr const char *pyClassName = "Float8E4M3Type";
173+
using PyConcreteType::PyConcreteType;
174+
175+
static void bindDerived(ClassTy &c) {
176+
c.def_static(
177+
"get",
178+
[](DefaultingPyMlirContext context) {
179+
MlirType t = mlirFloat8E4M3TypeGet(context->get());
180+
return PyFloat8E4M3Type(context->getRef(), t);
181+
},
182+
py::arg("context") = py::none(), "Create a float8_e4m3 type.");
183+
}
184+
};
185+
166186
/// Floating Point Type subclass - Float8E4M3FNUZ.
167187
class PyFloat8E4M3FNUZType
168188
: public PyConcreteType<PyFloat8E4M3FNUZType, PyFloatType> {
@@ -840,6 +860,7 @@ void mlir::python::populateIRTypes(py::module &m) {
840860
PyIndexType::bind(m);
841861
PyFloat8E4M3FNType::bind(m);
842862
PyFloat8E5M2Type::bind(m);
863+
PyFloat8E4M3Type::bind(m);
843864
PyFloat8E4M3FNUZType::bind(m);
844865
PyFloat8E4M3B11FNUZType::bind(m);
845866
PyFloat8E5M2FNUZType::bind(m);

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,18 @@ MlirType mlirFloat8E5M2TypeGet(MlirContext ctx) {
9898
return wrap(FloatType::getFloat8E5M2(unwrap(ctx)));
9999
}
100100

101+
MlirTypeID mlirFloat8E4M3TypeGetTypeID() {
102+
return wrap(Float8E4M3Type::getTypeID());
103+
}
104+
105+
bool mlirTypeIsAFloat8E4M3(MlirType type) {
106+
return unwrap(type).isFloat8E4M3();
107+
}
108+
109+
MlirType mlirFloat8E4M3TypeGet(MlirContext ctx) {
110+
return wrap(FloatType::getFloat8E4M3(unwrap(ctx)));
111+
}
112+
101113
MlirTypeID mlirFloat8E4M3FNTypeGetTypeID() {
102114
return wrap(Float8E4M3FNType::getTypeID());
103115
}

0 commit comments

Comments
 (0)