cubecl_hip/compute/storage/
gpu.rs

1use crate::compute::uninit_vec;
2use cubecl_common::backtrace::BackTrace;
3use cubecl_core::server::IoError;
4use cubecl_hip_sys::HIP_SUCCESS;
5use cubecl_runtime::storage::{ComputeStorage, StorageHandle, StorageId, StorageUtilization};
6use std::collections::HashMap;
7
8/// Buffer storage for AMD GPUs.
9///
10/// This struct manages memory resources for HIP kernels, allowing them to be used as bindings
11/// for launching kernels.
12pub struct GpuStorage {
13    mem_alignment: usize,
14    memory: HashMap<StorageId, cubecl_hip_sys::hipDeviceptr_t>,
15    deallocations: Vec<StorageId>,
16    ptr_bindings: PtrBindings,
17}
18
19/// A GPU memory resource allocated for HIP using [GpuStorage].
20#[derive(new, Debug)]
21pub struct GpuResource {
22    /// The GPU memory pointer.
23    pub ptr: cubecl_hip_sys::hipDeviceptr_t,
24    /// The HIP binding pointer.
25    pub binding: cubecl_hip_sys::hipDeviceptr_t,
26    /// The size of the resource.
27    pub size: u64,
28}
29
30impl GpuStorage {
31    /// Creates a new [GpuStorage] instance for the specified HIP stream.
32    ///
33    /// # Arguments
34    ///
35    /// * `mem_alignment` - The memory alignment requirement in bytes.
36    pub fn new(mem_alignment: usize) -> Self {
37        Self {
38            mem_alignment,
39            memory: HashMap::new(),
40            deallocations: Vec::new(),
41            ptr_bindings: PtrBindings::new(),
42        }
43    }
44
45    /// Deallocates buffers marked for deallocation.
46    ///
47    /// This method processes all pending deallocations by freeing the associated GPU memory.
48    pub fn perform_deallocations(&mut self) {
49        for id in self.deallocations.drain(..) {
50            if let Some(ptr) = self.memory.remove(&id) {
51                unsafe {
52                    cubecl_hip_sys::hipFree(ptr);
53                }
54            }
55        }
56    }
57}
58
59/// Manages active HIP buffer bindings in a ring buffer.
60///
61/// This ensures that pointers remain valid during kernel execution, preventing use-after-free errors.
62struct PtrBindings {
63    slots: Vec<u64>,
64    cursor: usize,
65}
66
67impl PtrBindings {
68    /// Creates a new [PtrBindings] instance with a fixed-size ring buffer.
69    fn new() -> Self {
70        Self {
71            slots: uninit_vec(crate::device::AMD_MAX_BINDINGS as usize),
72            cursor: 0,
73        }
74    }
75
76    /// Registers a new pointer in the ring buffer.
77    ///
78    /// # Arguments
79    ///
80    /// * `ptr` - The HIP device pointer to register.
81    ///
82    /// # Returns
83    ///
84    /// A reference to the registered pointer.
85    fn register(&mut self, ptr: u64) -> &u64 {
86        self.slots[self.cursor] = ptr;
87        let ptr_ref = self.slots.get(self.cursor).unwrap();
88
89        self.cursor += 1;
90
91        // Reset the cursor when the ring buffer is full.
92        if self.cursor >= self.slots.len() {
93            self.cursor = 0;
94        }
95
96        ptr_ref
97    }
98}
99
100impl ComputeStorage for GpuStorage {
101    type Resource = GpuResource;
102
103    fn alignment(&self) -> usize {
104        self.mem_alignment
105    }
106
107    fn get(&mut self, handle: &StorageHandle) -> Self::Resource {
108        let ptr = (*self.memory.get(&handle.id).unwrap()) as u64;
109
110        let offset = handle.offset();
111        let size = handle.size();
112        let ptr = self.ptr_bindings.register(ptr + offset);
113
114        GpuResource::new(
115            *ptr as cubecl_hip_sys::hipDeviceptr_t,
116            std::ptr::from_ref(ptr) as *mut std::ffi::c_void,
117            size,
118        )
119    }
120
121    #[cfg_attr(
122        feature = "tracing",
123        tracing::instrument(level = "trace", skip(self, size))
124    )]
125    fn alloc(&mut self, size: u64) -> Result<StorageHandle, IoError> {
126        let id = StorageId::new();
127        unsafe {
128            let mut dptr: *mut ::std::os::raw::c_void = std::ptr::null_mut();
129            let status = cubecl_hip_sys::hipMalloc(&mut dptr, size as usize);
130
131            match status {
132                HIP_SUCCESS => {}
133                other => {
134                    return Err(IoError::Unknown {
135                        description: format!("HIP allocation error: {other}"),
136                        backtrace: BackTrace::capture(),
137                    });
138                }
139            }
140            self.memory.insert(id, dptr);
141        };
142        Ok(StorageHandle::new(
143            id,
144            StorageUtilization { offset: 0, size },
145        ))
146    }
147
148    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
149    fn dealloc(&mut self, id: StorageId) {
150        self.deallocations.push(id);
151    }
152}
153
154unsafe impl Send for GpuStorage {}
155unsafe impl Send for GpuResource {}
156
157impl core::fmt::Debug for GpuStorage {
158    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
159        f.write_str("GpuStorage".to_string().as_str())
160    }
161}
cubecl_hip/compute/storage/gpu.rs

cubecl_hip/compute/storage/
gpu.rs
