1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
// SPDX-License-Identifier: MIT OR Apache-2.0
use std::os::raw::c_uint;
use crate::{context::CudaContext, error::CudaError, memory::CudaDeviceBuffer};
impl CudaContext {
fn validate_memset_target(
&self,
dst: &CudaDeviceBuffer,
required: usize,
) -> Result<bool, CudaError> {
if !dst.is_owned_by(self) {
return Err(CudaError::InvalidArgument {
message: "CUDA memset target must belong to the launch context".to_string(),
});
}
if required > dst.byte_len() {
return Err(CudaError::OutputTooSmall {
required,
have: dst.byte_len(),
});
}
Ok(required != 0)
}
/// Fill the first `bytes` of a device buffer with one byte value.
///
/// The destination must belong to this context. A zero-length fill is a
/// validated no-op.
///
/// # Errors
///
/// Returns [`CudaError::InvalidArgument`] when `dst` belongs to another
/// context, [`CudaError::OutputTooSmall`] when `bytes` exceeds its length,
/// or the underlying CUDA driver failure.
pub fn memset_d8(
&self,
dst: &CudaDeviceBuffer,
value: u8,
bytes: usize,
) -> Result<(), CudaError> {
if !self.validate_memset_target(dst, bytes)? {
return Ok(());
}
self.inner.with_current_resource_operation(|| {
// SAFETY: `dst` is a live CUDA allocation in this context, `bytes`
// was bounds-checked, and the context lifecycle gate is held.
self.inner.driver.check("cuMemsetD8_v2", unsafe {
(self.inner.driver.cu_memset_d8)(dst.device_ptr(), value, bytes)
})
})
}
/// Fill the first `words` 32-bit words of a device buffer synchronously.
///
/// This codec-neutral primitive is exposed for engine crates. The
/// destination must belong to this context.
#[doc(hidden)]
pub fn memset_d32(
&self,
dst: &CudaDeviceBuffer,
value: c_uint,
words: usize,
) -> Result<(), CudaError> {
let required = words
.checked_mul(std::mem::size_of::<u32>())
.ok_or(CudaError::LengthTooLarge { len: words })?;
if !self.validate_memset_target(dst, required)? {
return Ok(());
}
self.inner.with_current_resource_operation(|| {
// SAFETY: `dst` is a live CUDA allocation in this context, `words`
// was bounds-checked, and the context lifecycle gate is held.
self.inner.driver.check("cuMemsetD32_v2", unsafe {
(self.inner.driver.cu_memset_d32)(dst.device_ptr(), value, words)
})
})
}
/// Fill the first `words` 32-bit words of a device buffer on the default stream.
///
/// The destination must belong to this context. Submitted work is ordered
/// with later default-stream launches; callers retain the destination until
/// stream completion.
#[doc(hidden)]
pub fn memset_d32_async(
&self,
dst: &CudaDeviceBuffer,
value: c_uint,
words: usize,
) -> Result<(), CudaError> {
let required = words
.checked_mul(std::mem::size_of::<u32>())
.ok_or(CudaError::LengthTooLarge { len: words })?;
if !self.validate_memset_target(dst, required)? {
return Ok(());
}
self.inner.with_current_resource_operation(|| {
// SAFETY: the destination is live and bounds-checked. A null
// stream orders the memset with the codec's default-stream graph.
self.inner.driver.check("cuMemsetD32Async", unsafe {
(self.inner.driver.cu_memset_d32_async)(
dst.device_ptr(),
value,
words,
std::ptr::null_mut(),
)
})
})
}
}