gemini-code-assist[bot] commented on code in PR #81:
URL: https://github.com/apache/tvm-ffi/pull/81#discussion_r2395645394
##########
tests/python/test_load_inline.py:
##########
@@ -225,7 +225,7 @@ def test_load_inline_with_env_tensor_allocator() -> None:
namespace ffi = tvm::ffi;
ffi::Tensor return_add_one(ffi::Map<ffi::String,
ffi::Tuple<ffi::Tensor>> kwargs) {
- ffi::Tensor x = kwargs["x"].get<0>();
+ ffi::TensorView x = kwargs["x"].get<0>();
Review Comment:
This change introduces a critical use-after-free bug. `kwargs["x"].get<0>()`
returns a temporary `ffi::Tensor`. Initializing `ffi::TensorView x` from this
temporary results in a dangling view because the temporary is destroyed at the
end of the statement.
To fix this, `x` should be an `ffi::Tensor` to correctly manage the lifetime
of the tensor object.
```suggestion
ffi::Tensor x = kwargs["x"].get<0>();
```
##########
examples/quick_start/src/run_example.cc:
##########
@@ -36,7 +36,7 @@ int main() {
ffi::Module mod = ffi::Module::LoadFromFile("build/add_one_cpu.so");
// create an Tensor, alternatively, one can directly pass in a DLTensor*
- ffi::Tensor x = Empty({5}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU,
0}));
+ ffi::TensorView x = Empty({5}, DLDataType({kDLFloat, 32, 1}),
DLDevice({kDLCPU, 0}));
Review Comment:
This change introduces a critical use-after-free bug. `Empty()` returns a
temporary `ffi::Tensor` that owns the allocated memory. `ffi::TensorView x` is
a non-owning view initialized from this temporary. The temporary `ffi::Tensor`
is destroyed at the end of the statement, deallocating its memory. `x` then
becomes a dangling view, pointing to freed memory. Subsequent use of `x` is
undefined behavior.
To fix this, `x` should be an `ffi::Tensor` to correctly manage the memory's
lifetime. The conversion to `TensorView` will happen implicitly and safely at
the call site.
```suggestion
ffi::Tensor x = Empty({5}, DLDataType({kDLFloat, 32, 1}),
DLDevice({kDLCPU, 0}));
```
##########
include/tvm/ffi/container/tensor.h:
##########
@@ -458,6 +458,114 @@ class Tensor : public ObjectRef {
TensorObj* get_mutable() const { return const_cast<TensorObj*>(get()); }
};
+/*!
+ * \brief A non-owning view of a Tensor.
+ *
+ * This class stores a light-weight non-owning view of a Tensor.
+ * This is useful for accessing tensor data without retaining a strong
reference to the Tensor.
+ * Since the caller may not always be able to pass in a strong referenced
tensor.
+ *
+ * When exposing a function that expects only expects a TensorView, we
recommend using
+ * ffi::TensorView as the argument type instead of ffi::Tensor.
+ */
+class TensorView {
+ public:
+ TensorView(const Tensor& tensor) : tensor_(*tensor.operator->()) {} //
NOLINT(*)
Review Comment:
To improve API safety and prevent some misuse cases at compile time,
consider explicitly deleting the constructor that takes an rvalue reference of
`Tensor`. This would prevent direct initialization of a `TensorView` from a
temporary `Tensor` like `TensorView v(Empty(...));`.
While this doesn't prevent all forms of binding to temporaries (e.g.,
copy-initialization), it makes the API safer and the intent clearer.
```c
TensorView(const Tensor& tensor) : tensor_(*tensor.operator->()) {} //
NOLINT(*)
TensorView(Tensor&&) = delete;
```
##########
include/tvm/ffi/container/tensor.h:
##########
@@ -458,6 +458,114 @@ class Tensor : public ObjectRef {
TensorObj* get_mutable() const { return const_cast<TensorObj*>(get()); }
};
+/*!
+ * \brief A non-owning view of a Tensor.
+ *
+ * This class stores a light-weight non-owning view of a Tensor.
+ * This is useful for accessing tensor data without retaining a strong
reference to the Tensor.
+ * Since the caller may not always be able to pass in a strong referenced
tensor.
+ *
+ * When exposing a function that expects only expects a TensorView, we
recommend using
+ * ffi::TensorView as the argument type instead of ffi::Tensor.
+ */
Review Comment:
Given the non-owning nature of `TensorView`, it's crucial to warn users
about potential lifetime issues. Please add a warning to the class
documentation about the danger of creating a `TensorView` from a temporary
`Tensor`.
Here is a suggestion for an addition to the documentation:
```cpp
* \warning As `TensorView` is a non-owning view, it is the user's
responsibility to ensure
* that the underlying tensor data outlives the `TensorView`. Avoid creating
a `TensorView`
* from a temporary `Tensor` object, as this will lead to a dangling
reference.
*
* For example, this is incorrect usage:
* \code
* // BAD: `view` will be a dangling reference as the temporary Tensor is
destroyed.
* ffi::TensorView view = ffi::Tensor::FromDLPackAlloc(...);
*
* // GOOD: `tensor` owns the data, `view` is a valid view.
* ffi::Tensor tensor = ffi::Tensor::FromDLPackAlloc(...);
* ffi::TensorView view = tensor;
* \endcode
```
##########
examples/quick_start/src/run_example_cuda.cc:
##########
@@ -61,8 +61,8 @@ int main() {
constexpr int ARRAY_SIZE = 5;
- ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
- ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
+ ffi::TensorView x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
+ ffi::TensorView y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
Review Comment:
This change introduces a critical use-after-free bug, similar to the one in
`run_example.cc`. The `Empty()` function returns temporary `ffi::Tensor`
objects. Initializing `TensorView`s `x` and `y` from these temporaries results
in dangling views because the temporaries are destroyed immediately, freeing
the underlying memory.
To fix this, `x` and `y` must be `ffi::Tensor` to ensure their data remains
valid throughout their scope.
```suggestion
ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
```
--
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]
