use super::abi::DLDevice;
use super::codes::FDX_AFFINE_MAX_TERMS;
use core::ffi::c_void;
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXDTypeExt {
pub logical_dtype: u16,
pub bit_width: u16,
pub packing: u8,
pub lanes: u8,
pub sub_byte_bit_order: u8,
pub _pad: u8,
pub reserved: [u32; 2],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXQuant {
pub family: u16,
pub ggml_dtype: u16,
pub block_ndim: u8,
pub _pad0: [u8; 3],
pub block_shape: [u32; 4],
pub block_axes: [i32; 4],
pub pack_order: u8,
pub _pad1: [u8; 3],
pub scale_present: u8,
pub scale_dtype: u16,
pub scale_placement: u8,
pub scale_granularity: u8,
pub _pad2: [u8; 3],
pub scale_buffer: u32,
pub zp_present: u8,
pub zp_dtype: u16,
pub _pad3: u8,
pub zp_buffer: u32,
pub scale_pair_act: u8,
pub scale_pair_weight: u8,
pub role: u8,
pub _pad4: u8,
pub reserved: [u32; 6],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXAffineTerm {
pub coeff: i64,
pub sym_id: u32,
pub _pad: u32,
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXAffine {
pub c0: i64,
pub term_count: u8,
pub _pad: [u8; 7],
pub terms: [FDXAffineTerm; FDX_AFFINE_MAX_TERMS],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXExtent {
pub kind: u8,
pub _pad: [u8; 3],
pub min: u64,
pub capacity: u64,
pub sym_id: u32,
pub sym_scope: u8,
pub _pad2: [u8; 3], pub cap_kind: u8,
pub _pad3: [u8; 3], pub _pad4: u32, pub affine: FDXAffine,
pub reserved: [u32; 2],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXTiling {
pub alignment_bytes: u32,
pub access_granularity_bits: u32,
pub tile_ndim: u8,
pub _pad: [u8; 7],
pub tile_shape: [u32; 4],
pub reserved: [u32; 4],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXResidency {
pub tier: u8,
pub substrate: u8,
pub backend_id: u8,
pub _pad: u8,
pub device_index: u32,
pub is_mmap_view: u8,
pub _pad2: [u8; 7],
pub reserved: [u32; 4],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXStorage {
pub class: u8,
pub _pad: [u8; 3],
pub _pad_align: u32,
pub session_id: u64,
pub reserved: [u32; 4],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXOutputView {
pub byte_offset: u64,
pub len_elements: u64,
pub dtype: u16,
pub _pad: [u8; 2],
pub ndim: u32,
pub shape: [u64; 6], pub strides: [i64; 6],
pub name_hash: u64,
pub reserved: [u32; 4],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXBlockTable {
pub table_buffer: u32,
pub id_dtype: u16,
pub _pad0: u16,
pub max_blocks_per_seq: u32,
pub unmapped_sentinel: u32,
pub layout_flags: u32,
pub reserved: [u32; 4],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXIndexedResidency {
pub kind: u8,
pub _pad0: [u8; 3],
pub num_blocks: u64,
pub block_size: u64,
pub pool_buffer: u32,
pub _pad1: u32,
pub physical_ndim: u8,
pub _pad2: [u8; 7],
pub physical_shape: [u64; 6],
pub physical_strides: [i64; 6],
pub element_dtype: u16,
pub _pad3: [u8; 2],
pub block_table: FDXBlockTable,
pub num_sequences: u64,
pub max_seq_capacity: u64,
pub logical_ndim: u8,
pub seq_axis: u8,
pub _pad4: [u8; 6],
pub logical_shape: [u64; 6],
pub logical_extents_count: u8,
pub _pad5: [u8; 7],
pub logical_extents: [FDXExtent; 6],
pub context_lens_buffer: u32,
pub context_len_sym: u32,
pub context_len_scope: u8,
pub _pad6: [u8; 3],
pub reserved: [u32; 6],
}
#[repr(C)]
pub struct FDXBufferRef {
pub role: u8,
pub _pad: [u8; 1],
pub dtype: u16,
pub _pad2: u32,
pub data: *mut c_void,
pub device: DLDevice,
pub byte_offset: u64,
pub size_bytes: u64,
pub ndim: u32,
pub _pad3: u32,
pub shape: [u64; 6],
pub strides: [i64; 6],
pub reserved: [u32; 4],
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FDXSymBinding {
pub sym_id: u32,
pub _pad: u32,
pub value: u64,
}
#[repr(C)]
pub struct FDXSymEnv {
pub count: u32,
pub _pad: u32,
pub bindings: *const FDXSymBinding,
}
#[repr(C)]
pub struct FDXSidecar {
pub magic: u32,
pub version: u32,
pub struct_bytes: u32,
pub flags: u32,
pub dtype_ext: FDXDTypeExt,
pub quant: FDXQuant,
pub extents_count: u32,
pub _pad0: u32,
pub extents: *const FDXExtent,
pub tiling: FDXTiling,
pub residency: FDXResidency,
pub storage: FDXStorage,
pub buffers_count: u32,
pub _pad1: u32,
pub buffers: *const FDXBufferRef,
pub views_count: u32,
pub _pad2: u32,
pub views: *const FDXOutputView,
pub gather: FDXIndexedResidency,
pub reserved: [u64; 2],
}
#[cfg(test)]
mod tests {
use super::*;
use core::mem::{offset_of, size_of};
#[cfg(target_pointer_width = "64")]
#[test]
fn spec_stated_sizes() {
assert_eq!(
size_of::<FDXAffineTerm>(),
16,
"spec §6.4: FDXAffineTerm == 16"
);
assert_eq!(size_of::<FDXAffine>(), 80, "spec §6.4: FDXAffine == 80");
}
#[cfg(target_pointer_width = "64")]
#[test]
fn computed_sizes() {
assert_eq!(size_of::<FDXDTypeExt>(), 16);
assert_eq!(size_of::<FDXQuant>(), 100);
assert_eq!(size_of::<FDXExtent>(), 128);
assert_eq!(size_of::<FDXTiling>(), 48);
assert_eq!(size_of::<FDXResidency>(), 32);
assert_eq!(size_of::<FDXStorage>(), 32);
assert_eq!(size_of::<FDXOutputView>(), 144);
assert_eq!(size_of::<FDXBlockTable>(), 36);
assert_eq!(size_of::<FDXIndexedResidency>(), 1064);
assert_eq!(size_of::<FDXBufferRef>(), 160);
assert_eq!(size_of::<FDXSymBinding>(), 16);
assert_eq!(size_of::<FDXSymEnv>(), 16);
assert_eq!(size_of::<FDXSidecar>(), 1376);
}
#[cfg(target_pointer_width = "64")]
#[test]
fn fdx_extent_field_offsets_are_frozen() {
assert_eq!(offset_of!(FDXExtent, kind), 0);
assert_eq!(offset_of!(FDXExtent, _pad), 1);
assert_eq!(offset_of!(FDXExtent, min), 8);
assert_eq!(offset_of!(FDXExtent, capacity), 16);
assert_eq!(offset_of!(FDXExtent, sym_id), 24);
assert_eq!(offset_of!(FDXExtent, sym_scope), 28);
assert_eq!(offset_of!(FDXExtent, _pad2), 29);
assert_eq!(offset_of!(FDXExtent, cap_kind), 32);
assert_eq!(offset_of!(FDXExtent, _pad3), 33);
assert_eq!(offset_of!(FDXExtent, _pad4), 36);
assert_eq!(offset_of!(FDXExtent, affine), 40);
assert_eq!(offset_of!(FDXExtent, reserved), 120);
}
#[cfg(target_pointer_width = "64")]
#[test]
fn fdx_indexed_residency_trailing_offsets_are_pinned() {
assert_eq!(offset_of!(FDXIndexedResidency, logical_extents), 256);
assert_eq!(offset_of!(FDXIndexedResidency, context_lens_buffer), 1024);
assert_eq!(offset_of!(FDXIndexedResidency, context_len_sym), 1028);
assert_eq!(offset_of!(FDXIndexedResidency, context_len_scope), 1032);
assert_eq!(offset_of!(FDXIndexedResidency, reserved), 1036);
}
#[cfg(target_pointer_width = "64")]
#[test]
fn fdx_sidecar_gather_is_embedded_by_value() {
assert_eq!(
offset_of!(FDXSidecar, reserved),
offset_of!(FDXSidecar, gather) + size_of::<FDXIndexedResidency>()
);
assert_eq!(
size_of::<FDXSidecar>(),
offset_of!(FDXSidecar, reserved) + 16
);
}
}