aphelion-core 1.1.0

Core library for Aphelion AI pipeline framework - graphs, configs, pipelines, backends, diagnostics
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

//! Burn backend integration (feature-gated).

use crate::backend::{Backend, DeviceCapabilities};

/// Supported burn device selectors (placeholder).
#[derive(Debug, Clone)]
pub enum BurnDevice {
    Cpu,
    Cuda(u32),
    Metal(u32),
    Vulkan(u32),
}

impl BurnDevice {
    pub fn as_label(&self) -> String {
        match self {
            BurnDevice::Cpu => "cpu".to_string(),
            BurnDevice::Cuda(id) => format!("cuda:{}", id),
            BurnDevice::Metal(id) => format!("metal:{}", id),
            BurnDevice::Vulkan(id) => format!("vulkan:{}", id),
        }
    }
}

/// Burn backend configuration (placeholder).
#[derive(Debug, Clone)]
pub struct BurnBackendConfig {
    pub device: BurnDevice,
    pub allow_tf32: bool,
}

impl Default for BurnBackendConfig {
    fn default() -> Self {
        Self {
            device: BurnDevice::Cpu,
            allow_tf32: false,
        }
    }
}

/// Simple burn backend stub that carries device configuration.
#[derive(Debug, Clone)]
pub struct BurnBackend {
    config: BurnBackendConfig,
}

impl BurnBackend {
    pub fn new(config: BurnBackendConfig) -> Self {
        Self { config }
    }

    pub fn config(&self) -> &BurnBackendConfig {
        &self.config
    }
}

impl Backend for BurnBackend {
    fn name(&self) -> &str {
        "burn"
    }

    fn device(&self) -> &str {
        // NOTE: This returns a cached label for now.
        // Future: map to actual burn device handle.
        match &self.config.device {
            BurnDevice::Cpu => "cpu",
            BurnDevice::Cuda(_) => "cuda",
            BurnDevice::Metal(_) => "metal",
            BurnDevice::Vulkan(_) => "vulkan",
        }
    }

    fn capabilities(&self) -> DeviceCapabilities {
        match &self.config.device {
            BurnDevice::Cpu => DeviceCapabilities {
                supports_f16: false,
                supports_bf16: false,
                supports_tf32: false,
                max_memory_bytes: None,
                compute_units: None,
            },
            BurnDevice::Cuda(_) => DeviceCapabilities {
                supports_f16: true,
                supports_bf16: true,
                supports_tf32: self.config.allow_tf32,
                max_memory_bytes: None, // Would be queried from actual device
                compute_units: None,
            },
            BurnDevice::Metal(_) => DeviceCapabilities {
                supports_f16: true,
                supports_bf16: false,
                supports_tf32: false,
                max_memory_bytes: None,
                compute_units: None,
            },
            BurnDevice::Vulkan(_) => DeviceCapabilities {
                supports_f16: true,
                supports_bf16: false,
                supports_tf32: false,
                max_memory_bytes: None,
                compute_units: None,
            },
        }
    }

    fn is_available(&self) -> bool {
        // Placeholder: actual implementation would check device availability
        true
    }
}