atlas_core/
device.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Device-context wrapper + per-hardware constants.
4//!
5//! `sm121` is a pure-constants module (NUM_SMS, SMEM_PER_SM, …) that
6//! kernel-launch heuristics in spark-model consume regardless of the
7//! active backend — keeping it always-compiled lets metal-feature
8//! builds reuse the same dispatch tables and code paths.
9//!
10//! `AtlasDevice` is the cudarc-wrapped device handle, gated behind
11//! the `cuda` feature.
12
13/// SM121 hardware constants for DGX Spark GB10.
14pub mod sm121 {
15    /// Number of streaming multiprocessors
16    pub const NUM_SMS: u32 = 48;
17
18    /// Shared memory per SM (bytes)
19    pub const SMEM_PER_SM: usize = 99 * 1024; // 99 KB
20
21    /// Max registers per thread
22    pub const MAX_REGS_PER_THREAD: u32 = 255;
23
24    /// Max threads per block
25    pub const MAX_THREADS_PER_BLOCK: u32 = 1024;
26
27    /// Warp size
28    pub const WARP_SIZE: u32 = 32;
29
30    /// Memory bandwidth (GB/s) — LPDDR5X unified
31    pub const MEMORY_BW_GBS: f64 = 273.0;
32
33    /// Compute capability
34    pub const COMPUTE_MAJOR: u32 = 12;
35    pub const COMPUTE_MINOR: u32 = 1;
36}
37
38#[cfg(feature = "cuda")]
39mod cuda_impl {
40    use cudarc::driver::CudaContext;
41    use std::sync::Arc;
42
43    use crate::error::{AtlasError, Result};
44
45    /// Wrapper around a cudarc CudaContext with SM121-specific configuration.
46    #[derive(Clone)]
47    pub struct AtlasDevice {
48        pub ctx: Arc<CudaContext>,
49        pub ordinal: usize,
50    }
51
52    impl AtlasDevice {
53        /// Initialize an Atlas device on the given GPU ordinal.
54        pub fn new(ordinal: usize) -> Result<Self> {
55            let ctx = CudaContext::new(ordinal).map_err(AtlasError::CudaDriver)?;
56            Ok(Self { ctx, ordinal })
57        }
58    }
59}
60
61#[cfg(feature = "cuda")]
62pub use cuda_impl::AtlasDevice;
atlas_core/device.rs

atlas_core/
device.rs
