dlpark

A pure Rust implementation of dmlc/dlpack.

This implementation focuses on transferring tensor from Rust to Python and vice versa.

What is DLPack?

DLPack is a common in-memory tensor structure that enables sharing tensor data between different deep learning frameworks. It provides a standardized way to exchange tensor data without copying, making it efficient for framework interoperability.

Key features of DLPack: - Zero-copy tensor sharing between frameworks - Support for various data types and devices (CPU, GPU, etc.) - Memory management through deleter functions - Versioned ABI for compatibility

Implementation Details

Versioning

The library implements both legacy and versioned DLPack structures:

• SafeManagedTensor: Legacy implementation without versioning
• SafeManagedTensorVersioned: Versioned implementation (current standard) with:
• Major version: 1
• Minor version: 1
• Additional flags for tensor properties (read-only, copied, sub-byte type padding)

Safe Abstractions

The library provides safe Rust abstractions over the C-style DLPack structures:

1. SafeManagedTensor and SafeManagedTensorVersioned:
2. Safe wrappers around raw DLPack tensors
3. RAII-style memory management
4. Automatic cleanup through deleter functions

5. Safe conversion between different tensor types

6. Key Features:

7. Memory safety through Rust's ownership system
8. Zero-cost abstractions
9. Support for various tensor types (ndarray, image, standard containers)
10. Python interoperability through PyO3

Features

Quick Start

We provide a simple example of how to transfer image::RgbImage to Python and torch.Tensor to Rust.

Full code is here.

Usage Examples

Converting between Rust and Python

use dlpark::prelude::*;

// Rust to Python
#[pyfunction]
fn send() -> SafeManagedTensorVersioned {
    let v = vec![1i32, 2, 3];
    SafeManagedTensorVersioned::new(v).unwrap()
}

// Python to Rust
#[pyfunction]
fn receive(tensor: SafeManagedTensorVersioned) {
    let s: &[i32] = tensor.as_slice_contiguous().unwrap();
    // Do your work.
}

Working with ndarray

use dlpark::prelude::*;
use ndarray::ArrayD;

let arr = ArrayD::from_shape_vec(IxDyn(&[2, 3]), vec![1i32, 2, 3, 4, 5, 6])?;
let tensor = SafeManagedTensorVersioned::new(arr)?;
let view = ArrayViewD::<i32>::try_from(&tensor)?;

Image Processing

use dlpark::prelude::*;
use image::{ImageBuffer, Rgb};

let img = ImageBuffer::<Rgb<u8>, _>::from_vec(100, 100, vec![0; 100 * 100 * 3])?;
let tensor = SafeManagedTensorVersioned::new(img)?;
let img2 = ImageBuffer::<Rgb<u8>, _>::try_from(&tensor)?;

CUDA GPU data streams

use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaSlice, CudaStream, CudaView};
use dlpark::SafeManagedTensorVersioned;

let ctx: Arc<CudaContext> = CudaContext::new(0)?; // Set on GPU:0
let stream: Arc<CudaStream> = ctx.per_thread_stream();

let slice: CudaSlice<f32> = stream.alloc_zeros::<f32>(10)?;
let tensor = SafeManagedTensorVersioned::new(slice)?;
let view = CudaView::<f32>::try_from(&tensor)?;