tvm_ffi/collections/

tensor.rs

/*
 * 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 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};

//-----------------------------------------------------
// NDAllocator Trait
//-----------------------------------------------------
/// Trait for n-dimensional array allocators
pub unsafe trait NDAllocator: 'static {
    /// The minimum alignment of the data allocated by the allocator
    const MIN_ALIGN: usize;
    /// Allocate data for the given DLTensor
    ///
    /// # Arguments
    /// * `tensor` - The DLTensor to allocate data for
    ///
    /// This method should fill in the data pointer of the DLTensor.
    unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut core::ffi::c_void;

    /// Free data for the given DLTensor
    ///
    /// # Arguments
    /// * `tensor` - The DLTensor to free data for
    ///
    /// This method should free the data pointer of the DLTensor.
    unsafe fn free_data(&mut self, tensor: &DLTensor);
}

/// DLTensorExt trait
/// This trait provides methods to get the number of elements and the item size of a DLTensor
pub trait DLTensorExt {
    fn numel(&self) -> usize;
    fn item_size(&self) -> usize;
}

impl DLTensorExt for DLTensor {
    fn numel(&self) -> usize {
        unsafe {
            std::slice::from_raw_parts(self.shape, self.ndim as usize)
                .iter()
                .product::<i64>() as usize
        }
    }

    fn item_size(&self) -> usize {
        (self.dtype.bits as usize * self.dtype.lanes as usize + 7) / 8
    }
}

//-----------------------------------------------------
// Shape
//-----------------------------------------------------
// ShapeObj for heap-allocated shape
#[repr(C)]
#[derive(Object)]
#[type_key = "ffi.Tensor"]
#[type_index(TypeIndex::kTVMFFITensor)]
pub struct TensorObj {
    object: Object,
    dltensor: DLTensor,
}

/// ABI stable owned Shape for ffi
#[repr(C)]
#[derive(ObjectRef, Clone)]
pub struct Tensor {
    data: ObjectArc<TensorObj>,
}

impl Tensor {
    /// Get the data pointer of the Tensor
    ///
    /// # Returns
    /// * `*mut core::ffi::c_void` - The data pointer of the Tensor
    pub fn data_ptr(&self) -> *const core::ffi::c_void {
        self.data.dltensor.data
    }
    /// Get the data pointer of the Tensor
    ///
    /// # Returns
    /// * `*mut core::ffi::c_void` - The data pointer of the Tensor
    pub fn data_ptr_mut(&mut self) -> *mut core::ffi::c_void {
        self.data.dltensor.data
    }
    /// Check if the Tensor is contiguous
    ///
    /// # Returns
    /// * `bool` - True if the Tensor is contiguous, false otherwise
    pub fn is_contiguous(&self) -> bool {
        let strides = self.strides();
        let shape = self.shape();
        let mut expected_stride = 1;
        for i in (0..self.ndim()).rev() {
            if strides[i] != expected_stride {
                return false;
            }
            expected_stride *= shape[i];
        }
        true
    }

    pub fn data_as_slice<T: AsDLDataType>(&self) -> Result<&[T]> {
        let dtype = T::DL_DATA_TYPE;
        if self.dtype() != dtype {
            crate::bail!(
                crate::error::TYPE_ERROR,
                "Data type mismatch {} vs {}",
                self.dtype().to_string(),
                dtype.to_string()
            );
        }
        if self.device().device_type != DLDeviceType::kDLCPU {
            crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU");
        }
        crate::ensure!(
            self.is_contiguous(),
            crate::error::RUNTIME_ERROR,
            "Tensor is not contiguous"
        );

        unsafe {
            Ok(std::slice::from_raw_parts(
                self.data.dltensor.data as *const T,
                self.numel(),
            ))
        }
    }
    /// Returns the tensor data as a mutable slice.
    ///
    /// This method takes `&self` rather than `&mut self` by design: like
    /// `std::fs::File::write`, the *metadata* of a Tensor (shape, dtype,
    /// device) is governed by Rust's ownership rules, but writing to the
    /// underlying data buffer (CPU memory or a GPU pointer) is a side-effect
    /// outside Rust's aliasing model.  Most C/CUDA kernel APIs accept a
    /// non-mut Tensor and mutate its data content, so requiring `&mut self`
    /// here would force artificial mutability annotations throughout the
    /// deep-learning stack with no real safety benefit.
    ///
    /// # Safety contract (caller responsibility)
    /// If the `Tensor` has been cloned (via `ObjectArc`), the caller must
    /// ensure no other clone is concurrently reading the data.
    #[allow(clippy::wrong_self_convention)]
    pub fn data_as_slice_mut<T: AsDLDataType>(&self) -> Result<&mut [T]> {
        let dtype = T::DL_DATA_TYPE;
        if self.dtype() != dtype {
            crate::bail!(
                crate::error::TYPE_ERROR,
                "Data type mismatch: expected {}, got {}",
                dtype.to_string(),
                self.dtype().to_string()
            );
        }
        if self.device().device_type != DLDeviceType::kDLCPU {
            crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU");
        }
        crate::ensure!(
            self.is_contiguous(),
            crate::error::RUNTIME_ERROR,
            "Tensor is not contiguous"
        );
        unsafe {
            Ok(std::slice::from_raw_parts_mut(
                self.data.dltensor.data as *mut T,
                self.numel(),
            ))
        }
    }

    pub fn shape(&self) -> &[i64] {
        unsafe { std::slice::from_raw_parts(self.data.dltensor.shape, self.ndim()) }
    }

    pub fn ndim(&self) -> usize {
        self.data.dltensor.ndim as usize
    }

    pub fn numel(&self) -> usize {
        self.data.dltensor.numel()
    }

    pub fn strides(&self) -> &[i64] {
        unsafe { std::slice::from_raw_parts(self.data.dltensor.strides, self.ndim()) }
    }

    pub fn dtype(&self) -> DLDataType {
        self.data.dltensor.dtype
    }

    pub fn device(&self) -> DLDevice {
        self.data.dltensor.device
    }
}

struct TensorObjFromNDAlloc<TNDAlloc>
where
    TNDAlloc: NDAllocator,
{
    base: TensorObj,
    alloc: TNDAlloc,
}

unsafe impl<TNDAlloc: NDAllocator> ObjectCore for TensorObjFromNDAlloc<TNDAlloc> {
    const TYPE_KEY: &'static str = TensorObj::TYPE_KEY;
    const TYPE_DEPTH: i32 = TensorObj::TYPE_DEPTH;
    fn type_index() -> i32 {
        TensorObj::type_index()
    }
    unsafe fn object_header_mut(this: &mut Self) -> &mut tvm_ffi_sys::TVMFFIObject {
        TensorObj::object_header_mut(&mut this.base)
    }
}

unsafe impl<TNDAlloc: NDAllocator> ObjectCoreWithExtraItems for TensorObjFromNDAlloc<TNDAlloc> {
    type ExtraItem = i64;
    fn extra_items_count(this: &Self) -> usize {
        (this.base.dltensor.ndim * 2) as usize
    }
}

impl<TNDAlloc: NDAllocator> Drop for TensorObjFromNDAlloc<TNDAlloc> {
    fn drop(&mut self) {
        unsafe {
            self.alloc.free_data(&self.base.dltensor);
        }
    }
}

impl Tensor {
    // Create a Tensor from a NDAllocator
    ///
    /// # Arguments
    /// * `alloc` - The NDAllocator
    /// * `shape` - The shape of the Tensor
    /// * `dtype` - The data type of the Tensor
    /// * `device` - The device of the Tensor
    ///
    /// # Returns
    /// * `Tensor` - The created Tensor
    pub fn from_nd_alloc<TNDAlloc>(
        alloc: TNDAlloc,
        shape: &[i64],
        dtype: DLDataType,
        device: DLDevice,
    ) -> Self
    where
        TNDAlloc: NDAllocator,
    {
        let tensor_obj = TensorObjFromNDAlloc {
            base: TensorObj {
                object: Object::new(),
                dltensor: DLTensor {
                    data: std::ptr::null_mut(),
                    device: device,
                    ndim: shape.len() as i32,
                    dtype: dtype,
                    shape: std::ptr::null_mut(),
                    strides: std::ptr::null_mut(),
                    byte_offset: 0,
                },
            },
            alloc: alloc,
        };
        unsafe {
            let mut obj_arc = ObjectArc::new_with_extra_items(tensor_obj);
            // The fresh allocation is unique and has not been exposed.
            let object = &mut *ObjectArc::as_raw_mut(&mut obj_arc);
            object.base.dltensor.shape =
                TensorObjFromNDAlloc::extra_items(object).as_ptr() as *mut i64;
            object.base.dltensor.strides = object.base.dltensor.shape.add(shape.len());
            let extra_items = TensorObjFromNDAlloc::extra_items_mut(object);
            extra_items[..shape.len()].copy_from_slice(shape);
            Shape::fill_strides_from_shape(shape, &mut extra_items[shape.len()..]);
            let dltensor_ptr = &object.base.dltensor as *const DLTensor;
            object.base.dltensor.data = object.alloc.alloc_data(&*dltensor_ptr);
            Self {
                data: ObjectArc::from_raw(ObjectArc::into_raw(obj_arc) as *mut TensorObj),
            }
        }
    }
    /// 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
    /// * `slice` - The slice to create the Tensor from
    /// * `shape` - The shape of the Tensor
    ///
    /// # Returns
    /// * `Tensor` - The created Tensor
    pub fn from_slice<T: AsDLDataType>(slice: &[T], shape: &[i64]) -> Result<Self> {
        let dtype = T::DL_DATA_TYPE;
        let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
        let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, shape, dtype, device);
        if tensor.numel() != slice.len() {
            crate::bail!(crate::error::VALUE_ERROR, "Slice length mismatch");
        }
        tensor.data_as_slice_mut::<T>()?.copy_from_slice(slice);
        Ok(tensor)
    }
}

/// Example CPU NDAllocator
/// This allocator allocates data on the CPU
pub struct CPUNDAlloc {}

unsafe impl NDAllocator for CPUNDAlloc {
    const MIN_ALIGN: usize = 64;

    unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut core::ffi::c_void {
        let numel = prototype.numel() as usize;
        let item_size = prototype.item_size();
        let size = numel * item_size as usize;
        let layout = std::alloc::Layout::from_size_align(size, Self::MIN_ALIGN).unwrap();
        let ptr = std::alloc::alloc(layout);
        ptr as *mut core::ffi::c_void
    }

    unsafe fn free_data(&mut self, tensor: &DLTensor) {
        let numel = tensor.numel() as usize;
        let item_size = tensor.item_size();
        let size = numel * item_size;
        let layout = std::alloc::Layout::from_size_align(size, Self::MIN_ALIGN).unwrap();
        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(())
        }
    }
}