This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 68222aaa feat(dtype): expand parsing and introspection (#668)
68222aaa is described below
commit 68222aaa02129b4e090a28aabbece6a2e4724d38
Author: Junru Shao <[email protected]>
AuthorDate: Wed Jul 15 15:21:16 2026 -0700
feat(dtype): expand parsing and introspection (#668)
Summary:
- Add compact integer, floating-point, and low-precision dtype aliases
in the native parser.
- Add C++ dtype classification/abbreviation helpers and Python dtype
predicates/constants.
- Incorporate review feedback by caching Python classification sets,
removing dead returns, and force-inlining simple native predicates.
Architecture:
- Normalize compact dtype spellings before applying the canonical native
parsing rules.
- Keep reusable C++ classification helpers alongside dtype conversion
utilities.
- Express integer classification as direct code comparisons and
force-inline primitive predicates for downstream optimization.
- Store Python integer and floating-point classification codes once at
class scope.
Public Interfaces:
- Add `DTypeIsFloat`, `DTypeIsInt`, `DTypeIsBool`, and `DTypeAbbrev` to
the C++ API.
- Add Python dtype predicates, validity checking, and named float8,
float6, and float4 constants.
UI/UX:
- None.
Behavioral Changes:
- Accept compact `i`, `u`, `f`, `bf`, `fp8`, `fp6`, and `fp4` spellings,
including vector-lane suffixes.
- Reject partial low-precision floating-point variant matches
consistently.
- Avoid allocating type-code sets on every `is_integer` or `is_float`
property access.
- Expose simple force-inlined integer and Boolean checks to downstream
optimizers without changing results.
Docs:
- Add API docstrings and Doxygen comments for the new dtype helpers.
- No additional documentation is required for the review-only
performance cleanups.
Tests:
- `uv run --no-sync pytest -q tests/python/test_dtype.py` — 77 passed.
- Rebuilt `tvm_ffi_tests` and ran `--gtest_filter=*DType*:*DataType*` —
9 passed.
- Ruff, Ruff formatting, clang-format, ty, and all other applicable
pre-commit hooks passed.
Untested Edge Cases:
- Alias parsing for externally registered custom dtype codes.
- GCC and MSVC were not run locally; the C++ changes were compiled with
AppleClang.
- The MSVC-specific `TVM_FFI_INLINE` expansion remains covered by
cross-platform CI.
Commits:
- `a8834f73` feat(dtype): expand parsing and introspection
- `8cef0767` perf(dtype): avoid repeated classification allocations
- `d7a30b4c` perf(dtype): inline primitive type predicates
---
include/tvm/ffi/dtype.h | 146 ++++++++++++++++++++++++++++++++++++++++++
python/tvm_ffi/__init__.py | 6 ++
python/tvm_ffi/_dtype.py | 77 ++++++++++++++++++++++
src/ffi/dtype.cc | 155 +++++++++++++++++++++++++++++++++------------
tests/cpp/test_dtype.cc | 35 ++++++++++
tests/python/test_dtype.py | 92 +++++++++++++++++++++++++++
6 files changed, 471 insertions(+), 40 deletions(-)
diff --git a/include/tvm/ffi/dtype.h b/include/tvm/ffi/dtype.h
index 6aa2aa6c..f1713af1 100644
--- a/include/tvm/ffi/dtype.h
+++ b/include/tvm/ffi/dtype.h
@@ -138,6 +138,152 @@ inline String DLDataTypeToString(DLDataType dtype) {
return TypeTraits<String>::MoveFromAnyAfterCheck(&out);
}
+/*!
+ * \brief Check whether a DLDataType is a floating-point type.
+ * \param dtype The DLDataType to check.
+ * \return True if dtype is a floating-point type, false otherwise.
+ */
+inline bool DTypeIsFloat(DLDataType dtype) {
+ switch (static_cast<int>(dtype.code)) {
+ case kDLFloat:
+ case kDLBfloat:
+ case kDLFloat8_e3m4:
+ case kDLFloat8_e4m3:
+ case kDLFloat8_e4m3b11fnuz:
+ case kDLFloat8_e4m3fn:
+ case kDLFloat8_e4m3fnuz:
+ case kDLFloat8_e5m2:
+ case kDLFloat8_e5m2fnuz:
+ case kDLFloat8_e8m0fnu:
+ case kDLFloat6_e2m3fn:
+ case kDLFloat6_e3m2fn:
+ case kDLFloat4_e2m1fn:
+ return true;
+ default:
+ return false;
+ }
+}
+
+/*!
+ * \brief Check whether a DLDataType is an integer type.
+ * \param dtype The DLDataType to check.
+ * \return True if dtype is a signed or unsigned integer type, false otherwise.
+ */
+TVM_FFI_INLINE bool DTypeIsInt(DLDataType dtype) {
+ return dtype.code == kDLInt || dtype.code == kDLUInt;
+}
+
+/*!
+ * \brief Check whether a DLDataType is a boolean type.
+ * \param dtype The DLDataType to check.
+ * \return True if dtype is a boolean type, false otherwise.
+ */
+TVM_FFI_INLINE bool DTypeIsBool(DLDataType dtype) { return dtype.code ==
kDLBool; }
+
+/*!
+ * \brief Convert a DLDataType to a compact abbreviation for text formats.
+ * \param dtype The DLDataType to convert.
+ * \return The compact string, such as ``i32``, ``f16``, or ``bf16``.
+ */
+inline std::string DTypeAbbrev(DLDataType dtype) {
+ if (dtype.bits == 8 && dtype.lanes == 1 && dtype.code == kDLBool) {
+ return "bool";
+ }
+ if (dtype.code == kDLOpaqueHandle && dtype.lanes == 0 && dtype.bits == 0) {
+ return "void";
+ }
+
+ std::string result;
+ switch (static_cast<int>(dtype.code)) {
+ case kDLInt: {
+ result = "i";
+ break;
+ }
+ case kDLUInt: {
+ result = "u";
+ break;
+ }
+ case kDLFloat: {
+ result = "f";
+ break;
+ }
+ case kDLOpaqueHandle: {
+ return "handle";
+ }
+ case kDLBfloat: {
+ result = "bf";
+ break;
+ }
+ case kDLBool: {
+ result = "bool";
+ break;
+ }
+ case kDLFloat8_e3m4: {
+ result = "f8_e3m4";
+ break;
+ }
+ case kDLFloat8_e4m3: {
+ result = "f8_e4m3";
+ break;
+ }
+ case kDLFloat8_e4m3b11fnuz: {
+ result = "f8_e4m3b11fnuz";
+ break;
+ }
+ case kDLFloat8_e4m3fn: {
+ result = "f8_e4m3fn";
+ break;
+ }
+ case kDLFloat8_e4m3fnuz: {
+ result = "f8_e4m3fnuz";
+ break;
+ }
+ case kDLFloat8_e5m2: {
+ result = "f8_e5m2";
+ break;
+ }
+ case kDLFloat8_e5m2fnuz: {
+ result = "f8_e5m2fnuz";
+ break;
+ }
+ case kDLFloat8_e8m0fnu: {
+ result = "f8_e8m0fnu";
+ break;
+ }
+ case kDLFloat6_e2m3fn: {
+ result = "f6_e2m3fn";
+ break;
+ }
+ case kDLFloat6_e3m2fn: {
+ result = "f6_e3m2fn";
+ break;
+ }
+ case kDLFloat4_e2m1fn: {
+ result = "f4_e2m1fn";
+ break;
+ }
+ default: {
+ if (static_cast<int>(dtype.code) >=
static_cast<int>(DLExtDataTypeCode::kDLExtCustomBegin)) {
+ String full_name = DLDataTypeToString(dtype);
+ return std::string(full_name.data(), full_name.size());
+ }
+ TVM_FFI_THROW(ValueError) << "DLDataType contains unknown type_code="
+ << static_cast<int>(dtype.code);
+ }
+ }
+
+ int16_t lanes = static_cast<int16_t>(dtype.lanes);
+ if (dtype.code < kDLFloat8_e3m4 && (dtype.code != kDLBool || dtype.bits !=
8)) {
+ result += std::to_string(static_cast<int>(dtype.bits));
+ }
+ if (lanes > 1) {
+ result += "x" + std::to_string(lanes);
+ } else if (lanes < -1) {
+ result += "xvscalex" + std::to_string(-lanes);
+ }
+ return result;
+}
+
// DLDataType
template <>
struct TypeTraits<DLDataType> : public TypeTraitsBase {
diff --git a/python/tvm_ffi/__init__.py b/python/tvm_ffi/__init__.py
index 9676c718..9ce96656 100644
--- a/python/tvm_ffi/__init__.py
+++ b/python/tvm_ffi/__init__.py
@@ -111,11 +111,17 @@ if TYPE_CHECKING or not _is_config_mode():
float32,
float16,
bfloat16,
+ float8_e3m4,
+ float8_e4m3,
+ float8_e4m3b11fnuz,
float8_e4m3fn,
float8_e4m3fnuz,
float8_e5m2,
float8_e5m2fnuz,
float8_e8m0fnu,
+ float6_e2m3fn,
+ float6_e3m2fn,
+ float4_e2m1fn,
float4_e2m1fnx2,
)
elif sys.platform.startswith("win32"):
diff --git a/python/tvm_ffi/_dtype.py b/python/tvm_ffi/_dtype.py
index dcc4d073..6477c935 100644
--- a/python/tvm_ffi/_dtype.py
+++ b/python/tvm_ffi/_dtype.py
@@ -19,6 +19,7 @@
# pylint: disable=invalid-name
from __future__ import annotations
+import builtins
from enum import IntEnum
from typing import Any, ClassVar
@@ -89,6 +90,22 @@ class dtype(str):
_tvm_ffi_dtype: core.DataType
_NUMPY_DTYPE_TO_STR: ClassVar[dict[Any, str]] = {}
+ _INTEGER_CODES: ClassVar[set[int]] = {DataTypeCode.INT, DataTypeCode.UINT}
+ _FLOAT_CODES: ClassVar[set[int]] = {
+ DataTypeCode.FLOAT,
+ DataTypeCode.BFLOAT,
+ DataTypeCode.Float8E3M4,
+ DataTypeCode.Float8E4M3,
+ DataTypeCode.Float8E4M3B11FNUZ,
+ DataTypeCode.Float8E4M3FN,
+ DataTypeCode.Float8E4M3FNUZ,
+ DataTypeCode.Float8E5M2,
+ DataTypeCode.Float8E5M2FNUZ,
+ DataTypeCode.Float8E8M0FNU,
+ DataTypeCode.Float6E2M3FN,
+ DataTypeCode.Float6E3M2FN,
+ DataTypeCode.Float4E2M1FN,
+ }
def __new__(cls, content: Any) -> dtype:
content = str(content)
@@ -288,6 +305,58 @@ class dtype(str):
"""
return self._tvm_ffi_dtype.lanes
+ @property
+ def is_bool(self) -> builtins.bool:
+ """Whether this dtype stores boolean values."""
+ return self.type_code == DataTypeCode.BOOL
+
+ @property
+ def is_integer(self) -> builtins.bool:
+ """Whether this dtype stores signed or unsigned integer values."""
+ return self.type_code in self._INTEGER_CODES
+
+ @property
+ def is_float(self) -> builtins.bool:
+ """Whether this dtype stores floating-point values."""
+ return self.type_code in self._FLOAT_CODES
+
+ @property
+ def is_handle(self) -> builtins.bool:
+ """Whether this dtype stores opaque handle values."""
+ return self.type_code == DataTypeCode.HANDLE
+
+ @staticmethod
+ def is_dtype(arg: str | dtype) -> builtins.bool:
+ """Check if the given string is a valid dtype.
+
+ Parameters
+ ----------
+ arg
+ The string to check.
+
+ Returns
+ -------
+ bool
+ Whether the input string is a valid dtype.
+
+ Examples
+ --------
+ .. code-block:: python
+
+ import tvm_ffi
+
+ assert tvm_ffi.dtype.is_dtype("float32")
+ assert not tvm_ffi.dtype.is_dtype("not_a_dtype")
+
+ """
+ if isinstance(arg, dtype):
+ return True
+ try:
+ core.DataType(str(arg))
+ return True
+ except ValueError:
+ return False
+
try:
# this helps to make numpy as optional
@@ -342,11 +411,19 @@ float32 = dtype("float32")
float16 = dtype("float16")
bfloat16 = dtype("bfloat16")
# float8 dtypes
+float8_e3m4 = dtype("float8_e3m4")
+float8_e4m3 = dtype("float8_e4m3")
+float8_e4m3b11fnuz = dtype("float8_e4m3b11fnuz")
float8_e4m3fn = dtype("float8_e4m3fn")
float8_e4m3fnuz = dtype("float8_e4m3fnuz")
float8_e5m2 = dtype("float8_e5m2")
float8_e5m2fnuz = dtype("float8_e5m2fnuz")
float8_e8m0fnu = dtype("float8_e8m0fnu")
+# float6 dtypes
+float6_e2m3fn = dtype("float6_e2m3fn")
+float6_e3m2fn = dtype("float6_e3m2fn")
+# float4 dtypes
+float4_e2m1fn = dtype("float4_e2m1fn")
# float4x2 dtypes
float4_e2m1fnx2 = dtype("float4_e2m1fnx2")
# alias for torch naming pattern for f4x2
diff --git a/src/ffi/dtype.cc b/src/ffi/dtype.cc
index 14cfa5dc..16760783 100644
--- a/src/ffi/dtype.cc
+++ b/src/ffi/dtype.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/string.h>
+#include <string>
#include <string_view>
namespace tvm {
@@ -202,6 +203,31 @@ inline DLDataType StringViewToDLDataType_(std::string_view
str) {
dtype.lanes = 0;
return dtype;
}
+
+ auto is_digit = [](char ch) { return ch >= '0' && ch <= '9'; };
+
+ std::string normalized_dtype;
+ auto normalize_prefix = [&](std::string_view canonical_prefix, size_t
alias_prefix_size) {
+ normalized_dtype.assign(canonical_prefix);
+ normalized_dtype.append(str.substr(alias_prefix_size));
+ str = normalized_dtype;
+ };
+ if (str.size() >= 2 && str[0] == 'i' && is_digit(str[1])) {
+ normalize_prefix("int", 1);
+ } else if (str.size() >= 2 && str[0] == 'u' && is_digit(str[1])) {
+ normalize_prefix("uint", 1);
+ } else if (str.size() >= 3 && str[0] == 'b' && str[1] == 'f' &&
is_digit(str[2])) {
+ normalize_prefix("bfloat", 2);
+ } else if (str.compare(0, 4, "fp8_") == 0) {
+ normalize_prefix("float8_", 4);
+ } else if (str.compare(0, 4, "fp6_") == 0) {
+ normalize_prefix("float6_", 4);
+ } else if (str.compare(0, 4, "fp4_") == 0) {
+ normalize_prefix("float4_", 4);
+ } else if (str.size() >= 2 && str[0] == 'f' && is_digit(str[1])) {
+ normalize_prefix("float", 1);
+ }
+
// set the default values;
dtype.bits = 32;
dtype.lanes = 1;
@@ -253,10 +279,14 @@ inline DLDataType
StringViewToDLDataType_(std::string_view str) {
return static_cast<uint16_t>(multiplier * lanes_val);
};
- auto parse_float = [&](const std::string_view& str, int offset, int code,
int bits) {
+ auto parse_float = [&](const std::string_view& str, int offset, int code,
int bits,
+ bool allow_underscore_lanes = false) {
dtype.code = static_cast<uint8_t>(code);
dtype.bits = static_cast<uint8_t>(bits);
scan = str.data() + offset;
+ if (allow_underscore_lanes && scan < str_end && *scan == '_') {
+ ++scan;
+ }
const char* endpt = scan;
dtype.lanes = parse_lanes(&endpt, str_end, str);
scan = endpt;
@@ -266,63 +296,104 @@ inline DLDataType
StringViewToDLDataType_(std::string_view str) {
return dtype;
};
+ auto match_float_variant = [&](int offset, std::string_view pattern,
+ bool allow_underscore_lanes = false) -> int {
+ size_t end = static_cast<size_t>(offset) + pattern.size();
+ if (str.size() < end || str.compare(offset, pattern.size(), pattern) != 0)
{
+ return -1;
+ }
+ if (end == str.size() || str[end] == 'x' || (allow_underscore_lanes &&
str[end] == '_')) {
+ return static_cast<int>(end);
+ }
+ return -1;
+ };
+
+ auto parse_float8 = [&](int offset) {
+ if (int end = match_float_variant(offset, "e4m3b11fnuz"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e4m3b11fnuz, 8);
+ } else if (int end = match_float_variant(offset, "e4m3fnuz"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e4m3fnuz, 8);
+ } else if (int end = match_float_variant(offset, "e4m3fn"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e4m3fn, 8);
+ } else if (int end = match_float_variant(offset, "e5m2fnuz"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e5m2fnuz, 8);
+ } else if (int end = match_float_variant(offset, "e8m0fnu"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e8m0fnu, 8);
+ } else if (int end = match_float_variant(offset, "e3m4"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e3m4, 8);
+ } else if (int end = match_float_variant(offset, "e4m3"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e4m3, 8);
+ } else if (int end = match_float_variant(offset, "e5m2"); end >= 0) {
+ return parse_float(str, end, kDLFloat8_e5m2, 8);
+ } else {
+ TVM_FFI_THROW(ValueError) << "unknown float8 type `" << str << '`';
+ TVM_FFI_UNREACHABLE();
+ }
+ };
+
+ auto parse_float6 = [&](int offset) {
+ if (int end = match_float_variant(offset, "e2m3fn"); end >= 0) {
+ return parse_float(str, end, kDLFloat6_e2m3fn, 6);
+ } else if (int end = match_float_variant(offset, "e3m2fn"); end >= 0) {
+ return parse_float(str, end, kDLFloat6_e3m2fn, 6);
+ } else {
+ TVM_FFI_THROW(ValueError) << "unknown float6 type `" << str << '`';
+ TVM_FFI_UNREACHABLE();
+ }
+ };
+
+ auto parse_float4 = [&](int offset) {
+ if (int end = match_float_variant(offset, "e2m1fn",
/*allow_underscore_lanes=*/true);
+ end >= 0) {
+ return parse_float(str, end, kDLFloat4_e2m1fn, 4,
/*allow_underscore_lanes=*/true);
+ } else {
+ TVM_FFI_THROW(ValueError) << "unknown float4 type `" << str << '`';
+ TVM_FFI_UNREACHABLE();
+ }
+ };
+
if (str.compare(0, 3, "int") == 0) {
dtype.code = kDLInt;
scan = str.data() + 3;
+ } else if (str.size() >= 2 && str[0] == 'i' && is_digit(str[1])) {
+ dtype.code = kDLInt;
+ scan = str.data() + 1;
} else if (str.compare(0, 4, "uint") == 0) {
dtype.code = kDLUInt;
scan = str.data() + 4;
+ } else if (str.size() >= 2 && str[0] == 'u' && is_digit(str[1])) {
+ dtype.code = kDLUInt;
+ scan = str.data() + 1;
} else if (str.compare(0, 4, "bool") == 0) {
dtype.code = kDLBool;
dtype.bits = 8;
scan = str.data() + 4;
+ } else if (str.compare(0, 4, "fp8_") == 0) {
+ return parse_float8(4);
+ } else if (str.compare(0, 3, "f8_") == 0) {
+ return parse_float8(3);
+ } else if (str.compare(0, 4, "fp6_") == 0) {
+ return parse_float6(4);
+ } else if (str.compare(0, 3, "f6_") == 0) {
+ return parse_float6(3);
+ } else if (str.compare(0, 4, "fp4_") == 0) {
+ return parse_float4(4);
+ } else if (str.compare(0, 3, "f4_") == 0) {
+ return parse_float4(3);
} else if (str.compare(0, 5, "float") == 0) {
if (str.compare(5, 2, "8_") == 0) {
- if (str.compare(7, 4, "e3m4") == 0) {
- return parse_float(str, 11, kDLFloat8_e3m4, 8);
- } else if (str.compare(7, 4, "e4m3") == 0) {
- if (str.compare(11, 7, "b11fnuz") == 0) {
- return parse_float(str, 18, kDLFloat8_e4m3b11fnuz, 8);
- } else if (str.compare(11, 2, "fn") == 0) {
- if (str.compare(13, 2, "uz") == 0) {
- return parse_float(str, 15, kDLFloat8_e4m3fnuz, 8);
- } else {
- return parse_float(str, 13, kDLFloat8_e4m3fn, 8);
- }
- } else {
- return parse_float(str, 11, kDLFloat8_e4m3, 8);
- }
- } else if (str.compare(7, 8, "e5m2fnuz") == 0) {
- return parse_float(str, 15, kDLFloat8_e5m2fnuz, 8);
- } else if (str.compare(7, 4, "e5m2") == 0) {
- return parse_float(str, 11, kDLFloat8_e5m2, 8);
- } else if (str.compare(7, 7, "e8m0fnu") == 0) {
- return parse_float(str, 14, kDLFloat8_e8m0fnu, 8);
- } else {
- TVM_FFI_THROW(ValueError) << "unknown float8 type `" << str << '`';
- TVM_FFI_UNREACHABLE();
- }
+ return parse_float8(7);
} else if (str.compare(5, 2, "6_") == 0) {
- if (str.compare(7, 6, "e2m3fn") == 0) {
- return parse_float(str, 13, kDLFloat6_e2m3fn, 6);
- } else if (str.compare(7, 6, "e3m2fn") == 0) {
- return parse_float(str, 13, kDLFloat6_e3m2fn, 6);
- } else {
- TVM_FFI_THROW(ValueError) << "unknown float6 type `" << str << '`';
- TVM_FFI_UNREACHABLE();
- }
+ return parse_float6(7);
} else if (str.compare(5, 2, "4_") == 0) {
- // kFloat4_e2m1fn
- if (str.compare(7, 6, "e2m1fn") == 0) {
- return parse_float(str, 13, kDLFloat4_e2m1fn, 4);
- } else {
- TVM_FFI_THROW(ValueError) << "unknown float4 type `" << str << '`';
- TVM_FFI_UNREACHABLE();
- }
+ return parse_float4(7);
} else {
dtype.code = kDLFloat;
scan = str.data() + 5;
}
+ } else if (str.size() >= 2 && str[0] == 'f' && is_digit(str[1])) {
+ dtype.code = kDLFloat;
+ scan = str.data() + 1;
} else if (str.compare(0, 6, "handle") == 0) {
dtype.code = kDLOpaqueHandle;
dtype.bits = 64; // handle uses 64 bit by default.
@@ -331,6 +402,10 @@ inline DLDataType StringViewToDLDataType_(std::string_view
str) {
dtype.code = kDLBfloat;
dtype.bits = 16;
scan = str.data() + 6;
+ } else if (str.size() >= 3 && str[0] == 'b' && str[1] == 'f' &&
is_digit(str[2])) {
+ dtype.code = kDLBfloat;
+ dtype.bits = 16;
+ scan = str.data() + 2;
} else if (str.compare(0, 6, "custom") == 0) {
dtype.code = static_cast<uint8_t>(details::ParseCustomDataTypeCode(str,
&scan));
} else {
diff --git a/tests/cpp/test_dtype.cc b/tests/cpp/test_dtype.cc
index 67b74d8e..a2f24e34 100644
--- a/tests/cpp/test_dtype.cc
+++ b/tests/cpp/test_dtype.cc
@@ -75,6 +75,41 @@ TEST(DType, StringConversionAllDLPackTypes) {
}
}
+TEST(DType, StringConversionAliases) {
+ std::vector<std::pair<std::string, std::string>> test_cases = {
+ {"i32", "int32"},
+ {"u16", "uint16"},
+ {"f32", "float32"},
+ {"bf16", "bfloat16"},
+ {"f8_e3m4", "float8_e3m4"},
+ {"f8_e3m4x2", "float8_e3m4x2"},
+ {"f8_e4m3", "float8_e4m3"},
+ {"f8_e4m3b11fnuz", "float8_e4m3b11fnuz"},
+ {"f8_e4m3fn", "float8_e4m3fn"},
+ {"f8_e4m3fnuz", "float8_e4m3fnuz"},
+ {"f8_e4m3fnuzx2", "float8_e4m3fnuzx2"},
+ {"f8_e5m2", "float8_e5m2"},
+ {"f8_e5m2fnuz", "float8_e5m2fnuz"},
+ {"f8_e8m0fnu", "float8_e8m0fnu"},
+ {"fp8_e4m3", "float8_e4m3"},
+ {"fp8_e5m2x4", "float8_e5m2x4"},
+ {"f6_e2m3fn", "float6_e2m3fn"},
+ {"fp6_e3m2fn", "float6_e3m2fn"},
+ {"f4_e2m1fn", "float4_e2m1fn"},
+ {"fp4_e2m1fn", "float4_e2m1fn"},
+ {"f4_e2m1fnx2", "float4_e2m1fnx2"},
+ {"f4_e2m1fn_x2", "float4_e2m1fnx2"},
+ {"fp4_e2m1fn_x2", "float4_e2m1fnx2"},
+ {"float4_e2m1fn_x2", "float4_e2m1fnx2"},
+ };
+
+ for (const auto& [alias, expected] : test_cases) {
+ DLDataType parsed = StringToDLDataType(String(alias));
+ EXPECT_EQ(DLDataTypeToString(parsed), expected) << alias;
+ EXPECT_EQ(parsed, StringToDLDataType(String(expected))) << alias;
+ }
+}
+
TEST(DataType, AnyConversion) {
AnyView view0;
EXPECT_EQ(view0.CopyToTVMFFIAny().type_index, TypeIndex::kTVMFFINone);
diff --git a/tests/python/test_dtype.py b/tests/python/test_dtype.py
index 5864123c..a01d2c7c 100644
--- a/tests/python/test_dtype.py
+++ b/tests/python/test_dtype.py
@@ -32,6 +32,64 @@ def test_dtype() -> None:
assert x.dtype == float32
[email protected](
+ "alias, expected",
+ [
+ ("i8", "int8"),
+ ("i16", "int16"),
+ ("i32", "int32"),
+ ("i64", "int64"),
+ ("u8", "uint8"),
+ ("u16", "uint16"),
+ ("u32", "uint32"),
+ ("u64", "uint64"),
+ ("f16", "float16"),
+ ("f32", "float32"),
+ ("f64", "float64"),
+ ("bf16", "bfloat16"),
+ ("f8_e3m4", "float8_e3m4"),
+ ("f8_e4m3", "float8_e4m3"),
+ ("f8_e4m3b11fnuz", "float8_e4m3b11fnuz"),
+ ("f8_e4m3fn", "float8_e4m3fn"),
+ ("f8_e4m3fnuz", "float8_e4m3fnuz"),
+ ("f8_e5m2", "float8_e5m2"),
+ ("f8_e5m2fnuz", "float8_e5m2fnuz"),
+ ("f8_e8m0fnu", "float8_e8m0fnu"),
+ ("fp8_e3m4", "float8_e3m4"),
+ ("fp8_e4m3", "float8_e4m3"),
+ ("fp8_e4m3b11fnuz", "float8_e4m3b11fnuz"),
+ ("fp8_e4m3fn", "float8_e4m3fn"),
+ ("fp8_e4m3fnuz", "float8_e4m3fnuz"),
+ ("fp8_e5m2", "float8_e5m2"),
+ ("fp8_e5m2fnuz", "float8_e5m2fnuz"),
+ ("fp8_e8m0fnu", "float8_e8m0fnu"),
+ ("f6_e2m3fn", "float6_e2m3fn"),
+ ("f6_e3m2fn", "float6_e3m2fn"),
+ ("fp6_e2m3fn", "float6_e2m3fn"),
+ ("fp6_e3m2fn", "float6_e3m2fn"),
+ ("f4_e2m1fn", "float4_e2m1fn"),
+ ("fp4_e2m1fn", "float4_e2m1fn"),
+ ("f32x4", "float32x4"),
+ ("f8_e4m3x4", "float8_e4m3x4"),
+ ("f8_e4m3fnx4", "float8_e4m3fnx4"),
+ ("fp8_e5m2x4", "float8_e5m2x4"),
+ ("f6_e2m3fnx4", "float6_e2m3fnx4"),
+ ("f4_e2m1fnx2", "float4_e2m1fnx2"),
+ ("f4_e2m1fn_x2", "float4_e2m1fnx2"),
+ ("fp4_e2m1fnx2", "float4_e2m1fnx2"),
+ ("fp4_e2m1fn_x2", "float4_e2m1fnx2"),
+ ("float4_e2m1fn_x2", "float4_e2m1fnx2"),
+ ],
+)
+def test_dtype_aliases(alias: str, expected: str) -> None:
+ dtype = tvm_ffi.dtype(alias)
+ expected_dtype = tvm_ffi.dtype(expected)
+ assert str(dtype) == alias
+ assert dtype.type_code == expected_dtype.type_code
+ assert dtype.bits == expected_dtype.bits
+ assert dtype.lanes == expected_dtype.lanes
+
+
@pytest.mark.parametrize(
"dtype_str, expected_size",
[
@@ -182,11 +240,17 @@ def test_builtin_dtype_conversion() -> None:
_check_dtype(tvm_ffi.float32, 2, 32, 1)
_check_dtype(tvm_ffi.float64, 2, 64, 1)
_check_dtype(tvm_ffi.bfloat16, 4, 16, 1)
+ _check_dtype(tvm_ffi.float8_e3m4, 7, 8, 1)
+ _check_dtype(tvm_ffi.float8_e4m3, 8, 8, 1)
+ _check_dtype(tvm_ffi.float8_e4m3b11fnuz, 9, 8, 1)
_check_dtype(tvm_ffi.float8_e4m3fn, 10, 8, 1)
_check_dtype(tvm_ffi.float8_e4m3fnuz, 11, 8, 1)
_check_dtype(tvm_ffi.float8_e5m2, 12, 8, 1)
_check_dtype(tvm_ffi.float8_e5m2fnuz, 13, 8, 1)
_check_dtype(tvm_ffi.float8_e8m0fnu, 14, 8, 1)
+ _check_dtype(tvm_ffi.float6_e2m3fn, 15, 6, 1)
+ _check_dtype(tvm_ffi.float6_e3m2fn, 16, 6, 1)
+ _check_dtype(tvm_ffi.float4_e2m1fn, 17, 4, 1)
_check_dtype(tvm_ffi.float4_e2m1fnx2, 17, 4, 2)
@@ -208,3 +272,31 @@ def test_dtype_bool() -> None:
assert dtype_with_lanes.bits == 8
assert dtype_with_lanes.lanes == 4
assert dtype_with_lanes == "boolx4"
+
+
[email protected](
+ "dtype_str, is_bool, is_integer, is_float, is_handle",
+ [
+ ("bool", True, False, False, False),
+ ("int32", False, True, False, False),
+ ("uint8", False, True, False, False),
+ ("float32", False, False, True, False),
+ ("bfloat16", False, False, True, False),
+ ("float8_e4m3fn", False, False, True, False),
+ ("float6_e2m3fn", False, False, True, False),
+ ("float4_e2m1fn", False, False, True, False),
+ ("handle", False, False, False, True),
+ ],
+)
+def test_dtype_kind_properties(
+ dtype_str: str,
+ is_bool: bool,
+ is_integer: bool,
+ is_float: bool,
+ is_handle: bool,
+) -> None:
+ dtype = tvm_ffi.dtype(dtype_str)
+ assert dtype.is_bool is is_bool
+ assert dtype.is_integer is is_integer
+ assert dtype.is_float is is_float
+ assert dtype.is_handle is is_handle
