lucifer1004 opened a new pull request, #827:
URL: https://github.com/apache/tvm-ffi/pull/827
## Why
C++ kernel libraries allocate intermediate tensors with
`Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, ...)`, which `tensor.h` recommends:
the memory comes from the DLPack allocator the caller set with
`TVMFFIEnvSetDLPackManagedTensorAllocator`. The Python binding, for example,
sets the framework's allocator while it calls a function with that framework's
tensors. The Rust crates have no way to do this:
- `tvm-ffi-sys` does not bind `TVMFFIEnvTensorAlloc`.
- `TVMFFIEnvGetDLPackManagedTensorAllocator` is declared to return a
non-nullable `DLPackManagedTensorAllocator`, but the C API returns null when no
allocator is set. Calling it on a thread without an allocator produces an
invalid value. `TVMFFIEnvSetDLPackManagedTensorAllocator` likewise cannot unset
the allocator or report an unset original one.
## What
- `tvm-ffi-sys`:
- bind `TVMFFIEnvTensorAlloc`;
- the getter returns, and the setter takes and writes back,
`Option<DLPackManagedTensorAllocator>`. `Option` of a function pointer has the
same ABI as the nullable C pointer, so only the Rust types change.
- `tvm-ffi`: add `Tensor::from_env_alloc(shape, dtype, device) ->
Result<Tensor>`, the counterpart of C++ `Tensor::FromEnvAlloc`:
- it fills a prototype with the device, dtype and shape and calls
`TVMFFIEnvTensorAlloc`;
- it returns the allocator's error, which is a `RuntimeError` when no
allocator is set.
## Tests
The new tests in `rust/tvm-ffi/tests/test_tensor.rs` mirror
`tests/cpp/extra/test_c_env_api.cc`, using a DLPack allocator written in the
test:
- a tensor allocated through the allocator has the requested shape, compact
strides, dtype and device;
- an allocator that fails with `MemoryError` passes its kind and message
through;
- with the allocator unset on the thread, `from_env_alloc` fails with
`RuntimeError`.
Each test restores the thread's original allocator. `cargo test` in `rust/`
passes.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
