singe_cutensor/mp/
memory.rs

1use std::ptr;
2
3use singe_cuda::memory::DeviceMemory;
4
5use crate::{error::Result, mp::plan::Workspace, utility::to_usize};
6
7#[derive(Debug)]
8pub struct WorkspaceMemory {
9    device: DeviceMemory<u8>,
10    host: HostWorkspace,
11}
12
13#[derive(Debug)]
14struct HostWorkspace {
15    ptr: *mut (),
16}
17
18impl WorkspaceMemory {
19    pub fn create(workspace: Workspace) -> Result<Self> {
20        Ok(Self {
21            device: DeviceMemory::create(to_usize(workspace.device_size(), "device workspace")?)?,
22            host: HostWorkspace::create(to_usize(workspace.host_size(), "host workspace")?)?,
23        })
24    }
25
26    pub(crate) fn device_mut_ptr(&mut self) -> *mut () {
27        self.device.as_mut_ptr().cast()
28    }
29
30    pub(crate) fn host_mut_ptr(&self) -> *mut () {
31        self.host.ptr
32    }
33}
34
35impl HostWorkspace {
36    fn create(size: usize) -> singe_cuda::error::Result<Self> {
37        let ptr = if size == 0 {
38            ptr::null_mut()
39        } else {
40            unsafe { DeviceMemory::<u8>::alloc_host(size)? }
41        };
42        Ok(Self { ptr })
43    }
44}
45
46impl Drop for HostWorkspace {
47    fn drop(&mut self) {
48        if !self.ptr.is_null()
49            && let Err(err) = unsafe { DeviceMemory::<u8>::free_host(self.ptr) }
50        {
51            #[cfg(debug_assertions)]
52            eprintln!("failed to free cutensormp host workspace: {err}");
53        }
54    }
55}
singe_cutensor/mp/memory.rs

singe_cutensor/mp/
memory.rs
