[llvm] [mlir] [MLIR] Add f6E3M2FN type (PR #105573)
Sergey Kozub via llvm-commits
llvm-commits at lists.llvm.org
Mon Sep 9 00:03:45 PDT 2024
https://github.com/sergey-kozub updated https://github.com/llvm/llvm-project/pull/105573
>From 7afdded5575c9207f245202a1363c62fdadd9cbc Mon Sep 17 00:00:00 2001
From: Sergey Kozub <skozub at nvidia.com>
Date: Wed, 21 Aug 2024 13:27:01 +0200
Subject: [PATCH] [MLIR] Add f6E3M2FN type (#105573).
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
`f6E3M2FN` 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 S1E3M2. Unlike IEEE-754 types, there are no infinity or NaN values.
```c
f6E3M2FN
- Exponent bias: 3
- Maximum stored exponent value: 7 (binary 111)
- Maximum unbiased exponent value: 7 - 3 = 4
- Minimum stored exponent value: 1 (binary 001)
- Minimum unbiased exponent value: 1 − 3 = −2
- Has Positive and Negative zero
- Doesn't have infinity
- Doesn't have NaNs
Additional details:
- Zeros (+/-): S.000.00
- Max normal number: S.111.11 = ±2^(4) x (1 + 0.75) = ±28
- Min normal number: S.001.00 = ±2^(-2) = ±0.25
- Max subnormal number: S.000.11 = ±2^(-2) x 0.75 = ±0.1875
- Min subnormal number: S.000.01 = ±2^(-2) x 0.25 = ±0.0625
```
Related PRs:
- [PR-94735](https://github.com/llvm/llvm-project/pull/94735) [APFloat] Add APFloat support for FP6 data types
- [PR-97118](https://github.com/llvm/llvm-project/pull/97118) [MLIR] Add f8E4M3 type - was used as a template for this PR
---
llvm/unittests/ADT/APFloatTest.cpp | 7 +++---
mlir/include/mlir-c/BuiltinTypes.h | 10 +++++++++
mlir/include/mlir/IR/Builders.h | 1 +
mlir/include/mlir/IR/BuiltinTypes.h | 9 ++++++--
mlir/include/mlir/IR/BuiltinTypes.td | 21 ++++++++++++++++++
mlir/include/mlir/IR/CommonTypeConstraints.td | 2 ++
mlir/include/mlir/IR/Types.h | 1 +
mlir/lib/AsmParser/TokenKinds.def | 1 +
mlir/lib/AsmParser/TypeParser.cpp | 4 ++++
mlir/lib/Bindings/Python/IRTypes.cpp | 22 +++++++++++++++++++
mlir/lib/CAPI/IR/BuiltinTypes.cpp | 12 ++++++++++
.../Conversion/LLVMCommon/TypeConverter.cpp | 3 ++-
.../Transforms/EmulateUnsupportedFloats.cpp | 1 +
mlir/lib/IR/AsmPrinter.cpp | 1 +
mlir/lib/IR/Builders.cpp | 4 ++++
mlir/lib/IR/BuiltinTypes.cpp | 4 ++++
mlir/lib/IR/MLIRContext.cpp | 5 +++++
mlir/lib/IR/Types.cpp | 1 +
mlir/python/mlir/_mlir_libs/_mlir/ir.pyi | 14 ++++++++++++
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 +-
25 files changed, 137 insertions(+), 7 deletions(-)
diff --git a/llvm/unittests/ADT/APFloatTest.cpp b/llvm/unittests/ADT/APFloatTest.cpp
index be675bb7fe5a53..6c49d78e5c8ea9 100644
--- a/llvm/unittests/ADT/APFloatTest.cpp
+++ b/llvm/unittests/ADT/APFloatTest.cpp
@@ -2084,16 +2084,17 @@ TEST(APFloatTest, getSmallestNormalized) {
EXPECT_FALSE(test.isDenormal());
EXPECT_TRUE(test.bitwiseIsEqual(expected));
EXPECT_TRUE(test.isSmallestNormalized());
+
test = APFloat::getSmallestNormalized(APFloat::Float6E3M2FN(), false);
expected = APFloat(APFloat::Float6E3M2FN(), "0x1p-2");
-
- test = APFloat::getSmallestNormalized(APFloat::Float4E2M1FN(), false);
- expected = APFloat(APFloat::Float4E2M1FN(), "0x1p0");
EXPECT_FALSE(test.isNegative());
EXPECT_TRUE(test.isFiniteNonZero());
EXPECT_FALSE(test.isDenormal());
EXPECT_TRUE(test.bitwiseIsEqual(expected));
EXPECT_TRUE(test.isSmallestNormalized());
+
+ test = APFloat::getSmallestNormalized(APFloat::Float4E2M1FN(), false);
+ expected = APFloat(APFloat::Float4E2M1FN(), "0x1p0");
EXPECT_FALSE(test.isNegative());
EXPECT_TRUE(test.isFiniteNonZero());
EXPECT_FALSE(test.isDenormal());
diff --git a/mlir/include/mlir-c/BuiltinTypes.h b/mlir/include/mlir-c/BuiltinTypes.h
index d698bf4764568f..24531baecaa353 100644
--- a/mlir/include/mlir-c/BuiltinTypes.h
+++ b/mlir/include/mlir-c/BuiltinTypes.h
@@ -79,6 +79,16 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat(MlirType type);
/// Returns the bitwidth of a floating-point type.
MLIR_CAPI_EXPORTED unsigned mlirFloatTypeGetWidth(MlirType type);
+/// Returns the typeID of an Float6E3M2FN type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat6E3M2FNTypeGetTypeID(void);
+
+/// Checks whether the given type is an f6E3M2FN type.
+MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat6E3M2FN(MlirType type);
+
+/// Creates an f6E3M2FN type in the given context. The type is owned by the
+/// context.
+MLIR_CAPI_EXPORTED MlirType mlirFloat6E3M2FNTypeGet(MlirContext ctx);
+
/// Returns the typeID of an Float8E5M2 type.
MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E5M2TypeGetTypeID(void);
diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index b5962f3783924f..5ac3a04b1c26ba 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -60,6 +60,7 @@ class Builder {
Attribute metadata = Attribute());
// Types.
+ FloatType getFloat6E3M2FNType();
FloatType getFloat8E5M2Type();
FloatType getFloat8E4M3Type();
FloatType getFloat8E4M3FNType();
diff --git a/mlir/include/mlir/IR/BuiltinTypes.h b/mlir/include/mlir/IR/BuiltinTypes.h
index eefa4279df1a01..87ccc041f19758 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.h
+++ b/mlir/include/mlir/IR/BuiltinTypes.h
@@ -67,6 +67,7 @@ class FloatType : public Type {
static FloatType getFloat8E4M3FNUZ(MLIRContext *ctx);
static FloatType getFloat8E4M3B11FNUZ(MLIRContext *ctx);
static FloatType getFloat8E3M4(MLIRContext *ctx);
+ static FloatType getFloat6E3M2FN(MLIRContext *ctx);
/// Methods for support type inquiry through isa, cast, and dyn_cast.
static bool classof(Type type);
@@ -413,13 +414,17 @@ inline bool BaseMemRefType::isValidElementType(Type type) {
}
inline bool FloatType::classof(Type type) {
- return llvm::isa<Float8E5M2Type, Float8E4M3Type, Float8E4M3FNType,
- Float8E5M2FNUZType, Float8E4M3FNUZType,
+ return llvm::isa<Float6E3M2FNType, Float8E5M2Type, Float8E4M3Type,
+ Float8E4M3FNType, Float8E5M2FNUZType, Float8E4M3FNUZType,
Float8E4M3B11FNUZType, Float8E3M4Type, BFloat16Type,
Float16Type, FloatTF32Type, Float32Type, Float64Type,
Float80Type, Float128Type>(type);
}
+inline FloatType FloatType::getFloat6E3M2FN(MLIRContext *ctx) {
+ return Float6E3M2FNType::get(ctx);
+}
+
inline FloatType FloatType::getFloat8E5M2(MLIRContext *ctx) {
return Float8E5M2Type::get(ctx);
}
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 1ab1bbe9bfc9b2..b54d4ee4b7eb7a 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -233,6 +233,27 @@ def Builtin_Float8E3M4 : Builtin_FloatType<"Float8E3M4", "f8E3M4"> {
}];
}
+//===----------------------------------------------------------------------===//
+// Float6E3M2FNType
+
+def Builtin_Float6E3M2FN : Builtin_FloatType<"Float6E3M2FN", "f6E3M2FN"> {
+ let summary = "6-bit floating point with 3 bits exponent and 2 bit mantissa";
+ let description = [{
+ An 6-bit floating point type with 1 sign bit, 3 bits exponent and 2 bits
+ mantissa. This is not a standard type as defined by IEEE-754, but it
+ follows similar conventions with the following characteristics:
+
+ * bit encoding: S1E3M2
+ * exponent bias: 3
+ * infinities: Not supported
+ * NaNs: Not supported
+ * denormals when exponent is 0
+
+ Open Compute Project (OCP) microscaling formats (MX) specification:
+ https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
+ }];
+}
+
//===----------------------------------------------------------------------===//
// BFloat16Type
diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td
index 4536d781ef674f..09eab50f53a540 100644
--- a/mlir/include/mlir/IR/CommonTypeConstraints.td
+++ b/mlir/include/mlir/IR/CommonTypeConstraints.td
@@ -344,6 +344,8 @@ def F8E5M2FNUZ : Type<CPred<"$_self.isFloat8E5M2FNUZ()">, "f8E5M2FNUZ type">,
BuildableType<"$_builder.getFloat8E5M2FNUZType()">;
def F8E3M4 : Type<CPred<"$_self.isFloat8E3M4()">, "f8E3M4 type">,
BuildableType<"$_builder.getFloat8E3M4Type()">;
+def F6E3M2FN : Type<CPred<"$_self.isFloat6E3M2FN()">, "f6E3M2FN type">,
+ BuildableType<"$_builder.getFloat6E3M2FNType()">;
def AnyComplex : Type<CPred<"::llvm::isa<::mlir::ComplexType>($_self)">,
"complex-type", "::mlir::ComplexType">;
diff --git a/mlir/include/mlir/IR/Types.h b/mlir/include/mlir/IR/Types.h
index dfc7e69472c891..b6a307fd7cb0fe 100644
--- a/mlir/include/mlir/IR/Types.h
+++ b/mlir/include/mlir/IR/Types.h
@@ -125,6 +125,7 @@ class Type {
// Convenience predicates. This is only for floating point types,
// derived types should use isa/dyn_cast.
bool isIndex() const;
+ bool isFloat6E3M2FN() const;
bool isFloat8E5M2() const;
bool isFloat8E4M3() const;
bool isFloat8E4M3FN() const;
diff --git a/mlir/lib/AsmParser/TokenKinds.def b/mlir/lib/AsmParser/TokenKinds.def
index 4c1c1c21031c88..fa18cbe9e2b901 100644
--- a/mlir/lib/AsmParser/TokenKinds.def
+++ b/mlir/lib/AsmParser/TokenKinds.def
@@ -101,6 +101,7 @@ TOK_KEYWORD(f8E5M2FNUZ)
TOK_KEYWORD(f8E4M3FNUZ)
TOK_KEYWORD(f8E4M3B11FNUZ)
TOK_KEYWORD(f8E3M4)
+TOK_KEYWORD(f6E3M2FN)
TOK_KEYWORD(f128)
TOK_KEYWORD(false)
TOK_KEYWORD(floordiv)
diff --git a/mlir/lib/AsmParser/TypeParser.cpp b/mlir/lib/AsmParser/TypeParser.cpp
index f070c072c43296..05276031211fa9 100644
--- a/mlir/lib/AsmParser/TypeParser.cpp
+++ b/mlir/lib/AsmParser/TypeParser.cpp
@@ -39,6 +39,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) {
case Token::kw_tuple:
case Token::kw_vector:
case Token::inttype:
+ case Token::kw_f6E3M2FN:
case Token::kw_f8E5M2:
case Token::kw_f8E4M3:
case Token::kw_f8E4M3FN:
@@ -303,6 +304,9 @@ Type Parser::parseNonFunctionType() {
}
// float-type
+ case Token::kw_f6E3M2FN:
+ consumeToken(Token::kw_f6E3M2FN);
+ return builder.getFloat6E3M2FNType();
case Token::kw_f8E5M2:
consumeToken(Token::kw_f8E5M2);
return builder.getFloat8E5M2Type();
diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
index 2ee1d89c38884d..1cb429d9ca7b2d 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -124,6 +124,27 @@ class PyFloatType : public PyConcreteType<PyFloatType> {
}
};
+/// Floating Point Type subclass - Float6E3M2FNType.
+class PyFloat6E3M2FNType
+ : public PyConcreteType<PyFloat6E3M2FNType, PyFloatType> {
+public:
+ static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat6E3M2FN;
+ static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+ mlirFloat6E3M2FNTypeGetTypeID;
+ static constexpr const char *pyClassName = "Float6E3M2FNType";
+ using PyConcreteType::PyConcreteType;
+
+ static void bindDerived(ClassTy &c) {
+ c.def_static(
+ "get",
+ [](DefaultingPyMlirContext context) {
+ MlirType t = mlirFloat6E3M2FNTypeGet(context->get());
+ return PyFloat6E3M2FNType(context->getRef(), t);
+ },
+ py::arg("context") = py::none(), "Create a float6_e3m2fn type.");
+ }
+};
+
/// Floating Point Type subclass - Float8E4M3FNType.
class PyFloat8E4M3FNType
: public PyConcreteType<PyFloat8E4M3FNType, PyFloatType> {
@@ -880,6 +901,7 @@ void mlir::python::populateIRTypes(py::module &m) {
PyIntegerType::bind(m);
PyFloatType::bind(m);
PyIndexType::bind(m);
+ PyFloat6E3M2FNType::bind(m);
PyFloat8E4M3FNType::bind(m);
PyFloat8E5M2Type::bind(m);
PyFloat8E4M3Type::bind(m);
diff --git a/mlir/lib/CAPI/IR/BuiltinTypes.cpp b/mlir/lib/CAPI/IR/BuiltinTypes.cpp
index 2aa2e922f2abcc..254650d66a67e6 100644
--- a/mlir/lib/CAPI/IR/BuiltinTypes.cpp
+++ b/mlir/lib/CAPI/IR/BuiltinTypes.cpp
@@ -85,6 +85,18 @@ unsigned mlirFloatTypeGetWidth(MlirType type) {
return llvm::cast<FloatType>(unwrap(type)).getWidth();
}
+MlirTypeID mlirFloat6E3M2FNTypeGetTypeID() {
+ return wrap(Float6E3M2FNType::getTypeID());
+}
+
+bool mlirTypeIsAFloat6E3M2FN(MlirType type) {
+ return unwrap(type).isFloat6E3M2FN();
+}
+
+MlirType mlirFloat6E3M2FNTypeGet(MlirContext ctx) {
+ return wrap(FloatType::getFloat6E3M2FN(unwrap(ctx)));
+}
+
MlirTypeID mlirFloat8E5M2TypeGetTypeID() {
return wrap(Float8E5M2Type::getTypeID());
}
diff --git a/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp b/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
index 5313a64ed47e3a..b2c54bb3212edb 100644
--- a/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
+++ b/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
@@ -249,7 +249,8 @@ Type LLVMTypeConverter::convertIntegerType(IntegerType type) const {
Type LLVMTypeConverter::convertFloatType(FloatType type) const {
if (type.isFloat8E5M2() || type.isFloat8E4M3() || type.isFloat8E4M3FN() ||
type.isFloat8E5M2FNUZ() || type.isFloat8E4M3FNUZ() ||
- type.isFloat8E4M3B11FNUZ() || type.isFloat8E3M4())
+ type.isFloat8E4M3B11FNUZ() || type.isFloat8E3M4() ||
+ type.isFloat6E3M2FN())
return IntegerType::get(&getContext(), type.getWidth());
return type;
}
diff --git a/mlir/lib/Dialect/Arith/Transforms/EmulateUnsupportedFloats.cpp b/mlir/lib/Dialect/Arith/Transforms/EmulateUnsupportedFloats.cpp
index 51f229ef937c40..a5ee6edc6320d5 100644
--- a/mlir/lib/Dialect/Arith/Transforms/EmulateUnsupportedFloats.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/EmulateUnsupportedFloats.cpp
@@ -55,6 +55,7 @@ static std::optional<FloatType> parseFloatType(MLIRContext *ctx,
StringRef name) {
Builder b(ctx);
return llvm::StringSwitch<std::optional<FloatType>>(name)
+ .Case("f6E3M2FN", b.getFloat6E3M2FNType())
.Case("f8E5M2", b.getFloat8E5M2Type())
.Case("f8E4M3", b.getFloat8E4M3Type())
.Case("f8E4M3FN", b.getFloat8E4M3FNType())
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 02acc8c3f4659e..5142b462820786 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -2575,6 +2575,7 @@ void AsmPrinter::Impl::printTypeImpl(Type type) {
opaqueTy.getTypeData());
})
.Case<IndexType>([&](Type) { os << "index"; })
+ .Case<Float6E3M2FNType>([&](Type) { os << "f6E3M2FN"; })
.Case<Float8E5M2Type>([&](Type) { os << "f8E5M2"; })
.Case<Float8E4M3Type>([&](Type) { os << "f8E4M3"; })
.Case<Float8E4M3FNType>([&](Type) { os << "f8E4M3FN"; })
diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index e3d6d71fb61dfb..71f622b02adee0 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -34,6 +34,10 @@ Location Builder::getFusedLoc(ArrayRef<Location> locs, Attribute metadata) {
// Types.
//===----------------------------------------------------------------------===//
+FloatType Builder::getFloat6E3M2FNType() {
+ return FloatType::getFloat6E3M2FN(context);
+}
+
FloatType Builder::getFloat8E5M2Type() {
return FloatType::getFloat8E5M2(context);
}
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 16b53efa55fb80..4d91d98a9ba896 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/mlir/lib/IR/BuiltinTypes.cpp
@@ -91,6 +91,8 @@ IntegerType IntegerType::scaleElementBitwidth(unsigned scale) {
//===----------------------------------------------------------------------===//
unsigned FloatType::getWidth() {
+ if (llvm::isa<Float6E3M2FNType>(*this))
+ return 6;
if (llvm::isa<Float8E5M2Type, Float8E4M3Type, Float8E4M3FNType,
Float8E5M2FNUZType, Float8E4M3FNUZType, Float8E4M3B11FNUZType,
Float8E3M4Type>(*this))
@@ -110,6 +112,8 @@ unsigned FloatType::getWidth() {
/// Returns the floating semantics for the given type.
const llvm::fltSemantics &FloatType::getFloatSemantics() {
+ if (llvm::isa<Float6E3M2FNType>(*this))
+ return APFloat::Float6E3M2FN();
if (llvm::isa<Float8E5M2Type>(*this))
return APFloat::Float8E5M2();
if (llvm::isa<Float8E4M3Type>(*this))
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 5c93747438ecdb..2851e6457ea3cb 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -221,6 +221,7 @@ class MLIRContextImpl {
llvm::DenseMap<StringRef, AbstractType *> nameToType;
/// Cached Type Instances.
+ Float6E3M2FNType f6E3M2FNTy;
Float8E5M2Type f8E5M2Ty;
Float8E4M3Type f8E4M3Ty;
Float8E4M3FNType f8E4M3FNTy;
@@ -313,6 +314,7 @@ MLIRContext::MLIRContext(const DialectRegistry ®istry, Threading setting)
//// Types.
/// Floating-point Types.
+ impl->f6E3M2FNTy = TypeUniquer::get<Float6E3M2FNType>(this);
impl->f8E5M2Ty = TypeUniquer::get<Float8E5M2Type>(this);
impl->f8E4M3Ty = TypeUniquer::get<Float8E4M3Type>(this);
impl->f8E4M3FNTy = TypeUniquer::get<Float8E4M3FNType>(this);
@@ -1013,6 +1015,9 @@ AbstractType::lookup(StringRef name, MLIRContext *context) {
/// This should not be used directly.
StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; }
+Float6E3M2FNType Float6E3M2FNType::get(MLIRContext *context) {
+ return context->getImpl().f6E3M2FNTy;
+}
Float8E5M2Type Float8E5M2Type::get(MLIRContext *context) {
return context->getImpl().f8E5M2Ty;
}
diff --git a/mlir/lib/IR/Types.cpp b/mlir/lib/IR/Types.cpp
index 2bc26388b6218a..fa093664cf77f1 100644
--- a/mlir/lib/IR/Types.cpp
+++ b/mlir/lib/IR/Types.cpp
@@ -34,6 +34,7 @@ Type AbstractType::replaceImmediateSubElements(Type type,
MLIRContext *Type::getContext() const { return getDialect().getContext(); }
+bool Type::isFloat6E3M2FN() const { return llvm::isa<Float6E3M2FNType>(*this); }
bool Type::isFloat8E5M2() const { return llvm::isa<Float8E5M2Type>(*this); }
bool Type::isFloat8E4M3() const { return llvm::isa<Float8E4M3Type>(*this); }
bool Type::isFloat8E4M3FN() const { return llvm::isa<Float8E4M3FNType>(*this); }
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index 4a2d0e977ccf26..7b4fac7275bfc6 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -120,6 +120,7 @@ __all__ = [
"F32Type",
"F64Type",
"FlatSymbolRefAttr",
+ "Float6E3M2FNType",
"Float8E3M4Type",
"Float8E4M3B11FNUZType",
"Float8E4M3FNType",
@@ -1539,6 +1540,19 @@ class FlatSymbolRefAttr(Attribute):
Returns the value of the FlatSymbolRef attribute as a string
"""
+class Float6E3M2FNType(FloatType):
+ static_typeid: ClassVar[TypeID]
+ @staticmethod
+ def get(context: Optional[Context] = None) -> Float6E3M2FNType:
+ """
+ Create a float6_e3m2fn type.
+ """
+ @staticmethod
+ def isinstance(other: Type) -> bool: ...
+ def __init__(self, cast_from_type: Type) -> None: ...
+ @property
+ def typeid(self) -> TypeID: ...
+
class Float8E3M4Type(FloatType):
static_typeid: ClassVar[TypeID]
@staticmethod
diff --git a/mlir/python/mlir/extras/types.py b/mlir/python/mlir/extras/types.py
index fe7c3e25d16906..0c6ece91d8b94a 100644
--- a/mlir/python/mlir/extras/types.py
+++ b/mlir/python/mlir/extras/types.py
@@ -12,6 +12,7 @@
F16Type,
F32Type,
F64Type,
+ Float6E3M2FNType,
Float8E3M4Type,
Float8E4M3B11FNUZType,
Float8E4M3FNType,
@@ -74,6 +75,7 @@ def ui(width):
f8E4M3FN = lambda: Float8E4M3FNType.get()
f8E4M3B11FNUZ = lambda: Float8E4M3B11FNUZType.get()
f8E3M4 = lambda: Float8E3M4Type.get()
+f6E3M2FN = lambda: Float6E3M2FNType.get()
none = lambda: NoneType.get()
diff --git a/mlir/test/IR/attribute.mlir b/mlir/test/IR/attribute.mlir
index ac0aec113add17..38cbf9d5d2b579 100644
--- a/mlir/test/IR/attribute.mlir
+++ b/mlir/test/IR/attribute.mlir
@@ -36,6 +36,10 @@ func.func @any_attr_of_fail() {
//===----------------------------------------------------------------------===//
func.func @float_attrs_pass() {
+ "test.float_attrs"() {
+ // CHECK: float_attr = 2.000000e+00 : f6E3M2FN
+ float_attr = 2. : f6E3M2FN
+ } : () -> ()
"test.float_attrs"() {
// CHECK: float_attr = 2.000000e+00 : f8E5M2
float_attr = 2. : f8E5M2
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index df61fef605fde0..b53583e8e67da9 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -39,6 +39,9 @@ llvm.mlir.global internal constant @string_const("foobar") : !llvm.array<6 x i8>
// CHECK: @int_global_undef = internal global i64 undef
llvm.mlir.global internal @int_global_undef() : i64
+// CHECK: @f6E3M2FN_global_as_i6 = internal global i6 14
+llvm.mlir.global internal @f6E3M2FN_global_as_i6(1.5 : f6E3M2FN) : i6
+
// CHECK: @f8E3M4_global_as_i8 = internal global i8 56
llvm.mlir.global internal @f8E3M4_global_as_i8(1.5 : f8E3M4) : i8
diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py
index f95cccc54105ed..b72ef4de0bd6dd 100644
--- a/mlir/test/python/ir/builtin_types.py
+++ b/mlir/test/python/ir/builtin_types.py
@@ -113,6 +113,8 @@ def testTypeIsInstance():
def testFloatTypeSubclasses():
ctx = Context()
# CHECK: True
+ print(isinstance(Type.parse("f6E3M2FN", ctx), FloatType))
+ # CHECK: True
print(isinstance(Type.parse("f8E3M4", ctx), FloatType))
# CHECK: True
print(isinstance(Type.parse("f8E4M3", ctx), FloatType))
@@ -233,6 +235,8 @@ def testIndexType():
@run
def testFloatType():
with Context():
+ # CHECK: float: f6E3M2FN
+ print("float:", Float6E3M2FNType.get())
# CHECK: float: f8E3M4
print("float:", Float8E3M4Type.get())
# CHECK: float: f8E4M3
@@ -609,6 +613,7 @@ def testTypeIDs():
types = [
(IntegerType, IntegerType.get_signless(16)),
(IndexType, IndexType.get()),
+ (Float6E3M2FNType, Float6E3M2FNType.get()),
(Float8E3M4Type, Float8E3M4Type.get()),
(Float8E4M3Type, Float8E4M3Type.get()),
(Float8E4M3FNType, Float8E4M3FNType.get()),
@@ -634,6 +639,7 @@ def testTypeIDs():
# CHECK: IntegerType(i16)
# CHECK: IndexType(index)
+ # CHECK: Float6E3M2FNType(f6E3M2FN)
# CHECK: Float8E3M4Type(f8E3M4)
# CHECK: Float8E4M3Type(f8E4M3)
# CHECK: Float8E4M3FNType(f8E4M3FN)
@@ -713,6 +719,9 @@ def print_downcasted(typ):
# CHECK: F64Type
# CHECK: F64Type(f64)
print_downcasted(F64Type.get())
+ # CHECK: Float6E3M2FNType
+ # CHECK: Float6E3M2FNType(f6E3M2FN)
+ print_downcasted(Float6E3M2FNType.get())
# CHECK: Float8E3M4Type
# CHECK: Float8E3M4Type(f8E3M4)
print_downcasted(Float8E3M4Type.get())
diff --git a/mlir/utils/lldb-scripts/mlirDataFormatters.py b/mlir/utils/lldb-scripts/mlirDataFormatters.py
index e7c526842439b9..fed149d03ecf31 100644
--- a/mlir/utils/lldb-scripts/mlirDataFormatters.py
+++ b/mlir/utils/lldb-scripts/mlirDataFormatters.py
@@ -50,6 +50,7 @@ def build_ptr_str_from_addr(addrValue: lldb.SBValue, type: lldb.SBType):
"mlir::CallSiteLoc": '"loc(callsite(...))"',
"mlir::FusedLoc": '"loc(fused<...>[...])"',
"mlir::UnknownLoc": '"loc(unknown)"',
+ "mlir::Float6E3M2FNType": '"f6E3M2FN"',
"mlir::Float8E5M2Type": '"f8E5M2"',
"mlir::Float8E4M3Type": '"f8E4M3"',
"mlir::Float8E4M3FNType": '"f8E4M3FN"',
diff --git a/mlir/utils/tree-sitter-mlir/grammar.js b/mlir/utils/tree-sitter-mlir/grammar.js
index b5926d75da4f25..d2c66714b4b118 100644
--- a/mlir/utils/tree-sitter-mlir/grammar.js
+++ b/mlir/utils/tree-sitter-mlir/grammar.js
@@ -231,7 +231,7 @@ const common = {
token(seq(choice('si', 'ui', 'i'), /[1-9]/, repeat(/[0-9]/))),
float_type : $ => token(
choice('f16', 'f32', 'f64', 'f80', 'f128', 'bf16', 'f8E3M4', 'f8E4M3FN',
- 'f8E4M3', 'f8E5M2')),
+ 'f8E4M3', 'f8E5M2', 'f6E3M2FN')),
index_type : $ => token('index'),
none_type : $ => token('none'),
complex_type : $ => seq(token('complex'), '<', $._prim_type, '>'),
More information about the llvm-commits
mailing list