diff --git a/Cargo.lock b/Cargo.lock index ac1796a..6d1491e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14,6 +14,12 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" +[[package]] +name = "bytemuck" +version = "1.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6aedf8ae72766347502cf3cb4f41cf5e9cc37d28bee90f1fdaaae15f9cf9424" + [[package]] name = "cc" version = "1.2.65" @@ -177,9 +183,12 @@ name = "oneapi-rs" version = "0.1.0" dependencies = [ "allocator-api2", + "bytemuck", "cxx", "oneapi-rs-sys", + "pin-project", "thiserror", + "tokio", ] [[package]] @@ -191,6 +200,32 @@ dependencies = [ "which", ] +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "proc-macro2" version = "1.0.106" @@ -297,6 +332,27 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "pin-project-lite", + "tokio-macros", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/oneapi-rs-sys/build.rs b/oneapi-rs-sys/build.rs index dae79ad..c57aa82 100644 --- a/oneapi-rs-sys/build.rs +++ b/oneapi-rs-sys/build.rs @@ -19,6 +19,7 @@ fn main() { "src/device-sys.rs", "src/queue-sys.rs", "src/usm-sys.rs", + "src/event-sys.rs", ]; let cpp_sources = [ @@ -26,6 +27,7 @@ fn main() { "src/device.cpp", "src/queue.cpp", "src/usm.cpp", + "src/event.cpp", ]; let cpp_headers = [ @@ -34,6 +36,7 @@ fn main() { "include/device.hpp", "include/queue.hpp", "include/usm.hpp", + "include/event.hpp", ]; cxx_build::bridges(&rust_sources) diff --git a/oneapi-rs-sys/include/event.hpp b/oneapi-rs-sys/include/event.hpp new file mode 100644 index 0000000..e799771 --- /dev/null +++ b/oneapi-rs-sys/include/event.hpp @@ -0,0 +1,26 @@ +// +// Copyright (C) 2026 Intel Corporation +// +// Under the MIT License or the Apache License v2.0. +// See LICENSE-MIT and LICENSE-APACHE for license information. +// SPDX-License-Identifier: MIT OR Apache-2.0 +// + +#pragma once + +#include "oneapi-rs-sys/include/types.hpp" +#include "oneapi-rs-sys/src/event-sys.rs.h" +#include "rust/cxx.h" + +#include + +namespace sycl_shims { +enum class EventCommandStatus : std::uint8_t; +} // namespace sycl_shims + +namespace sycl_shims::event { +void wait(std::unique_ptr &); +void register_callback(std::unique_ptr &, Event const &, rust::Box); +EventCommandStatus get_command_execution_status(Event const &); +std::unique_ptr clone(Event const &); +} // namespace sycl_shims::event diff --git a/oneapi-rs-sys/include/queue.hpp b/oneapi-rs-sys/include/queue.hpp index 55a6877..344102e 100644 --- a/oneapi-rs-sys/include/queue.hpp +++ b/oneapi-rs-sys/include/queue.hpp @@ -13,7 +13,22 @@ #include +namespace sycl_shims { +struct EventPtr; +} // namespace sycl_shims + namespace sycl_shims::queue { std::unique_ptr new_queue(); +std::unique_ptr new_queue_immediate(); std::unique_ptr new_queue_from_device(Device const &); +std::unique_ptr clone(Queue const &); +std::unique_ptr memset( + std::unique_ptr &, + std::uint8_t * ptr, + int value, + std::size_t num_bytes, + rust::Vec +); +std::unique_ptr barrier(std::unique_ptr &, rust::Vec); +void wait(std::unique_ptr &); } // namespace sycl_shims::queue diff --git a/oneapi-rs-sys/include/types.hpp b/oneapi-rs-sys/include/types.hpp index f97d0af..82e9bbb 100644 --- a/oneapi-rs-sys/include/types.hpp +++ b/oneapi-rs-sys/include/types.hpp @@ -14,4 +14,5 @@ namespace sycl_shims { using Device = sycl::device; using Platform = sycl::platform; using Queue = sycl::queue; +using Event = sycl::event; } diff --git a/oneapi-rs-sys/src/event-sys.rs b/oneapi-rs-sys/src/event-sys.rs new file mode 100644 index 0000000..c94f99c --- /dev/null +++ b/oneapi-rs-sys/src/event-sys.rs @@ -0,0 +1,42 @@ +// +// Copyright (C) 2026 Intel Corporation +// +// Under the MIT License or the Apache License v2.0. +// See LICENSE-MIT and LICENSE-APACHE for license information. +// SPDX-License-Identifier: MIT OR Apache-2.0 +// + +use crate::types::Waker; + +#[cxx::bridge(namespace = "sycl_shims::event")] +pub mod ffi { + #[namespace = "sycl_shims"] + extern "C++" { + include!("oneapi-rs-sys/src/types-sys.rs.h"); + type EventCommandStatus = crate::types::ffi::EventCommandStatus; + } + + unsafe extern "C++" { + include!("oneapi-rs-sys/include/event.hpp"); + + #[namespace = "sycl_shims"] + type Event = crate::types::ffi::Event; + + #[namespace = "sycl_shims"] + type Queue = crate::types::ffi::Queue; + + fn wait(event: &mut UniquePtr); + fn register_callback(queue: &mut UniquePtr, event: &Event, waker: Box); + fn get_command_execution_status(event: &Event) -> EventCommandStatus; + fn clone(event: &Event) -> UniquePtr; + } + + extern "Rust" { + type Waker; + fn wake(waker: &Box); + } +} + +fn wake(waker: &Box) { + waker.0.wake_by_ref(); +} diff --git a/oneapi-rs-sys/src/event.cpp b/oneapi-rs-sys/src/event.cpp new file mode 100644 index 0000000..b3725c0 --- /dev/null +++ b/oneapi-rs-sys/src/event.cpp @@ -0,0 +1,41 @@ +// +// Copyright (C) 2026 Intel Corporation +// +// Under the MIT License or the Apache License v2.0. +// See LICENSE-MIT and LICENSE-APACHE for license information. +// SPDX-License-Identifier: MIT OR Apache-2.0 +// + +#include "oneapi-rs-sys/include/event.hpp" + +using sycl::info::event_command_status; + +namespace syclintel = sycl::ext::intel; + +namespace sycl_shims::event { +void wait(std::unique_ptr & event) { + event->wait(); +} +std::unique_ptr clone(Event const & event) { + return std::make_unique(sycl::event(event)); +} +EventCommandStatus get_command_execution_status(Event const & event) { + auto status = event.get_info(); + switch (status) { + case event_command_status::submitted: + return EventCommandStatus::Submitted; + case event_command_status::running: + return EventCommandStatus::Running; + case event_command_status::complete: + return EventCommandStatus::Complete; + default: + return EventCommandStatus::Unknown; + } +} +void register_callback(std::unique_ptr & queue, Event const & event, rust::Box waker) { + queue->submit([waker = std::move(waker), event](sycl::handler& cgh) { + cgh.depends_on(event); + cgh.host_task([&]() { wake(waker); }); + }); +} +} // namespace sycl_shims::event diff --git a/oneapi-rs-sys/src/lib.rs b/oneapi-rs-sys/src/lib.rs index 9f71b49..f2f723d 100644 --- a/oneapi-rs-sys/src/lib.rs +++ b/oneapi-rs-sys/src/lib.rs @@ -20,3 +20,6 @@ pub mod queue; #[path = "usm-sys.rs"] pub mod usm; + +#[path = "event-sys.rs"] +pub mod event; diff --git a/oneapi-rs-sys/src/queue-sys.rs b/oneapi-rs-sys/src/queue-sys.rs index b7b1c65..d9f7a1f 100644 --- a/oneapi-rs-sys/src/queue-sys.rs +++ b/oneapi-rs-sys/src/queue-sys.rs @@ -8,6 +8,12 @@ #[cxx::bridge(namespace = "sycl_shims::queue")] pub mod ffi { + #[namespace = "sycl_shims"] + extern "C++" { + include!("oneapi-rs-sys/src/types-sys.rs.h"); + type EventPtr = crate::types::ffi::EventPtr; + } + unsafe extern "C++" { include!("oneapi-rs-sys/include/queue.hpp"); @@ -15,8 +21,21 @@ pub mod ffi { type Queue = crate::types::ffi::Queue; #[namespace = "sycl_shims"] type Device = crate::types::ffi::Device; + #[namespace = "sycl_shims"] + type Event = crate::types::ffi::Event; fn new_queue() -> UniquePtr; + fn new_queue_immediate() -> UniquePtr; fn new_queue_from_device(device: &Device) -> UniquePtr; + fn clone(queue: &Queue) -> UniquePtr; + unsafe fn memset( + queue: &mut UniquePtr, + ptr: *mut u8, + value: i32, + num_bytes: usize, + dep_events: Vec + ) -> UniquePtr; + fn barrier(queue: &mut UniquePtr, dep_events: Vec) -> UniquePtr; + fn wait(queue: &mut UniquePtr); } } diff --git a/oneapi-rs-sys/src/queue.cpp b/oneapi-rs-sys/src/queue.cpp index 22e012c..13695bd 100644 --- a/oneapi-rs-sys/src/queue.cpp +++ b/oneapi-rs-sys/src/queue.cpp @@ -9,11 +9,46 @@ #include "oneapi-rs-sys/include/queue.hpp" #include "oneapi-rs-sys/src/queue-sys.rs.h" +using sycl::property::queue::in_order; +using sycl::ext::intel::property::queue::immediate_command_list; + namespace sycl_shims::queue { std::unique_ptr new_queue() { - return std::make_unique(sycl::queue()); + return std::make_unique(sycl::queue({ + in_order() + })); +} +std::unique_ptr new_queue_immediate() { + return std::make_unique(sycl::queue({ + in_order(), + immediate_command_list() + })); } std::unique_ptr new_queue_from_device(Device const & device) { - return std::make_unique(sycl::queue(device)); + return std::make_unique(sycl::queue(device, {in_order()})); +} +std::unique_ptr clone(Queue const & queue) { + return std::make_unique(sycl::queue(queue)); +} +std::unique_ptr memset( + std::unique_ptr & queue, + std::uint8_t * ptr, + int value, + std::size_t num_bytes, + rust::Vec dep_events +) { + std::vector deps; + for (auto&& e: dep_events) + deps.push_back(std::move(*e.ptr.release())); + return std::make_unique(queue->memset(ptr, value, num_bytes, deps)); +} +std::unique_ptr barrier(std::unique_ptr & queue, rust::Vec dep_events) { + std::vector deps; + for (auto&& e: dep_events) + deps.push_back(std::move(*e.ptr.release())); + return std::make_unique(queue->ext_oneapi_submit_barrier(deps)); +} +void wait(std::unique_ptr & queue) { + queue->wait(); } } // namespace sycl_shims::queue diff --git a/oneapi-rs-sys/src/types-sys.rs b/oneapi-rs-sys/src/types-sys.rs index 5075deb..1b485e0 100644 --- a/oneapi-rs-sys/src/types-sys.rs +++ b/oneapi-rs-sys/src/types-sys.rs @@ -6,6 +6,14 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 // +// cxx requires Rust opaque types to be defined in the current crate. +pub struct Waker(pub std::task::Waker); +impl From for Waker { + fn from(value: std::task::Waker) -> Self { + Waker(value) + } +} + #[cxx::bridge(namespace = "sycl_shims")] pub mod ffi { unsafe extern "C++" { @@ -13,6 +21,7 @@ pub mod ffi { type Device; type Platform; type Queue; + type Event; } // This is a workaround - cxx currently doesn't support passing @@ -26,6 +35,10 @@ pub mod ffi { struct PlatformPtr { ptr: UniquePtr } + + struct EventPtr { + ptr: UniquePtr + } #[derive(Debug)] enum DeviceType { @@ -38,10 +51,20 @@ pub mod ffi { Unimplemented } + #[derive(Debug)] + enum EventCommandStatus { + Submitted, + Running, + Complete, + Unknown + } + impl UniquePtr {} impl UniquePtr {} impl UniquePtr {} + impl UniquePtr {} impl Vec {} impl Vec {} + impl Vec {} } diff --git a/oneapi-rs/Cargo.toml b/oneapi-rs/Cargo.toml index 3bb6be3..4a9876e 100644 --- a/oneapi-rs/Cargo.toml +++ b/oneapi-rs/Cargo.toml @@ -6,6 +6,11 @@ license = "MIT OR Apache-2.0" [dependencies] allocator-api2 = "0.4.0" +bytemuck = "1.25.1" cxx = "1.0.194" oneapi-rs-sys = { path = "../oneapi-rs-sys" } +pin-project = "1.1.13" thiserror = "2.0.18" + +[dev-dependencies] +tokio = { version = "1.52.3", features = ["macros", "rt", "rt-multi-thread"] } diff --git a/oneapi-rs/examples/alloc.rs b/oneapi-rs/examples/alloc.rs index 5cecde4..2e8476b 100644 --- a/oneapi-rs/examples/alloc.rs +++ b/oneapi-rs/examples/alloc.rs @@ -9,8 +9,14 @@ use oneapi_rs::queue::Queue; fn main() { - let queue = Queue::new(); - let mut buffer = unsafe { queue.alloc_uninit_shared::(10) }; + let mut queue = Queue::new(); + let mut buffer = queue.alloc_shared::(10).wait(); + + for e in buffer.iter() { + print!("{e} ") + } + + println!(); for i in 0..buffer.len() { buffer[i] = i as u32; diff --git a/oneapi-rs/examples/async.rs b/oneapi-rs/examples/async.rs new file mode 100644 index 0000000..ad54f51 --- /dev/null +++ b/oneapi-rs/examples/async.rs @@ -0,0 +1,31 @@ +// +// Copyright (C) 2026 Intel Corporation +// +// Under the MIT License or the Apache License v2.0. +// See LICENSE-MIT and LICENSE-APACHE for license information. +// SPDX-License-Identifier: MIT OR Apache-2.0 +// + +use oneapi_rs::queue::Queue; + +#[tokio::main] +async fn main() { + let mut queue = Queue::new(); + let mut buffer = queue.alloc_shared::(10).await; + + for e in buffer.iter() { + print!("{e} ") + } + + println!(); + + for i in 0..buffer.len() { + buffer[i] = i as u32; + } + + for e in buffer.iter() { + print!("{e} ") + } + + println!(); +} diff --git a/oneapi-rs/src/buffer.rs b/oneapi-rs/src/buffer.rs index c174815..e1fa270 100644 --- a/oneapi-rs/src/buffer.rs +++ b/oneapi-rs/src/buffer.rs @@ -6,9 +6,11 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 // -use std::{alloc::{Layout, handle_alloc_error}, ops::{Deref, DerefMut}, ptr::NonNull, slice}; +use std::{alloc::{Layout, handle_alloc_error}, ops::{Deref, DerefMut}, pin::Pin, ptr::NonNull, slice, task::{Context, Poll}}; -use crate::usm::UsmAlloc; +use pin_project::pin_project; + +use crate::{event::{Event, EventFuture}, usm::UsmAlloc}; /// The Buffer struct defines a shared array of one, two or three dimensions that can be used /// by the SYCL kernel. Buffers are templated on the type of their data, and the number of @@ -46,6 +48,14 @@ impl Buffer { allocator, } } + + pub(crate) fn get_byte_ptr(&mut self) -> *mut u8 { + self.data.as_ptr().cast() + } + + pub(crate) fn get_byte_size(&self) -> usize { + self.layout.size() + } } impl Deref for Buffer { @@ -70,3 +80,54 @@ impl Drop for Buffer { unsafe { self.allocator.deallocate(self.data.cast(), self.layout); } } } + +/// A [`Buffer`] whose initialization has been enqueued. You need to wait/await it. +pub struct EnqueuedBuffer { + buffer: Buffer, + event: Event +} + +impl EnqueuedBuffer { + pub(crate) fn new(buffer: Buffer, event: Event) -> Self { + Self { + buffer, + event + } + } +} + +impl EnqueuedBuffer { + /// Waits for [`Buffer`] initialization to finish. + pub fn wait(mut self) -> Buffer { + self.event.wait(); + self.buffer + } +} + +#[pin_project] +/// A [`Future`] which represents a pending [`Buffer`] allocation. +pub struct BufferFuture { + buffer: Option>, + #[pin] + event_future: EventFuture +} + +impl Future for BufferFuture { + type Output = Buffer; + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + this.event_future.poll(cx).map(|_| this.buffer.take().unwrap()) + } +} + +impl IntoFuture for EnqueuedBuffer { + type Output = Buffer; + type IntoFuture = BufferFuture; + + fn into_future(self) -> Self::IntoFuture { + Self::IntoFuture { + buffer: Some(self.buffer), + event_future: self.event.into_future() + } + } +} diff --git a/oneapi-rs/src/event.rs b/oneapi-rs/src/event.rs new file mode 100644 index 0000000..206dfd6 --- /dev/null +++ b/oneapi-rs/src/event.rs @@ -0,0 +1,78 @@ +// +// Copyright (C) 2026 Intel Corporation +// +// Under the MIT License or the Apache License v2.0. +// See LICENSE-MIT and LICENSE-APACHE for license information. +// SPDX-License-Identifier: MIT OR Apache-2.0 +// + +use std::{pin::Pin, task::{Context, Poll}}; + +use oneapi_rs_sys::event::ffi; + +use pin_project::pin_project; + +use crate::{info::{EventCommandStatus, event::{CommandExecutionStatus, EventInfo}}, queue::Queue}; + +pub struct Event(pub(crate) cxx::UniquePtr); + +impl Event { + pub fn wait(&mut self) { + ffi::wait(&mut self.0); + } + + pub fn get_info(&self) -> T::Item { + T::get_item(self) + } +} + +impl From> for Event { + fn from(value: cxx::UniquePtr) -> Self { + Self(value) + } +} + +impl Clone for Event { + fn clone(&self) -> Self { + ffi::clone(&self.0).into() + } +} + +#[pin_project] +pub struct EventFuture { + event: Event, + set_callback: bool, + queue: Queue, +} + +impl Future for EventFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + if self.event.get_info::() == EventCommandStatus::Complete { + Poll::Ready(()) + } + else { + if self.set_callback == false { + let this = self.project(); + *this.set_callback = true; + let waker = Box::new(cx.waker().clone().into()); + ffi::register_callback(&mut this.queue.0, &this.event.0, waker); + } + Poll::Pending + } + } +} + +impl IntoFuture for Event { + type Output = (); + type IntoFuture = EventFuture; + + fn into_future(self) -> Self::IntoFuture { + EventFuture { + queue: Queue::new_immediate(), + event: self, + set_callback: false + } + } +} diff --git a/oneapi-rs/src/info.rs b/oneapi-rs/src/info.rs index b7a1883..9811994 100644 --- a/oneapi-rs/src/info.rs +++ b/oneapi-rs/src/info.rs @@ -12,5 +12,11 @@ pub mod platform; #[path = "./info/device-info.rs"] pub mod device; +#[path = "./info/event-info.rs"] +pub mod event; + /// The type of the SYCL device. pub use oneapi_rs_sys::device::ffi::DeviceType; + +/// Event status of the contained action associated with this event. +pub use oneapi_rs_sys::event::ffi::EventCommandStatus; diff --git a/oneapi-rs/src/info/event-info.rs b/oneapi-rs/src/info/event-info.rs new file mode 100644 index 0000000..1a8e6bf --- /dev/null +++ b/oneapi-rs/src/info/event-info.rs @@ -0,0 +1,24 @@ +// +// Copyright (C) 2026 Intel Corporation +// +// Under the MIT License or the Apache License v2.0. +// See LICENSE-MIT and LICENSE-APACHE for license information. +// SPDX-License-Identifier: MIT OR Apache-2.0 +// + +use crate::event::Event; +use oneapi_rs_sys::event::ffi; + +pub trait EventInfo { + type Item; + fn get_item(event: &Event) -> Self::Item; +} + +/// Returns the event status of the action associated with this event. +pub struct CommandExecutionStatus; +impl EventInfo for CommandExecutionStatus { + type Item = crate::info::EventCommandStatus; + fn get_item(event: &Event) -> Self::Item { + ffi::get_command_execution_status(&event.0) + } +} diff --git a/oneapi-rs/src/lib.rs b/oneapi-rs/src/lib.rs index 82cd93a..1f7a5b7 100644 --- a/oneapi-rs/src/lib.rs +++ b/oneapi-rs/src/lib.rs @@ -12,3 +12,4 @@ pub mod device; pub mod queue; pub mod usm; pub mod buffer; +pub mod event; diff --git a/oneapi-rs/src/queue.rs b/oneapi-rs/src/queue.rs index d1000e9..a7bd77a 100644 --- a/oneapi-rs/src/queue.rs +++ b/oneapi-rs/src/queue.rs @@ -6,9 +6,10 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 // -use oneapi_rs_sys::queue::ffi; +use bytemuck::Pod; +use oneapi_rs_sys::{queue::ffi, types::ffi::EventPtr}; -use crate::{buffer::Buffer, device::Device, usm::{HostAllocator, SharedAllocator, UsmAllocator}}; +use crate::{buffer::{Buffer, EnqueuedBuffer}, device::Device, event::Event, usm::{HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator}}; /// The `Queue` connects a host program to a single device. Programs submit tasks to a device via the /// `Queue` and may monitor the `Queue` for completion. A program initiates the task by submitting @@ -21,19 +22,84 @@ impl Queue { Self(ffi::new_queue()) } + /// Construct an immediate `Queue` based on the device returned from the default selector. + pub fn new_immediate() -> Self { + Self(ffi::new_queue_immediate()) + } + + /// Allocates zeroed memory and creates a host-side [`Buffer`] that can store an array of T. + pub fn alloc_host(&mut self, len: usize) -> EnqueuedBuffer> { + unsafe { + let mut buffer = self.alloc_uninit_host(len); + let event = self.memset(&mut buffer, 0); + EnqueuedBuffer::new(buffer, event) + } + } + + /// Allocates zeroed memory and creates a shared [`Buffer`] that can store an array of T. + pub fn alloc_shared(&mut self, len: usize) -> EnqueuedBuffer> { + unsafe { + let mut buffer = self.alloc_uninit_shared(len); + let event = self.memset(&mut buffer, 0); + EnqueuedBuffer::new(buffer, event) + } + } + /// Allocates memory and creates a host-side [`Buffer`] that can store an array of T. /// Safety: the buffer contents are uninitialized. - pub unsafe fn alloc_uninit_host(&self, len: usize) -> Buffer> { + pub unsafe fn alloc_uninit_host(&self, len: usize) -> Buffer> { let allocator = UsmAllocator::from(self); unsafe { Buffer::new(allocator, len) } } /// Allocates memory and creates a shared [`Buffer`] that can store an array of T. /// Safety: the buffer contents are uninitialized. - pub unsafe fn alloc_uninit_shared(&self, len: usize) -> Buffer> { + pub unsafe fn alloc_uninit_shared(&self, len: usize) -> Buffer> { let allocator = UsmAllocator::from(self); unsafe { Buffer::new(allocator, len) } } + + /// Sets memory allocated with USM allocations. + /// Safety: the caller must make sure the underlying memory isn't being aliased somewhere else. + pub unsafe fn memset(&mut self, buffer: &mut Buffer, value: i32) -> Event { + unsafe { self.memset_with_deps(buffer, value, &[]) } + } + + /// Sets memory allocated with USM allocations after all specified events finish. + /// Safety: the caller must make sure the underlying memory isn't being aliased somewhere else. + pub unsafe fn memset_with_deps( + &mut self, + buffer: &mut Buffer, + value: i32, + dep_events: &[&Event] + ) -> Event { + let ptr = buffer.get_byte_ptr(); + let num_bytes = buffer.get_byte_size(); + let dep_events = dep_events + .iter() + .map(|e| EventPtr { ptr: (*e).clone().0 }) + .collect::>(); + unsafe { ffi::memset(&mut self.0, ptr, value, num_bytes, dep_events) }.into() + } + + /// Submits a barrier to the queue. + pub fn barrier(&mut self) -> Event { + self.barrier_with_deps(&[]) + } + + /// Submits a barrier to the queue after all specified events finish. + pub fn barrier_with_deps(&mut self, dep_events: &[&Event]) -> Event { + let dep_events = dep_events + .iter() + .map(|e| EventPtr { ptr: (*e).clone().0 }) + .collect::>(); + ffi::barrier(&mut self.0, dep_events).into() + } + + /// Performs a blocking wait for the completion of all enqueued tasks in the queue. + pub fn wait(&mut self) { + ffi::wait(&mut self.0); + } } impl From<&Device> for Queue { @@ -41,3 +107,15 @@ impl From<&Device> for Queue { Self(ffi::new_queue_from_device(&value.0)) } } + +impl From> for Queue { + fn from(value: cxx::UniquePtr) -> Self { + Self(value) + } +} + +impl Clone for Queue { + fn clone(&self) -> Self { + ffi::clone(&self.0).into() + } +} diff --git a/oneapi-rs/src/usm.rs b/oneapi-rs/src/usm.rs index a4d0060..39b41fd 100644 --- a/oneapi-rs/src/usm.rs +++ b/oneapi-rs/src/usm.rs @@ -16,30 +16,30 @@ use std::{alloc::Layout, marker::PhantomData, ptr::NonNull}; type CxxResult = cxx::core::result::Result; /// An instance of a USM allocator. -pub struct UsmAllocator<'a, T: UsmAllocatorKind> { - queue: &'a Queue, +pub struct UsmAllocator { + queue: Queue, _kind: PhantomData } /// A marker trait for USM allocators. pub unsafe trait UsmAlloc : Allocator {} -unsafe impl<'a, T: UsmAllocatorKind> UsmAlloc for UsmAllocator<'a, T> {} +unsafe impl UsmAlloc for UsmAllocator {} pub trait UsmAllocatorKind { unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8>; } -impl<'a, T: UsmAllocatorKind> From<&'a Queue> for UsmAllocator<'a, T> { - fn from(queue: &'a Queue) -> Self { +impl From<&Queue> for UsmAllocator { + fn from(queue: &Queue) -> Self { Self { - queue, + queue: queue.clone(), _kind: PhantomData } } } -unsafe impl Allocator for UsmAllocator<'_, T> { +unsafe impl Allocator for UsmAllocator { fn allocate(&self, layout: Layout) -> Result, AllocError> { let ptr = unsafe { T::alloc(layout.align(), layout.size(), &self.queue) } .map_err(|_e| AllocError)?;