This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch codex/remove-uint1-getdatasize-special-case in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
commit da7cbb186e895cf3d6b988a5b07c66e80543aa03 Author: tqchen <[email protected]> AuthorDate: Tue Jul 21 08:07:13 2026 +0000 [FIX] Pack uint1 tensors by data type width --- 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();
