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 d6d5e20f [FIX] Preserve bit widths when formatting custom dtypes (#913)
d6d5e20f is described below
commit d6d5e20f20d2d1c0ec68798302f5a3a8751b6a2b
Author: Lei Wang <[email protected]>
AuthorDate: Wed Nov 4 20:11:12 2026 +0800
[FIX] Preserve bit widths when formatting custom dtypes (#913)
Custom dtype bit widths are lost when formatting native dtype values,
which can change their types after a JSON round trip.
`DLDataTypeToString_` only appends `dtype.bits` for type codes below
`kDLFloat8_e3m4`, excluding custom type codes at or above
`kDLExtCustomBegin`.
For example, TileLang uses `custom[float4_e2m1_unpacked]8` to describe
unpacked FP4 shared-memory storage. The formatter drops the trailing
`8`, and the parser then uses its default width of 32 bits:
```text
Before:
(code=131, bits=8, lanes=1)
-> "custom[float4_e2m1_unpacked]"
-> (code=131, bits=32, lanes=1)
After:
(code=131, bits=8, lanes=1)
-> "custom[float4_e2m1_unpacked]8"
-> (code=131, bits=8, lanes=1)
```
Include custom type codes in the condition that emits the bit-width
suffix, matching the existing parser's `custom[name]<bits>` syntax.
---
src/ffi/dtype.cc | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/src/ffi/dtype.cc b/src/ffi/dtype.cc
index aa4bc01e..0780910a 100644
--- a/src/ffi/dtype.cc
+++ b/src/ffi/dtype.cc
@@ -178,7 +178,8 @@ inline std::string DLDataTypeToString_(DLDataType dtype) {
// NOLINT(*)
}
if (dtype.code == kDLOpaqueHandle) return os.str();
int16_t lanes = static_cast<int16_t>(dtype.lanes);
- if (dtype.code < kDLFloat8_e3m4 && (dtype.code != kDLBool || dtype.bits !=
8)) {
+ if ((dtype.code < kDLFloat8_e3m4 || dtype.code >= kDLExtCustomBegin) &&
+ (dtype.code != kDLBool || dtype.bits != 8)) {
os << static_cast<int>(dtype.bits);
}
if (lanes > 1) {
