use std::ffi::{CStr, CString, c_void};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
pub use crate::sys::ExportInfo;
use crate::{Error, Result, TARGET_FAMILY, Target, sys};
pub(crate) fn check(status: sys::Status, what: impl std::fmt::Display) -> Result<()> {
if sys::is_ok(status) {
return Ok(());
}
let code = unsafe { sys::hrx_status_code(status) };
let mut message: *mut std::ffi::c_char = std::ptr::null_mut();
let mut length: usize = 0;
let text = unsafe {
let to_string = sys::hrx_status_to_string(status, &mut message, &mut length);
let text = if sys::is_ok(to_string) && !message.is_null() {
let owned = CStr::from_ptr(message).to_string_lossy().into_owned();
sys::hrx_status_free_message(message);
owned
} else {
sys::hrx_status_ignore(to_string);
"unknown error".to_string()
};
sys::hrx_status_ignore(status);
text
};
Err(Error::Runtime {
context: what.to_string(),
code,
message: text,
})
}
struct Inner {
device: sys::Device,
target: Target,
stream: sys::Stream,
}
impl Drop for Inner {
fn drop(&mut self) {
unsafe { sys::hrx_stream_release(self.stream) };
}
}
unsafe impl Send for Inner {}
unsafe impl Sync for Inner {}
pub(crate) fn initialize_runtime() -> Result<()> {
sys::load()?;
static READY: std::sync::Mutex<Option<()>> = std::sync::Mutex::new(None);
crate::cached_init(&READY, || {
let _lock = sys::runtime_lock()?;
unsafe {
let status = sys::hrx_gpu_initialize(0);
if sys::hrx_status_code(status) == sys::STATUS_ALREADY_EXISTS {
sys::hrx_status_ignore(status);
} else {
check(status, "hrx_gpu_initialize")?;
}
}
Ok(())
})
}
pub(crate) fn open_device(index: i32) -> Result<(sys::Device, Target)> {
initialize_runtime()?;
unsafe {
let mut count = 0;
check(
sys::hrx_gpu_device_count(&mut count),
"hrx_gpu_device_count",
)?;
if index < 0 || index >= count {
return Err(Error::Message(
"no GPU device (check GPU access and the README provisioning instructions)".into(),
));
}
let mut device = std::ptr::null_mut();
check(
sys::hrx_gpu_device_get(index, &mut device),
"hrx_gpu_device_get",
)?;
let mut architecture = [0u8; 64];
check(
sys::hrx_device_get_property(
device,
sys::DEVICE_PROPERTY_ARCHITECTURE,
architecture.as_mut_ptr().cast(),
architecture.len(),
),
"device architecture",
)?;
let key = std::str::from_utf8(architecture.split(|b| *b == 0).next().unwrap())
.map_err(|_| Error::Message("invalid device architecture".into()))?;
Ok((device, Target::from_device_architecture(key)?))
}
}
impl Stream {
fn new(index: i32) -> Result<Self> {
let (device, target) = open_device(index)?;
unsafe {
let mut stream = std::ptr::null_mut();
check(
sys::hrx_stream_create(device, 0, &mut stream),
"hrx_stream_create",
)?;
Ok(Self {
inner: std::sync::Arc::new(Inner {
device,
target,
stream,
}),
_not_sync: std::marker::PhantomData,
staging: Vec::new(),
staging_pool: Vec::new(),
scratch: std::collections::BTreeMap::new(),
scratch_bytes: 0,
scratch_limit: 256 * 1024 * 1024,
})
}
}
fn owns(&self, buffer: &Buffer) -> Result<()> {
owns(&self.inner, buffer)
}
fn synchronize_native(&self) -> Result<()> {
unsafe {
check(
sys::hrx_stream_synchronize(self.inner.stream),
"hrx_stream_synchronize",
)
}
}
unsafe fn loaded_export(&self, executable: sys::Executable, symbol: &str) -> Result<Kernel> {
let kernel = (|| {
let symbol_c =
CString::new(symbol).map_err(|_| Error::Message("export contains a NUL".into()))?;
let mut ordinal = 0;
let mut info = sys::ExportInfo::default();
unsafe {
check(
sys::hrx_executable_lookup_export_by_name(
executable,
symbol_c.as_ptr(),
&mut ordinal,
),
"looking up export",
)?;
check(
sys::hrx_executable_export_info(executable, ordinal, &mut info),
"export metadata",
)?;
}
Ok(Kernel {
executable,
ordinal,
info,
symbol: symbol.into(),
_device: self.inner.clone(),
})
})();
if kernel.is_err() {
unsafe { sys::hrx_executable_release(executable) }
}
kernel
}
}
pub struct Buffer {
raw: sys::Buffer,
bytes: usize,
shared: AtomicBool,
_device: std::sync::Arc<Inner>,
}
unsafe impl Send for Buffer {}
unsafe impl Sync for Buffer {}
#[derive(Clone, Copy, Debug)]
pub struct View<'a> {
raw: sys::BufferRef,
owner: &'a Buffer,
}
impl<'a> View<'a> {
fn new(raw: sys::BufferRef, owner: &'a Buffer) -> Self {
Self { raw, owner }
}
pub fn slice(self, offset: usize, length: usize) -> Result<Self> {
checked_span(offset, length, self.len())?;
Ok(Self::new(
sys::BufferRef {
offset: self.raw.offset + offset,
length,
..self.raw
},
self.owner,
))
}
#[must_use]
pub fn offset(self) -> usize {
self.raw.offset
}
#[must_use]
pub fn owner(self) -> &'a Buffer {
self.owner
}
#[must_use]
pub fn len(self) -> usize {
self.raw.length
}
#[must_use]
pub fn is_empty(self) -> bool {
self.raw.length == 0
}
}
impl Buffer {
pub unsafe fn share_on(&self, stream: &Stream) -> Result<Buffer> {
if self._device.device != stream.inner.device {
return Err(Error::Message("buffer belongs to another device".into()));
}
unsafe { sys::hrx_buffer_retain(self.raw) };
self.shared.store(true, Ordering::Relaxed);
Ok(Buffer {
raw: self.raw,
bytes: self.bytes,
shared: AtomicBool::new(true),
_device: stream.inner.clone(),
})
}
pub fn bytes(&self) -> usize {
self.bytes
}
pub fn try_slice(&self, offset: usize, length: usize) -> Result<View<'_>> {
self.binding().slice(offset, length)
}
pub fn binding(&self) -> View<'_> {
View::new(
sys::BufferRef {
buffer: self.raw,
offset: 0,
length: self.bytes,
},
self,
)
}
pub fn slice(&self, offset: usize, length: usize) -> View<'_> {
self.try_slice(offset, length)
.expect("slice past the allocation or overflow")
}
}
impl Drop for Buffer {
fn drop(&mut self) {
unsafe { sys::hrx_buffer_release(self.raw) }
}
}
pub struct Kernel {
executable: sys::Executable,
ordinal: u32,
info: sys::ExportInfo,
symbol: String,
_device: std::sync::Arc<Inner>,
}
unsafe impl Send for Kernel {}
unsafe impl Sync for Kernel {}
impl Kernel {
pub fn info(&self) -> &ExportInfo {
&self.info
}
pub fn symbol(&self) -> &str {
&self.symbol
}
}
impl Drop for Kernel {
fn drop(&mut self) {
unsafe { sys::hrx_executable_release(self.executable) }
}
}
#[cfg(test)]
mod tests {
#[test]
fn compiled_workgroup_dimensions_are_enforced() {
let info = super::ExportInfo {
workgroup_size: [64, 2, 1],
..Default::default()
};
assert!(super::validate_export_launch(&info, [1; 3], [64, 2, 1]).is_ok());
assert!(super::validate_export_launch(&info, [1; 3], [128, 1, 1]).is_err());
assert!(super::validate_export_launch(&info, [0, 1, 1], [64, 2, 1]).is_err());
let dynamic = super::ExportInfo {
workgroup_size: [0, 2, 1],
..Default::default()
};
assert!(super::validate_export_launch(&dynamic, [1; 3], [32, 2, 1]).is_ok());
assert!(super::validate_export_launch(&dynamic, [1; 3], [32, 1, 1]).is_err());
}
#[test]
fn a_slice_past_the_allocation_is_refused_even_when_the_sum_wraps() {
let checked = |offset, length, bytes| super::checked_span(offset, length, bytes).is_ok();
assert!(checked(0, 8, 8));
assert!(checked(4, 4, 8));
assert!(!checked(4, 5, 8));
assert!(!checked(usize::MAX, 2, 8));
assert!(!checked(usize::MAX - 1, 4, 8));
}
}
fn owns(inner: &std::sync::Arc<Inner>, buffer: &Buffer) -> Result<()> {
if std::sync::Arc::ptr_eq(inner, &buffer._device) {
Ok(())
} else {
Err(Error::Message("buffer belongs to another stream".into()))
}
}
fn raw_bindings<'a>(
views: &[View<'_>],
stack: &'a mut [std::mem::MaybeUninit<sys::BufferRef>; 32],
) -> std::borrow::Cow<'a, [sys::BufferRef]> {
if views.len() <= stack.len() {
for (out, view) in stack.iter_mut().zip(views) {
out.write(view.raw);
}
std::borrow::Cow::Borrowed(unsafe {
std::slice::from_raw_parts(stack.as_ptr().cast(), views.len())
})
} else {
std::borrow::Cow::Owned(views.iter().map(|v| v.raw).collect())
}
}
pub(crate) fn checked_span(offset: usize, length: usize, capacity: usize) -> Result<()> {
if offset.checked_add(length).is_none_or(|end| end > capacity) {
Err(Error::Message(format!(
"span {offset}+{length} exceeds {capacity} bytes"
)))
} else {
Ok(())
}
}
pub(crate) fn validate_launch(grid: [u32; 3], block: [u32; 3]) -> Result<()> {
if grid.contains(&0)
|| block.contains(&0)
|| block
.iter()
.try_fold(1u32, |n, b| n.checked_mul(*b))
.is_none_or(|n| n > 1024)
{
Err(Error::Message("invalid GPU launch dimensions".into()))
} else {
Ok(())
}
}
pub(crate) fn validate_export_launch(
info: &sys::ExportInfo,
grid: [u32; 3],
block: [u32; 3],
) -> Result<()> {
validate_launch(grid, block)?;
if info
.workgroup_size
.iter()
.zip(block)
.any(|(&expected, actual)| expected != 0 && expected != actual)
{
return Err(Error::Message(format!(
"workgroup size {block:?} disagrees with compiled size {:?}",
info.workgroup_size
)));
}
Ok(())
}
#[derive(Clone, Debug)]
pub struct Device {
index: i32,
target: Target,
}
impl Device {
pub fn open(index: i32) -> Result<Self> {
let (_, target) = open_device(index)?;
Ok(Self { index, target })
}
pub fn target(&self) -> &Target {
&self.target
}
pub fn stream(&self) -> Result<Stream> {
Stream::new(self.index)
}
}
pub struct Stream {
inner: std::sync::Arc<Inner>,
_not_sync: std::marker::PhantomData<std::cell::Cell<()>>,
staging: Vec<Buffer>,
staging_pool: Vec<Buffer>,
scratch: std::collections::BTreeMap<usize, Vec<Buffer>>,
scratch_bytes: usize,
scratch_limit: usize,
}
impl Stream {
pub fn target(&self) -> &Target {
&self.inner.target
}
pub fn record_event(&mut self) -> Result<Event> {
let mut raw = std::ptr::null_mut();
unsafe {
check(
sys::hrx_event_create(self.inner.device, sys::EVENT_FLAG_DISABLE_TIMING, &mut raw),
"create event",
)?;
let event = Event {
raw,
inner: self.inner.clone(),
};
check(
sys::hrx_event_record(raw, self.inner.stream),
"record event",
)?;
Ok(event)
}
}
pub fn wait_event(&mut self, event: &Event) -> Result<()> {
if self.inner.device != event.inner.device {
return Err(Error::Message("event belongs to another device".into()));
}
unsafe {
check(
sys::hrx_stream_wait_event(self.inner.stream, event.raw),
"wait event",
)
}
}
pub fn open() -> Result<Self> {
Device::open(0)?.stream()
}
pub fn allocate(&self, bytes: usize) -> Result<Buffer> {
let mut buffer = std::ptr::null_mut();
unsafe {
check(
sys::hrx_allocator_allocate_buffer(
sys::hrx_device_allocator(self.inner.device),
sys::BufferParams {
memory_type: sys::MEMORY_TYPE_DEVICE_LOCAL,
access: sys::MEMORY_ACCESS_ALL,
usage: sys::BUFFER_USAGE_DEFAULT,
queue_affinity: u64::MAX,
},
bytes.max(1),
&mut buffer,
),
"hrx_buffer_allocate",
)?;
}
Ok(Buffer {
raw: buffer,
bytes: bytes.max(1),
shared: AtomicBool::new(false),
_device: self.inner.clone(),
})
}
pub fn synchronize(&mut self) -> Result<()> {
self.synchronize_native()?;
self.reclaim_staging();
Ok(())
}
pub fn upload_blocking(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()> {
self.owns(dst.owner)?;
checked_span(0, bytes.len(), dst.len())?;
self.synchronize()?;
if bytes.is_empty() {
return Ok(());
}
unsafe {
check(
sys::hrx_synchronous_h2d(
self.inner.device,
bytes.as_ptr().cast(),
dst.raw.buffer,
dst.raw.offset,
bytes.len(),
),
"hrx_synchronous_h2d",
)
}
}
pub fn read(&mut self, src: View<'_>, bytes: &mut [u8]) -> Result<()> {
self.owns(src.owner)?;
checked_span(0, bytes.len(), src.len())?;
self.synchronize()?;
if bytes.is_empty() {
return Ok(());
}
unsafe {
check(
sys::hrx_synchronous_d2h(
self.inner.device,
src.raw.buffer,
src.raw.offset,
bytes.as_mut_ptr().cast(),
bytes.len(),
),
"hrx_synchronous_d2h",
)
}
}
pub fn fill(&mut self, dst: View<'_>, value: u8) -> Result<()> {
self.owns(dst.owner)?;
if dst.is_empty() {
return Err(Error::Message("empty stream fill".into()));
}
unsafe {
check(
sys::hrx_stream_fill_buffer(
self.inner.stream,
dst.raw.buffer,
dst.raw.offset,
dst.len(),
&value as *const u8 as *const c_void,
1,
),
"hrx_stream_fill_buffer",
)
}
}
pub fn copy(&mut self, dst: View<'_>, src: View<'_>) -> Result<()> {
self.owns(dst.owner)?;
self.owns(src.owner)?;
if dst.len() != src.len() || dst.is_empty() {
return Err(Error::Message(
"stream copy requires equal nonempty spans".into(),
));
}
unsafe {
check(
sys::hrx_stream_copy_buffer(
self.inner.stream,
src.raw.buffer,
src.raw.offset,
dst.raw.buffer,
dst.raw.offset,
src.len(),
),
"hrx_stream_copy_buffer",
)
}
}
pub fn upload(&mut self, dst: View<'_>, bytes: &[u8]) -> Result<()> {
self.owns(dst.owner)?;
checked_span(0, bytes.len(), dst.len())?;
if bytes.is_empty() {
return Ok(());
}
let next_capacity = self
.staging_pool
.iter()
.filter(|b| b.bytes >= bytes.len())
.map(|b| b.bytes)
.min()
.unwrap_or(bytes.len());
if !self.staging.is_empty()
&& (self.staging.len() >= 8
|| self
.staging
.iter()
.map(|b| b.bytes)
.sum::<usize>()
.saturating_add(next_capacity)
> STAGING_LIMIT)
&& !self.submit()?.is_complete()?
{
self.synchronize()?;
}
let staging = if let Some(i) = self
.staging_pool
.iter()
.enumerate()
.filter(|(_, b)| b.bytes >= bytes.len())
.min_by_key(|(_, b)| b.bytes)
.map(|(i, _)| i)
{
self.staging_pool.swap_remove(i)
} else {
self.allocate_host(bytes.len())?
};
let raw = staging.raw;
unsafe {
let mut pointer = std::ptr::null_mut();
check(
sys::hrx_buffer_get_device_ptr(raw, &mut pointer),
"map staging",
)?;
std::ptr::copy_nonoverlapping(bytes.as_ptr(), pointer.cast::<u8>(), bytes.len());
self.staging.push(staging);
check(
sys::hrx_stream_copy_buffer(
self.inner.stream,
raw,
0,
dst.raw.buffer,
dst.raw.offset,
bytes.len(),
),
"enqueue upload",
)
}
}
fn allocate_host(&self, bytes: usize) -> Result<Buffer> {
let mut raw = std::ptr::null_mut();
unsafe {
check(
sys::hrx_allocator_allocate_buffer(
sys::hrx_device_allocator(self.inner.device),
sys::BufferParams {
memory_type: sys::MEMORY_TYPE_HOST_LOCAL
| sys::MEMORY_TYPE_DEVICE_VISIBLE
| sys::MEMORY_TYPE_HOST_COHERENT,
access: sys::MEMORY_ACCESS_ALL,
usage: sys::BUFFER_USAGE_DEFAULT | sys::BUFFER_USAGE_MAPPING_SCOPED,
queue_affinity: u64::MAX,
},
bytes.max(1),
&mut raw,
),
"allocate staging",
)?;
}
Ok(Buffer {
raw,
bytes,
shared: AtomicBool::new(false),
_device: self.inner.clone(),
})
}
fn reclaim_staging(&mut self) {
let mut cached = self.staging_pool.iter().map(|b| b.bytes).sum::<usize>();
for buffer in self.staging.drain(..) {
if self.staging_pool.len() < 8 && buffer.bytes <= STAGING_LIMIT.saturating_sub(cached) {
cached += buffer.bytes;
self.staging_pool.push(buffer);
}
}
}
pub fn submit(&mut self) -> Result<Submission<'_>> {
unsafe {
check(sys::hrx_stream_flush(self.inner.stream), "submit")?;
}
Ok(Submission { stream: self })
}
pub unsafe fn load(&self, path: &Path, symbol: &str) -> Result<Kernel> {
let c_path = CString::new(path.as_os_str().as_encoded_bytes())
.map_err(|_| Error::Message(format!("{} contains a NUL", path.display())))?;
let c_family = TARGET_FAMILY;
let c_key = self.inner.target.as_c_str();
unsafe {
let mut executable = std::ptr::null_mut();
check(
sys::hrx_executable_load_file(
self.inner.device,
c_path.as_ptr(),
c_family.as_ptr(),
c_key.as_ptr(),
&mut executable,
),
format_args!("loading {}", path.display()),
)?;
self.loaded_export(executable, symbol)
}
}
#[cfg(feature = "loom")]
pub unsafe fn load_artifact(&self, artifact: &crate::loom::Artifact) -> Result<Kernel> {
if artifact.target() != self.inner.target.as_str() {
return Err(Error::Message(
"artifact target does not match this runtime".into(),
));
}
let bytes = artifact.bytes();
let mut executable = std::ptr::null_mut();
unsafe {
check(
sys::hrx_executable_load_data(
self.inner.device,
bytes.as_ptr().cast(),
bytes.len(),
TARGET_FAMILY.as_ptr(),
self.inner.target.as_c_str().as_ptr(),
&mut executable,
),
"loading compiled artifact",
)?;
self.loaded_export(executable, artifact.symbol())
}
}
pub unsafe fn dispatch(
&mut self,
kernel: &Kernel,
grid: [u32; 3],
block: [u32; 3],
constants: &Constants,
bindings: &[View<'_>],
) -> Result<()> {
if kernel._device.device != self.inner.device {
return Err(Error::Message("kernel belongs to another device".into()));
}
for view in bindings {
owns(&self.inner, view.owner)?;
}
let mut binding_storage = [std::mem::MaybeUninit::uninit(); 32];
let raw_bindings = raw_bindings(bindings, &mut binding_storage);
validate_export_launch(&kernel.info, grid, block)?;
if kernel.info.binding_count as usize != bindings.len()
|| kernel.info.constant_byte_length as usize != constants.len
{
return Err(Error::Message(
"kernel binding or constant byte count mismatch".into(),
));
}
let config = sys::DispatchConfig {
workgroup_count: grid,
workgroup_size: block,
subgroup_size: sys::SUBGROUP_SIZE_FROM_EXECUTABLE,
};
unsafe {
check(
sys::hrx_stream_dispatch(
self.inner.stream,
kernel.executable,
kernel.ordinal,
&config,
constants.bytes.as_ptr().cast(),
constants.len,
raw_bindings.as_ptr(),
bindings.len(),
0,
),
"dispatch",
)
}
}
pub fn scratch(&mut self, bytes: usize) -> Result<Buffer> {
let choice = self
.scratch
.range(bytes.max(1)..=bytes.saturating_mul(2).max(1))
.next()
.map(|(&size, _)| size);
if let Some(size) = choice {
let buffers = self.scratch.get_mut(&size).unwrap();
let buffer = buffers.pop().unwrap();
if buffers.is_empty() {
self.scratch.remove(&size);
}
self.scratch_bytes -= buffer.bytes;
Ok(buffer)
} else {
self.allocate(bytes)
}
}
pub fn recycle(&mut self, buffer: Buffer) -> Result<()> {
if !std::sync::Arc::ptr_eq(&buffer._device, &self.inner) {
return Err(Error::Message("scratch belongs to another stream".into()));
}
if buffer.shared.load(Ordering::Relaxed) {
return Err(Error::Message(
"shared allocations cannot enter the scratch pool".into(),
));
}
if buffer.bytes <= self.scratch_limit.saturating_sub(self.scratch_bytes) {
self.scratch_bytes += buffer.bytes;
self.scratch.entry(buffer.bytes).or_default().push(buffer);
}
Ok(())
}
pub fn set_scratch_limit(&mut self, bytes: usize) {
self.scratch_limit = bytes;
while self.scratch_bytes > bytes {
if let Some((_, buffers)) = self.scratch.pop_last() {
self.scratch_bytes -= buffers.iter().map(|b| b.bytes).sum::<usize>();
} else {
break;
}
}
}
}
pub struct Event {
raw: sys::Event,
inner: std::sync::Arc<Inner>,
}
unsafe impl Send for Event {}
unsafe impl Sync for Event {}
impl Event {
pub fn is_complete(&self) -> Result<bool> {
let mut complete = false;
unsafe {
check(sys::hrx_event_query(self.raw, &mut complete), "query event")?;
}
Ok(complete)
}
pub fn synchronize(&self) -> Result<()> {
unsafe { check(sys::hrx_event_synchronize(self.raw), "synchronize event") }
}
}
impl Drop for Event {
fn drop(&mut self) {
unsafe { sys::hrx_event_release(self.raw) }
}
}
#[must_use]
pub struct Submission<'a> {
stream: &'a mut Stream,
}
impl Submission<'_> {
pub fn is_complete(&mut self) -> Result<bool> {
let mut done = false;
unsafe {
check(
sys::hrx_stream_query(self.stream.inner.stream, &mut done),
"query submission",
)?;
}
if done {
self.stream.reclaim_staging();
}
Ok(done)
}
pub fn wait(self) -> Result<()> {
self.stream.synchronize()
}
}
#[derive(Clone, Debug)]
pub struct Constants {
bytes: [u8; 256],
len: usize,
}
impl Default for Constants {
fn default() -> Self {
Self {
bytes: [0; 256],
len: 0,
}
}
}
impl Constants {
pub fn indices(kernel: &Kernel, values: &[u32]) -> Result<Self> {
let size = kernel.info.constant_byte_length as usize;
let mut constants = Self::new();
if values.is_empty() {
if size != 0 {
return Err(Error::Message("missing index constants".into()));
}
return Ok(constants);
}
if size == values.len() * 4 {
for &value in values {
constants.push(value)?;
}
} else if size == values.len() * 8 {
for &value in values {
constants.push(u64::from(value))?;
}
} else {
return Err(Error::Message(
"export is not a homogeneous 32/64-bit index list".into(),
));
}
Ok(constants)
}
pub fn new() -> Self {
Self::default()
}
#[must_use = "constant packing can fail when capacity is exceeded"]
pub fn push<T: Scalar>(&mut self, value: T) -> Result<()> {
let bytes = value.bytes();
checked_span(self.len, bytes.as_ref().len(), self.bytes.len())?;
self.bytes[self.len..self.len + bytes.as_ref().len()].copy_from_slice(bytes.as_ref());
self.len += bytes.as_ref().len();
Ok(())
}
pub fn as_bytes(&self) -> &[u8] {
&self.bytes[..self.len]
}
}
mod sealed {
pub trait Sealed {}
}
pub trait Scalar: sealed::Sealed {
type Bytes: AsRef<[u8]>;
fn bytes(self) -> Self::Bytes;
}
macro_rules! scalars {
($($ty:ty => $n:literal),*) => {
$(
impl sealed::Sealed for $ty {}
impl Scalar for $ty {
type Bytes = [u8; $n];
fn bytes(self) -> Self::Bytes { self.to_le_bytes() }
}
)*
};
}
scalars!(u32 => 4, i32 => 4, f32 => 4, u64 => 8, i64 => 8, f64 => 8);
impl Drop for Stream {
fn drop(&mut self) {
if !self.staging.is_empty() && self.synchronize_native().is_err() {
std::mem::forget(std::mem::take(&mut self.staging));
}
}
}
pub struct SequenceBuilder<'a> {
graph: sys::Graph,
last: sys::GraphNode,
inner: std::sync::Arc<Inner>,
_resources: std::marker::PhantomData<(&'a Buffer, &'a Kernel)>,
}
impl Stream {
pub fn sequence(&self) -> Result<SequenceBuilder<'_>> {
let mut graph = std::ptr::null_mut();
unsafe {
check(
sys::hrx_graph_create(self.inner.device, 0, &mut graph),
"create sequence",
)?;
}
Ok(SequenceBuilder {
graph,
last: std::ptr::null_mut(),
inner: self.inner.clone(),
_resources: std::marker::PhantomData,
})
}
pub fn launch_sequence(&mut self, sequence: &mut FixedSequence) -> Result<()> {
if !std::sync::Arc::ptr_eq(&sequence.inner, &self.inner) {
return Err(Error::Message("sequence belongs to another stream".into()));
}
unsafe {
check(
sys::hrx_graph_exec_launch(sequence.raw, self.inner.stream),
"launch sequence",
)
}
}
}
impl<'a> SequenceBuilder<'a> {
fn deps(&self) -> (*const sys::GraphNode, usize) {
(&self.last, usize::from(!self.last.is_null()))
}
pub fn fill(&mut self, dst: View<'a>, pattern: u8) -> Result<&mut Self> {
owns(&self.inner, dst.owner)?;
if dst.is_empty() {
return Err(Error::Message("empty sequence fill".into()));
}
let attrs = sys::GraphFill {
dst: dst.raw,
pattern: pattern.into(),
pattern_size: 1,
};
let (deps, count) = self.deps();
let mut next = std::ptr::null_mut();
unsafe {
check(
sys::hrx_graph_add_fill_buffer_node(self.graph, deps, count, &attrs, &mut next),
"record fill",
)?;
}
self.last = next;
Ok(self)
}
pub fn copy(&mut self, dst: View<'a>, src: View<'a>) -> Result<&mut Self> {
owns(&self.inner, dst.owner)?;
owns(&self.inner, src.owner)?;
if dst.len() != src.len() || dst.is_empty() {
return Err(Error::Message(
"sequence copy requires equal nonempty spans".into(),
));
}
let attrs = sys::GraphCopy {
src: src.raw,
dst: dst.raw,
};
let (deps, count) = self.deps();
let mut next = std::ptr::null_mut();
unsafe {
check(
sys::hrx_graph_add_copy_buffer_node(self.graph, deps, count, &attrs, &mut next),
"record copy",
)?;
}
self.last = next;
Ok(self)
}
pub unsafe fn dispatch(
&mut self,
kernel: &'a Kernel,
grid: [u32; 3],
block: [u32; 3],
constants: &Constants,
bindings: &[View<'a>],
) -> Result<&mut Self> {
if kernel._device.device != self.inner.device {
return Err(Error::Message("kernel belongs to another device".into()));
}
for view in bindings {
owns(&self.inner, view.owner)?;
}
let mut binding_storage = [std::mem::MaybeUninit::uninit(); 32];
let raw_bindings = raw_bindings(bindings, &mut binding_storage);
validate_export_launch(&kernel.info, grid, block)?;
if kernel.info.binding_count as usize != bindings.len()
|| kernel.info.constant_byte_length as usize != constants.len
{
return Err(Error::Message(
"sequence binding or constant byte count mismatch".into(),
));
}
let attrs = sys::GraphKernel {
executable: kernel.executable,
ordinal: kernel.ordinal,
config: sys::DispatchConfig {
workgroup_count: grid,
workgroup_size: block,
subgroup_size: sys::SUBGROUP_SIZE_FROM_EXECUTABLE,
},
constants: constants.bytes.as_ptr().cast(),
constants_size: constants.len,
bindings: raw_bindings.as_ptr(),
binding_count: bindings.len(),
flags: 0,
};
let (deps, count) = self.deps();
let mut next = std::ptr::null_mut();
unsafe {
check(
sys::hrx_graph_add_kernel_node(self.graph, deps, count, &attrs, &mut next),
"record kernel",
)?;
}
self.last = next;
Ok(self)
}
pub fn finish(self) -> Result<FixedSequence> {
let mut raw = std::ptr::null_mut();
unsafe {
check(
sys::hrx_graph_instantiate(self.graph, 0, &mut raw),
"instantiate sequence",
)?;
}
Ok(FixedSequence {
raw,
inner: self.inner.clone(),
})
}
}
impl Drop for SequenceBuilder<'_> {
fn drop(&mut self) {
unsafe {
sys::hrx_graph_release(self.graph);
}
}
}
pub struct FixedSequence {
raw: sys::GraphExec,
inner: std::sync::Arc<Inner>,
}
unsafe impl Send for FixedSequence {}
impl Drop for FixedSequence {
fn drop(&mut self) {
unsafe {
sys::hrx_graph_exec_release(self.raw);
}
}
}
#[must_use]
pub struct Readback {
buffer: Buffer,
}
impl Stream {
pub fn read_queued(&mut self, source: View<'_>) -> Result<Readback> {
self.owns(source.owner)?;
let buffer = self.allocate_host(source.len())?;
let raw = buffer.raw;
unsafe {
if !source.is_empty() {
check(
sys::hrx_stream_copy_buffer(
self.inner.stream,
source.raw.buffer,
source.raw.offset,
raw,
0,
source.len(),
),
"enqueue readback",
)?;
}
Ok(Readback { buffer })
}
}
}
impl Readback {
pub fn wait(self, stream: &mut Stream) -> Result<Vec<u8>> {
if !std::sync::Arc::ptr_eq(&self.buffer._device, &stream.inner) {
return Err(Error::Message("readback belongs to another stream".into()));
}
stream.synchronize()?;
let mut bytes = Vec::with_capacity(self.buffer.bytes);
if self.buffer.bytes != 0 {
let mut pointer = std::ptr::null_mut();
unsafe {
check(
sys::hrx_buffer_get_device_ptr(self.buffer.raw, &mut pointer),
"map completed readback",
)?;
std::ptr::copy_nonoverlapping(
pointer.cast::<u8>(),
bytes.as_mut_ptr(),
self.buffer.bytes,
);
bytes.set_len(self.buffer.bytes);
}
}
Ok(bytes)
}
}
impl std::fmt::Debug for Buffer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Buffer")
.field("bytes", &self.bytes)
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for Kernel {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Kernel")
.field("symbol", &self.symbol)
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for Stream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Stream")
.field("scratch_bytes", &self.scratch_bytes)
.finish_non_exhaustive()
}
}
impl std::fmt::Debug for SequenceBuilder<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SequenceBuilder").finish_non_exhaustive()
}
}
impl std::fmt::Debug for FixedSequence {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FixedSequence").finish_non_exhaustive()
}
}
impl std::fmt::Debug for Readback {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Readback").finish_non_exhaustive()
}
}
impl std::fmt::Debug for Submission<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Submission").finish_non_exhaustive()
}
}
const STAGING_LIMIT: usize = 64 * 1024 * 1024;
#[cfg(test)]
mod staging_tests {
use super::*;
#[test]
#[ignore = "requires gfx1151"]
fn scratch_pools_reject_shared_handles_and_their_original() -> Result<()> {
let device = Device::open(0)?;
let mut source_stream = device.stream()?;
let mut compute = device.stream()?;
let mut third = device.stream()?;
let private = compute.scratch(4096)?;
let private_raw = private.raw;
compute.recycle(private)?;
let source = source_stream.allocate(4096)?;
let shared = unsafe { source.share_on(&compute)? };
let chained = unsafe { shared.share_on(&third)? };
let same_stream = unsafe { source.share_on(&source_stream)? };
for result in [
compute.recycle(shared),
source_stream.recycle(source),
third.recycle(chained),
source_stream.recycle(same_stream),
] {
assert!(
result
.unwrap_err()
.to_string()
.contains("shared allocations")
);
}
assert!(source_stream.scratch.is_empty());
assert!(third.scratch.is_empty());
assert_eq!(compute.scratch_bytes, 4096);
assert_eq!(compute.scratch(4096)?.raw, private_raw);
let source = source_stream.allocate(4096)?;
drop(unsafe { source.share_on(&compute)? });
assert!(source_stream.recycle(source).is_err());
assert!(source_stream.scratch.is_empty());
Ok(())
}
#[test]
#[ignore = "requires gfx1151"]
fn completed_staging_is_reclaimed_and_reused_on_every_completion_path() -> Result<()> {
let mut stream = Stream::open()?;
let buffer = stream.allocate(1024)?;
stream.upload(buffer.binding(), &[7; 1024])?;
let original = stream.staging[0].raw;
let mut output = [0; 1024];
stream.read(buffer.binding(), &mut output)?;
assert_eq!(output, [7; 1024]);
assert!(stream.staging.is_empty());
stream.upload(buffer.binding(), &[9; 1024])?;
assert_eq!(stream.staging[0].raw, original);
stream.upload_blocking(buffer.binding(), &[11; 1024])?;
assert!(stream.staging.is_empty());
stream.upload(buffer.binding(), &[12; 512])?;
assert_eq!(stream.staging[0].raw, original);
stream.synchronize()?;
stream.upload(buffer.binding(), &[13; 1024])?;
stream.synchronize_native()?; assert!(stream.submit()?.is_complete()?);
assert!(stream.staging.is_empty());
for _ in 0..32 {
stream.upload(buffer.binding(), &[1; 1024])?;
assert!(stream.staging.len() <= 8);
}
stream.synchronize()?;
assert!(stream.staging_pool.len() <= 8);
assert!(stream.staging_pool.iter().map(|b| b.bytes).sum::<usize>() <= STAGING_LIMIT);
Ok(())
}
#[test]
#[ignore = "requires gfx1151"]
fn uploads_batch_until_pressure_and_reuse_the_smallest_capacity() -> Result<()> {
let mut stream = Stream::open()?;
let buffer = stream.allocate(4096)?;
for size in [1024, 2048, 4096] {
stream.upload(buffer.binding(), &vec![7; size])?;
}
assert_eq!(stream.staging.len(), 3);
let medium = stream.staging[1].raw;
stream.synchronize()?;
stream.upload(buffer.binding(), &[9; 1500])?;
assert_eq!(stream.staging[0].raw, medium);
for _ in 0..7 {
stream.upload(buffer.binding(), &[11; 1024])?;
}
assert_eq!(stream.staging.len(), 8);
stream.upload(buffer.binding(), &[13; 1024])?;
assert_eq!(stream.staging.len(), 1);
stream.synchronize()?;
Ok(())
}
}