singe_cutensor/mg/
memory.rs

1use std::ptr;
2
3use singe_cuda::{device::Device, memory::DeviceMemory};
4
5use crate::{
6    error::{Error, Result},
7    mg::{context::ContextRef, tensor::TensorDevice, workspace::Workspace},
8};
9
10#[derive(Debug)]
11pub struct DistributedDeviceMemory<T> {
12    buffers: Vec<DeviceMemory<T>>,
13    devices: Vec<TensorDevice>,
14    len_per_pointer: usize,
15}
16
17#[derive(Debug)]
18pub struct WorkspaceMemory {
19    device: Vec<DeviceMemory<u8>>,
20    host: HostWorkspace,
21}
22
23#[derive(Debug)]
24struct HostWorkspace {
25    ptr: *mut (),
26}
27
28impl<T> DistributedDeviceMemory<T> {
29    pub fn create(devices: &[TensorDevice], len_per_pointer: usize) -> Result<Self> {
30        let mut buffers = Vec::with_capacity(devices.len());
31        for device in devices {
32            let TensorDevice::Device(device_id) = *device else {
33                return Err(Error::MgHostTensorMemoryUnsupported {
34                    name: "distributed device memory".into(),
35                });
36            };
37            Device::new(device_id).set_current()?;
38            buffers.push(DeviceMemory::<T>::create(len_per_pointer)?);
39        }
40
41        Ok(Self {
42            buffers,
43            devices: devices.to_vec(),
44            len_per_pointer,
45        })
46    }
47
48    pub fn pointer_count(&self) -> usize {
49        self.buffers.len()
50    }
51
52    pub fn len_per_pointer(&self) -> usize {
53        self.len_per_pointer
54    }
55
56    pub fn devices(&self) -> &[TensorDevice] {
57        &self.devices
58    }
59
60    pub(crate) fn const_ptrs(&self) -> Vec<*const ()> {
61        self.buffers
62            .iter()
63            .map(|buffer| buffer.as_ptr().cast())
64            .collect()
65    }
66
67    pub(crate) fn mut_ptrs(&mut self) -> Vec<*mut ()> {
68        self.buffers
69            .iter()
70            .map(|buffer| buffer.as_mut_ptr().cast())
71            .collect()
72    }
73}
74
75impl WorkspaceMemory {
76    pub fn create(context: &ContextRef, workspace: &Workspace) -> Result<Self> {
77        workspace.validate_for_context(context)?;
78
79        let device_sizes = workspace.device_sizes_bytes()?;
80        let mut device = Vec::with_capacity(context.device_count());
81        for (&device_id, &size) in context.devices().iter().zip(&device_sizes) {
82            Device::new(device_id).set_current()?;
83            device.push(DeviceMemory::<u8>::create(size)?);
84        }
85
86        Ok(Self {
87            device,
88            host: HostWorkspace::create(workspace.host_size_bytes()?)?,
89        })
90    }
91
92    pub(crate) fn device_mut_ptrs(&mut self) -> Vec<*mut ()> {
93        self.device
94            .iter()
95            .map(|buffer| buffer.as_mut_ptr().cast())
96            .collect()
97    }
98
99    pub(crate) fn device_pointer_count(&self) -> usize {
100        self.device.len()
101    }
102
103    pub(crate) fn host_mut_ptr(&self) -> *mut () {
104        self.host.ptr
105    }
106}
107
108impl HostWorkspace {
109    fn create(size: usize) -> singe_cuda::error::Result<Self> {
110        let ptr = if size == 0 {
111            ptr::null_mut()
112        } else {
113            unsafe { DeviceMemory::<u8>::alloc_host(size)? }
114        };
115        Ok(Self { ptr })
116    }
117}
118
119impl Drop for HostWorkspace {
120    fn drop(&mut self) {
121        if !self.ptr.is_null()
122            && let Err(err) = unsafe { DeviceMemory::<u8>::free_host(self.ptr) }
123        {
124            #[cfg(debug_assertions)]
125            eprintln!("failed to free cutensormg host workspace: {err}");
126        }
127    }
128}
singe_cutensor/mg/memory.rs

singe_cutensor/mg/
memory.rs
