use core::ffi::c_void;
pub mod device_type {
pub const K_DL_CPU: i32 = 1;
pub const K_DL_CUDA: i32 = 2;
pub const K_DL_CUDA_HOST: i32 = 3;
pub const K_DL_VULKAN: i32 = 7;
pub const K_DL_METAL: i32 = 8;
}
pub mod dtype_code {
pub const K_DL_INT: u8 = 0;
pub const K_DL_UINT: u8 = 1;
pub const K_DL_FLOAT: u8 = 2;
pub const K_DL_BFLOAT: u8 = 4;
pub const K_DL_COMPLEX: u8 = 5;
pub const K_DL_BOOL: u8 = 6;
}
pub const DLPACK_FLAG_BITMASK_READ_ONLY: u64 = 1 << 0;
pub const DLPACK_FLAG_BITMASK_IS_COPIED: u64 = 1 << 1;
pub const DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED: u64 = 1 << 2;
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DLDevice {
pub device_type: i32,
pub device_id: i32,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DLDataType {
pub code: u8,
pub bits: u8,
pub lanes: u16,
}
#[repr(C)]
pub struct DLTensor {
pub data: *mut c_void,
pub device: DLDevice,
pub ndim: i32,
pub dtype: DLDataType,
pub shape: *mut i64,
pub strides: *mut i64,
pub byte_offset: u64,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DLPackVersion {
pub major: u32,
pub minor: u32,
}
#[repr(C)]
pub struct DLManagedTensorVersioned {
pub version: DLPackVersion,
pub manager_ctx: *mut c_void,
pub deleter: Option<unsafe extern "C" fn(*mut DLManagedTensorVersioned)>,
pub flags: u64,
pub dl_tensor: DLTensor,
}
#[cfg(test)]
mod tests {
use super::*;
use core::mem::size_of;
#[cfg(target_pointer_width = "64")]
#[test]
fn dlpack_struct_sizes_match_c_header() {
assert_eq!(size_of::<DLDevice>(), 8);
assert_eq!(size_of::<DLDataType>(), 4);
assert_eq!(size_of::<DLPackVersion>(), 8);
assert_eq!(size_of::<DLTensor>(), 48);
assert_eq!(size_of::<DLManagedTensorVersioned>(), 80);
}
#[test]
fn dlpack_standard_flag_bits() {
assert_eq!(DLPACK_FLAG_BITMASK_READ_ONLY, 1);
assert_eq!(DLPACK_FLAG_BITMASK_IS_COPIED, 2);
assert_eq!(DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED, 4);
}
}