ai-hwaccel 1.1.0

Universal AI hardware accelerator detection, capability querying, and workload planning for Rust
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

//! Cerebras Wafer-Scale Engine (WSE) detection via `/dev` and `cerebras_cli`.

use tracing::debug;

use crate::error::DetectionError;
use crate::hardware::AcceleratorType;
use crate::profile::AcceleratorProfile;

use super::command::{DEFAULT_TIMEOUT, run_tool};

/// Default memory: 44 GB SRAM (WSE-3).
const DEFAULT_MEMORY_BYTES: u64 = 44 * 1024 * 1024 * 1024;

pub(crate) fn detect_cerebras_wse(
    profiles: &mut Vec<AcceleratorProfile>,
    warnings: &mut Vec<DetectionError>,
) {
    // Try CLI first for richer info.
    match run_tool("cerebras_cli", &["system-info"], DEFAULT_TIMEOUT) {
        Ok(output) => {
            debug!(
                device_id = 0,
                memory_gb = 44,
                "Cerebras WSE detected via cerebras_cli"
            );
            profiles.push(AcceleratorProfile {
                accelerator: AcceleratorType::CerebrasWse { device_id: 0 },
                available: true,
                memory_bytes: parse_memory_from_cli(&output.stdout).unwrap_or(DEFAULT_MEMORY_BYTES),
                compute_capability: Some("WSE".into()),
                ..Default::default()
            });
            return;
        }
        Err(DetectionError::ToolNotFound { .. }) => {
            debug!("cerebras_cli not found on $PATH, skipping Cerebras CLI detection");
        }
        Err(_) => {}
    }

    // Fallback: check /dev/cerebras* devices.
    if super::has_dev_device("cerebras") {
        debug!(
            device_id = 0,
            memory_gb = 44,
            "Cerebras WSE detected via /dev"
        );
        profiles.push(AcceleratorProfile {
            accelerator: AcceleratorType::CerebrasWse { device_id: 0 },
            available: true,
            memory_bytes: DEFAULT_MEMORY_BYTES,
            compute_capability: Some("WSE".into()),
            ..Default::default()
        });
        return;
    }

    let _ = warnings;
}

/// Best-effort parse of memory from `cerebras_cli system-info` output.
pub(crate) fn parse_memory_from_cli(stdout: &str) -> Option<u64> {
    for line in stdout.lines() {
        let lower = line.to_lowercase();
        if lower.contains("memory") || lower.contains("sram") {
            // Look for a number followed by "gb" or "GB".
            for word in line.split_whitespace() {
                if let Some(num_str) = word.strip_suffix("GB").or_else(|| word.strip_suffix("gb"))
                    && let Ok(gb) = num_str.parse::<u64>()
                {
                    return Some(gb.saturating_mul(1024 * 1024 * 1024));
                }
            }
        }
    }
    None
}