Skip to main content

sp1_gpu_sys/
runtime.rs

1use std::ffi::{c_char, c_void};
2
3#[derive(Clone, Copy, PartialEq, Eq)]
4#[repr(C)]
5pub struct CudaRustError {
6    pub message: *const c_char,
7}
8
9extern "C" {
10    pub static CUDA_SUCCESS_CSL: CudaRustError;
11
12    pub static CUDA_OUT_OF_MEMORY: CudaRustError;
13
14    pub static CUDA_ERROR_NOT_READY_SLOP: CudaRustError;
15
16    pub fn cuda_malloc(ptr: *mut *mut c_void, count: usize) -> CudaRustError;
17
18    pub fn cuda_free(ptr: *const c_void) -> CudaRustError;
19
20    pub fn cuda_mem_get_info(free: *mut usize, total: *mut usize) -> CudaRustError;
21
22    pub fn cuda_get_device_name(name: *mut c_char, len: usize) -> CudaRustError;
23
24    pub fn cuda_malloc_host(ptr: *mut *mut c_void, count: usize) -> CudaRustError;
25    pub fn cuda_host_register(ptr: *const c_void, count: usize) -> CudaRustError;
26    pub fn cuda_free_host(ptr: *const c_void) -> CudaRustError;
27    pub fn cuda_host_unregister(ptr: *const c_void) -> CudaRustError;
28
29    pub fn cuda_mem_set(dst: *mut c_void, value: u8, size: usize) -> CudaRustError;
30
31    pub fn cuda_mem_copy_host_to_device(
32        dst: *mut c_void,
33        src: *const c_void,
34        count: usize,
35    ) -> CudaRustError;
36
37    pub fn cuda_mem_copy_device_to_host(
38        dst: *mut c_void,
39        src: *const c_void,
40        count: usize,
41    ) -> CudaRustError;
42
43    pub fn cuda_mem_copy_device_to_device(
44        dst: *const c_void,
45        src: *const c_void,
46        count: usize,
47    ) -> CudaRustError;
48}
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
51#[repr(transparent)]
52pub struct CudaStreamHandle(pub *mut c_void);
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
55#[repr(transparent)]
56pub struct CudaEventHandle(pub *mut c_void);
57
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
59#[repr(C)]
60pub struct Dim3 {
61    pub x: u32,
62    pub y: u32,
63    pub z: u32,
64}
65
66#[repr(transparent)]
67pub struct KernelPtr(pub *const c_void);
68
69#[repr(transparent)]
70pub struct CudaMemPool(pub *mut c_void);
71
72#[repr(transparent)]
73pub struct CudaDevice(pub i32);
74
75extern "C" {
76
77    pub static DEFAULT_STREAM: CudaStreamHandle;
78
79    pub fn cuda_device_synchronize() -> CudaRustError;
80    pub fn cuda_event_create(event: *mut CudaEventHandle) -> CudaRustError;
81    pub fn cuda_event_destroy(event: CudaEventHandle) -> CudaRustError;
82    pub fn cuda_event_record(event: CudaEventHandle, stream: CudaStreamHandle) -> CudaRustError;
83    pub fn cuda_event_synchronize(event: CudaEventHandle) -> CudaRustError;
84    pub fn cuda_event_elapsed_time(
85        ms: *mut f32,
86        start: CudaEventHandle,
87        end: CudaEventHandle,
88    ) -> CudaRustError;
89
90    pub fn cuda_stream_create(stream: *mut CudaStreamHandle) -> CudaRustError;
91    pub fn cuda_stream_destroy(stream: CudaStreamHandle) -> CudaRustError;
92    pub fn cuda_stream_synchronize(stream: CudaStreamHandle) -> CudaRustError;
93
94    pub fn cuda_stream_wait_event(
95        stream: CudaStreamHandle,
96        event: CudaEventHandle,
97    ) -> CudaRustError;
98
99    // Async memory operations.
100
101    pub fn cuda_malloc_async(
102        devPtr: *mut *mut c_void,
103        size: usize,
104        stream: CudaStreamHandle,
105    ) -> CudaRustError;
106
107    pub fn cuda_mem_set_async(
108        dst: *mut c_void,
109        value: u8,
110        size: usize,
111        stream: CudaStreamHandle,
112    ) -> CudaRustError;
113
114    pub fn cuda_free_async(devPtr: *mut c_void, stream: CudaStreamHandle) -> CudaRustError;
115
116    pub fn cuda_mem_copy_device_to_device_async(
117        dst: *mut c_void,
118        src: *const c_void,
119        count: usize,
120        stream: CudaStreamHandle,
121    ) -> CudaRustError;
122    pub fn cuda_mem_copy_host_to_device_async(
123        dst: *mut c_void,
124        src: *const c_void,
125        count: usize,
126        stream: CudaStreamHandle,
127    ) -> CudaRustError;
128    pub fn cuda_mem_copy_device_to_host_async(
129        dst: *mut c_void,
130        src: *const c_void,
131        count: usize,
132        stream: CudaStreamHandle,
133    ) -> CudaRustError;
134    pub fn cuda_mem_copy_host_to_host_async(
135        dst: *mut c_void,
136        src: *const c_void,
137        count: usize,
138        stream: CudaStreamHandle,
139    ) -> CudaRustError;
140
141    pub fn cuda_stream_query(stream: CudaStreamHandle) -> CudaRustError;
142
143    pub fn cuda_event_query(event: CudaEventHandle) -> CudaRustError;
144
145    pub fn cuda_launch_host_function(
146        stream: CudaStreamHandle,
147        host_fn: Option<unsafe extern "C" fn(*mut c_void)>,
148        data: *const c_void,
149    ) -> CudaRustError;
150
151    pub fn cuda_launch_kernel(
152        kernel: KernelPtr,
153        grid: Dim3,
154        block: Dim3,
155        args: *mut *mut c_void,
156        shared_mem: usize,
157        stream: CudaStreamHandle,
158    ) -> CudaRustError;
159
160    pub fn cuda_device_get_default_mem_pool(
161        memPool: *mut CudaMemPool,
162        device: CudaDevice,
163    ) -> CudaRustError;
164
165    pub fn cuda_device_get_mem_pool(memPool: *mut CudaMemPool, device: CudaDevice)
166        -> CudaRustError;
167
168    pub fn cuda_mem_pool_set_release_threshold(
169        memPool: CudaMemPool,
170        threshold: u64,
171    ) -> CudaRustError;
172}
173
174#[derive(Debug, Clone, Copy)]
175#[repr(transparent)]
176pub struct NvtxRangeId(u64);
177
178extern "C" {
179    pub fn nvtx_range_start(name: *const c_char) -> NvtxRangeId;
180
181    pub fn nvtx_range_end(domain: NvtxRangeId);
182}
183
184impl Dim3 {
185    pub fn new(x: u32, y: u32, z: u32) -> Self {
186        Self { x, y, z }
187    }
188
189    pub fn x(num_elements: u32) -> Self {
190        Self { x: num_elements, y: 1, z: 1 }
191    }
192}
193
194impl From<u32> for Dim3 {
195    fn from(x: u32) -> Self {
196        Self { x, y: 1, z: 1 }
197    }
198}
199
200impl From<u64> for Dim3 {
201    fn from(x: u64) -> Self {
202        Self { x: x as u32, y: 1, z: 1 }
203    }
204}
205
206impl From<i32> for Dim3 {
207    fn from(x: i32) -> Self {
208        Self { x: x as u32, y: 1, z: 1 }
209    }
210}
211
212impl From<i64> for Dim3 {
213    fn from(x: i64) -> Self {
214        Self { x: x as u32, y: 1, z: 1 }
215    }
216}
217
218impl From<usize> for Dim3 {
219    fn from(x: usize) -> Self {
220        Self { x: x as u32, y: 1, z: 1 }
221    }
222}
223
224impl From<(u32, u32, u32)> for Dim3 {
225    fn from((x, y, z): (u32, u32, u32)) -> Self {
226        Self { x, y, z }
227    }
228}
229
230impl From<(u64, u64, u64)> for Dim3 {
231    fn from((x, y, z): (u64, u64, u64)) -> Self {
232        Self { x: x as u32, y: y as u32, z: z as u32 }
233    }
234}
235
236impl From<(usize, usize, usize)> for Dim3 {
237    fn from((x, y, z): (usize, usize, usize)) -> Self {
238        Self { x: x as u32, y: y as u32, z: z as u32 }
239    }
240}