Merge pull request #2 from szymon-zadworny/usm-alloc

bratpiorka · web-flow · commit 0a536495cda2 · 2026-07-14T11:58:56.000+02:00
Add USM allocations
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/oneapi-rs-sys/build.rs b/oneapi-rs-sys/build.rs
@@ -17,17 +17,23 @@ fn main() {
         "src/types-sys.rs",
         "src/platform-sys.rs",
         "src/device-sys.rs",
+        "src/queue-sys.rs",
+        "src/usm-sys.rs",
     ];
 
     let cpp_sources = [
         "src/platform.cpp",
         "src/device.cpp",
+        "src/queue.cpp",
+        "src/usm.cpp",
     ];
 
     let cpp_headers = [
         "include/types.hpp",
         "include/platform.hpp",
         "include/device.hpp",
+        "include/queue.hpp",
+        "include/usm.hpp",
     ];
 
     cxx_build::bridges(&rust_sources)
diff --git a/oneapi-rs-sys/include/queue.hpp b/oneapi-rs-sys/include/queue.hpp
@@ -0,0 +1,19 @@
+//
+// 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 "rust/cxx.h"
+
+#include <memory>
+
+namespace sycl_shims::queue {
+std::unique_ptr<Queue> new_queue();
+std::unique_ptr<Queue> new_queue_from_device(Device const &);
+} // namespace sycl_shims::queue
diff --git a/oneapi-rs-sys/include/types.hpp b/oneapi-rs-sys/include/types.hpp
@@ -13,4 +13,5 @@
 namespace sycl_shims {
 using Device = sycl::device;
 using Platform = sycl::platform;
+using Queue = sycl::queue;
 }
diff --git a/oneapi-rs-sys/include/usm.hpp b/oneapi-rs-sys/include/usm.hpp
@@ -0,0 +1,21 @@
+//
+// 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 "rust/cxx.h"
+
+#include <memory>
+
+namespace sycl_shims::usm {
+std::uint8_t* aligned_alloc_device(std::size_t alignment, std::size_t num_bytes, Queue const &);
+std::uint8_t* aligned_alloc_host(std::size_t alignment, std::size_t num_bytes, Queue const &);
+std::uint8_t* aligned_alloc_shared(std::size_t alignment, std::size_t num_bytes, Queue const &);
+void free(std::uint8_t*, Queue const &);
+} // namespace sycl_shims::usm
diff --git a/oneapi-rs-sys/src/lib.rs b/oneapi-rs-sys/src/lib.rs
@@ -14,3 +14,9 @@ pub mod platform;
 
 #[path = "device-sys.rs"]
 pub mod device;
+
+#[path = "queue-sys.rs"]
+pub mod queue;
+
+#[path = "usm-sys.rs"]
+pub mod usm;
diff --git a/oneapi-rs-sys/src/queue-sys.rs b/oneapi-rs-sys/src/queue-sys.rs
@@ -0,0 +1,22 @@
+//
+// 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
+//
+
+#[cxx::bridge(namespace = "sycl_shims::queue")]
+pub mod ffi {
+    unsafe extern "C++" {
+        include!("oneapi-rs-sys/include/queue.hpp");
+
+        #[namespace = "sycl_shims"]
+        type Queue = crate::types::ffi::Queue;
+        #[namespace = "sycl_shims"]
+        type Device = crate::types::ffi::Device;
+
+        fn new_queue() -> UniquePtr<Queue>;
+        fn new_queue_from_device(device: &Device) -> UniquePtr<Queue>;
+    }
+}
diff --git a/oneapi-rs-sys/src/queue.cpp b/oneapi-rs-sys/src/queue.cpp
@@ -0,0 +1,19 @@
+//
+// 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/queue.hpp"
+#include "oneapi-rs-sys/src/queue-sys.rs.h"
+
+namespace sycl_shims::queue {
+std::unique_ptr<Queue> new_queue() {
+  return std::make_unique<Queue>(sycl::queue());
+}
+std::unique_ptr<Queue> new_queue_from_device(Device const & device) {
+  return std::make_unique<Queue>(sycl::queue(device));
+}
+} // namespace sycl_shims::queue
diff --git a/oneapi-rs-sys/src/types-sys.rs b/oneapi-rs-sys/src/types-sys.rs
@@ -12,6 +12,7 @@ pub mod ffi {
         include!("oneapi-rs-sys/include/types.hpp");
         type Device;
         type Platform;
+        type Queue;
     }
 
     // This is a workaround - cxx currently doesn't support passing
@@ -39,6 +40,7 @@ pub mod ffi {
 
     impl UniquePtr<Device> {}
     impl UniquePtr<Platform> {}
+    impl UniquePtr<Queue> {}
 
     impl Vec<DevicePtr> {}
     impl Vec<PlatformPtr> {}
diff --git a/oneapi-rs-sys/src/usm-sys.rs b/oneapi-rs-sys/src/usm-sys.rs
@@ -0,0 +1,23 @@
+//
+// 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
+//
+
+#[cxx::bridge(namespace = "sycl_shims::usm")]
+pub mod ffi {
+    unsafe extern "C++" {
+        #[namespace = "sycl_shims"]
+        type Queue = crate::types::ffi::Queue;
+    }
+
+    extern "C++" {
+        include!("oneapi-rs-sys/include/usm.hpp");
+        unsafe fn aligned_alloc_device(alignment: usize, num_bytes: usize, queue: &Queue) -> Result<*mut u8>;
+        unsafe fn aligned_alloc_host(alignment: usize, num_bytes: usize, queue: &Queue) -> Result<*mut u8>;
+        unsafe fn aligned_alloc_shared(alignment: usize, num_bytes: usize, queue: &Queue) -> Result<*mut u8>;
+        unsafe fn free(ptr: *mut u8, queue: &Queue);
+    }
+}
diff --git a/oneapi-rs-sys/src/usm.cpp b/oneapi-rs-sys/src/usm.cpp
@@ -0,0 +1,28 @@
+//
+// 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/usm.hpp"
+#include "oneapi-rs-sys/src/usm-sys.rs.h"
+
+namespace sycl_shims::usm {
+std::uint8_t* aligned_alloc_device(std::size_t alignment, std::size_t num_bytes, Queue const & queue) {
+  return static_cast<std::uint8_t*>(sycl::aligned_alloc_device(alignment, num_bytes, queue));
+}
+
+std::uint8_t* aligned_alloc_host(std::size_t alignment, std::size_t num_bytes, Queue const & queue) {
+  return static_cast<std::uint8_t*>(sycl::aligned_alloc_host(alignment, num_bytes, queue));
+}
+
+std::uint8_t* aligned_alloc_shared(std::size_t alignment, std::size_t num_bytes, Queue const & queue) {
+  return static_cast<std::uint8_t*>(sycl::aligned_alloc_shared(alignment, num_bytes, queue));
+}
+
+void free(std::uint8_t* memory, Queue const & queue) {
+  sycl::free(static_cast<void*>(memory), queue);
+}
+} // namespace sycl_shims::usm
diff --git a/oneapi-rs/Cargo.toml b/oneapi-rs/Cargo.toml
@@ -5,5 +5,7 @@ edition = "2024"
 license = "MIT OR Apache-2.0"
 
 [dependencies]
+allocator-api2 = "0.4.0"
 cxx = "1.0.194"
 oneapi-rs-sys = { path = "../oneapi-rs-sys" }
+thiserror = "2.0.18"
diff --git a/oneapi-rs/examples/alloc.rs b/oneapi-rs/examples/alloc.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 oneapi_rs::queue::Queue;
+
+fn main() {
+    let queue = Queue::new();
+    let mut buffer = unsafe { queue.alloc_uninit_shared::<u32>(10) };
+
+    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
@@ -0,0 +1,72 @@
+//
+// 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::{alloc::{Layout, handle_alloc_error}, ops::{Deref, DerefMut}, ptr::NonNull, slice};
+
+use crate::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
+/// dimensions that the data is stored and accessed through.
+///
+/// A Buffer does not map to only one underlying backend object, and all SYCL backend memory objects
+/// may be temporary for use on a specific device.
+///
+/// Buffers can be constructed by methods provided by the [`Queue`](`crate::queue::Queue`) class.
+///
+/// The Buffer struct template takes a template parameter [`UsmAlloc`](`crate::usm::UsmAlloc`) for
+/// specifying an allocator which is used by the SYCL runtime when allocating temporary memory on
+/// the host.
+pub struct Buffer<T, A: UsmAlloc> {
+    data: NonNull<T>,
+    len: usize,
+    layout: Layout,
+    allocator: A,
+}
+
+impl<T, A: UsmAlloc> Buffer<T, A> {
+    /// Creates a new buffer given an allocator.
+    /// Safety: returns uninitialized memory.
+    pub(crate) unsafe fn new(allocator: A, len: usize) -> Self {
+        let layout = Layout::array::<T>(len).unwrap();
+        let ptr = match allocator.allocate(layout.clone()) {
+            Ok(ptr) => ptr,
+            _ => handle_alloc_error(layout)
+        };
+
+        Self {
+            data: ptr.cast(),
+            len,
+            layout,
+            allocator,
+        }
+    }
+}
+
+impl<T, A: UsmAlloc> Deref for Buffer<T, A> {
+    type Target = [T];
+    fn deref(&self) -> &Self::Target {
+        unsafe {
+            slice::from_raw_parts(self.data.as_ptr(), self.len)
+        }
+    }
+}
+
+impl<T, A: UsmAlloc> DerefMut for Buffer<T, A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe {
+            slice::from_raw_parts_mut(self.data.as_ptr(), self.len)
+        }
+    }
+}
+
+impl<T, A: UsmAlloc> Drop for Buffer<T, A> {
+    fn drop(&mut self) {
+        unsafe { self.allocator.deallocate(self.data.cast(), self.layout); }
+    }
+}
diff --git a/oneapi-rs/src/lib.rs b/oneapi-rs/src/lib.rs
@@ -9,3 +9,6 @@
 pub mod info;
 pub mod platform;
 pub mod device;
+pub mod queue;
+pub mod usm;
+pub mod buffer;
diff --git a/oneapi-rs/src/queue.rs b/oneapi-rs/src/queue.rs
diff --git a/oneapi-rs/src/usm.rs b/oneapi-rs/src/usm.rs

Original file line number	Diff line number	Diff line change
`@@ -13,4 +13,5 @@`
`13`	`13`	`namespace sycl_shims {`
`14`	`14`	`using Device = sycl::device;`
`15`	`15`	`using Platform = sycl::platform;`
	`16`	`+using Queue = sycl::queue;`
`16`	`17`	`}`