Trait AsDlTensor

Source

pub trait AsDlTensor {
    // Required method
    fn as_dl_tensor(&self) -> Result<DLTensorView<'_>, DLPackError>;
}

Expand description

Borrows a tensor as a read-only DLTensorView for tensor inputs.

Implement this for your own tensor type by calling DLTensorView::from_raw_parts inside a small unsafe block, upholding its safety contract.

§Examples

A minimal adapter for a row-major matrix in device memory:

use cuvs::dlpack::{AsDlTensor, DLDevice, DLDeviceType, DLPackError, DLTensorView, DType};

struct GpuMatrix<T> {
    ptr: *mut T,
    rows: usize,
    cols: usize,
}

impl<T: DType> AsDlTensor for GpuMatrix<T> {
    fn as_dl_tensor(&self) -> Result<DLTensorView<'_>, DLPackError> {
        let shape = [self.rows as i64, self.cols as i64];
        // SAFETY: `ptr` points to `rows * cols` initialized elements of `T`
        // in device 0's memory, valid while `self` is borrowed, and is
        // row-major contiguous.
        unsafe {
            DLTensorView::from_raw_parts(
                self.ptr.cast(),
                DLDevice { device_type: DLDeviceType::kDLCUDA, device_id: 0 },
                &shape,
                None,
                T::dl_dtype(),
            )
        }
    }
}