[Mlir-commits] [mlir] 9de8463 - Extending UniformQuantizedType with interface-based support for new storage types in Quant dialect (#152966)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Feb 10 01:56:29 PST 2026
Author: Roman-Pevnyi
Date: 2026-02-10T09:56:24Z
New Revision: 9de84638b97425db97f5d86879fd071593702eca
URL: https://github.com/llvm/llvm-project/commit/9de84638b97425db97f5d86879fd071593702eca
DIFF: https://github.com/llvm/llvm-project/commit/9de84638b97425db97f5d86879fd071593702eca.diff
LOG: Extending UniformQuantizedType with interface-based support for new storage types in Quant dialect (#152966)
Currently, UniformQuantizedType only supports built-in MLIR storage
types such as Integer. LLM quantization research introducing feature of
using NF4 as a low precision datatype (see
https://arxiv.org/pdf/2305.14314). There is a growing need to make the
system extensible and maintainable as more types are added. Ensuring
that MLIR can natively support NF4 through a clean, extensible interface
is essential for both current and future quantization workflows.
**Current Approach and Its Limitations:**
- The present implementation relies on dynamic checks (e.g., type
switches or if-else chains) to determine the storage type and retrieve
type-specific information for legality checks.
- This approach works for a small, fixed set of types, but as the number
of supported types grows, the code becomes harder to read, maintain, and
extend.
**Proposed Interface-Based Approach:**
- Define a StorageTypeInterface that specifies the required methods any
storage type must implement to be used in UniformQuantizedType.
- Each storage type (Integer, Float8E5M2, Float8E4M3FN, and new types
like NF4) would implement this interface, encapsulating their
type-specific logic.
- When UniformQuantizedType needs to check legality or retrieve
information, it can use MLIR’s dyn_cast mechanism to check if the type
implements the interface and then call the required methods.
- This design decouples UniformQuantizedType from the specifics of each
storage type, making it easy to add new types (such as NF4) without
modifying the core logic or introducing more type checks.
**Benefits:**
- Extensibility: New storage types can be added by simply implementing
the interface, without touching the core UniformQuantizedType logic.
- Readability: The code is cleaner, as it avoids large switch statements
or if-else chains.
- Maintainability: Type-specific logic is encapsulated within each type,
reducing the risk of errors and making the codebase easier to understand
and update.
Added:
mlir/include/mlir/IR/QuantStorageTypeInterface.h
mlir/include/mlir/IR/QuantStorageTypeInterface.td
mlir/lib/IR/QuantStorageTypeInterface.cpp
Modified:
mlir/cmake/modules/AddMLIR.cmake
mlir/include/mlir/IR/BuiltinTypes.h
mlir/include/mlir/IR/BuiltinTypes.td
mlir/include/mlir/IR/CMakeLists.txt
mlir/lib/Dialect/Quant/IR/QuantTypes.cpp
mlir/lib/Dialect/Quant/IR/TypeParser.cpp
mlir/lib/IR/CMakeLists.txt
mlir/test/Dialect/Quant/parse-uniform-invalid.mlir
mlir/test/Dialect/Quant/parse-uniform.mlir
Removed:
################################################################################
diff --git a/mlir/cmake/modules/AddMLIR.cmake b/mlir/cmake/modules/AddMLIR.cmake
index 6589458ab7894..be28a2ab900c9 100644
--- a/mlir/cmake/modules/AddMLIR.cmake
+++ b/mlir/cmake/modules/AddMLIR.cmake
@@ -216,6 +216,15 @@ macro(add_mlir_generic_tablegen_target target)
add_dependencies(mlir-generic-headers ${target})
endmacro()
+# Declare a dialect in the include directory
+function(add_mlir_type_interface interface)
+ set(LLVM_TARGET_DEFINITIONS ${interface}.td)
+ mlir_tablegen(${interface}.h.inc -gen-type-interface-decls)
+ mlir_tablegen(${interface}.cpp.inc -gen-type-interface-defs)
+ add_public_tablegen_target(MLIR${interface}IncGen)
+ add_dependencies(mlir-generic-headers MLIR${interface}IncGen)
+endfunction()
+
# Generate Documentation
function(add_mlir_doc doc_filename output_file output_directory command)
set(LLVM_TARGET_DEFINITIONS ${doc_filename}.td)
diff --git a/mlir/include/mlir/IR/BuiltinTypes.h b/mlir/include/mlir/IR/BuiltinTypes.h
index 86ec5c43970b1..d30cba29c9814 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.h
+++ b/mlir/include/mlir/IR/BuiltinTypes.h
@@ -167,6 +167,8 @@ class BaseMemRefType : public Type,
// Tablegen Type Declarations
//===----------------------------------------------------------------------===//
+#include "mlir/IR/QuantStorageTypeInterface.h"
+
#define GET_TYPEDEF_CLASSES
#include "mlir/IR/BuiltinTypes.h.inc"
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 08847dd11c685..806064faeda00 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -17,6 +17,7 @@
include "mlir/IR/AttrTypeBase.td"
include "mlir/IR/BuiltinDialect.td"
include "mlir/IR/BuiltinTypeInterfaces.td"
+include "mlir/IR/QuantStorageTypeInterface.td"
include "mlir/IR/CommonTypeConstraints.td"
// TODO: Currently the types defined in this file are prefixed with `Builtin_`.
@@ -80,17 +81,19 @@ def Builtin_Complex : Builtin_Type<"Complex", "complex"> {
// Base class for Builtin dialect float types.
class Builtin_FloatType<string name, string mnemonic,
+ list<Trait> traits = [],
list<string> declaredInterfaceMethods = []>
- : Builtin_Type<name, mnemonic, /*traits=*/[
- DeclareTypeInterfaceMethods<
+ : Builtin_Type<name, mnemonic, /*traits=*/
+ traits # [DeclareTypeInterfaceMethods<
FloatTypeInterface,
["getFloatSemantics"] # declaredInterfaceMethods>]> {
}
// Float types that are cached in MLIRContext.
class Builtin_CachedFloatType<string name, string mnemonic,
+ list<Trait> traits = [],
list<string> declaredInterfaceMethods = []>
- : Builtin_FloatType<name, mnemonic, declaredInterfaceMethods> {
+ : Builtin_FloatType<name, mnemonic, traits, declaredInterfaceMethods> {
let extraClassDeclaration = [{
static }] # name # [{Type get(MLIRContext *context);
}];
@@ -100,7 +103,8 @@ class Builtin_CachedFloatType<string name, string mnemonic,
// Float8E5M2Type
//===----------------------------------------------------------------------===//
-def Builtin_Float8E5M2 : Builtin_FloatType<"Float8E5M2", "f8E5M2"> {
+def Builtin_Float8E5M2 : Builtin_FloatType<"Float8E5M2", "f8E5M2",
+ [QuantStorageTypeInterface]> {
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
@@ -116,6 +120,34 @@ def Builtin_Float8E5M2 : Builtin_FloatType<"Float8E5M2", "f8E5M2"> {
Described in: https://arxiv.org/abs/2209.05433
}];
+
+ let extraClassDeclaration = [{
+ /// QuantStorageTypeInterface method implementations
+ /// Returns true if the the storage type should default to signed
+ /// when used in quantization.
+ bool shouldDefaultToSigned() const { return true; }
+ /// Get the bit width of this 8-bit floating point type.
+ unsigned getStorageWidth() const { return 8; }
+
+ /// Get default min, max value for this 8-bit floating point type.
+ int64_t getDefaultMaximum(bool isSigned) const { return 57344; }
+ int64_t getDefaultMinimum(bool isSigned) const { return -getDefaultMaximum(isSigned); }
+
+ /// Get the storage type as a string.
+ std::string getStorageTypeName(bool isSigned) const { return "f8E5M2"; }
+
+ /// Check if this 8-bit floating point type uses packed representation.
+ bool isPacked() const { return false; }
+
+ /// Get the logical bit width per value for this 8-bit floating point type.
+ unsigned getLogicalBitWidth() const { return 8; }
+
+ /// Returns how many values of this type fit in one byte
+ unsigned getElementsPerByte() const { return 1; }
+
+ /// Get the preferred alignment in bytes for this 8-bit floating point type.
+ std::optional<unsigned> getPreferredAlignmentBytes() const { return std::nullopt; }
+ }];
}
//===----------------------------------------------------------------------===//
@@ -142,7 +174,8 @@ def Builtin_Float8E4M3 : Builtin_FloatType<"Float8E4M3", "f8E4M3"> {
// Float8E4M3FNType
//===----------------------------------------------------------------------===//
-def Builtin_Float8E4M3FN : Builtin_FloatType<"Float8E4M3FN", "f8E4M3FN"> {
+def Builtin_Float8E4M3FN : Builtin_FloatType<"Float8E4M3FN", "f8E4M3FN",
+ [QuantStorageTypeInterface]> {
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
@@ -159,6 +192,34 @@ def Builtin_Float8E4M3FN : Builtin_FloatType<"Float8E4M3FN", "f8E4M3FN"> {
Described in: https://arxiv.org/abs/2209.05433
}];
+
+ let extraClassDeclaration = [{
+ /// QuantStorageTypeInterface method implementations
+ /// Returns true if the the storage type should default to signed
+ /// when used in quantization.
+ bool shouldDefaultToSigned() const { return true; }
+ /// Get the bit width of this 8-bit floating point type.
+ unsigned getStorageWidth() const { return 8; }
+
+ /// Get default min, max value for this 8-bit floating point type.
+ int64_t getDefaultMaximum(bool isSigned) const { return 448; }
+ int64_t getDefaultMinimum(bool isSigned) const { return -getDefaultMaximum(isSigned); }
+
+ /// Get the storage type as a string.
+ std::string getStorageTypeName(bool isSigned) const { return "f8E4M3FN"; }
+
+ /// Check if this 8-bit floating point type uses packed representation.
+ bool isPacked() const { return false; }
+
+ /// Get the logical bit width per value for this 8-bit floating point type.
+ unsigned getLogicalBitWidth() const { return 8; }
+
+ /// Returns how many values of this type fit in one byte
+ unsigned getElementsPerByte() const { return 1; }
+
+ /// Get the preferred alignment in bytes for this 8-bit floating point type.
+ std::optional<unsigned> getPreferredAlignmentBytes() const { return std::nullopt; }
+ }];
}
//===----------------------------------------------------------------------===//
@@ -254,7 +315,8 @@ def Builtin_Float8E3M4 : Builtin_FloatType<"Float8E3M4", "f8E3M4"> {
// Float4E2M1FNType
//===----------------------------------------------------------------------===//
-def Builtin_Float4E2M1FN : Builtin_FloatType<"Float4E2M1FN", "f4E2M1FN"> {
+def Builtin_Float4E2M1FN : Builtin_FloatType<"Float4E2M1FN", "f4E2M1FN",
+ [QuantStorageTypeInterface]> {
let summary = "4-bit floating point with 2-bit exponent and 1-bit mantissa";
let description = [{
An 4-bit floating point type with 1 sign bit, 2 bits exponent and 1 bit
@@ -270,6 +332,34 @@ def Builtin_Float4E2M1FN : Builtin_FloatType<"Float4E2M1FN", "f4E2M1FN"> {
Open Compute Project (OCP) microscaling formats (MX) specification:
https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
}];
+
+ let extraClassDeclaration = [{
+ /// QuantStorageTypeInterface method implementations
+ /// Returns true if the the storage type should default to signed
+ /// when used in quantization.
+ bool shouldDefaultToSigned() const { return true; }
+ /// Get the bit width of this 4-bit floating point type.
+ unsigned getStorageWidth() const { return 4; }
+
+ /// Get default min, max value for this 4-bit floating point type.
+ int64_t getDefaultMaximum(bool isSigned) const { return 6; }
+ int64_t getDefaultMinimum(bool isSigned) const { return -getDefaultMaximum(isSigned); }
+
+ /// Get the storage type as a string.
+ std::string getStorageTypeName(bool isSigned) const { return "f4E2M1FN"; }
+
+ /// Check if this 4-bit floating point type uses packed representation.
+ bool isPacked() const { return true; }
+
+ /// Get the logical bit width per value for this 4-bit floating point type.
+ unsigned getLogicalBitWidth() const { return 4; }
+
+ /// Returns how many values of this type fit in one byte
+ unsigned getElementsPerByte() const { return 2; }
+
+ /// Get the preferred alignment in bytes for this 4-bit floating point type.
+ std::optional<unsigned> getPreferredAlignmentBytes() const { return std::nullopt; }
+ }];
}
//===----------------------------------------------------------------------===//
@@ -345,7 +435,7 @@ def Builtin_Float8E8M0FNU : Builtin_FloatType<"Float8E8M0FNU", "f8E8M0FNU"> {
//===----------------------------------------------------------------------===//
def Builtin_BFloat16 : Builtin_CachedFloatType<"BFloat16", "bf16",
- /*declaredInterfaceMethods=*/["scaleElementBitwidth"]> {
+ /*traits=*/[], /*declaredInterfaceMethods=*/["scaleElementBitwidth"]> {
let summary = "bfloat16 floating-point type";
}
@@ -354,7 +444,7 @@ def Builtin_BFloat16 : Builtin_CachedFloatType<"BFloat16", "bf16",
//===----------------------------------------------------------------------===//
def Builtin_Float16 : Builtin_CachedFloatType<"Float16", "f16",
- /*declaredInterfaceMethods=*/["scaleElementBitwidth"]> {
+ /*traits=*/[], /*declaredInterfaceMethods=*/["scaleElementBitwidth"]> {
let summary = "16-bit floating-point type";
}
@@ -371,7 +461,7 @@ def Builtin_FloatTF32 : Builtin_CachedFloatType<"FloatTF32", "tf32"> {
//===----------------------------------------------------------------------===//
def Builtin_Float32 : Builtin_CachedFloatType<"Float32", "f32",
- /*declaredInterfaceMethods=*/["scaleElementBitwidth"]> {
+ /*traits=*/[], /*declaredInterfaceMethods=*/["scaleElementBitwidth"]> {
let summary = "32-bit floating-point type";
}
@@ -501,7 +591,7 @@ def Builtin_Index : Builtin_Type<"Index", "index",
//===----------------------------------------------------------------------===//
def Builtin_Integer : Builtin_Type<"Integer", "integer",
- [VectorElementTypeInterface]> {
+ [VectorElementTypeInterface, QuantStorageTypeInterface]> {
let summary = "Integer type with arbitrary precision up to a fixed limit";
let description = [{
Syntax:
@@ -558,6 +648,44 @@ def Builtin_Integer : Builtin_Type<"Integer", "integer",
/// Integer representation maximal bitwidth.
/// Note: This is aligned with the maximum width of llvm::IntegerType.
static constexpr unsigned kMaxWidth = (1 << 24) - 1;
+
+ /// QuantStorageTypeInterface method implementations
+ /// Returns true if the storage type should default to signed when used
+ /// in quantization. Returns true for signed or signless integer types.
+ bool shouldDefaultToSigned() const { return !isUnsigned(); }
+ /// Get the bit width of this integer type.
+ unsigned getStorageWidth() const { return getWidth(); }
+
+ /// Get default min, max value for this integer type.
+ int64_t getDefaultMaximum(bool isSigned) const {
+ if (isSigned) {
+ return llvm::maxIntN(getWidth());
+ }
+ return llvm::maxUIntN(getWidth());
+ }
+ int64_t getDefaultMinimum(bool isSigned) const {
+ if (isSigned) {
+ return llvm::minIntN(getWidth());
+ }
+ return 0;
+ }
+
+ /// Get the storage type as a string.
+ std::string getStorageTypeName(bool isSigned) const {
+ return (isSigned ? "i" : "u") + std::to_string(getWidth());
+ }
+
+ /// Check if this integer type uses packed representation.
+ bool isPacked() const { return getStorageWidth() < 8; }
+
+ /// Get the logical bit width per value for this integer type.
+ unsigned getLogicalBitWidth() const { return getWidth(); }
+
+ /// Returns how many values of this type fit in one byte
+ unsigned getElementsPerByte() const { return 8 / getLogicalBitWidth(); }
+
+ /// Get the preferred alignment in bytes for this integer type.
+ std::optional<unsigned> getPreferredAlignmentBytes() const { return std::nullopt; }
}];
}
diff --git a/mlir/include/mlir/IR/CMakeLists.txt b/mlir/include/mlir/IR/CMakeLists.txt
index 3d30d92ed6ec4..15518901b901a 100644
--- a/mlir/include/mlir/IR/CMakeLists.txt
+++ b/mlir/include/mlir/IR/CMakeLists.txt
@@ -4,6 +4,8 @@ mlir_tablegen(SymbolInterfacesAttrInterface.h.inc -gen-attr-interface-decls)
mlir_tablegen(SymbolInterfacesAttrInterface.cpp.inc -gen-attr-interface-defs)
add_mlir_interface(RegionKindInterface)
+add_mlir_type_interface(QuantStorageTypeInterface)
+
set(LLVM_TARGET_DEFINITIONS OpAsmInterface.td)
mlir_tablegen(OpAsmAttrInterface.h.inc -gen-attr-interface-decls)
mlir_tablegen(OpAsmAttrInterface.cpp.inc -gen-attr-interface-defs)
diff --git a/mlir/include/mlir/IR/QuantStorageTypeInterface.h b/mlir/include/mlir/IR/QuantStorageTypeInterface.h
new file mode 100644
index 0000000000000..4f5e102feea85
--- /dev/null
+++ b/mlir/include/mlir/IR/QuantStorageTypeInterface.h
@@ -0,0 +1,20 @@
+//===- QuantStorageTypeInterface.h - Quantzation Interfaces -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_QuantStorageTypeInterface_H
+#define MLIR_IR_QuantStorageTypeInterface_H
+
+#include "mlir/IR/Types.h"
+
+//===----------------------------------------------------------------------===//
+// Tablegen Interface Declarations
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/QuantStorageTypeInterface.h.inc"
+
+#endif // MLIR_IR_QuantStorageTypeInterface_H
diff --git a/mlir/include/mlir/IR/QuantStorageTypeInterface.td b/mlir/include/mlir/IR/QuantStorageTypeInterface.td
new file mode 100644
index 0000000000000..22dabc58fca54
--- /dev/null
+++ b/mlir/include/mlir/IR/QuantStorageTypeInterface.td
@@ -0,0 +1,70 @@
+#ifndef MLIR_IR_QUANTSTORAGETYPEINTERFACE
+#define MLIR_IR_QUANTSTORAGETYPEINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def QuantStorageTypeInterface : TypeInterface<"QuantStorageTypeInterface"> {
+ let description = [{
+ Interface for types that can be used as storage types in Quant dialect.
+ This interface provides methods to determine storage characteristics for quantization purposes,
+ including packing behavior, and alignment requirements.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ InterfaceMethod<[{
+ Whether the storage type should default to signed when used in quantization.
+ Returns true if the type defaults to signed (e.g., si8, i8 or float types),
+ false if it defaults to unsigned.
+ }],
+ "bool", "shouldDefaultToSigned", (ins)>,
+
+ InterfaceMethod<[{
+ Get the bit width of this type.
+ Returns the number of bits used to store values of this type.
+ }],
+ "unsigned", "getStorageWidth", (ins)>,
+
+ InterfaceMethod<[{
+ Get default minimum value for this type.
+ }],
+ "int64_t", "getDefaultMinimum", (ins "bool":$isSigned)>,
+
+ InterfaceMethod<[{
+ Get default maximum value for this type.
+ }],
+ "int64_t", "getDefaultMaximum", (ins "bool":$isSigned)>,
+
+ InterfaceMethod<[{
+ Get the storage type as a string.
+ }],
+ "std::string", "getStorageTypeName", (ins "bool":$isSigned)>,
+
+ InterfaceMethod<[{
+ Check if the storage type uses packed representation.
+ Returns true if multiple values are packed into one byte (e.g., sub-byte types),
+ false if value uses full byte.
+ }],
+ "bool", "isPacked", (ins)>,
+
+ InterfaceMethod<[{
+ Get the logical bit width per value.
+ For packed sub-byte types, this may
diff er from getStorageWidth().
+ }],
+ "unsigned", "getLogicalBitWidth", (ins)>,
+
+ InterfaceMethod<[{
+ Get the number of logical elements that fit in one byte.
+ For packed sub-byte types, this returns how many values can be stored per byte.
+ }],
+ "unsigned", "getElementsPerByte", (ins)>,
+
+ InterfaceMethod<[{
+ Returns the preferred alignment for this type, in bytes.
+ }],
+ "std::optional<unsigned>", "getPreferredAlignmentBytes", (ins)>
+ ];
+
+}
+
+#endif // MLIR_IR_QUANTSTORAGETYPEINTERFACE
diff --git a/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp b/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp
index b2227792f32ca..c5a36f7106ad3 100644
--- a/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp
+++ b/mlir/lib/Dialect/Quant/IR/QuantTypes.cpp
@@ -9,6 +9,7 @@
#include "mlir/Dialect/Quant/IR/QuantTypes.h"
#include "TypeDetail.h"
#include "mlir/Dialect/Quant/IR/Quant.h"
+#include "mlir/IR/QuantStorageTypeInterface.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
@@ -46,32 +47,28 @@ QuantizedType::verifyInvariants(function_ref<InFlightDiagnostic()> emitError,
unsigned flags, Type storageType,
Type expressedType, int64_t storageTypeMin,
int64_t storageTypeMax) {
- // Verify that the storage type is integral.
- // This restriction may be lifted at some point in favor of using bf16
- // or f16 as exact representations on hardware where that is advantageous.
- auto intStorageType = llvm::dyn_cast<IntegerType>(storageType);
- if (!intStorageType)
- return emitError() << "storage type must be integral";
- unsigned integralWidth = intStorageType.getWidth();
-
- // Verify storage width.
- if (integralWidth == 0 || integralWidth > MaxStorageBits)
- return emitError() << "illegal storage type size: " << integralWidth;
-
- // Verify storageTypeMin and storageTypeMax.
- bool isSigned =
- (flags & QuantizationFlags::Signed) == QuantizationFlags::Signed;
- int64_t defaultIntegerMin =
- getDefaultMinimumForInteger(isSigned, integralWidth);
- int64_t defaultIntegerMax =
- getDefaultMaximumForInteger(isSigned, integralWidth);
- if (storageTypeMax - storageTypeMin <= 0 ||
- storageTypeMin < defaultIntegerMin ||
- storageTypeMax > defaultIntegerMax) {
- return emitError() << "illegal storage min and storage max: ("
- << storageTypeMin << ":" << storageTypeMax << ")";
+ if (auto quantStorageTypeInterface =
+ llvm::dyn_cast<QuantStorageTypeInterface>(storageType)) {
+ unsigned integralWidth = quantStorageTypeInterface.getStorageWidth();
+
+ // Verify storage width.
+ if (integralWidth == 0 || integralWidth > MaxStorageBits)
+ return emitError() << "illegal storage type size: " << integralWidth;
+
+ bool isSigned = flags & QuantizationFlags::Signed;
+ int64_t defaultMin = quantStorageTypeInterface.getDefaultMinimum(isSigned);
+ int64_t defaultMax = quantStorageTypeInterface.getDefaultMaximum(isSigned);
+
+ if (storageTypeMax - storageTypeMin <= 0 || storageTypeMin < defaultMin ||
+ storageTypeMax > defaultMax) {
+ return emitError() << "illegal storage min and storage max: ("
+ << storageTypeMin << ":" << storageTypeMax << ")";
+ }
+
+ return success();
}
- return success();
+
+ return emitError() << "storage type must implement QuantStorageTypeInterface";
}
Type QuantizedType::getStorageType() const {
@@ -87,20 +84,22 @@ int64_t QuantizedType::getStorageTypeMax() const {
}
bool QuantizedType::hasStorageTypeBounds() const {
- unsigned int integralWidth = getStorageTypeIntegralWidth();
- bool isSignedInteger = isSigned();
- int64_t defaultIntegerMin =
- getDefaultMinimumForInteger(isSignedInteger, integralWidth);
- int64_t defaultIntegerMax =
- getDefaultMaximumForInteger(isSignedInteger, integralWidth);
- return defaultIntegerMin != getStorageTypeMin() ||
- defaultIntegerMax != getStorageTypeMax();
+ Type storageType = static_cast<ImplType *>(impl)->storageType;
+ auto quantStorageTypeInterface =
+ llvm::dyn_cast<QuantStorageTypeInterface>(storageType);
+
+ int64_t defaultMin = quantStorageTypeInterface.getDefaultMinimum(isSigned());
+ int64_t defaultMax = quantStorageTypeInterface.getDefaultMaximum(isSigned());
+
+ return defaultMin != getStorageTypeMin() || defaultMax != getStorageTypeMax();
}
unsigned QuantizedType::getStorageTypeIntegralWidth() const {
- // NOTE: If ever supporting non-integral storage types, some other scheme
- // for determining the width will be needed.
- return static_cast<ImplType *>(impl)->storageType.getIntOrFloatBitWidth();
+ Type storageType = static_cast<ImplType *>(impl)->storageType;
+ auto quantStorageTypeInterface =
+ llvm::dyn_cast<QuantStorageTypeInterface>(storageType);
+
+ return quantStorageTypeInterface.getStorageWidth();
}
Type QuantizedType::getExpressedType() const {
diff --git a/mlir/lib/Dialect/Quant/IR/TypeParser.cpp b/mlir/lib/Dialect/Quant/IR/TypeParser.cpp
index 7c3840abbf91c..1a42b90ac31e2 100644
--- a/mlir/lib/Dialect/Quant/IR/TypeParser.cpp
+++ b/mlir/lib/Dialect/Quant/IR/TypeParser.cpp
@@ -10,6 +10,7 @@
#include "mlir/Dialect/Quant/IR/QuantTypes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/QuantStorageTypeInterface.h"
#include "mlir/IR/Types.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/SmallVectorExtras.h"
@@ -17,9 +18,9 @@
using namespace mlir;
using namespace quant;
-static IntegerType parseStorageType(DialectAsmParser &parser, bool &isSigned) {
+static Type parseStorageType(DialectAsmParser &parser, bool &isSigned) {
auto typeLoc = parser.getCurrentLocation();
- IntegerType type;
+ Type type;
// Parse storage type (alpha_ident, integer_literal).
StringRef identifier;
@@ -28,20 +29,30 @@ static IntegerType parseStorageType(DialectAsmParser &parser, bool &isSigned) {
if (result.has_value()) {
if (!succeeded(*result))
return nullptr;
- isSigned = !type.isUnsigned();
- storageTypeWidth = type.getWidth();
- } else if (succeeded(parser.parseKeyword(&identifier))) {
- // Otherwise, this must be an unsigned integer (`u` integer-literal).
- if (!identifier.consume_front("u")) {
+
+ if (auto quantStorageTypeInterface =
+ llvm::dyn_cast<QuantStorageTypeInterface>(type)) {
+ // Returns true if the type defaults to signed (e.g., si8, i8 or float
+ // types), false if it defaults to unsigned.
+ isSigned = quantStorageTypeInterface.shouldDefaultToSigned();
+ storageTypeWidth = quantStorageTypeInterface.getStorageWidth();
+ } else {
parser.emitError(typeLoc, "illegal storage type prefix");
return nullptr;
}
- if (identifier.getAsInteger(10, storageTypeWidth)) {
- parser.emitError(typeLoc, "expected storage type width");
+ } else if (succeeded(parser.parseKeyword(&identifier))) {
+ // Otherwise, this must be an unsigned integer (`u` integer-literal)
+ if (identifier.consume_front("u")) {
+ if (identifier.getAsInteger(10, storageTypeWidth)) {
+ parser.emitError(typeLoc, "expected storage type width");
+ return nullptr;
+ }
+ isSigned = false;
+ type = parser.getBuilder().getIntegerType(storageTypeWidth);
+ } else {
+ parser.emitError(typeLoc, "illegal storage type prefix");
return nullptr;
}
- isSigned = false;
- type = parser.getBuilder().getIntegerType(storageTypeWidth);
} else {
return nullptr;
}
@@ -56,17 +67,18 @@ static IntegerType parseStorageType(DialectAsmParser &parser, bool &isSigned) {
return type;
}
-static ParseResult parseStorageRange(DialectAsmParser &parser,
- IntegerType storageType, bool isSigned,
- int64_t &storageTypeMin,
+static ParseResult parseStorageRange(DialectAsmParser &parser, Type storageType,
+ bool isSigned, int64_t &storageTypeMin,
int64_t &storageTypeMax) {
- int64_t defaultIntegerMin = QuantizedType::getDefaultMinimumForInteger(
- isSigned, storageType.getWidth());
- int64_t defaultIntegerMax = QuantizedType::getDefaultMaximumForInteger(
- isSigned, storageType.getWidth());
+ auto quantStorageTypeInterface =
+ llvm::dyn_cast<QuantStorageTypeInterface>(storageType);
+
+ int64_t defaultMin = quantStorageTypeInterface.getDefaultMinimum(isSigned);
+ int64_t defaultMax = quantStorageTypeInterface.getDefaultMaximum(isSigned);
+
if (failed(parser.parseOptionalLess())) {
- storageTypeMin = defaultIntegerMin;
- storageTypeMax = defaultIntegerMax;
+ storageTypeMin = defaultMin;
+ storageTypeMax = defaultMax;
return success();
}
@@ -76,11 +88,11 @@ static ParseResult parseStorageRange(DialectAsmParser &parser,
parser.getCurrentLocation(&maxLoc) ||
parser.parseInteger(storageTypeMax) || parser.parseGreater())
return failure();
- if (storageTypeMin < defaultIntegerMin) {
+ if (storageTypeMin < defaultMin) {
return parser.emitError(minLoc, "illegal storage type minimum: ")
<< storageTypeMin;
}
- if (storageTypeMax > defaultIntegerMax) {
+ if (storageTypeMax > defaultMax) {
return parser.emitError(maxLoc, "illegal storage type maximum: ")
<< storageTypeMax;
}
@@ -114,7 +126,7 @@ static FloatType parseExpressedTypeAndRange(DialectAsmParser &parser,
/// storage-type ::= (`i` | `u`) integer-literal
/// expressed-type-spec ::= `:` `f` integer-literal
static Type parseAnyType(DialectAsmParser &parser) {
- IntegerType storageType;
+ Type storageType;
FloatType expressedType;
unsigned typeFlags = 0;
int64_t storageTypeMin;
@@ -323,7 +335,7 @@ parseQuantParamListUntilRBrace(DialectAsmParser &parser, Type expressedType,
/// scale-zero-tensor (`,` scale-zero-tensor)*
/// `}`
static Type parseUniformType(DialectAsmParser &parser) {
- IntegerType storageType;
+ Type storageType;
FloatType expressedType;
unsigned typeFlags = 0;
int64_t storageTypeMin;
@@ -488,13 +500,10 @@ Type QuantDialect::parseType(DialectAsmParser &parser) const {
static void printStorageType(QuantizedType type, DialectAsmPrinter &out) {
// storage type
- unsigned storageWidth = type.getStorageTypeIntegralWidth();
- bool isSigned = type.isSigned();
- if (isSigned) {
- out << "i" << storageWidth;
- } else {
- out << "u" << storageWidth;
- }
+ auto quantStorageTypeInterface =
+ llvm::dyn_cast<QuantStorageTypeInterface>(type.getStorageType());
+
+ out << quantStorageTypeInterface.getStorageTypeName(type.isSigned());
// storageTypeMin and storageTypeMax if not default.
if (type.hasStorageTypeBounds()) {
diff --git a/mlir/lib/IR/CMakeLists.txt b/mlir/lib/IR/CMakeLists.txt
index 563c8c6285ef3..632b06e4378f2 100644
--- a/mlir/lib/IR/CMakeLists.txt
+++ b/mlir/lib/IR/CMakeLists.txt
@@ -31,6 +31,7 @@ add_mlir_library(MLIRIR
OperationSupport.cpp
PatternLoggingListener.cpp
PatternMatch.cpp
+ QuantStorageTypeInterface.cpp
Region.cpp
RegionKindInterface.cpp
Remarks.cpp
@@ -63,6 +64,7 @@ add_mlir_library(MLIRIR
MLIRCastInterfacesIncGen
MLIRDataLayoutInterfacesIncGen
MLIROpAsmInterfaceIncGen
+ MLIRQuantStorageTypeInterfaceIncGen
MLIRRegionKindInterfaceIncGen
MLIRSideEffectInterfacesIncGen
MLIRSymbolInterfacesIncGen
@@ -72,4 +74,3 @@ add_mlir_library(MLIRIR
LINK_LIBS PUBLIC
MLIRSupport
)
-
diff --git a/mlir/lib/IR/QuantStorageTypeInterface.cpp b/mlir/lib/IR/QuantStorageTypeInterface.cpp
new file mode 100644
index 0000000000000..e9c27187f3f33
--- /dev/null
+++ b/mlir/lib/IR/QuantStorageTypeInterface.cpp
@@ -0,0 +1,17 @@
+//===- QuantStorageTypeInterface.cpp --------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/QuantStorageTypeInterface.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+/// Tablegen Interface Definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/QuantStorageTypeInterface.cpp.inc"
diff --git a/mlir/test/Dialect/Quant/parse-uniform-invalid.mlir b/mlir/test/Dialect/Quant/parse-uniform-invalid.mlir
index 3b358443e43f2..6dbc86263bd71 100644
--- a/mlir/test/Dialect/Quant/parse-uniform-invalid.mlir
+++ b/mlir/test/Dialect/Quant/parse-uniform-invalid.mlir
@@ -235,3 +235,32 @@
!qalias = !quant.uniform<i8:f16:{0:1,1:2},
{{6.6e4:120,9.987200e-01:127}, {2.000000e+02:256,9.987200e-01}}>
+// -----
+// Illegal storage min/max: max > defaultMax
+// expected-error at +1 {{illegal storage type maximum: 60000}}
+!qalias = !quant.uniform<f8E5M2<-57344:60000>:f32, 0.99872:127>
+
+// -----
+// Illegal storage min/max: min < defaultMin
+// expected-error at +1 {{illegal storage type minimum: -60000}}
+!qalias = !quant.uniform<f8E5M2<-60000:57344>:f32, 0.99872:127>
+
+// -----
+// Illegal storage min/max: max > defaultMax
+// expected-error at +1 {{illegal storage type maximum: 500}}
+!qalias = !quant.uniform<f8E4M3FN<-448:500>:f32, 0.99872:127>
+
+// -----
+// Illegal storage min/max: min < defaultMin
+// expected-error at +1 {{illegal storage type minimum: -500}}
+!qalias = !quant.uniform<f8E4M3FN<-500:448>:f32, 0.99872:127>
+
+// -----
+// Illegal storage min/max: max > defaultMax
+// expected-error at +1 {{illegal storage type maximum: 10}}
+!qalias = !quant.uniform<f4E2M1FN<-6:10>:f32, 0.99872:127>
+
+// -----
+// Illegal storage min/max: min < defaultMin
+// expected-error at +1 {{illegal storage type minimum: -10}}
+!qalias = !quant.uniform<f4E2M1FN<-10:6>:f32, 0.99872:127>
diff --git a/mlir/test/Dialect/Quant/parse-uniform.mlir b/mlir/test/Dialect/Quant/parse-uniform.mlir
index 80a6621ed6979..1431403aece41 100644
--- a/mlir/test/Dialect/Quant/parse-uniform.mlir
+++ b/mlir/test/Dialect/Quant/parse-uniform.mlir
@@ -172,3 +172,58 @@ func.func @parse() -> !qalias {
%0 = "foo"() : () -> !qalias
return %0 : !qalias
}
+
+// -----
+// Default min/max value optimization for f8E5M2.
+// CHECK: !quant.uniform<f8E5M2:f32, 9.987200e-01:127>
+!qalias = !quant.uniform<f8E5M2<-57344:57344>:f32, 0.99872:127 >
+func.func @parse() -> !qalias {
+ %0 = "foo"() : () -> !qalias
+ return %0 : !qalias
+}
+
+// -----
+// Storage type: f8E5M2
+// CHECK: !quant.uniform<f8E5M2:f32, 2.000000e+02>
+!qalias = !quant.uniform<f8E5M2:f32, 2.0e+2>
+func.func @parse() -> !qalias {
+ %0 = "foo"() : () -> !qalias
+ return %0 : !qalias
+}
+
+// -----
+// Default min/max value optimization for f8E4M3FN.
+// CHECK: !quant.uniform<f8E4M3FN:f32, 9.987200e-01:127>
+!qalias = !quant.uniform<f8E4M3FN<-448:448>:f32, 0.99872:127 >
+func.func @parse() -> !qalias {
+ %0 = "foo"() : () -> !qalias
+ return %0 : !qalias
+}
+
+// -----
+// Storage type: f8E4M3FN
+// CHECK: !quant.uniform<f8E4M3FN:f32, 2.000000e+02>
+!qalias = !quant.uniform<f8E4M3FN:f32, 2.0e+2>
+func.func @parse() -> !qalias {
+ %0 = "foo"() : () -> !qalias
+ return %0 : !qalias
+}
+
+
+// -----
+// Default min/max value optimization for f4E2M1FN.
+// CHECK: !quant.uniform<f4E2M1FN:f32, 9.987200e-01:127>
+!qalias = !quant.uniform<f4E2M1FN<-6:6>:f32, 0.99872:127 >
+func.func @parse() -> !qalias {
+ %0 = "foo"() : () -> !qalias
+ return %0 : !qalias
+}
+
+// -----
+// Storage type: f4E2M1FN
+// CHECK: !quant.uniform<f4E2M1FN:f32, 2.000000e+02>
+!qalias = !quant.uniform<f4E2M1FN:f32, 2.0e+2>
+func.func @parse() -> !qalias {
+ %0 = "foo"() : () -> !qalias
+ return %0 : !qalias
+}
More information about the Mlir-commits
mailing list