Seven-Streams commented on code in PR #813:
URL: https://github.com/apache/tvm-ffi/pull/813#discussion_r4106095235
##########
rust/tvm-ffi/src/collections/tensor.rs:
##########
@@ -342,3 +344,156 @@ 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.
+#[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.
+ ///
+ /// # Panics
+ /// If the `DLTensor` has null strides; see [`strides`](Self::strides).
+ pub fn is_contiguous(&self) -> bool {
+ let mut expected_stride = 1;
+ for (size, stride) in self.shape().iter().zip(self.strides()).rev() {
+ if *size != 1 && *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,
+ }
+ }
+}
+
+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) {
Review Comment:
according to include/tvm/ffi/container/tensor.h:928, on the C++ side,
`move_to_any` is not supported. Not sure if we need to align the behavior.
##########
rust/tvm-ffi/src/collections/tensor.rs:
##########
@@ -342,3 +344,156 @@ 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.
+#[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.
+ ///
+ /// # Panics
+ /// If the `DLTensor` has null strides; see [`strides`](Self::strides).
+ pub fn is_contiguous(&self) -> bool {
Review Comment:
It seems the behavior between the C++ and Rust sides is not aligned.
--
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]
