oneliner_executor/
lib.rs

1#![no_std]
2
3use oneliner_iree_abi::{
4    iree_hal_executable_dispatch_state_v0_t, iree_hal_executable_environment_v0_t,
5    iree_hal_executable_workgroup_state_v0_t, DispatchFn,
6};
7use portable_atomic::{AtomicI32, Ordering};
8
9#[derive(Clone, Copy)]
10pub struct WorkItem {
11    kind: WorkItemKind,
12}
13
14#[derive(Clone, Copy)]
15enum WorkItemKind {
16    IreeWorkload {
17        dispatch_fn: DispatchFn,
18        environment: *mut iree_hal_executable_environment_v0_t,
19        dispatch_state: *mut iree_hal_executable_dispatch_state_v0_t,
20        workgroup_state: iree_hal_executable_workgroup_state_v0_t,
21        status: *const AtomicI32,
22    },
23}
24
25impl WorkItem {
26    /// Creates an IREE work item borrowing dispatch state owned by the caller.
27    ///
28    /// # Safety
29    ///
30    /// All pointers must remain valid until the executor reports completion.
31    #[doc(hidden)]
32    pub unsafe fn iree(
33        dispatch_fn: DispatchFn,
34        environment: *mut iree_hal_executable_environment_v0_t,
35        dispatch_state: *mut iree_hal_executable_dispatch_state_v0_t,
36        workgroup_state: iree_hal_executable_workgroup_state_v0_t,
37        status: *const AtomicI32,
38    ) -> Self {
39        Self {
40            kind: WorkItemKind::IreeWorkload {
41                dispatch_fn,
42                environment,
43                dispatch_state,
44                workgroup_state,
45                status,
46            },
47        }
48    }
49
50    #[doc(hidden)]
51    pub fn run(self) {
52        match self.kind {
53            WorkItemKind::IreeWorkload {
54                dispatch_fn,
55                environment,
56                dispatch_state,
57                workgroup_state,
58                status,
59            } => {
60                let dispatch_status =
61                    unsafe { dispatch_fn(environment, dispatch_state, &workgroup_state) };
62                if dispatch_status != 0 {
63                    let status = unsafe { &*status };
64                    let _ = status.compare_exchange(
65                        0,
66                        dispatch_status,
67                        Ordering::Release,
68                        Ordering::Relaxed,
69                    );
70                }
71            }
72        }
73    }
74}
75
76/// Schedules work items and provides a completion barrier.
77pub trait Executor {
78    /// Schedules one work item for execution.
79    fn schedule(&mut self, item: WorkItem);
80
81    /// Waits until all previously scheduled work has finished.
82    fn wait_job_completion(&mut self);
83}
84
85/// Executor that runs work items immediately in submission order.
86#[derive(Debug, Default, Clone, Copy)]
87pub struct SequentialExecutor;
88
89impl SequentialExecutor {
90    pub const fn new() -> Self {
91        Self
92    }
93}
94
95impl Executor for SequentialExecutor {
96    fn schedule(&mut self, item: WorkItem) {
97        item.run();
98    }
99
100    fn wait_job_completion(&mut self) {}
101}
oneliner_executor/lib.rs

oneliner_executor/
lib.rs
