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
21pub struct Context {
22    l: &'static Loaded,
23    pub driver: ZeDriverHandle,
24    pub device: ZeDeviceHandle,
25    pub context: ZeContextHandle,
26    pub queue: ZeCommandQueueHandle,
27    pub list: ZeCommandListHandle,
28    /// Command-queue-group ordinal chosen for compute. Kept for later
29    /// `zeCommandListAppendLaunchKernel` dispatches.
30    pub queue_ordinal: u32,
31}
32
33impl Context {
34    /// Walk drivers + devices, pick the `global_ordinal`-th device in
35    /// the same order [`drv::enumerate`] produces, and bring up a
36    /// context + compute queue + list on it.
37    pub fn new(global_ordinal: u32) -> Option<Self> {
38        let l = drv::loaded()?;
39        unsafe {
40            let (driver, device) = locate_device(l, global_ordinal)?;
41
42            let mut context: ZeContextHandle = core::ptr::null_mut();
43            let ctx_desc = ZeContextDesc {
44                stype: ZE_STRUCTURE_TYPE_CONTEXT_DESC,
45                p_next: core::ptr::null(),
46                flags: 0,
47            };
48            if (l.ze_context_create)(driver, &ctx_desc, &mut context) != ZE_RESULT_SUCCESS {
49                return None;
50            }
51
52            // Ordinal 0 is the default compute group across every Level
53            // Zero driver shipped to date. A fuller implementation would
54            // query `zeDeviceGetCommandQueueGroupProperties` and pick
55            // the first group whose flags advertise COMPUTE.
56            let queue_ordinal = 0u32;
57
58            let mut queue: ZeCommandQueueHandle = core::ptr::null_mut();
59            let q_desc = ZeCommandQueueDesc {
60                stype: ZE_STRUCTURE_TYPE_COMMAND_QUEUE_DESC,
61                p_next: core::ptr::null(),
62                ordinal: queue_ordinal,
63                index: 0,
64                flags: 0,
65                mode: ZE_COMMAND_QUEUE_MODE_DEFAULT,
66                priority: ZE_COMMAND_QUEUE_PRIORITY_NORMAL,
67            };
68            if (l.ze_command_queue_create)(context, device, &q_desc, &mut queue)
69                != ZE_RESULT_SUCCESS
70            {
71                (l.ze_context_destroy)(context);
72                return None;
73            }
74
75            let mut list: ZeCommandListHandle = core::ptr::null_mut();
76            let l_desc = ZeCommandListDesc {
77                stype: ZE_STRUCTURE_TYPE_COMMAND_LIST_DESC,
78                p_next: core::ptr::null(),
79                command_queue_group_ordinal: queue_ordinal,
80                flags: 0,
81            };
82            if (l.ze_command_list_create)(context, device, &l_desc, &mut list) != ZE_RESULT_SUCCESS
83            {
84                (l.ze_command_queue_destroy)(queue);
85                (l.ze_context_destroy)(context);
86                return None;
87            }
88
89            Some(Context {
90                l,
91                driver,
92                device,
93                context,
94                queue,
95                list,
96                queue_ordinal,
97            })
98        }
99    }
100}
101
102impl Context {
103    /// Allocate `size` bytes of device-local memory on this context's
104    /// device. Returned pointer is an unmapped USM address valid for
105    /// `zeCommandListAppendMemoryCopy` and kernel arg binding.
106    pub fn alloc_device(&self, size: usize, alignment: usize) -> Option<DeviceBuffer> {
107        let desc = ZeDeviceMemAllocDesc {
108            stype: ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
109            p_next: core::ptr::null(),
110            flags: 0,
111            ordinal: self.queue_ordinal,
112        };
113        let mut ptr: *mut c_void = core::ptr::null_mut();
114        unsafe {
115            if (self.l.ze_mem_alloc_device)(
116                self.context,
117                &desc,
118                size,
119                alignment,
120                self.device,
121                &mut ptr,
122            ) != ZE_RESULT_SUCCESS
123            {
124                return None;
125            }
126        }
127        Some(DeviceBuffer {
128            l: self.l,
129            context: self.context,
130            ptr,
131            size,
132        })
133    }
134
135    /// Allocate `size` bytes of shared USM (host + device accessible).
136    pub fn alloc_shared(&self, size: usize, alignment: usize) -> Option<DeviceBuffer> {
137        let ddesc = ZeDeviceMemAllocDesc {
138            stype: ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
139            p_next: core::ptr::null(),
140            flags: 0,
141            ordinal: self.queue_ordinal,
142        };
143        let hdesc = ZeHostMemAllocDesc {
144            stype: ZE_STRUCTURE_TYPE_HOST_MEM_ALLOC_DESC,
145            p_next: core::ptr::null(),
146            flags: 0,
147        };
148        let mut ptr: *mut c_void = core::ptr::null_mut();
149        unsafe {
150            if (self.l.ze_mem_alloc_shared)(
151                self.context,
152                &ddesc,
153                &hdesc,
154                size,
155                alignment,
156                self.device,
157                &mut ptr,
158            ) != ZE_RESULT_SUCCESS
159            {
160                return None;
161            }
162        }
163        Some(DeviceBuffer {
164            l: self.l,
165            context: self.context,
166            ptr,
167            size,
168        })
169    }
170
171    /// Load a SPIR-V module onto the device.
172    pub fn load_spirv(&self, spirv: &[u8]) -> Option<Module> {
173        let desc = ZeModuleDesc {
174            stype: ZE_STRUCTURE_TYPE_MODULE_DESC,
175            p_next: core::ptr::null(),
176            format: ZE_MODULE_FORMAT_IL_SPIRV,
177            input_size: spirv.len(),
178            p_input_module: spirv.as_ptr(),
179            p_build_flags: core::ptr::null(),
180            p_constants: core::ptr::null(),
181        };
182        let mut module: ZeModuleHandle = core::ptr::null_mut();
183        unsafe {
184            if (self.l.ze_module_create)(
185                self.context,
186                self.device,
187                &desc,
188                &mut module,
189                core::ptr::null_mut(),
190            ) != ZE_RESULT_SUCCESS
191            {
192                return None;
193            }
194        }
195        Some(Module { l: self.l, module })
196    }
197
198    /// Append a kernel launch to `self.list`, close + execute the list,
199    /// and wait for the queue to drain. One-shot pattern — higher layers
200    /// will want to reuse command lists.
201    pub fn launch(&self, kernel: &Kernel, group_count: [u32; 3]) -> Result<(), u32> {
202        let gc = ZeGroupCount {
203            group_count_x: group_count[0],
204            group_count_y: group_count[1],
205            group_count_z: group_count[2],
206        };
207        unsafe {
208            let r = (self.l.ze_command_list_append_launch_kernel)(
209                self.list,
210                kernel.kernel,
211                &gc,
212                core::ptr::null_mut(),
213                0,
214                core::ptr::null_mut(),
215            );
216            if r != ZE_RESULT_SUCCESS {
217                return Err(r);
218            }
219            let r = (self.l.ze_command_list_close)(self.list);
220            if r != ZE_RESULT_SUCCESS {
221                return Err(r);
222            }
223            let lists = [self.list];
224            let r = (self.l.ze_command_queue_execute_command_lists)(
225                self.queue,
226                1,
227                lists.as_ptr(),
228                core::ptr::null_mut(),
229            );
230            if r != ZE_RESULT_SUCCESS {
231                return Err(r);
232            }
233            let r = (self.l.ze_command_queue_synchronize)(self.queue, u64::MAX);
234            if r != ZE_RESULT_SUCCESS {
235                return Err(r);
236            }
237            let _ = (self.l.ze_command_list_reset)(self.list);
238        }
239        Ok(())
240    }
241}
242
243impl Drop for Context {
244    fn drop(&mut self) {
245        unsafe {
246            (self.l.ze_command_list_destroy)(self.list);
247            (self.l.ze_command_queue_destroy)(self.queue);
248            (self.l.ze_context_destroy)(self.context);
249        }
250    }
251}
252
253pub struct DeviceBuffer {
254    l: &'static Loaded,
255    context: ZeContextHandle,
256    pub ptr: *mut c_void,
257    pub size: usize,
258}
259
260impl Drop for DeviceBuffer {
261    fn drop(&mut self) {
262        unsafe {
263            (self.l.ze_mem_free)(self.context, self.ptr);
264        }
265    }
266}
267
268pub struct Module {
269    l: &'static Loaded,
270    pub module: ZeModuleHandle,
271}
272
273impl Module {
274    /// Create a kernel object bound to `name` inside this module.
275    pub fn kernel(&self, name: &str) -> Option<Kernel> {
276        let cname = std::ffi::CString::new(name).ok()?;
277        let desc = ZeKernelDesc {
278            stype: ZE_STRUCTURE_TYPE_KERNEL_DESC,
279            p_next: core::ptr::null(),
280            flags: 0,
281            p_kernel_name: cname.as_ptr(),
282        };
283        let mut k: ZeKernelHandle = core::ptr::null_mut();
284        unsafe {
285            if (self.l.ze_kernel_create)(self.module, &desc, &mut k) != ZE_RESULT_SUCCESS {
286                return None;
287            }
288        }
289        Some(Kernel {
290            l: self.l,
291            kernel: k,
292        })
293    }
294}
295
296impl Drop for Module {
297    fn drop(&mut self) {
298        unsafe {
299            (self.l.ze_module_destroy)(self.module);
300        }
301    }
302}
303
304pub struct Kernel {
305    l: &'static Loaded,
306    pub kernel: ZeKernelHandle,
307}
308
309impl Kernel {
310    pub fn set_group_size(&self, gx: u32, gy: u32, gz: u32) -> Result<(), u32> {
311        unsafe {
312            match (self.l.ze_kernel_set_group_size)(self.kernel, gx, gy, gz) {
313                ZE_RESULT_SUCCESS => Ok(()),
314                e => Err(e),
315            }
316        }
317    }
318
319    /// Bind argument `index` to `value` (bytewise — for pointer args,
320    /// pass `&buf.ptr`; for scalar args, pass `&scalar`).
321    ///
322    /// # Safety
323    /// `value` must remain valid for the call's duration and match the
324    /// kernel's declared argument layout.
325    pub unsafe fn set_arg<T>(&self, index: u32, value: &T) -> Result<(), u32> {
326        match (self.l.ze_kernel_set_argument_value)(
327            self.kernel,
328            index,
329            core::mem::size_of::<T>(),
330            value as *const T as *const c_void,
331        ) {
332            ZE_RESULT_SUCCESS => Ok(()),
333            e => Err(e),
334        }
335    }
336}
337
338impl Drop for Kernel {
339    fn drop(&mut self) {
340        unsafe {
341            (self.l.ze_kernel_destroy)(self.kernel);
342        }
343    }
344}
345
346unsafe fn locate_device(l: &Loaded, target: u32) -> Option<(ZeDriverHandle, ZeDeviceHandle)> {
347    let mut driver_count: u32 = 0;
348    if (l.ze_driver_get)(&mut driver_count, core::ptr::null_mut()) != ZE_RESULT_SUCCESS
349        || driver_count == 0
350    {
351        return None;
352    }
353    let mut drivers = vec![core::ptr::null_mut::<c_void>(); driver_count as usize];
354    if (l.ze_driver_get)(&mut driver_count, drivers.as_mut_ptr()) != ZE_RESULT_SUCCESS {
355        return None;
356    }
357    let mut seen = 0u32;
358    for driver in drivers.into_iter().take(driver_count as usize) {
359        let mut dev_count: u32 = 0;
360        if (l.ze_device_get)(driver, &mut dev_count, core::ptr::null_mut()) != ZE_RESULT_SUCCESS
361            || dev_count == 0
362        {
363            continue;
364        }
365        let mut devs = vec![core::ptr::null_mut::<c_void>(); dev_count as usize];
366        if (l.ze_device_get)(driver, &mut dev_count, devs.as_mut_ptr()) != ZE_RESULT_SUCCESS {
367            continue;
368        }
369        for dev in devs.into_iter().take(dev_count as usize) {
370            if seen == target {
371                return Some((driver, dev));
372            }
373            seen += 1;
374        }
375    }
376    None
377}
ironaccelerator_levelzero/compute.rs

ironaccelerator_levelzero/
compute.rs
