This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new aedb1ede [FEAT][Rust] Allocate tensors with the environment allocator
(#827)
aedb1ede is described below
commit aedb1ede7a19f4052bd6afede9ca8c8ad0126ae6
Author: Gabriel Wu <[email protected]>
AuthorDate: Mon Sep 28 22:22:55 2026 +0800
[FEAT][Rust] Allocate tensors with the environment allocator (#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.
---
rust/tvm-ffi-sys/src/c_env_api.rs | 15 +++-
rust/tvm-ffi/src/collections/tensor.rs | 37 +++++++++
rust/tvm-ffi/tests/test_tensor.rs | 144 +++++++++++++++++++++++++++++++++
3 files changed, 193 insertions(+), 3 deletions(-)
diff --git a/rust/tvm-ffi-sys/src/c_env_api.rs
b/rust/tvm-ffi-sys/src/c_env_api.rs
index 94feec12..25fed43c 100644
--- a/rust/tvm-ffi-sys/src/c_env_api.rs
+++ b/rust/tvm-ffi-sys/src/c_env_api.rs
@@ -53,13 +53,22 @@ unsafe extern "C" {
pub fn TVMFFIEnvGetStream(device_type: i32, device_id: i32) ->
TVMFFIStreamHandle;
+ /// Set the allocator of this thread, or unset it with `None`; the
+ /// original allocator, possibly `None`, is written to
+ /// `opt_out_original_allocator` unless it is null.
pub fn TVMFFIEnvSetDLPackManagedTensorAllocator(
- allocator: DLPackManagedTensorAllocator,
+ allocator: Option<DLPackManagedTensorAllocator>,
write_to_global_context: i32,
- opt_out_original_allocator: *mut DLPackManagedTensorAllocator,
+ opt_out_original_allocator: *mut Option<DLPackManagedTensorAllocator>,
) -> i32;
- pub fn TVMFFIEnvGetDLPackManagedTensorAllocator() ->
DLPackManagedTensorAllocator;
+ /// The current allocator, or `None` when none is set.
+ pub fn TVMFFIEnvGetDLPackManagedTensorAllocator() ->
Option<DLPackManagedTensorAllocator>;
+
+ /// Allocate a tensor with the current allocator. Only the device, dtype,
+ /// ndim and shape of `prototype` are used; `out` receives an owned
+ /// `ffi.Tensor` handle.
+ pub fn TVMFFIEnvTensorAlloc(prototype: *mut DLTensor, out: *mut
TVMFFIObjectHandle) -> i32;
pub fn TVMFFIEnvCheckSignals() -> i32;
diff --git a/rust/tvm-ffi/src/collections/tensor.rs
b/rust/tvm-ffi/src/collections/tensor.rs
index 78928049..ebf184dd 100644
--- a/rust/tvm-ffi/src/collections/tensor.rs
+++ b/rust/tvm-ffi/src/collections/tensor.rs
@@ -24,6 +24,7 @@ use crate::error::Result;
use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems};
use tvm_ffi_sys::dlpack::{DLDataType, DLDevice, DLDeviceType, DLTensor};
use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+use tvm_ffi_sys::{TVMFFIEnvTensorAlloc, TVMFFIObjectHandle};
//-----------------------------------------------------
// NDAllocator Trait
@@ -298,6 +299,42 @@ impl Tensor {
}
}
}
+ /// Create a Tensor with the environment allocator of this thread, as C++
+ /// `Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, ...)` does.
+ ///
+ /// Kernel libraries allocate intermediate tensors this way, so that they
+ /// come from the caller's allocator, such as the one a framework sets
+ /// with `TVMFFIEnvSetDLPackManagedTensorAllocator`.
+ ///
+ /// # Arguments
+ /// * `shape` - The shape of the Tensor
+ /// * `dtype` - The data type of the Tensor
+ /// * `device` - The device of the Tensor
+ ///
+ /// # Returns
+ /// * `Result<Tensor>` - The created Tensor, or the error the allocator
+ /// raised, which is a `RuntimeError` when no allocator is set
+ pub fn from_env_alloc(shape: &[i64], dtype: DLDataType, device: DLDevice)
-> Result<Self> {
+ let mut prototype = DLTensor {
+ data: std::ptr::null_mut(),
+ device,
+ ndim: shape.len() as i32,
+ dtype,
+ shape: shape.as_ptr() as *mut i64,
+ strides: std::ptr::null_mut(),
+ byte_offset: 0,
+ };
+ let mut out: TVMFFIObjectHandle = std::ptr::null_mut();
+ unsafe {
+ if TVMFFIEnvTensorAlloc(&mut prototype, &mut out) != 0 {
+ return Err(crate::error::Error::from_raised());
+ }
+ Ok(Self {
+ data: ObjectArc::from_raw(out as *const TensorObj),
+ })
+ }
+ }
+
/// Create a Tensor from a slice
///
/// # Arguments
diff --git a/rust/tvm-ffi/tests/test_tensor.rs
b/rust/tvm-ffi/tests/test_tensor.rs
index d7e8f9bd..1a6207ca 100644
--- a/rust/tvm-ffi/tests/test_tensor.rs
+++ b/rust/tvm-ffi/tests/test_tensor.rs
@@ -110,3 +110,147 @@ fn test_any_tensor() {
assert_eq!(converted_view.shape(), &shape);
assert_eq!(converted_view.dtype().code, DLDataTypeCode::kDLFloat as u8);
}
+
+// ============================================================================
+// Environment allocator tests
+// ============================================================================
+
+mod env_alloc {
+ use std::ffi::{c_char, c_void, CStr};
+ use tvm_ffi::tvm_ffi_sys::{
+ dlpack::DLTensor, DLPackManagedTensorAllocator,
TVMFFIEnvSetDLPackManagedTensorAllocator,
+ };
+ use tvm_ffi::*;
+
+ /// DLPack's `DLManagedTensorVersioned`.
+ #[repr(C)]
+ struct ManagedTensor {
+ version: [u32; 2],
+ manager_ctx: *mut c_void,
+ deleter: Option<unsafe extern "C" fn(*mut ManagedTensor)>,
+ flags: u64,
+ dl_tensor: DLTensor,
+ }
+
+ /// The memory of a tensor made by `cpu_allocator`.
+ struct Owned {
+ shape: Vec<i64>,
+ data: Vec<u8>,
+ }
+
+ unsafe extern "C" fn delete(tensor: *mut ManagedTensor) {
+ let tensor = Box::from_raw(tensor);
+ drop(Box::from_raw(tensor.manager_ctx as *mut Owned));
+ }
+
+ unsafe extern "C" fn cpu_allocator(
+ prototype: *mut DLTensor,
+ out: *mut *mut c_void,
+ _error_ctx: *mut c_void,
+ _set_error: unsafe extern "C" fn(*mut c_void, *const c_char, *const
c_char),
+ ) -> i32 {
+ let prototype = &*prototype;
+ let shape = std::slice::from_raw_parts(prototype.shape, prototype.ndim
as usize).to_vec();
+ let bytes = shape.iter().product::<i64>() as usize *
(prototype.dtype.bits as usize / 8);
+ let mut owned = Box::new(Owned {
+ shape,
+ data: vec![0; bytes],
+ });
+ let dl_tensor = DLTensor {
+ data: owned.data.as_mut_ptr() as *mut c_void,
+ device: prototype.device,
+ ndim: prototype.ndim,
+ dtype: prototype.dtype,
+ shape: owned.shape.as_mut_ptr(),
+ strides: std::ptr::null_mut(),
+ byte_offset: 0,
+ };
+ let tensor = Box::new(ManagedTensor {
+ version: [1, 1],
+ manager_ctx: Box::into_raw(owned) as *mut c_void,
+ deleter: Some(delete),
+ flags: 0,
+ dl_tensor,
+ });
+ *out = Box::into_raw(tensor) as *mut c_void;
+ 0
+ }
+
+ unsafe extern "C" fn failing_allocator(
+ _prototype: *mut DLTensor,
+ _out: *mut *mut c_void,
+ error_ctx: *mut c_void,
+ set_error: unsafe extern "C" fn(*mut c_void, *const c_char, *const
c_char),
+ ) -> i32 {
+ let kind = CStr::from_bytes_with_nul(b"MemoryError\0").unwrap();
+ let message = CStr::from_bytes_with_nul(b"out of test
memory\0").unwrap();
+ set_error(error_ctx, kind.as_ptr(), message.as_ptr());
+ -1
+ }
+
+ /// Runs `f` with `allocator` set on this thread, then restores the
+ /// original allocator.
+ fn with_allocator<R>(
+ allocator: Option<DLPackManagedTensorAllocator>,
+ f: impl FnOnce() -> R,
+ ) -> R {
+ let mut original = None;
+ unsafe {
+ assert_eq!(
+ TVMFFIEnvSetDLPackManagedTensorAllocator(allocator, 0, &mut
original),
+ 0
+ );
+ }
+ let result = f();
+ unsafe {
+ assert_eq!(
+ TVMFFIEnvSetDLPackManagedTensorAllocator(original, 0,
std::ptr::null_mut()),
+ 0
+ );
+ }
+ result
+ }
+
+ fn f32_on_cpu() -> (DLDataType, DLDevice) {
+ (
+ DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1),
+ DLDevice::new(DLDeviceType::kDLCPU, 0),
+ )
+ }
+
+ #[test]
+ fn allocates_with_the_environment_allocator() {
+ let (dtype, device) = f32_on_cpu();
+ let tensor = with_allocator(Some(cpu_allocator), || {
+ Tensor::from_env_alloc(&[1, 2, 3], dtype, device)
+ })
+ .unwrap();
+ assert_eq!(tensor.shape(), &[1, 2, 3]);
+ assert_eq!(tensor.strides(), &[6, 3, 1]);
+ assert_eq!(tensor.dtype(), dtype);
+ assert_eq!(tensor.device(), device);
+ assert!(!tensor.data_ptr().is_null());
+ tensor.data_as_slice::<f32>().unwrap();
+ }
+
+ #[test]
+ fn raises_the_allocator_error() {
+ let (dtype, device) = f32_on_cpu();
+ let error = with_allocator(Some(failing_allocator), || {
+ Tensor::from_env_alloc(&[4], dtype, device)
+ })
+ .err()
+ .unwrap();
+ assert_eq!(error.kind().as_str(), "MemoryError");
+ assert_eq!(error.message(), "out of test memory");
+ }
+
+ #[test]
+ fn fails_without_an_allocator() {
+ let (dtype, device) = f32_on_cpu();
+ let error = with_allocator(None, || Tensor::from_env_alloc(&[4],
dtype, device))
+ .err()
+ .unwrap();
+ assert_eq!(error.kind(), RUNTIME_ERROR);
+ }
+}
