(tvm-ffi) branch main updated: [FIX] Pack uint1 tensors by data type width (#680)

[yaxingcai]

This is an automated email from the ASF dual-hosted git repository.

cyx-6 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 d6c18a66 [FIX] Pack uint1 tensors by data type width (#680)
d6c18a66 is described below

commit d6c18a666a79d6261239e134861c749c168c8e6a
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Wed Jul 22 11:41:32 2026 +0800

    [FIX] Pack uint1 tensors by data type width (#680)
    
    Remove the legacy scalar `uint1` byte-per-element special case from
    `GetDataSize`, so all sub-byte integer types use the packed size
    calculation. Add focused coverage for packed `uint1` and byte-sized
    `bool8` tensors.
---
 include/tvm/ffi/container/tensor.h |  7 +------
 tests/cpp/test_tensor.cc           | 10 ++++++++++
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/include/tvm/ffi/container/tensor.h 
b/include/tvm/ffi/container/tensor.h
index b20ec022..f6660e0c 100644
--- a/include/tvm/ffi/container/tensor.h
+++ b/include/tvm/ffi/container/tensor.h
@@ -103,12 +103,7 @@ inline bool IsAligned(const DLTensor& arr, size_t 
alignment) {
  * \return the total number of bytes needed to store packed data
  */
 inline size_t GetDataSize(size_t numel, DLDataType dtype) {
-  // compatible handling sub-byte uint1(bool), which usually stored as uint8_t
-  // TODO(tqchen): revisit and switch to kDLBool
-  if (dtype.code == kDLUInt && dtype.bits == 1 && dtype.lanes == 1) {
-    return numel;
-  }
-  // for other sub-byte types, packing is preferred
+  // Sub-byte types are stored packed.
   // Use uint64_t to avoid overflow on 32-bit platforms (WASM) for large 
allocations.
   return static_cast<size_t>((static_cast<uint64_t>(numel) * dtype.bits * 
dtype.lanes + 7) / 8);
 }
diff --git a/tests/cpp/test_tensor.cc b/tests/cpp/test_tensor.cc
index dcb9b69a..b3877ffb 100644
--- a/tests/cpp/test_tensor.cc
+++ b/tests/cpp/test_tensor.cc
@@ -49,6 +49,16 @@ int TestEnvTensorAllocatorError(DLTensor* prototype, 
TVMFFIObjectHandle* out) {
   return -1;
 }
 
+TEST(Tensor, GetDataSize) {
+  DLDataType uint1 = DLDataType({kDLUInt, 1, 1});
+  EXPECT_EQ(GetDataSize(1, uint1), 1);
+  EXPECT_EQ(GetDataSize(8, uint1), 1);
+  EXPECT_EQ(GetDataSize(9, uint1), 2);
+
+  DLDataType bool8 = DLDataType({kDLBool, 8, 1});
+  EXPECT_EQ(GetDataSize(9, bool8), 9);
+}
+
 TEST(Tensor, Basic) {
   Tensor nd = Empty({1, 2, 3}, DLDataType({kDLFloat, 32, 1}), 
DLDevice({kDLCPU, 0}));
   Shape shape = nd.shape();