ironaccelerator_levelzero/
compute.rs

1//! Level Zero context + command queue / list scaffold.
2//!
3//! Resolves one driver + device by ordinal, creates a `ze_context` on
4//! it, a compute `ze_command_queue`, and a default `ze_command_list`.
5//! Higher layers will allocate memory and kernels on top; this module
6//! is intentionally the minimum to reach a dispatchable state.
7
8use core::ffi::c_void;
9
10use crate::drv::{
11    self, Loaded, ZeCommandListDesc, ZeCommandListHandle, ZeCommandQueueDesc, ZeCommandQueueHandle,
12    ZeContextDesc, ZeContextHandle, ZeDeviceHandle, ZeDeviceMemAllocDesc, ZeDriverHandle,
13    ZeGroupCount, ZeHostMemAllocDesc, ZeKernelDesc, ZeKernelHandle, ZeModuleDesc, ZeModuleHandle,
14    ZE_COMMAND_QUEUE_MODE_DEFAULT, ZE_COMMAND_QUEUE_PRIORITY_NORMAL, ZE_MODULE_FORMAT_IL_SPIRV,
15    ZE_RESULT_SUCCESS, ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC, ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC,
16    ZE_STRUCTURE_TYPE_CONTEXT_DESC, ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
17    ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC, ZE_STRUCTURE_TYPE_KERNEL_DESC,
18    ZE_STRUCTURE_TYPE_MODULE_DESC,
19};
20
21/// Compute context bound to one Level Zero device: driver, device, context, command queue and command list.
22pub struct Context {
23    l: &'static Loaded,
24    /// Driver handle (`ze_driver_handle_t`).
25    pub driver: ZeDriverHandle,
26    /// Device handle (`ze_device_handle_t`).
27    pub device: ZeDeviceHandle,
28    /// Context handle (`ze_context_handle_t`).
29    pub context: ZeContextHandle,
30    /// Command queue handle (`ze_command_queue_handle_t`).
31    pub queue: ZeCommandQueueHandle,
32    /// Command list handle (`ze_command_list_handle_t`).
33    pub list: ZeCommandListHandle,
34    /// Command-queue-group ordinal chosen for compute. Kept for later
35    /// `zeCommandListAppendLaunchKernel` dispatches.
36    pub queue_ordinal: u32,
37}
38
39impl Context {
40    /// Walk drivers + devices, pick the `global_ordinal`-th device in
41    /// the same order [`drv::enumerate`] produces, and bring up a
42    /// context + compute queue + list on it.
43    pub fn new(global_ordinal: u32) -> Option<Self> {
44        let l = drv::loaded()?;
45        unsafe {
46            let (driver, device) = locate_device(l, global_ordinal)?;
47
48            let mut context: ZeContextHandle = core::ptr::null_mut();
49            let ctx_desc = ZeContextDesc {
50                stype: ZE_STRUCTURE_TYPE_CONTEXT_DESC,
51                p_next: core::ptr::null(),
52                flags: 0,
53            };
54            if (l.ze_context_create)(driver, &ctx_desc, &mut context) != ZE_RESULT_SUCCESS {
55                return None;
56            }
57
58            // Ordinal 0 is the default compute group across every Level
59            // Zero driver shipped to date. A fuller implementation would
60            // query `zeDeviceGetCommandQueueGroupProperties` and pick
61            // the first group whose flags advertise COMPUTE.
62            let queue_ordinal = 0u32;
63
64            let mut queue: ZeCommandQueueHandle = core::ptr::null_mut();
65            let q_desc = ZeCommandQueueDesc {
66                stype: ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC,
67                p_next: core::ptr::null(),
68                ordinal: queue_ordinal,
69                index: 0,
70                flags: 0,
71                mode: ZE_COMMAND_QUEUE_MODE_DEFAULT,
72                priority: ZE_COMMAND_QUEUE_PRIORITY_NORMAL,
73            };
74            if (l.ze_command_queue_create)(context, device, &q_desc, &mut queue)
75                != ZE_RESULT_SUCCESS
76            {
77                (l.ze_context_destroy)(context);
78                return None;
79            }
80
81            let mut list: ZeCommandListHandle = core::ptr::null_mut();
82            let l_desc = ZeCommandListDesc {
83                stype: ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC,
84                p_next: core::ptr::null(),
85                command_queue_group_ordinal: queue_ordinal,
86                flags: 0,
87            };
88            if (l.ze_command_list_create)(context, device, &l_desc, &mut list) != ZE_RESULT_SUCCESS
89            {
90                (l.ze_command_queue_destroy)(queue);
91                (l.ze_context_destroy)(context);
92                return None;
93            }
94
95            Some(Context {
96                l,
97                driver,
98                device,
99                context,
100                queue,
101                list,
102                queue_ordinal,
103            })
104        }
105    }
106}
107
108impl Context {
109    /// Allocate `size` bytes of device-local memory on this context's
110    /// device. Returned pointer is an unmapped USM address valid for
111    /// `zeCommandListAppendMemoryCopy` and kernel arg binding.
112    pub fn alloc_device(&self, size: usize, alignment: usize) -> Option<DeviceBuffer> {
113        let desc = ZeDeviceMemAllocDesc {
114            stype: ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
115            p_next: core::ptr::null(),
116            flags: 0,
117            ordinal: self.queue_ordinal,
118        };
119        let mut ptr: *mut c_void = core::ptr::null_mut();
120        unsafe {
121            if (self.l.ze_mem_alloc_device)(
122                self.context,
123                &desc,
124                size,
125                alignment,
126                self.device,
127                &mut ptr,
128            ) != ZE_RESULT_SUCCESS
129            {
130                return None;
131            }
132        }
133        Some(DeviceBuffer {
134            l: self.l,
135            context: self.context,
136            ptr,
137            size,
138        })
139    }
140
141    /// Allocate `size` bytes of shared USM (host + device accessible).
142    pub fn alloc_shared(&self, size: usize, alignment: usize) -> Option<DeviceBuffer> {
143        let ddesc = ZeDeviceMemAllocDesc {
144            stype: ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
145            p_next: core::ptr::null(),
146            flags: 0,
147            ordinal: self.queue_ordinal,
148        };
149        let hdesc = ZeHostMemAllocDesc {
150            stype: ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC,
151            p_next: core::ptr::null(),
152            flags: 0,
153        };
154        let mut ptr: *mut c_void = core::ptr::null_mut();
155        unsafe {
156            if (self.l.ze_mem_alloc_shared)(
157                self.context,
158                &ddesc,
159                &hdesc,
160                size,
161                alignment,
162                self.device,
163                &mut ptr,
164            ) != ZE_RESULT_SUCCESS
165            {
166                return None;
167            }
168        }
169        Some(DeviceBuffer {
170            l: self.l,
171            context: self.context,
172            ptr,
173            size,
174        })
175    }
176
177    /// Load a SPIR-V module onto the device.
178    pub fn load_spirv(&self, spirv: &[u8]) -> Option<Module> {
179        let desc = ZeModuleDesc {
180            stype: ZE_STRUCTURE_TYPE_MODULE_DESC,
181            p_next: core::ptr::null(),
182            format: ZE_MODULE_FORMAT_IL_SPIRV,
183            input_size: spirv.len(),
184            p_input_module: spirv.as_ptr(),
185            p_build_flags: core::ptr::null(),
186            p_constants: core::ptr::null(),
187        };
188        let mut module: ZeModuleHandle = core::ptr::null_mut();
189        unsafe {
190            if (self.l.ze_module_create)(
191                self.context,
192                self.device,
193                &desc,
194                &mut module,
195                core::ptr::null_mut(),
196            ) != ZE_RESULT_SUCCESS
197            {
198                return None;
199            }
200        }
201        Some(Module { l: self.l, module })
202    }
203
204    /// Append a kernel launch to `self.list`, close + execute the list,
205    /// and wait for the queue to drain. One-shot pattern — higher layers
206    /// will want to reuse command lists.
207    pub fn launch(&self, kernel: &Kernel, group_count: [u32; 3]) -> Result<(), u32> {
208        let gc = ZeGroupCount {
209            group_count_x: group_count[0],
210            group_count_y: group_count[1],
211            group_count_z: group_count[2],
212        };
213        unsafe {
214            let r = (self.l.ze_command_list_append_launch_kernel)(
215                self.list,
216                kernel.kernel,
217                &gc,
218                core::ptr::null_mut(),
219                0,
220                core::ptr::null_mut(),
221            );
222            if r != ZE_RESULT_SUCCESS {
223                return Err(r);
224            }
225            let r = (self.l.ze_command_list_close)(self.list);
226            if r != ZE_RESULT_SUCCESS {
227                return Err(r);
228            }
229            let lists = [self.list];
230            let r = (self.l.ze_command_queue_execute_command_lists)(
231                self.queue,
232                1,
233                lists.as_ptr(),
234                core::ptr::null_mut(),
235            );
236            if r != ZE_RESULT_SUCCESS {
237                return Err(r);
238            }
239            let r = (self.l.ze_command_queue_synchronize)(self.queue, u64::MAX);
240            if r != ZE_RESULT_SUCCESS {
241                return Err(r);
242            }
243            let _ = (self.l.ze_command_list_reset)(self.list);
244        }
245        Ok(())
246    }
247}
248
249impl Drop for Context {
250    fn drop(&mut self) {
251        unsafe {
252            (self.l.ze_command_list_destroy)(self.list);
253            (self.l.ze_command_queue_destroy)(self.queue);
254            (self.l.ze_context_destroy)(self.context);
255        }
256    }
257}
258
259/// Device or shared allocation from `zeMemAllocDevice` or `zeMemAllocShared`, freed with `zeMemFree` on drop.
260pub struct DeviceBuffer {
261    l: &'static Loaded,
262    context: ZeContextHandle,
263    /// Base pointer of the allocation.
264    pub ptr: *mut c_void,
265    /// Allocation size in bytes.
266    pub size: usize,
267}
268
269impl Drop for DeviceBuffer {
270    fn drop(&mut self) {
271        unsafe {
272            (self.l.ze_mem_free)(self.context, self.ptr);
273        }
274    }
275}
276
277/// Module created with `zeModuleCreate`, destroyed on drop.
278pub struct Module {
279    l: &'static Loaded,
280    /// Module handle (`ze_module_handle_t`).
281    pub module: ZeModuleHandle,
282}
283
284impl Module {
285    /// Create a kernel object bound to `name` inside this module.
286    pub fn kernel(&self, name: &str) -> Option<Kernel> {
287        let cname = std::ffi::CString::new(name).ok()?;
288        let desc = ZeKernelDesc {
289            stype: ZE_STRUCTURE_TYPE_KERNEL_DESC,
290            p_next: core::ptr::null(),
291            flags: 0,
292            p_kernel_name: cname.as_ptr(),
293        };
294        let mut k: ZeKernelHandle = core::ptr::null_mut();
295        unsafe {
296            if (self.l.ze_kernel_create)(self.module, &desc, &mut k) != ZE_RESULT_SUCCESS {
297                return None;
298            }
299        }
300        Some(Kernel {
301            l: self.l,
302            kernel: k,
303        })
304    }
305}
306
307impl Drop for Module {
308    fn drop(&mut self) {
309        unsafe {
310            (self.l.ze_module_destroy)(self.module);
311        }
312    }
313}
314
315/// Kernel created with `zeKernelCreate`, destroyed on drop.
316pub struct Kernel {
317    l: &'static Loaded,
318    /// Kernel handle (`ze_kernel_handle_t`).
319    pub kernel: ZeKernelHandle,
320}
321
322impl Kernel {
323    /// Sets the group size via `zeKernelSetGroupSize`.
324    pub fn set_group_size(&self, gx: u32, gy: u32, gz: u32) -> Result<(), u32> {
325        unsafe {
326            match (self.l.ze_kernel_set_group_size)(self.kernel, gx, gy, gz) {
327                ZE_RESULT_SUCCESS => Ok(()),
328                e => Err(e),
329            }
330        }
331    }
332
333    /// Bind argument `index` to `value` (bytewise — for pointer args,
334    /// pass `&buf.ptr`; for scalar args, pass `&scalar`).
335    ///
336    /// # Safety
337    /// `value` must remain valid for the call's duration and match the
338    /// kernel's declared argument layout.
339    pub unsafe fn set_arg<T>(&self, index: u32, value: &T) -> Result<(), u32> {
340        match (self.l.ze_kernel_set_argument_value)(
341            self.kernel,
342            index,
343            core::mem::size_of::<T>(),
344            value as *const T as *const c_void,
345        ) {
346            ZE_RESULT_SUCCESS => Ok(()),
347            e => Err(e),
348        }
349    }
350}
351
352impl Drop for Kernel {
353    fn drop(&mut self) {
354        unsafe {
355            (self.l.ze_kernel_destroy)(self.kernel);
356        }
357    }
358}
359
360/// A SPIR-V module and one kernel from it, kept together — Level Zero destroys
361/// the kernel before the module, and the unified
362/// [`ComputeDevice`](ironaccelerator_core::ComputeDevice) trait hands back a
363/// single pipeline object. `kernel` is declared first so it drops first.
364pub struct Pipeline {
365    kernel: Kernel,
366    _module: Module,
367}
368
369/// Unified cross-backend compute surface. `code` is **OpenCL/SYCL-flavored**
370/// SPIR-V — the `Kernel` execution model, where the entry point takes its
371/// buffers as `__global` pointer arguments. This is *not* the descriptor-bound
372/// `GLCompute` SPIR-V that Vulkan's GLSL produces: `dispatch` binds each buffer
373/// as a pointer argument (`zeKernelSetArgumentValue`), matching how oneAPI,
374/// `ocloc`, and `clang -target spir64` emit compute kernels. The entry point is
375/// assumed to be `main`. Buffers are shared USM, so `upload`/`download` are a
376/// `memcpy` through the pointer with no staging.
377///
378/// Like Metal, Level Zero sets the group size at dispatch rather than in the
379/// shader, so [`dispatch`](ironaccelerator_core::ComputeDevice::dispatch)
380/// assumes a 1-D group of 64. Use [`Kernel::set_group_size`] +
381/// [`Context::launch`] directly for other geometries.
382impl ironaccelerator_core::ComputeDevice for Context {
383    type Buffer = DeviceBuffer;
384    type Pipeline = Pipeline;
385    type Error = String;
386
387    fn device_buffer(&self, bytes: u64) -> Result<DeviceBuffer, String> {
388        self.alloc_shared(bytes as usize, 64)
389            .ok_or_else(|| "level-zero: shared USM allocation failed".to_string())
390    }
391
392    fn upload(&self, data: &[u8]) -> Result<DeviceBuffer, String> {
393        let buf = self.device_buffer(data.len() as u64)?;
394        // SAFETY: shared USM is host-accessible; `buf.ptr` is valid for
395        // `data.len()` bytes (the buffer was sized to it).
396        unsafe {
397            core::ptr::copy_nonoverlapping(data.as_ptr(), buf.ptr as *mut u8, data.len());
398        }
399        Ok(buf)
400    }
401
402    fn download(&self, buffer: &DeviceBuffer, out: &mut [u8]) -> Result<(), String> {
403        let n = out.len().min(buffer.size);
404        // SAFETY: shared USM is host-accessible; reading `n` bytes stays within
405        // the buffer's length.
406        unsafe {
407            core::ptr::copy_nonoverlapping(buffer.ptr as *const u8, out.as_mut_ptr(), n);
408        }
409        Ok(())
410    }
411
412    fn pipeline(&self, code: &[u8], _bindings: u32) -> Result<Pipeline, String> {
413        let module = self
414            .load_spirv(code)
415            .ok_or_else(|| "level-zero: SPIR-V module build failed".to_string())?;
416        let kernel = module
417            .kernel("main")
418            .ok_or_else(|| "level-zero: kernel `main` not found in module".to_string())?;
419        Ok(Pipeline {
420            kernel,
421            _module: module,
422        })
423    }
424
425    fn dispatch(
426        &self,
427        pipeline: &Pipeline,
428        buffers: &[&DeviceBuffer],
429        groups: [u32; 3],
430    ) -> Result<(), String> {
431        let fmt = |e: u32| format!("level-zero error {e:#010x}");
432        pipeline.kernel.set_group_size(64, 1, 1).map_err(fmt)?;
433        for (i, b) in buffers.iter().enumerate() {
434            // Pointer kernel argument: pass the USM address by value.
435            // SAFETY: `b.ptr` outlives the launch; the kernel's arg `i` is a
436            // pointer, matching `size_of::<*mut c_void>()`.
437            unsafe {
438                pipeline.kernel.set_arg(i as u32, &b.ptr).map_err(fmt)?;
439            }
440        }
441        self.launch(&pipeline.kernel, groups).map_err(fmt)
442    }
443
444    fn buffer_len(&self, buffer: &DeviceBuffer) -> u64 {
445        buffer.size as u64
446    }
447}
448
449unsafe fn locate_device(l: &Loaded, target: u32) -> Option<(ZeDriverHandle, ZeDeviceHandle)> {
450    let mut driver_count: u32 = 0;
451    if (l.ze_driver_get)(&mut driver_count, core::ptr::null_mut()) != ZE_RESULT_SUCCESS
452        || driver_count == 0
453    {
454        return None;
455    }
456    let mut drivers = vec![core::ptr::null_mut::<c_void>(); driver_count as usize];
457    if (l.ze_driver_get)(&mut driver_count, drivers.as_mut_ptr()) != ZE_RESULT_SUCCESS {
458        return None;
459    }
460    let mut seen = 0u32;
461    for driver in drivers.into_iter().take(driver_count as usize) {
462        let mut dev_count: u32 = 0;
463        if (l.ze_device_get)(driver, &mut dev_count, core::ptr::null_mut()) != ZE_RESULT_SUCCESS
464            || dev_count == 0
465        {
466            continue;
467        }
468        let mut devs = vec![core::ptr::null_mut::<c_void>(); dev_count as usize];
469        if (l.ze_device_get)(driver, &mut dev_count, devs.as_mut_ptr()) != ZE_RESULT_SUCCESS {
470            continue;
471        }
472        for dev in devs.into_iter().take(dev_count as usize) {
473            if seen == target {
474                return Some((driver, dev));
475            }
476            seen += 1;
477        }
478    }
479    None
480}
ironaccelerator_levelzero/compute.rs

ironaccelerator_levelzero/
compute.rs
