Skip to main content

get_device_cpu_set

Function get_device_cpu_set 

Source
pub fn get_device_cpu_set(device_id: u32) -> Option<Vec<usize>>
Expand description

Get a deterministic CPU subset for a CUDA device, subdivided among ALL GPUs sharing the same NUMA node (including those hidden by CUDA_VISIBLE_DEVICES).

§Algorithm

  1. Get PCI address + NUMA node for target device (CUDA driver API)
  2. Enumerate ALL GPUs on the system:
    • Try NVML first (sees all GPUs, ignores CUDA_VISIBLE_DEVICES)
    • Fall back to CUDA driver API (only sees visible devices)
  3. For each GPU, get its NUMA node via sysfs (PCI address → /sys/…/numa_node)
  4. Group GPUs by NUMA node
  5. Sort by PCI address within each group (deterministic)
  6. Get full CPU set for the node via topology
  7. Divide into N equal slices (N = GPUs on same node)
  8. Return the slice for the target device’s position

§Example

System: 8 GPUs, 2 NUMA nodes, 4 GPUs per node. CUDA_VISIBLE_DEVICES=0,1 (only 2 visible). NVML sees all 8 → correctly subdivides into 4 slices per node.

Returns None if NUMA node can’t be determined.