use core::ffi::c_void;
use core::marker::PhantomData;
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use fuel_ir::dlpack::abi::{DLDataType, DLDevice, DLTensor, device_type};
use fuel_ir::dlpack::codes::*;
use fuel_ir::dlpack::convert::{dl_dtype, extent_to_fdx};
use fuel_ir::dlpack::sidecar::{
FDXAffine, FDXAffineTerm, FDXBufferRef, FDXDTypeExt, FDXExtent, FDXIndexedResidency,
FDXOutputView, FDXQuant, FDXResidency, FDXSidecar, FDXStorage, FDXSymBinding, FDXSymEnv,
FDXTiling,
};
use fuel_ir::dlpack::validate::{self, FdxValidationError};
use fuel_ir::shape::Extent;
use fuel_ir::storage::OutputView;
use fuel_ir::{DType, DeviceLocation, Error, Layout, Result, SymEnv};
use crate::Storage;
const MAX_RANK: usize = 6;
pub struct DlpackView<'a> {
pub dl: DLTensor,
pub sidecar: Option<FDXSidecar>,
_shape: [i64; MAX_RANK],
_strides: [i64; MAX_RANK],
_extents: Vec<FDXExtent>,
_buffers: Vec<FDXBufferRef>,
_views: Vec<FDXOutputView>,
_marker: PhantomData<(&'a Storage, &'a Layout)>,
}
impl<'a> DlpackView<'a> {
pub fn ndim(&self) -> usize {
self.dl.ndim.max(0) as usize
}
pub fn shape(&self) -> &[i64] {
&self._shape[..self.ndim()]
}
pub fn strides(&self) -> &[i64] {
&self._strides[..self.ndim()]
}
pub fn dl_tensor(&self) -> DLTensor {
let (shape, strides) = if self.dl.ndim != 0 {
(
self._shape.as_ptr() as *mut i64,
self._strides.as_ptr() as *mut i64,
)
} else {
(core::ptr::null_mut(), core::ptr::null_mut())
};
DLTensor {
data: self.dl.data,
device: self.dl.device,
ndim: self.dl.ndim,
dtype: self.dl.dtype,
shape,
strides,
byte_offset: self.dl.byte_offset,
}
}
pub fn buffers(&self) -> &[FDXBufferRef] {
&self._buffers
}
pub fn validate(&self) -> std::result::Result<(), FdxValidationError> {
match &self.sidecar {
Some(sc) => validate::validate(sc, &self.dl_tensor(), &self._buffers),
None => Ok(()),
}
}
pub fn validate_realize(&self, env: &FDXSymEnv) -> std::result::Result<(), FdxValidationError> {
match &self.sidecar {
Some(sc) => validate::validate_realize(sc, &self.dl_tensor(), &self._buffers, env),
None => Ok(()),
}
}
}
fn dl_device(loc: DeviceLocation) -> DLDevice {
let (device_type, device_id) = match loc {
DeviceLocation::Cpu => (device_type::K_DL_CPU, 0),
DeviceLocation::Cuda { gpu_id } => (device_type::K_DL_CUDA, gpu_id as i32),
DeviceLocation::Vulkan { gpu_id } => (device_type::K_DL_VULKAN, gpu_id as i32),
DeviceLocation::Metal { gpu_id } => (device_type::K_DL_METAL, gpu_id as i32),
};
DLDevice {
device_type,
device_id,
}
}
fn base_ptr_and_device(storage: &Storage) -> Result<(*mut c_void, DeviceLocation)> {
match &storage.inner {
crate::BackendStorage::Cpu(cpu) => {
Ok((cpu.bytes().as_ptr() as *mut c_void, DeviceLocation::Cpu))
}
#[cfg(any(feature = "cuda", feature = "vulkan", feature = "metal"))]
other => {
let _ = crate::dispatch_storage!(other, _inner => _inner.len_bytes());
Err(Error::Msg(
"DlpackView: device-pointer extraction for GPU storage is a later \
comm-layer slice (CPU is wired); [consumer-ahead]"
.into(),
)
.bt())
}
}
}
fn is_sub_byte(d: DType) -> bool {
d.size_in_bytes() == 0
}
fn sub_byte_bits_and_packing(d: DType) -> Result<(u16, u8)> {
match d {
DType::F4 => Ok((4, FDX_PACKING_DENSE_SUBBYTE)),
DType::F6E2M3 | DType::F6E3M2 => Ok((6, FDX_PACKING_DENSE_SUBBYTE)),
other => Err(Error::Msg(format!(
"DlpackView: no authored sub-byte bit-width/packing for dtype {other:?}; \
declining rather than emitting an invented FDX wire value a foreign \
runtime would stride over (GAP-075)"
))
.bt()),
}
}
fn physical_elem_bytes(d: DType) -> usize {
let s = d.size_in_bytes();
if s == 0 { 1 } else { s }
}
fn dtype_ext_none() -> FDXDTypeExt {
FDXDTypeExt {
logical_dtype: FDX_DTYPE_NONE,
bit_width: 0,
packing: 0,
lanes: 0,
sub_byte_bit_order: 0,
_pad: 0,
reserved: [0; 2],
}
}
fn quant_none() -> FDXQuant {
FDXQuant {
family: FDX_QUANT_NONE,
ggml_dtype: FDX_DTYPE_NONE,
block_ndim: 0,
_pad0: [0; 3],
block_shape: [0; 4],
block_axes: [-1; 4],
pack_order: 0,
_pad1: [0; 3],
scale_present: 0,
scale_dtype: FDX_DTYPE_NONE,
scale_placement: 0,
scale_granularity: 0,
_pad2: [0; 3],
scale_buffer: FDX_BUFFER_NONE,
zp_present: 0,
zp_dtype: FDX_DTYPE_NONE,
_pad3: 0,
zp_buffer: FDX_BUFFER_NONE,
scale_pair_act: 0,
scale_pair_weight: 0,
role: 0,
_pad4: 0,
reserved: [0; 6],
}
}
fn tiling_none() -> FDXTiling {
FDXTiling {
alignment_bytes: 0,
access_granularity_bits: 0,
tile_ndim: 0,
_pad: [0; 7],
tile_shape: [0; 4],
reserved: [0; 4],
}
}
fn residency_for(loc: DeviceLocation) -> FDXResidency {
let (tier, substrate, backend_id, device_index) = match loc {
DeviceLocation::Cpu => (FDX_TIER_HOST, FDX_SUBSTRATE_HOST_BYTES, FDX_BACKEND_CPU, 0),
DeviceLocation::Cuda { gpu_id } => (
FDX_TIER_DEVICE,
FDX_SUBSTRATE_CUDA_UNTYPED,
FDX_BACKEND_CUDA,
gpu_id as u32,
),
DeviceLocation::Vulkan { gpu_id } => (
FDX_TIER_DEVICE,
FDX_SUBSTRATE_VULKAN_BUFFER,
FDX_BACKEND_VULKAN,
gpu_id as u32,
),
DeviceLocation::Metal { gpu_id } => (
FDX_TIER_DEVICE,
FDX_SUBSTRATE_METAL_BUFFER,
FDX_BACKEND_METAL,
gpu_id as u32,
),
};
FDXResidency {
tier,
substrate,
backend_id,
_pad: 0,
device_index,
is_mmap_view: 0,
_pad2: [0; 7],
reserved: [0; 4],
}
}
fn storage_none() -> FDXStorage {
FDXStorage {
class: FDX_STORAGE_TRANSIENT,
_pad: [0; 3],
_pad_align: 0,
session_id: FDX_SESSION_NONE as u64,
reserved: [0; 4],
}
}
fn gather_none() -> FDXIndexedResidency {
FDXIndexedResidency {
kind: FDX_GATHER_NONE as u8,
_pad0: [0; 3],
num_blocks: 0,
block_size: 0,
pool_buffer: 0,
_pad1: 0,
physical_ndim: 0,
_pad2: [0; 7],
physical_shape: [0; 6],
physical_strides: [0; 6],
element_dtype: FDX_DTYPE_NONE,
_pad3: [0; 2],
block_table: fuel_ir::dlpack::sidecar::FDXBlockTable {
table_buffer: 0,
id_dtype: FDX_DTYPE_NONE,
_pad0: 0,
max_blocks_per_seq: 0,
unmapped_sentinel: 0,
layout_flags: 0,
reserved: [0; 4],
},
num_sequences: 0,
max_seq_capacity: 0,
logical_ndim: 0,
seq_axis: 0xFF,
_pad4: [0; 6],
logical_shape: [0; 6],
logical_extents_count: 0,
_pad5: [0; 7],
logical_extents: [zero_extent(); 6],
context_lens_buffer: FDX_BUFFER_NONE,
context_len_sym: FDX_SYM_NONE,
context_len_scope: 0,
_pad6: [0; 3],
reserved: [0; 6],
}
}
const fn zero_extent() -> FDXExtent {
FDXExtent {
kind: FDX_EXTENT_SCALAR as u8,
_pad: [0; 3],
min: 0,
capacity: 0,
sym_id: FDX_SYM_NONE,
sym_scope: 0,
_pad2: [0; 3],
cap_kind: FDX_CAP_KIND_EXPLICIT as u8,
_pad3: [0; 3],
_pad4: 0,
affine: FDXAffine {
c0: 0,
term_count: 0,
_pad: [0; 7],
terms: [FDXAffineTerm {
coeff: 0,
sym_id: FDX_SYM_NONE,
_pad: 0,
}; FDX_AFFINE_MAX_TERMS],
},
reserved: [0; 2],
}
}
static FDX_NAME_TABLE: OnceLock<Mutex<HashMap<u64, String>>> = OnceLock::new();
fn name_table() -> &'static Mutex<HashMap<u64, String>> {
FDX_NAME_TABLE.get_or_init(|| Mutex::new(HashMap::new()))
}
pub fn fdx_slot_name(hash: u64) -> Option<String> {
if hash == 0 {
return None;
}
name_table().lock().ok()?.get(&hash).cloned()
}
fn output_view_to_fdx(ov: &OutputView) -> Result<FDXOutputView> {
let dims = ov.shape.dims();
let ndim = dims.len();
if ndim > MAX_RANK {
return Err(Error::Msg(format!(
"DlpackView bundle slot rank {ndim} exceeds 6 (BundleRankExceeds6)"
))
.bt());
}
let mut shape = [0u64; MAX_RANK];
let mut strides = [0i64; MAX_RANK];
let st = ov.layout.stride();
for i in 0..ndim {
shape[i] = dims[i] as u64;
strides[i] = st[i] as i64;
}
let name_hash = match ov.name {
Some(n) => {
let h = fnv1a(n);
if let Ok(mut table) = name_table().lock() {
table.insert(h, n.to_string());
}
h
}
None => 0,
};
Ok(FDXOutputView {
byte_offset: ov.byte_offset as u64,
len_elements: ov.len_elements as u64,
dtype: dtype_to_fdx_code(ov.dtype)?,
_pad: [0; 2],
ndim: ndim as u32,
shape,
strides,
name_hash,
reserved: [0; 4],
})
}
fn fnv1a(s: &str) -> u64 {
let mut h: u64 = 0xcbf2_9ce4_8422_2325;
for b in s.as_bytes() {
h ^= *b as u64;
h = h.wrapping_mul(0x0000_0100_0000_01b3);
}
h
}
fn dtype_to_fdx_code(d: DType) -> fuel_ir::Result<u16> {
fuel_ir::dlpack::convert::dtype_to_fdx(d).map(|disp| disp.namespace_code())
}
pub fn view<'a>(
storage: &'a Storage,
layout: &'a Layout,
env: Option<&SymEnv>,
) -> Result<DlpackView<'a>> {
let dtype = storage.dtype;
let dims = layout.dims();
let ndim = dims.len();
if ndim > MAX_RANK {
return Err(Error::Msg(format!(
"DlpackView: layout rank {ndim} exceeds 6 (RankExceeds6)"
))
.bt());
}
let (data, loc) = base_ptr_and_device(storage)?;
let device = dl_device(loc);
let sub_byte = is_sub_byte(dtype);
let bundled = storage.is_bundled();
let symbolic = layout.has_dynamic();
let quant_proj = storage.stype.to_fdx(None)?;
let has_quant = quant_proj.is_some();
let need_sidecar = sub_byte || symbolic || bundled || has_quant;
let physical_byte_base = sub_byte || bundled || has_quant;
let mut shape = [0i64; MAX_RANK];
let mut strides = [0i64; MAX_RANK];
let base_ndim: usize;
let dl_dt: DLDataType;
let byte_offset: u64;
if physical_byte_base {
base_ndim = 1;
shape[0] = storage.len_bytes() as i64;
strides[0] = 1;
dl_dt = DLDataType {
code: fuel_ir::dlpack::abi::dtype_code::K_DL_UINT,
bits: 8,
lanes: 1,
};
byte_offset =
(layout.start_offset() as u64).saturating_mul(physical_elem_bytes(dtype) as u64);
} else {
base_ndim = ndim;
let st = layout.stride();
for i in 0..ndim {
shape[i] = dims[i] as i64;
strides[i] = st[i] as i64;
}
dl_dt = dl_dtype(dtype)?;
byte_offset = (layout.start_offset() as u64).saturating_mul(dtype.size_in_bytes() as u64);
}
let dl = DLTensor {
data,
device,
ndim: base_ndim as i32,
dtype: dl_dt,
shape: core::ptr::null_mut(),
strides: core::ptr::null_mut(),
byte_offset,
};
if !need_sidecar {
return Ok(DlpackView {
dl,
sidecar: None,
_shape: shape,
_strides: strides,
_extents: Vec::new(),
_buffers: Vec::new(),
_views: Vec::new(),
_marker: PhantomData,
});
}
let mut flags: u32 = 0;
let mut dtype_ext = dtype_ext_none();
let mut extents: Vec<FDXExtent> = Vec::new();
let mut views: Vec<FDXOutputView> = Vec::new();
let quant = match quant_proj {
Some(q) => {
flags |= FDX_FLAG_HAS_QUANT | FDX_FLAG_MEANING_REQUIRES_EXT;
q
}
None => quant_none(),
};
if sub_byte {
flags |= FDX_FLAG_HAS_DTYPE_EXT | FDX_FLAG_MEANING_REQUIRES_EXT;
let (bit_width, packing) = sub_byte_bits_and_packing(dtype)?;
dtype_ext = FDXDTypeExt {
logical_dtype: dtype_to_fdx_code(dtype)?,
bit_width,
packing,
lanes: 1,
sub_byte_bit_order: 0, _pad: 0,
reserved: [0; 2],
};
}
if symbolic {
flags |= FDX_FLAG_HAS_SYMBOLIC;
for e in layout.shape().extents() {
extents.push(extent_to_fdx(e));
}
}
if bundled {
flags |= FDX_FLAG_IS_BUNDLE;
if let Some(slots) = storage.bundle() {
for ov in slots {
views.push(output_view_to_fdx(ov)?);
}
}
}
let buf0_size = storage.len_bytes() as u64;
let mut buf0_shape = [0u64; MAX_RANK];
let mut buf0_strides = [0i64; MAX_RANK];
for i in 0..base_ndim {
buf0_shape[i] = shape[i] as u64;
buf0_strides[i] = strides[i];
}
let buffers = vec![FDXBufferRef {
role: FDX_BUFFER_ROLE_DATA,
_pad: [0; 1],
dtype: if physical_byte_base {
FDX_DTYPE_U8
} else {
dtype_to_fdx_code(dtype)?
},
_pad2: 0,
data,
device,
byte_offset,
size_bytes: buf0_size,
ndim: base_ndim as u32,
_pad3: 0,
shape: buf0_shape,
strides: buf0_strides,
reserved: [0; 4],
}];
let extents_count = extents.len() as u32;
let views_count = views.len() as u32;
let sidecar = FDXSidecar {
magic: FDX_MAGIC,
version: FDX_VERSION_1,
struct_bytes: core::mem::size_of::<FDXSidecar>() as u32,
flags,
dtype_ext,
quant,
extents_count,
_pad0: 0,
extents: core::ptr::null(),
tiling: tiling_none(),
residency: residency_for(loc),
storage: storage_none(),
buffers_count: 1,
_pad1: 0,
buffers: core::ptr::null(),
views_count,
_pad2: 0,
views: core::ptr::null(),
gather: gather_none(),
reserved: [0; 2],
};
let mut v = DlpackView {
dl,
sidecar: Some(sidecar),
_shape: shape,
_strides: strides,
_extents: extents,
_buffers: buffers,
_views: views,
_marker: PhantomData,
};
if let Some(sc) = v.sidecar.as_mut() {
sc.extents = if v._extents.is_empty() {
core::ptr::null()
} else {
v._extents.as_ptr()
};
sc.buffers = if v._buffers.is_empty() {
core::ptr::null()
} else {
v._buffers.as_ptr()
};
sc.views = if v._views.is_empty() {
core::ptr::null()
} else {
v._views.as_ptr()
};
}
if symbolic {
if let Some(env) = env {
let bindings = collect_bindings(layout, env);
let fdx_env = FDXSymEnv {
count: bindings.len() as u32,
_pad: 0,
bindings: if bindings.is_empty() {
core::ptr::null()
} else {
bindings.as_ptr()
},
};
v.validate_realize(&fdx_env).map_err(|e| {
Error::Msg(format!(
"DlpackView: FDX realize-time validation failed: {e}"
))
.bt()
})?;
}
}
Ok(v)
}
pub fn view_with_quant<'a>(
storage: &'a Storage,
layout: &'a Layout,
env: Option<&SymEnv>,
scale: (&'a Storage, &'a Layout),
) -> Result<DlpackView<'a>> {
let mut v = view(storage, layout, env)?;
if !storage.stype.requires_scale_sibling() {
return Ok(v);
}
let (scale_storage, scale_layout) = scale;
let scale_ref = scale_buffer_ref(scale_storage, scale_layout)?;
let scale_index = v._buffers.len() as u32;
v._buffers.push(scale_ref);
let buffers_ptr = v._buffers.as_ptr();
let buffers_count = v._buffers.len() as u32;
if let Some(sc) = v.sidecar.as_mut() {
sc.quant.scale_buffer = scale_index;
sc.buffers_count = buffers_count;
sc.buffers = buffers_ptr;
}
Ok(v)
}
fn scale_buffer_ref(storage: &Storage, layout: &Layout) -> Result<FDXBufferRef> {
let dtype = storage.dtype;
let dims = layout.dims();
let ndim = dims.len();
if ndim > MAX_RANK {
return Err(Error::Msg(format!(
"view_with_quant: scale rank {ndim} exceeds 6 (RankExceeds6)"
))
.bt());
}
let (data, loc) = base_ptr_and_device(storage)?;
let device = dl_device(loc);
let st = layout.stride();
let mut shape = [0u64; MAX_RANK];
let mut strides = [0i64; MAX_RANK];
for i in 0..ndim {
shape[i] = dims[i] as u64;
strides[i] = st[i] as i64;
}
let byte_offset = (layout.start_offset() as u64).saturating_mul(dtype.size_in_bytes() as u64);
Ok(FDXBufferRef {
role: FDX_BUFFER_ROLE_SCALE,
_pad: [0; 1],
dtype: dtype_to_fdx_code(dtype)?,
_pad2: 0,
data,
device,
byte_offset,
size_bytes: storage.len_bytes() as u64,
ndim: ndim as u32,
_pad3: 0,
shape,
strides,
reserved: [0; 4],
})
}
fn collect_bindings(layout: &Layout, env: &SymEnv) -> Vec<FDXSymBinding> {
let mut out: Vec<FDXSymBinding> = Vec::new();
for e in layout.shape().extents() {
if let Extent::Range { sym, .. } = e {
if let Some(v) = env.get(sym) {
if !out.iter().any(|b| b.sym_id == sym.0) {
out.push(FDXSymBinding {
sym_id: sym.0,
_pad: 0,
value: v as u64,
});
}
}
}
}
out.sort_unstable_by_key(|b| b.sym_id);
out
}
#[cfg(test)]
mod tests;