ferrum_kernels/backend/reference/
composition.rs

1use std::collections::BTreeSet;
2use std::sync::Arc;
3
4use ferrum_interfaces::vnext::{
5    dense_linear_contract, CapabilityCatalog, CapabilityId, ContractVersion, DeviceClass,
6    DeviceDescriptor, DeviceId, DeviceRuntime, DynamicStorageAllocator, DynamicStorageProfile,
7    DynamicStorageView, EngineProviderDescriptor, OperationContract, OperationProvider,
8    OperationRuntimeRegistry, ProviderId, VNextError, WeightMaterializerId,
9    WeightMaterializerRegistry, DENSE_LINEAR_F16_CAPABILITY_ID, IDENTITY_WEIGHT_MATERIALIZER_ID,
10};
11
12use super::dense_linear::{implementation_fingerprint, ReferenceDenseLinearProvider};
13use super::runtime::{
14    ReferenceDeviceRuntime, ReferenceDeviceRuntimeConfig, ReferenceDeviceRuntimeError,
15};
16
17const REFERENCE_ENGINE_PROVIDER_ID: &str = "provider.engine.reference.vnext";
18pub const REFERENCE_DENSE_SAFETENSORS_FORMAT_ID: &str = "weight-format.safetensors.dense";
19const REFERENCE_MEMORY_BYTES: u64 = 64 * 1024 * 1024;
20
21pub fn reference_vnext_capabilities() -> Result<BTreeSet<CapabilityId>, VNextError> {
22    Ok(BTreeSet::from([CapabilityId::new(
23        DENSE_LINEAR_F16_CAPABILITY_ID,
24    )?]))
25}
26
27pub(super) fn reference_vnext_runtime_config(
28    device_id: DeviceId,
29) -> Result<ReferenceDeviceRuntimeConfig, VNextError> {
30    let descriptor = DeviceDescriptor {
31        id: device_id,
32        class: DeviceClass::Reference,
33        ordinal: 0,
34        total_memory_bytes: REFERENCE_MEMORY_BYTES,
35        runtime_implementation_fingerprint: implementation_fingerprint(&[
36            include_str!("runtime.rs").as_bytes(),
37            include_str!("dense_linear.rs").as_bytes(),
38            include_str!("composition.rs").as_bytes(),
39        ]),
40        capabilities: reference_vnext_capabilities()?,
41        dynamic_storage_profiles: BTreeSet::from([DynamicStorageProfile::new(
42            DynamicStorageAllocator::LinearArena,
43            DynamicStorageView::Contiguous,
44        )?]),
45    };
46    descriptor.validate()?;
47    Ok(ReferenceDeviceRuntimeConfig { descriptor })
48}
49
50pub fn reference_vnext_operation_registry(
51    runtime: &ReferenceDeviceRuntime,
52) -> Result<OperationRuntimeRegistry<ReferenceDeviceRuntime>, ReferenceDeviceRuntimeError> {
53    let contracts: Vec<Box<dyn OperationContract>> =
54        vec![Box::new(dense_linear_contract().map_err(contract_error)?)];
55    let providers: Vec<Box<dyn OperationProvider<ReferenceDeviceRuntime>>> =
56        vec![Box::new(ReferenceDenseLinearProvider::new(runtime)?)];
57    OperationRuntimeRegistry::new(contracts, providers).map_err(contract_error)
58}
59
60/// Single composition authority for reference planning, allocation, weight
61/// initialization, provider binding, and execution.
62pub struct ReferenceVNextComposition {
63    runtime: Arc<ReferenceDeviceRuntime>,
64    registry: OperationRuntimeRegistry<ReferenceDeviceRuntime>,
65    weight_materializers: WeightMaterializerRegistry,
66    weight_materializer_id: WeightMaterializerId,
67    catalog: CapabilityCatalog,
68}
69
70impl ReferenceVNextComposition {
71    pub fn create(device_id: DeviceId) -> Result<Self, ReferenceDeviceRuntimeError> {
72        let config = reference_vnext_runtime_config(device_id).map_err(contract_error)?;
73        let runtime = Arc::new(ReferenceDeviceRuntime::new(config)?);
74        let registry = reference_vnext_operation_registry(&runtime)?;
75        let weight_materializers =
76            WeightMaterializerRegistry::identity_only().map_err(contract_error)?;
77        let weight_materializer_id =
78            WeightMaterializerId::new(IDENTITY_WEIGHT_MATERIALIZER_ID).map_err(contract_error)?;
79        let engine = EngineProviderDescriptor::new(
80            ProviderId::new(REFERENCE_ENGINE_PROVIDER_ID).map_err(contract_error)?,
81            ContractVersion::new(1, 0),
82            implementation_fingerprint(&[
83                include_str!("composition.rs").as_bytes(),
84                REFERENCE_ENGINE_PROVIDER_ID.as_bytes(),
85            ]),
86            runtime.descriptor().id.clone(),
87            runtime.descriptor().capabilities.clone(),
88        )
89        .map_err(contract_error)?;
90        let catalog = registry
91            .capability_catalog(runtime.descriptor().clone(), vec![engine])
92            .map_err(contract_error)?;
93        let catalog = weight_materializers
94            .augment_catalog(catalog)
95            .map_err(contract_error)?;
96        Ok(Self {
97            runtime,
98            registry,
99            weight_materializers,
100            weight_materializer_id,
101            catalog,
102        })
103    }
104
105    pub fn runtime(&self) -> &Arc<ReferenceDeviceRuntime> {
106        &self.runtime
107    }
108
109    pub fn registry(&self) -> &OperationRuntimeRegistry<ReferenceDeviceRuntime> {
110        &self.registry
111    }
112
113    pub fn catalog(&self) -> &CapabilityCatalog {
114        &self.catalog
115    }
116
117    pub fn weight_materializers(&self) -> &WeightMaterializerRegistry {
118        &self.weight_materializers
119    }
120
121    pub fn weight_materializer_id(&self) -> &WeightMaterializerId {
122        &self.weight_materializer_id
123    }
124}
125
126fn contract_error(error: VNextError) -> ReferenceDeviceRuntimeError {
127    ReferenceDeviceRuntimeError::contract(error.to_string())
128}
ferrum_kernels/backend/reference/composition.rs

ferrum_kernels/backend/reference/
composition.rs
