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 2eb546fd [FEAT][Rust] Add TensorView for DLTensor* arguments (#813)
2eb546fd is described below
commit 2eb546fd93d70b4424266e174a3e7aad81c58f76
Author: Gabriel Wu <[email protected]>
AuthorDate: Wed Sep 30 07:09:26 2026 +0800
[FEAT][Rust] Add TensorView for DLTensor* arguments (#813)
## Why
C and C++ callers pass tensors to tvm-ffi functions as borrowed
`DLTensor*` (`kTVMFFIDLTensorPtr`), which is also how C++ `TensorView`
arguments are passed on. The Rust `Tensor` only accepts `ffi.Tensor`
objects, so a Rust function exported with
`tvm_ffi_dll_export_typed_func!` fails for those callers with `Cannot
convert from type 'DLTensor*' to 'ffi.Tensor'`. Python callers are not
affected, because they pass `ffi.Tensor`.
## What
`tvm_ffi::TensorView` mirrors C++ `tvm::ffi::TensorView`:
- it accepts `DLTensor*` and `ffi.Tensor` arguments and is passed on as
a `DLTensor*`;
- it is non-owning: it holds a pointer to the `DLTensor`, and the caller
keeps the tensor alive, as in C++;
- its accessors follow `Tensor`'s: `data_ptr`, `device`, `dtype`,
`ndim`, `shape`, `numel`, `strides`, `is_contiguous`. `strides` panics
on null strides, which DLPack forbids since v1.2, as C++
`TensorView::strides` checks. `is_contiguous` follows C++
`tvm::ffi::IsContiguous`: null strides and tensors without elements are
contiguous, and a dimension of extent 1 may have any stride.
- as with C++ `TypeTraits<DLTensor*>` / `TypeTraits<TensorView>`, a view
can be borrowed as an `AnyView` but not moved into or out of an owned
`Any`, since it does not retain ownership; doing so panics with the C++
error message.
- `TryFrom<AnyView>` converts a `DLTensor*` or a tensor object to a
view, as C++ `AnyView::cast<TensorView>` / `try_cast` do.
`try_as::<TensorView>()` stays aligned with C++ `as` and accepts only a
`DLTensor*`. There is no `TryFrom<Any>`, since the view would point into
a value the conversion drops.
Typed functions require `'static` argument types, so the view carries no
lifetime. The alternative, borrow-checked views of call arguments, would
need changes to typed-function argument conversion and is left for a
follow-up discussion.
## Tests
`rust/tvm-ffi/tests/test_tensor_view.rs`: the same typed function called
with a `DLTensor*` built as a C caller would and with an `ffi.Tensor`,
rejection of other types, contiguity (including null strides, empty
tensors and extent-1 dimensions), the null-strides panic of `strides`,
the panic on moving a view into `Any`, and `TryFrom<AnyView>` from both
a tensor object and a `DLTensor*`, with other values rejected. `cargo
test` passes against libtvm_ffi built from `main`.
Also checked end to end: a Rust `cdylib` exporting a CUDA kernel entry
with `TensorView` arguments, called from C with `DLTensor*` and from
Python with torch tensors.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 5.5 (1M context) <[email protected]>
---
rust/tvm-ffi/src/collections/tensor.rs | 185 +++++++++++++++++++++++++++++++++
rust/tvm-ffi/src/lib.rs | 2 +-
rust/tvm-ffi/tests/test_tensor_view.rs | 155 +++++++++++++++++++++++++++
3 files changed, 341 insertions(+), 1 deletion(-)
diff --git a/rust/tvm-ffi/src/collections/tensor.rs
b/rust/tvm-ffi/src/collections/tensor.rs
index ebf184dd..bb92f3ed 100644
--- a/rust/tvm-ffi/src/collections/tensor.rs
+++ b/rust/tvm-ffi/src/collections/tensor.rs
@@ -16,13 +16,16 @@
* specific language governing permissions and limitations
* under the License.
*/
+use crate::any::{AnyView, TryFromTemp};
use crate::collections::shape::Shape;
use crate::derive::{Object, ObjectRef};
use crate::dtype::AsDLDataType;
use crate::dtype::DLDataTypeExt;
use crate::error::Result;
use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems};
+use crate::type_traits::AnyCompatible;
use tvm_ffi_sys::dlpack::{DLDataType, DLDevice, DLDeviceType, DLTensor};
+use tvm_ffi_sys::TVMFFIAny;
use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
use tvm_ffi_sys::{TVMFFIEnvTensorAlloc, TVMFFIObjectHandle};
@@ -379,3 +382,185 @@ unsafe impl NDAllocator for CPUNDAlloc {
std::alloc::dealloc(tensor.data as *mut u8, layout);
}
}
+
+/// A non-owning view of a tensor: a pointer to a `DLTensor`.
+///
+/// This mirrors `tvm::ffi::TensorView` in C++. A function argument of this
+/// type accepts both a borrowed `DLTensor*` (type index `kTVMFFIDLTensorPtr`,
+/// how C and C++ callers pass tensors) and a `Tensor` object (type index
+/// `kTVMFFITensor`), and is passed on as a `DLTensor*`. Prefer it over
+/// [`Tensor`] for arguments of exported functions that only read a tensor's
+/// metadata and data.
+///
+/// The view does not keep anything alive. The caller must ensure that the
+/// `DLTensor` it points to, and the memory that the `DLTensor` points to
+/// (data, shape and strides), outlive every use of the view. For an argument
+/// of an exported function, that is the duration of the call. For the same
+/// reason, as in C++, a view cannot be moved into an owned
[`Any`](crate::Any);
+/// use [`Tensor`] instead.
+#[derive(Clone, Copy, Debug)]
+pub struct TensorView {
+ tensor: *const DLTensor,
+}
+
+impl TensorView {
+ /// Creates a view of a `DLTensor`.
+ ///
+ /// # Safety
+ /// `tensor` must be non-null and valid for every use of the view.
+ pub unsafe fn from_raw(tensor: *const DLTensor) -> Self {
+ assert!(!tensor.is_null(), "TensorView of a null DLTensor");
+ Self { tensor }
+ }
+
+ /// The viewed `DLTensor`.
+ pub fn as_raw(&self) -> *const DLTensor {
+ self.tensor
+ }
+
+ fn dltensor(&self) -> &DLTensor {
+ unsafe { &*self.tensor }
+ }
+
+ /// The data pointer, as in the `DLTensor`; `byte_offset` is not applied.
+ pub fn data_ptr(&self) -> *mut core::ffi::c_void {
+ self.dltensor().data
+ }
+
+ pub fn device(&self) -> DLDevice {
+ self.dltensor().device
+ }
+
+ pub fn dtype(&self) -> DLDataType {
+ self.dltensor().dtype
+ }
+
+ pub fn ndim(&self) -> usize {
+ self.dltensor().ndim as usize
+ }
+
+ pub fn shape(&self) -> &[i64] {
+ let t = self.dltensor();
+ if t.ndim == 0 {
+ return &[];
+ }
+ unsafe { std::slice::from_raw_parts(t.shape, t.ndim as usize) }
+ }
+
+ pub fn numel(&self) -> usize {
+ self.shape().iter().product::<i64>() as usize
+ }
+
+ /// The strides, in elements.
+ ///
+ /// # Panics
+ /// If the `DLTensor` has null strides and at least one dimension, which
+ /// DLPack forbids since v1.2, as C++ `TensorView::strides` checks.
+ pub fn strides(&self) -> &[i64] {
+ let t = self.dltensor();
+ if t.ndim == 0 {
+ return &[];
+ }
+ assert!(
+ !t.strides.is_null(),
+ "TensorView::strides of a DLTensor with null strides"
+ );
+ unsafe { std::slice::from_raw_parts(t.strides, t.ndim as usize) }
+ }
+
+ /// Whether the elements are laid out in row-major order without gaps, as
+ /// C++ `tvm::ffi::IsContiguous` decides: a `DLTensor` with null strides
+ /// or no elements is contiguous, and a dimension of extent 1 may have any
+ /// stride.
+ pub fn is_contiguous(&self) -> bool {
+ if self.dltensor().strides.is_null() {
+ return true;
+ }
+ let shape = self.shape();
+ if shape.contains(&0) {
+ return true;
+ }
+ let mut expected_stride = 1;
+ for (size, stride) in shape.iter().zip(self.strides()).rev() {
+ if *size == 1 {
+ continue;
+ }
+ if *stride != expected_stride {
+ return false;
+ }
+ expected_stride *= size;
+ }
+ true
+ }
+}
+
+impl From<&Tensor> for TensorView {
+ fn from(tensor: &Tensor) -> Self {
+ Self {
+ tensor: &tensor.data.dltensor as *const DLTensor,
+ }
+ }
+}
+
+// A borrowed `AnyView` converts as C++ `AnyView::cast<TensorView>` does: from
+// a `DLTensor*` or from a tensor object. `try_as`, like C++ `as`, accepts only
+// a `DLTensor*`. There is no conversion from an owned `Any`, which would leave
+// the view pointing into a value the conversion drops.
+impl<'a> TryFrom<AnyView<'a>> for TensorView {
+ type Error = crate::error::Error;
+ #[inline]
+ fn try_from(value: AnyView<'a>) -> Result<Self> {
+ TryFromTemp::<Self>::try_from(value).map(TryFromTemp::into_value)
+ }
+}
+
+// As in C++, where `TypeTraits<DLTensor*>::MoveToAny` throws this, and
+// `TypeTraits<TensorView>` does not support moves.
+const NOT_OWNED: &str =
+ "DLTensor* cannot be held in Any as it does not retain ownership, use
Tensor instead";
+
+unsafe impl AnyCompatible for TensorView {
+ fn type_str() -> String {
+ // make it consistent with c++ representation
+ "DLTensor*".to_string()
+ }
+
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIDLTensorPtr as i32;
+ data.small_str_len = 0;
+ data.data_union.v_uint64 = 0;
+ data.data_union.v_ptr = src.tensor as *mut core::ffi::c_void;
+ }
+
+ unsafe fn move_to_any(_src: Self, _data: &mut TVMFFIAny) {
+ panic!("{}", NOT_OWNED);
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ data.type_index == TypeIndex::kTVMFFIDLTensorPtr as i32
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ Self {
+ tensor: data.data_union.v_ptr as *const DLTensor,
+ }
+ }
+
+ unsafe fn move_from_any_after_check(_data: &mut TVMFFIAny) -> Self {
+ panic!("{}", NOT_OWNED);
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) ->
std::result::Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIDLTensorPtr as i32 {
+ Ok(Self::copy_from_any_view_after_check(data))
+ } else if data.type_index == TypeIndex::kTVMFFITensor as i32 {
+ // A tensor object's DLTensor follows its object header.
+ let obj = data.data_union.v_obj as *const TensorObj;
+ Ok(Self {
+ tensor: &(*obj).dltensor as *const DLTensor,
+ })
+ } else {
+ Err(())
+ }
+ }
+}
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index 5d6425c0..8b3cbfd1 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -42,7 +42,7 @@ pub use crate::big_int::BigInt;
pub use crate::collections::array::Array;
pub use crate::collections::map::Map;
pub use crate::collections::shape::Shape;
-pub use crate::collections::tensor::{CPUNDAlloc, NDAllocator, Tensor};
+pub use crate::collections::tensor::{CPUNDAlloc, NDAllocator, Tensor,
TensorView};
pub use crate::device::{current_stream, with_stream};
pub use crate::dtype::DLDataTypeExt;
pub use crate::error::{Error, ErrorKind, Result};
diff --git a/rust/tvm-ffi/tests/test_tensor_view.rs
b/rust/tvm-ffi/tests/test_tensor_view.rs
new file mode 100644
index 00000000..c0fe1409
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_tensor_view.rs
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::tvm_ffi_sys::dlpack::{DLDataType, DLDataTypeCode, DLDevice,
DLDeviceType, DLTensor};
+use tvm_ffi::*;
+
+/// A 2x3 f32 DLTensor over `data`, with the given strides, as a C caller
+/// would build it.
+fn dltensor(data: &mut [f32; 6], shape: &mut [i64; 2], strides: *mut i64) ->
DLTensor {
+ DLTensor {
+ data: data.as_mut_ptr() as *mut core::ffi::c_void,
+ device: DLDevice::new(DLDeviceType::kDLCPU, 0),
+ ndim: 2,
+ dtype: DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1),
+ shape: shape.as_mut_ptr(),
+ strides,
+ byte_offset: 0,
+ }
+}
+
+/// `(ndim, shape[1], strides[0], first element)` of the argument.
+fn describe() -> Function {
+ Function::from_typed(|t: TensorView| -> Result<Array<i64>> {
+ let first = unsafe { *(t.data_ptr() as *const f32) } as i64;
+ Ok(Array::new(vec![
+ t.ndim() as i64,
+ t.shape()[1],
+ t.strides()[0],
+ first,
+ ]))
+ })
+}
+
+fn call(f: &Function, arg: AnyView) -> Result<Vec<i64>> {
+ let out = Array::<i64>::try_from(f.call_packed(&[arg])?)?;
+ (0..out.len()).map(|i| Ok(i64::try_from(out[i])?)).collect()
+}
+
+#[test]
+fn test_tensor_view_from_dltensor_ptr() {
+ let (mut data, mut shape, mut strides) = ([7.0f32, 1., 2., 3., 4., 5.],
[2, 3], [3i64, 1]);
+ let raw = dltensor(&mut data, &mut shape, strides.as_mut_ptr());
+ let view = unsafe { TensorView::from_raw(&raw) };
+ let arg = AnyView::from(&view);
+ assert_eq!(arg.type_index(), TypeIndex::kTVMFFIDLTensorPtr as i32);
+ assert_eq!(call(&describe(), arg).unwrap(), vec![2, 3, 3, 7]);
+}
+
+#[test]
+fn test_tensor_view_from_tensor_object() {
+ let tensor = Tensor::from_slice(&[9.0f32, 1., 2., 3., 4., 5.], &[2,
3]).unwrap();
+ let arg = AnyView::from(&tensor);
+ assert_eq!(arg.type_index(), TypeIndex::kTVMFFITensor as i32);
+ assert_eq!(call(&describe(), arg).unwrap(), vec![2, 3, 3, 9]);
+ let view = TensorView::from(&tensor);
+ assert_eq!(view.shape(), tensor.shape());
+ assert_eq!(view.strides(), tensor.strides());
+ assert_eq!(
+ view.data_ptr() as *const core::ffi::c_void,
+ tensor.data_ptr()
+ );
+}
+
+#[test]
+fn test_tensor_view_rejects_other_types() {
+ let err = call(&describe(), AnyView::from(&1i32)).unwrap_err();
+ assert!(err.message().contains("DLTensor*"), "{}", err.message());
+}
+
+#[test]
+fn test_tensor_view_contiguity() {
+ let (mut data, mut shape) = ([0f32; 6], [2i64, 3]);
+ let mut compact_strides = [3i64, 1];
+ let compact = dltensor(&mut data, &mut shape,
compact_strides.as_mut_ptr());
+ let view = unsafe { TensorView::from_raw(&compact) };
+ assert!(view.is_contiguous());
+ assert_eq!(view.numel(), 6);
+ let mut padded_strides = [4i64, 1];
+ let padded = dltensor(&mut data, &mut shape, padded_strides.as_mut_ptr());
+ assert!(!unsafe { TensorView::from_raw(&padded) }.is_contiguous());
+ // As in C++ `IsContiguous`: null strides and empty tensors are
+ // contiguous, and a dimension of extent 1 may have any stride.
+ let unstrided = dltensor(&mut data, &mut shape, std::ptr::null_mut());
+ assert!(unsafe { TensorView::from_raw(&unstrided) }.is_contiguous());
+ let mut row_shape = [1i64, 3];
+ let mut row_strides = [7i64, 1];
+ let row = dltensor(&mut data, &mut row_shape, row_strides.as_mut_ptr());
+ assert!(unsafe { TensorView::from_raw(&row) }.is_contiguous());
+ let mut empty_shape = [0i64, 3];
+ let empty = dltensor(&mut data, &mut empty_shape,
padded_strides.as_mut_ptr());
+ assert!(unsafe { TensorView::from_raw(&empty) }.is_contiguous());
+}
+
+#[test]
+#[should_panic(expected = "null strides")]
+fn test_tensor_view_null_strides() {
+ let (mut data, mut shape) = ([0f32; 6], [2i64, 3]);
+ let compact = dltensor(&mut data, &mut shape, std::ptr::null_mut());
+ unsafe { TensorView::from_raw(&compact) }.strides();
+}
+
+#[test]
+#[should_panic(expected = "does not retain ownership")]
+fn test_tensor_view_cannot_move_to_any() {
+ let tensor = Tensor::from_slice(&[0.0f32; 6], &[2, 3]).unwrap();
+ let _ = Any::from(TensorView::from(&tensor));
+}
+
+#[test]
+fn test_tensor_view_try_from_any_view() {
+ // From a tensor object, as C++ `AnyView::cast<TensorView>` converts it;
the
+ // strict `try_as`, like C++ `as`, does not.
+ let tensor = Tensor::from_slice(&[5.0f32, 1., 2., 3., 4., 5.], &[2,
3]).unwrap();
+ let object = AnyView::from(&tensor);
+ let view = TensorView::try_from(object).unwrap();
+ assert_eq!(
+ view.data_ptr() as *const core::ffi::c_void,
+ tensor.data_ptr()
+ );
+ assert_eq!(view.shape(), tensor.shape());
+ assert!(object.try_as::<TensorView>().is_none());
+
+ // From a DLTensor*, which both accept.
+ let (mut data, mut shape, mut strides) = ([7.0f32, 1., 2., 3., 4., 5.],
[2, 3], [3i64, 1]);
+ let raw = dltensor(&mut data, &mut shape, strides.as_mut_ptr());
+ let borrowed = unsafe { TensorView::from_raw(&raw) };
+ let pointer = AnyView::from(&borrowed);
+ assert_eq!(
+ TensorView::try_from(pointer).unwrap().as_raw(),
+ &raw as *const DLTensor
+ );
+ assert_eq!(
+ pointer.try_as::<TensorView>().unwrap().as_raw(),
+ &raw as *const DLTensor
+ );
+
+ // Other values are rejected with the expected type.
+ let err = TensorView::try_from(AnyView::from(&1i32)).unwrap_err();
+ assert!(err.message().contains("DLTensor*"), "{}", err.message());
+}
