gemini-code-assist[bot] commented on code in PR #679:
URL: https://github.com/apache/tvm-ffi/pull/679#discussion_r3620000145
##########
src/ffi/tensor.cc:
##########
@@ -25,9 +25,70 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
namespace tvm {
namespace ffi {
+namespace {
+
+size_t GetValidatedDataSize(const Shape& shape, DLDataType dtype) {
+ int16_t lanes = static_cast<int16_t>(dtype.lanes);
+ if (dtype.bits == 0 || lanes <= 0) {
+ TVM_FFI_THROW(ValueError) << "Cannot determine byte size for dtype "
+ << DLDataTypeToString(dtype);
+ }
+
+ size_t numel = 1;
+ for (int64_t dimension : shape) {
+ if (dimension < 0) {
+ TVM_FFI_THROW(ValueError) << "Tensor shape cannot contain negative
dimensions";
+ }
+ if (dimension != 0 &&
+ numel > std::numeric_limits<size_t>::max() /
static_cast<size_t>(dimension)) {
+ TVM_FFI_THROW(OverflowError) << "Tensor shape is too large";
+ }
+ numel *= static_cast<size_t>(dimension);
+ }
Review Comment:
On 32-bit systems, `size_t` is 32-bit. If `dimension` is a multiple of
`2^32` (e.g., `4294967296LL`), `static_cast<size_t>(dimension)` truncates to
`0`. Since `dimension != 0` is true, the expression `numel >
std::numeric_limits<size_t>::max() / static_cast<size_t>(dimension)` results in
a division by zero, causing a crash.
To prevent this, we should validate that `dimension` fits within the limits
of `size_t` before casting.
```c
size_t numel = 1;
for (int64_t dimension : shape) {
if (dimension < 0) {
TVM_FFI_THROW(ValueError) << "Tensor shape cannot contain negative
dimensions";
}
if (static_cast<uint64_t>(dimension) >
std::numeric_limits<size_t>::max()) {
TVM_FFI_THROW(OverflowError) << "Tensor shape dimension is too large";
}
if (dimension != 0 &&
numel > std::numeric_limits<size_t>::max() /
static_cast<size_t>(dimension)) {
TVM_FFI_THROW(OverflowError) << "Tensor shape is too large";
}
numel *= static_cast<size_t>(dimension);
}
```
##########
src/ffi/tensor.cc:
##########
@@ -25,9 +25,70 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
namespace tvm {
namespace ffi {
+namespace {
+
+size_t GetValidatedDataSize(const Shape& shape, DLDataType dtype) {
+ int16_t lanes = static_cast<int16_t>(dtype.lanes);
+ if (dtype.bits == 0 || lanes <= 0) {
+ TVM_FFI_THROW(ValueError) << "Cannot determine byte size for dtype "
+ << DLDataTypeToString(dtype);
+ }
+
+ size_t numel = 1;
+ for (int64_t dimension : shape) {
+ if (dimension < 0) {
+ TVM_FFI_THROW(ValueError) << "Tensor shape cannot contain negative
dimensions";
+ }
+ if (dimension != 0 &&
+ numel > std::numeric_limits<size_t>::max() /
static_cast<size_t>(dimension)) {
+ TVM_FFI_THROW(OverflowError) << "Tensor shape is too large";
+ }
+ numel *= static_cast<size_t>(dimension);
+ }
+
+ if (dtype.code == kDLUInt && dtype.bits == 1 && dtype.lanes == 1) {
+ return numel;
+ }
+ size_t bits_per_element = static_cast<size_t>(dtype.bits) *
static_cast<size_t>(lanes);
+ if (numel > (std::numeric_limits<size_t>::max() - 7) / bits_per_element) {
+ TVM_FFI_THROW(OverflowError) << "Tensor byte size is too large";
+ }
+ return (numel * bits_per_element + 7) / 8;
Review Comment:
`GetValidatedDataSize` has a special case for `uint1` (bits=1, lanes=1)
where it returns `numel` (assuming 1 byte per element). However, the standard
`GetDataSize` function in TVM (and DLPack) calculates the size as `(numel *
bits * lanes + 7) / 8`, which for `uint1` returns `(numel + 7) / 8` (packed 1
bit per element). This inconsistency means that a `uint1` tensor created via
`from_bytes` will expect `numel` bytes, but calling `tobytes()` on it will
return only `(numel + 7) / 8` bytes, leading to data truncation or size
mismatch errors.
Removing this special case ensures that `uint1` is consistently treated as
packed 1-bit per element, matching `GetDataSize`.
```c
size_t bits_per_element = static_cast<size_t>(dtype.bits) *
static_cast<size_t>(lanes);
if (numel > (std::numeric_limits<size_t>::max() - 7) / bits_per_element) {
TVM_FFI_THROW(OverflowError) << "Tensor byte size is too large";
}
return (numel * bits_per_element + 7) / 8;
```
##########
src/ffi/tensor.cc:
##########
@@ -25,9 +25,70 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
namespace tvm {
namespace ffi {
+namespace {
+
+size_t GetValidatedDataSize(const Shape& shape, DLDataType dtype) {
+ int16_t lanes = static_cast<int16_t>(dtype.lanes);
+ if (dtype.bits == 0 || lanes <= 0) {
+ TVM_FFI_THROW(ValueError) << "Cannot determine byte size for dtype "
+ << DLDataTypeToString(dtype);
+ }
+
+ size_t numel = 1;
+ for (int64_t dimension : shape) {
+ if (dimension < 0) {
+ TVM_FFI_THROW(ValueError) << "Tensor shape cannot contain negative
dimensions";
+ }
+ if (dimension != 0 &&
+ numel > std::numeric_limits<size_t>::max() /
static_cast<size_t>(dimension)) {
+ TVM_FFI_THROW(OverflowError) << "Tensor shape is too large";
+ }
+ numel *= static_cast<size_t>(dimension);
+ }
+
+ if (dtype.code == kDLUInt && dtype.bits == 1 && dtype.lanes == 1) {
+ return numel;
+ }
+ size_t bits_per_element = static_cast<size_t>(dtype.bits) *
static_cast<size_t>(lanes);
+ if (numel > (std::numeric_limits<size_t>::max() - 7) / bits_per_element) {
+ TVM_FFI_THROW(OverflowError) << "Tensor byte size is too large";
+ }
+ return (numel * bits_per_element + 7) / 8;
+}
+
+class CPUBytesNDAlloc {
+ public:
+ void AllocData(DLTensor* tensor, const Bytes& source, size_t data_size) {
+ uint8_t* data = new uint8_t[std::max<size_t>(data_size, 1)];
Review Comment:
Using standard `new uint8_t[]` does not guarantee 64-byte alignment (or even
32-byte alignment), which is often required or highly recommended for CPU
tensors in TVM to avoid performance degradation or crashes when executing
vectorized (SIMD) instructions. Consider using an aligned allocation mechanism.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
