themis-topology 0.11.1

Typed placement law for NUMA, HBM, memory-tier, and worker locality contracts
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

//! TPU topology query types.

use super::types::TpuDeviceProperties;
use crate::law::TopologyEpoch;

/// TPU device topology snapshot.
///
/// Provider-fed: themis stays stateless law, so there is no `detect()` here.
/// TPU backends construct this from PJRT or runtime-reported device
/// attributes via [`TpuTopology::from_provider`]. Unreported fields are zero,
/// never inferred.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TpuTopology {
    epoch: TopologyEpoch,
    properties: TpuDeviceProperties,
}

impl TpuTopology {
    /// Construct a snapshot from provider-reported TPU properties.
    #[must_use]
    pub const fn from_provider(properties: TpuDeviceProperties) -> Self {
        Self {
            epoch: TopologyEpoch::INITIAL,
            properties,
        }
    }

    /// Snapshot epoch.
    #[must_use]
    #[inline]
    pub const fn epoch(&self) -> TopologyEpoch {
        self.epoch
    }

    /// TPU core count.
    #[must_use]
    #[inline]
    pub const fn core_count(&self) -> u32 {
        self.properties.core_count
    }

    /// HBM capacity per TPU core in bytes.
    #[must_use]
    #[inline]
    pub const fn hbm_bytes_per_core(&self) -> u64 {
        self.properties.hbm_bytes_per_core
    }

    /// Total provider-reported HBM capacity in bytes.
    ///
    /// Uses saturating arithmetic so malformed provider input cannot wrap the
    /// capacity value.
    #[must_use]
    #[inline]
    pub const fn total_hbm_bytes(&self) -> u64 {
        self.properties
            .hbm_bytes_per_core
            .saturating_mul(self.properties.core_count as u64)
    }
}