gemini-code-assist[bot] commented on code in PR #335:
URL: https://github.com/apache/tvm-ffi/pull/335#discussion_r2612664094
##########
include/tvm/ffi/container/tensor.h:
##########
@@ -417,6 +462,39 @@ class Tensor : public ObjectRef {
num_extra_i64_at_tail, alloc, shape, dtype, device,
std::forward<ExtraArgs>(extra_args)...));
}
+
+ /*!
+ * \brief A variant of FromNDAlloc that allows explicit passing a strides.
+ *
+ * \note This function needs to ensure that strides are well-defined
+ * with respect to the allocated compact shape.
+ *
+ * \param alloc The NDAllocator.
+ * \param shape The shape of the Tensor.
+ * \param strides The strides of the Tensor.
+ * \param dtype The data type of the Tensor.
+ * \param device The device of the Tensor.
+ * \param extra_args Extra arguments to be forwarded to TNDAlloc.
+ * \return The created Tensor.
+ * \tparam TNDAlloc The type of the NDAllocator, impelments Alloc and Free.
+ * \tparam ExtraArgs Extra arguments to be passed to Alloc.
+ */
+ template <typename TNDAlloc, typename... ExtraArgs>
+ static Tensor FromNDAllocStrided(TNDAlloc alloc, ffi::ShapeView shape,
ffi::ShapeView strides,
+ DLDataType dtype, DLDevice device,
ExtraArgs&&... extra_args) {
+ // inplace alloc shape and strides after data structure (as a result why
multiply 2)
Review Comment:
The function doesn't check if `shape` and `strides` have the same size. If
`strides.size() < shape.size()`, it will lead to a buffer over-read when
copying the strides, as `prototype.ndim` (which is `shape.size()`) is used for
the copy length. Please add a check to ensure their sizes are equal.
```c
TVM_FFI_ICHECK(shape.size() == strides.size())
<< "ValueError: shape and strides must have the same size.";
// inplace alloc shape and strides after data structure (as a result why
multiply 2)
```
##########
include/tvm/ffi/container/tensor.h:
##########
@@ -348,6 +361,38 @@ class Tensor : public ObjectRef {
* \return True if the Tensor data is aligned to the given alignment, false
otherwise.
*/
bool IsAligned(size_t alignment) const { return tvm::ffi::IsAligned(*get(),
alignment); }
+
+ /*!
+ * \brief Create a new Tensor as a strided view of the current Tensor.
+ * \param shape The shape of the new Tensor.
+ * \param strides The strides of the new Tensor.
+ * \param element_offset The element offset of the new Tensor in the unit of
dtype elements.
+ * \return The new Tensor.
+ * \note element_offset is in the unit of dtype elements not bytes.
+ */
+ Tensor as_strided(ShapeView shape, ShapeView strides,
+ std::optional<int64_t> element_offset = std::nullopt)
const {
+ DLTensor prototype;
+ prototype = *static_cast<const DLTensor*>(get());
+ prototype.shape = const_cast<int64_t*>(shape.data());
+ prototype.ndim = static_cast<int>(shape.size());
+ prototype.strides = const_cast<int64_t*>(strides.data());
+ prototype.byte_offset +=
+ GetDataSize(static_cast<size_t>(element_offset.value_or(0)),
prototype.dtype);
Review Comment:
The `element_offset` is an `int64_t` and can be negative. Casting a negative
value to `size_t` results in a large positive number, which can lead to
incorrect behavior or memory access errors. It's safer to add a check to ensure
`element_offset` is non-negative.
```c
const int64_t offset = element_offset.value_or(0);
TVM_FFI_ICHECK(offset >= 0) << "ValueError: element_offset must be
non-negative, but got " << offset;
prototype.byte_offset += GetDataSize(static_cast<size_t>(offset),
prototype.dtype);
```
--
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]
