Skip to main content

baracuda_driver/
pointer.rs

1//! Pointer attribute queries (`cuPointerGetAttribute`).
2
3use baracuda_cuda_sys::driver;
4use baracuda_cuda_sys::types::CUpointer_attribute;
5use baracuda_cuda_sys::CUdeviceptr;
6
7use crate::error::{check, Result};
8
9/// Raw pointer-attribute query. `attribute` is one of
10/// [`baracuda_cuda_sys::types::CUpointer_attribute`]; the caller must
11/// provide a writable `out` of the correct size for that attribute.
12///
13/// # Safety
14///
15/// `out` must point to a buffer large enough to receive the attribute's
16/// value type (typically 4 or 8 bytes). See NVIDIA's driver API reference
17/// for per-attribute size.
18pub unsafe fn raw_attribute(
19    attribute: i32,
20    ptr: CUdeviceptr,
21    out: *mut core::ffi::c_void,
22) -> Result<()> { unsafe {
23    let d = driver()?;
24    let cu = d.cu_pointer_get_attribute()?;
25    check(cu(out, attribute, ptr))
26}}
27
28/// Memory "kind" returned by `CUpointer_attribute::MEMORY_TYPE`.
29#[derive(Copy, Clone, Debug, Eq, PartialEq)]
30pub enum MemoryType {
31    Host,
32    Device,
33    Array,
34    Unified,
35    /// Unrecognized value; includes the raw code for forward compatibility.
36    Unknown(u32),
37}
38
39impl MemoryType {
40    #[inline]
41    fn from_raw(raw: u32) -> Self {
42        use baracuda_cuda_sys::types::CUmemorytype;
43        match raw {
44            CUmemorytype::HOST => MemoryType::Host,
45            CUmemorytype::DEVICE => MemoryType::Device,
46            CUmemorytype::ARRAY => MemoryType::Array,
47            CUmemorytype::UNIFIED => MemoryType::Unified,
48            other => MemoryType::Unknown(other),
49        }
50    }
51}
52
53/// Query the memory type of a device pointer.
54pub fn memory_type(ptr: CUdeviceptr) -> Result<MemoryType> {
55    let mut raw: u32 = 0;
56    // SAFETY: `raw` is a `u32` (4 bytes) which matches the attribute's size.
57    unsafe {
58        raw_attribute(
59            CUpointer_attribute::MEMORY_TYPE,
60            ptr,
61            &mut raw as *mut u32 as *mut core::ffi::c_void,
62        )?;
63    }
64    Ok(MemoryType::from_raw(raw))
65}
66
67/// Query whether a pointer refers to managed (unified) memory.
68pub fn is_managed(ptr: CUdeviceptr) -> Result<bool> {
69    let mut raw: u32 = 0;
70    unsafe {
71        raw_attribute(
72            CUpointer_attribute::IS_MANAGED,
73            ptr,
74            &mut raw as *mut u32 as *mut core::ffi::c_void,
75        )?;
76    }
77    Ok(raw != 0)
78}
79
80/// Query the device ordinal this allocation was created on.
81pub fn device_ordinal(ptr: CUdeviceptr) -> Result<i32> {
82    let mut raw: i32 = 0;
83    unsafe {
84        raw_attribute(
85            CUpointer_attribute::DEVICE_ORDINAL,
86            ptr,
87            &mut raw as *mut i32 as *mut core::ffi::c_void,
88        )?;
89    }
90    Ok(raw)
91}
92
93/// Query the size (bytes) of the range this pointer sits inside.
94pub fn range_size(ptr: CUdeviceptr) -> Result<usize> {
95    let mut raw: usize = 0;
96    unsafe {
97        raw_attribute(
98            CUpointer_attribute::RANGE_SIZE,
99            ptr,
100            &mut raw as *mut usize as *mut core::ffi::c_void,
101        )?;
102    }
103    Ok(raw)
104}
105
106/// Batched pointer-attribute query. For each `attribute` in `attributes`
107/// there must be a matching writable slot in `data` sized for that
108/// attribute's value type.
109///
110/// # Safety
111///
112/// Each `data[i]` must point to a buffer large enough to receive
113/// `attributes[i]`'s value type.
114pub unsafe fn raw_attributes_batched(
115    attributes: &mut [i32],
116    data: &mut [*mut core::ffi::c_void],
117    ptr: CUdeviceptr,
118) -> Result<()> { unsafe {
119    assert_eq!(
120        attributes.len(),
121        data.len(),
122        "attributes / data length mismatch"
123    );
124    let d = driver()?;
125    let cu = d.cu_pointer_get_attributes()?;
126    check(cu(
127        attributes.len() as core::ffi::c_uint,
128        attributes.as_mut_ptr(),
129        data.as_mut_ptr(),
130        ptr,
131    ))
132}}
133
134/// Query a single attribute on a managed-memory range.
135///
136/// # Safety
137///
138/// `out` must point to a buffer of `data_size` bytes sized for the
139/// attribute — consult the NVIDIA docs for
140/// [`CUmem_range_attribute`](baracuda_cuda_sys::types).
141pub unsafe fn range_attribute_raw(
142    attribute: i32,
143    ptr: CUdeviceptr,
144    count: usize,
145    out: *mut core::ffi::c_void,
146    data_size: usize,
147) -> Result<()> { unsafe {
148    let d = driver()?;
149    let cu = d.cu_mem_range_get_attribute()?;
150    check(cu(out, data_size, attribute, ptr, count))
151}}
152
153/// Batched range-attribute query. `data[i]` has size `data_sizes[i]`.
154///
155/// # Safety
156///
157/// Caller guarantees slot sizes match the attribute types.
158pub unsafe fn range_attributes_batched(
159    attributes: &mut [i32],
160    data: &mut [*mut core::ffi::c_void],
161    data_sizes: &mut [usize],
162    ptr: CUdeviceptr,
163    count: usize,
164) -> Result<()> { unsafe {
165    assert_eq!(attributes.len(), data.len());
166    assert_eq!(attributes.len(), data_sizes.len());
167    let d = driver()?;
168    let cu = d.cu_mem_range_get_attributes()?;
169    check(cu(
170        data.as_mut_ptr(),
171        data_sizes.as_mut_ptr(),
172        attributes.as_mut_ptr(),
173        attributes.len(),
174        ptr,
175        count,
176    ))
177}}
178
179/// Set a single pointer attribute (typically `SYNC_MEMOPS`).
180///
181/// # Safety
182///
183/// `value` must point to a valid payload matching the attribute.
184pub unsafe fn set_attribute_raw(
185    value: *const core::ffi::c_void,
186    attribute: i32,
187    ptr: CUdeviceptr,
188) -> Result<()> { unsafe {
189    let d = driver()?;
190    let cu = d.cu_pointer_set_attribute()?;
191    check(cu(value, attribute, ptr))
192}}