hrx/execution/
mod.rs

1//! Coordinated execution for GPU and NPU kernels.
2//!
3//! Device placement is explicit. Contracts describe memory access; prepared
4//! graphs infer dependencies and retain all allocations through completion.
5//! The low-level GPU API remains available in [`crate::gpu`].
6mod buffer;
7mod completion;
8mod contract;
9mod dependencies;
10mod graph;
11mod handoff;
12mod native_session;
13mod scheduler;
14mod statistics;
15mod trace;
16use crate::{Error, Result};
17pub use buffer::{Buffer, BufferView, MemoryPlacement, ReadGuard, WriteGuard};
18use buffer::{Engine, HostState, Storage, Visibility};
19pub use completion::{Completion, CompletionProfile};
20pub use contract::{Access, BindingContract, KernelContract};
21pub use graph::{ExecutableGraph, GpuKernel, GpuLane, Graph, Node};
22pub use native_session::NativeSession;
23use scheduler::Core;
24pub use statistics::Statistics;
25use std::sync::{Arc, Mutex};
26pub use trace::{ExecutionTrace, TraceEvent};
27
28/// A buffer region and the access performed during an external GPU handoff.
29#[derive(Clone)]
30pub struct GpuAccess {
31    /// Region retained and reserved until the supplied stream completes.
32    pub view: BufferView,
33    /// Reads and writes performed by the callback.
34    pub access: Access,
35}
36
37pub(super) struct RuntimeOwner {
38    core: Arc<Core>,
39    workers: Mutex<Vec<std::thread::JoinHandle<()>>>,
40    gpu_index: i32,
41    copy_streams: [Mutex<Option<Arc<Mutex<graph::CopyStream>>>>; 3],
42}
43impl Drop for RuntimeOwner {
44    fn drop(&mut self) {
45        self.core
46            .state
47            .lock()
48            .unwrap_or_else(|e| e.into_inner())
49            .shutdown = true;
50        self.core.changed.notify_all();
51        for worker in self
52            .workers
53            .get_mut()
54            .unwrap_or_else(|e| e.into_inner())
55            .drain(..)
56        {
57            if worker.thread().id() != std::thread::current().id() {
58                let _ = worker.join();
59            }
60        }
61    }
62}
63/// Runtime limits fixed before workers or prepared run slots are created.
64#[derive(Clone, Debug)]
65pub struct RuntimeOptions {
66    /// Optional shared byte ceiling for tracked allocations, including weights,
67    /// private scratch and transfer staging. Native allocations made outside
68    /// this runtime are not covered; adoption charges them at the boundary.
69    pub memory_budget: Option<crate::residency::MemoryBudget>,
70    /// Physical GPU index; zero selects the integrated GPU on the tested host.
71    pub gpu_index: i32,
72    /// Maximum simultaneous submissions. Exhaustion returns `Busy`.
73    /// This bounds queued work; it does not increase per-engine concurrency.
74    pub max_submissions: usize,
75    /// Preallocated completion slots per prepared graph.
76    pub graph_slots: usize,
77}
78impl Default for RuntimeOptions {
79    fn default() -> Self {
80        Self {
81            memory_budget: None,
82            gpu_index: 0,
83            max_submissions: 64,
84            graph_slots: 2,
85        }
86    }
87}
88/// Shared scheduler and allocation domain. Clones use the same dependency state.
89///
90/// Upload, compute, download, and NPU regions have independent scheduling lanes.
91/// Each lane runs one region at a time; conflicting memory accesses are ordered
92/// across all lanes and submissions. Actual device overlap depends on hardware.
93#[derive(Clone)]
94pub struct Runtime {
95    pub(super) inner: Arc<RuntimeOwner>,
96    options: RuntimeOptions,
97}
98/// A selected GPU and its verified compilation target.
99#[derive(Clone)]
100pub struct GpuDevice {
101    index: i32,
102    target: crate::Target,
103}
104impl GpuDevice {
105    /// Physical GPU ordinal.
106    pub fn index(&self) -> i32 {
107        self.index
108    }
109    /// Target for compiling a kernel on this GPU.
110    pub fn target(&self) -> &crate::Target {
111        &self.target
112    }
113}
114/// An explicitly selected NPU ordinal.
115#[cfg(feature = "npu")]
116#[derive(Clone)]
117pub struct NpuDevice {
118    index: i32,
119    budget: Option<crate::residency::MemoryBudget>,
120}
121#[cfg(feature = "npu")]
122impl NpuDevice {
123    /// Physical NPU ordinal.
124    pub fn index(&self) -> i32 {
125        self.index
126    }
127    /// Load trusted code on this device.
128    /// # Safety
129    /// The image must satisfy [`crate::npu::NpuProgram::load`]'s contract.
130    pub unsafe fn load_program(
131        &self,
132        path: impl AsRef<std::path::Path>,
133    ) -> Result<crate::npu::NpuProgram> {
134        let mut program = unsafe { crate::npu::NpuProgram::load(self.index, path) }?;
135        program.budget = self.budget.clone();
136        Ok(program)
137    }
138}
139impl Runtime {
140    /// Begin an opt-in bounded trace of host-observed native-region latencies.
141    pub fn start_trace(&self, event_capacity: usize) -> Result<()> {
142        self.inner.core.tracer.start(event_capacity)
143    }
144    /// Stop capture without waiting for pending work. Wait for its completions
145    /// first when the capture should include the entire request.
146    pub fn finish_trace(&self) -> Option<ExecutionTrace> {
147        self.inner.core.tracer.finish()
148    }
149    /// Whether a view belongs to this runtime's allocation and scheduling domain.
150    pub fn owns(&self, view: &BufferView) -> bool {
151        Arc::ptr_eq(&self.inner, &view.buffer.storage.runtime)
152    }
153
154    /// Whether two handles share the same allocation and scheduling domain.
155    pub fn same_domain(&self, other: &Self) -> bool {
156        Arc::ptr_eq(&self.inner, &other.inner)
157    }
158
159    /// Shared allocation ceiling, if selected at construction. Native clients
160    /// can attach it to their streams before loading weights or workspace.
161    pub fn memory_budget(&self) -> Option<&crate::residency::MemoryBudget> {
162        self.options.memory_budget.as_ref()
163    }
164
165    /// Create a runtime. Native devices are opened lazily when requested.
166    pub fn new() -> Result<Self> {
167        Self::with_options(RuntimeOptions::default())
168    }
169    /// Create a runtime with bounded pending submission and graph capacities.
170    pub fn with_options(options: RuntimeOptions) -> Result<Self> {
171        if options.max_submissions == 0 || options.graph_slots == 0 || options.gpu_index < 0 {
172            return Err(Error::Message(
173                "runtime capacities must be nonzero and GPU index nonnegative".into(),
174            ));
175        }
176        let core = Arc::new(Core::new(options.max_submissions));
177        let owner = Arc::new(RuntimeOwner {
178            core: core.clone(),
179            workers: Mutex::new(Vec::new()),
180            gpu_index: options.gpu_index,
181            copy_streams: std::array::from_fn(|_| Mutex::new(None)),
182        });
183        for name in ["hrx-upload", "hrx-compute", "hrx-download", "hrx-npu"] {
184            let core = core.clone();
185            let worker = std::thread::Builder::new()
186                .name(name.into())
187                .spawn(move || scheduler::worker(core))?;
188            owner
189                .workers
190                .lock()
191                .unwrap_or_else(|e| e.into_inner())
192                .push(worker);
193        }
194        Ok(Self {
195            inner: owner,
196            options,
197        })
198    }
199    /// Open the selected GPU and return its actual target.
200    pub fn gpu(&self) -> Result<GpuDevice> {
201        let device = crate::gpu::Device::open(self.inner.gpu_index)?;
202        Ok(GpuDevice {
203            index: self.inner.gpu_index,
204            target: device.target().clone(),
205        })
206    }
207    /// Select an NPU. Opening its program validates actual device access.
208    #[cfg(feature = "npu")]
209    pub fn npu(&self, index: i32) -> Result<NpuDevice> {
210        if index < 0 {
211            return Err(Error::Message("NPU index must be nonnegative".into()));
212        }
213        Ok(NpuDevice {
214            index,
215            budget: self.options.memory_budget.clone(),
216        })
217    }
218    /// Allocate initialized storage. Shared memory is exported and imported once.
219    pub fn allocate(&self, bytes: usize, placement: MemoryPlacement) -> Result<Buffer> {
220        self.allocate_inner(bytes, placement, None)
221    }
222
223    /// Transfer an initialized allocation into this runtime, optionally importing
224    /// it into an NPU program once. Only GpuLocal and Shared placement are accepted.
225    ///
226    /// # Safety
227    /// All bytes must be initialized and prior uses complete. No old pointer,
228    /// recorded graph or other external alias may access the allocation after
229    /// adoption, except through this runtime's scoped GPU handoff.
230    pub unsafe fn adopt_gpu_buffer(
231        &self,
232        buffer: crate::gpu::Buffer,
233        placement: MemoryPlacement,
234    ) -> Result<Buffer> {
235        let stream = crate::gpu::Device::open(self.inner.gpu_index)?.stream()?;
236        if buffer.device_id() != stream.device_id() {
237            return Err(Error::Message("buffer belongs to another device".into()));
238        }
239        if matches!(placement, MemoryPlacement::HostVisible) || matches_npu_local(&placement) {
240            return Err(Error::Unsupported(
241                "adoption requires GpuLocal or Shared placement".into(),
242            ));
243        }
244        self.allocate_inner(buffer.bytes(), placement, Some(buffer))
245    }
246
247    fn allocate_inner(
248        &self,
249        bytes: usize,
250        placement: MemoryPlacement,
251        mut adopted: Option<crate::gpu::Buffer>,
252    ) -> Result<Buffer> {
253        if bytes == 0 || bytes > isize::MAX as usize {
254            return Err(Error::Message(
255                "allocation size must be in 1..=isize::MAX".into(),
256            ));
257        }
258        let mut storage = Storage {
259            _reservation: self
260                .options
261                .memory_budget
262                .as_ref()
263                .filter(|budget| {
264                    !adopted
265                        .as_ref()
266                        .is_some_and(|buffer| buffer.charged_to(budget))
267                })
268                .map(|budget| budget.reserve(bytes))
269                .transpose()?,
270            #[cfg(feature = "npu")]
271            bo: None,
272            #[cfg(feature = "npu")]
273            descriptor: None,
274            accounted: false,
275            gpu: None,
276            pointer: std::ptr::null_mut(),
277            bytes,
278            #[cfg(feature = "npu")]
279            npu_device: None,
280            #[cfg(feature = "npu")]
281            group: None,
282            #[cfg(test)]
283            _test_memory: None,
284            shared: false,
285            host: Mutex::new(HostState::default()),
286            visibility: Mutex::new(Visibility::new()),
287            runtime: self.inner.clone(),
288        };
289        match &placement {
290            MemoryPlacement::GpuLocal => {
291                let mut stream = crate::gpu::Device::open(self.inner.gpu_index)?.stream()?;
292                let buffer = if let Some(buffer) = adopted.take() {
293                    buffer
294                } else {
295                    let buffer = stream.allocate(bytes)?;
296                    stream.fill(buffer.binding(), 0)?;
297                    stream.synchronize()?;
298                    buffer
299                };
300                storage.gpu = Some(buffer);
301                storage.visibility.get_mut().unwrap().wrote(Engine::Gpu);
302            }
303            MemoryPlacement::HostVisible => {
304                let stream = crate::gpu::Device::open(self.inner.gpu_index)?.stream()?;
305                let buffer = stream.allocate_shared(bytes)?;
306                let pointer = buffer.device_ptr()?.cast::<u8>();
307                // Coherent host-local allocation, with no aliases or device uses.
308                unsafe {
309                    std::ptr::write_bytes(pointer, 0, bytes);
310                }
311                storage.pointer = pointer;
312                storage.gpu = Some(buffer);
313            }
314            #[cfg(feature = "npu")]
315            MemoryPlacement::Shared(program) | MemoryPlacement::NpuLocal(program) => {
316                let context = &program.inner.context;
317                let group = context.group_id(3).map_err(Error::Message)?;
318                let bo = if matches!(placement, MemoryPlacement::Shared(_)) {
319                    use std::os::fd::AsRawFd;
320                    let mut stream = crate::gpu::Device::open(self.inner.gpu_index)?.stream()?;
321                    let gpu = if let Some(buffer) = adopted.take() {
322                        buffer
323                    } else {
324                        let buffer = stream.allocate(bytes)?;
325                        stream.fill(buffer.binding(), 0)?;
326                        stream.synchronize()?;
327                        buffer
328                    };
329                    let (fd, offset) = gpu.export_dmabuf()?;
330                    let offset = usize::try_from(offset).map_err(|_| {
331                        Error::Unsupported("dma-buf offset exceeds address space".into())
332                    })?;
333                    // Only the owned subregion is initialized; an HSA pool may
334                    // export a larger root containing unrelated allocations.
335                    let bo = unsafe { context.import_dmabuf_region(fd.as_raw_fd(), offset, bytes) }
336                        .map_err(Error::Message)?;
337                    storage.gpu = Some(gpu);
338                    storage.descriptor = Some(fd);
339                    storage.shared = true;
340                    bo
341                } else {
342                    context
343                        .alloc_bo(bytes, crate::npu::raw::BoKind::HostOnly, group)
344                        .map_err(Error::Message)?
345                };
346                let pointer = bo.map().map_err(Error::Message)?;
347                // No aliases or device submissions exist during initialization.
348                // Shared GPU allocations are already initialized, including
349                // adopted data. Only freshly allocated NPU-local bytes need zeroing.
350                if !storage.shared {
351                    unsafe {
352                        std::ptr::write_bytes(pointer, 0, bytes);
353                    }
354                } else {
355                    storage.visibility.get_mut().unwrap().wrote(Engine::Gpu);
356                }
357                storage.pointer = pointer;
358                storage.bo = Some(bo);
359                storage.npu_device = Some(program.inner.device);
360                storage.group = Some(group);
361            }
362        }
363        use std::sync::atomic::Ordering;
364        let counters = &self.inner.core.counters;
365        counters.allocations.fetch_add(1, Ordering::Relaxed);
366        if storage.shared {
367            counters.imports.fetch_add(1, Ordering::Relaxed);
368        }
369        let live = counters
370            .live_bytes
371            .fetch_add(bytes as u64, Ordering::Relaxed)
372            + bytes as u64;
373        counters.peak_bytes.fetch_max(live, Ordering::Relaxed);
374        storage.accounted = true;
375        Ok(Buffer {
376            storage: Arc::new(storage),
377        })
378    }
379    /// Observe allocation, transfer, and execution counters.
380    pub fn statistics(&self) -> Statistics {
381        self.inner.core.counters.snapshot()
382    }
383
384    /// Start describing a reusable pipeline in this allocation domain.
385    pub fn graph(&self) -> Graph {
386        Graph::new(self.clone())
387    }
388    /// Load a fixed GPU specialization with its trusted memory contract.
389    /// # Safety
390    /// The code must obey the contract for this grid/block and every accepted
391    /// binding. Kernel memory access is not sandboxed by native binding lengths.
392    pub unsafe fn load_gpu_kernel(
393        &self,
394        path: impl AsRef<std::path::Path>,
395        symbol: &str,
396        grid: [u32; 3],
397        block: [u32; 3],
398        contract: KernelContract,
399    ) -> Result<GpuKernel> {
400        let stream = crate::gpu::Device::open(self.inner.gpu_index)?.stream()?;
401        let raw = unsafe { stream.load(path.as_ref(), symbol) }?;
402        unsafe { self.adopt_gpu_kernel(raw, grid, block, contract) }
403    }
404
405    /// Retain a loaded executable without loading or copying its code again.
406    /// # Safety
407    /// The kernel must obey the declared memory contract for this launch and
408    /// every accepted binding, as for [`Self::load_gpu_kernel`].
409    pub unsafe fn adopt_gpu_kernel(
410        &self,
411        raw: crate::gpu::Kernel,
412        grid: [u32; 3],
413        block: [u32; 3],
414        contract: KernelContract,
415    ) -> Result<GpuKernel> {
416        contract.validate()?;
417        crate::runtime::validate_export_launch(raw.info(), grid, block)?;
418        let stream = crate::gpu::Device::open(self.inner.gpu_index)?.stream()?;
419        if raw.device_id() != stream.device_id() {
420            return Err(Error::Message("kernel belongs to another device".into()));
421        }
422        if raw.info().binding_count as usize != contract.bindings.len()
423            || raw.info().constant_byte_length as usize != contract.constants.len()
424        {
425            return Err(Error::Message(
426                "GPU metadata does not match contract".into(),
427            ));
428        }
429        Ok(GpuKernel {
430            raw,
431            contract,
432            grid,
433            block,
434            runtime: self.inner.clone(),
435        })
436    }
437}
438
439fn matches_npu_local(placement: &MemoryPlacement) -> bool {
440    #[cfg(feature = "npu")]
441    {
442        matches!(placement, MemoryPlacement::NpuLocal(_))
443    }
444    #[cfg(not(feature = "npu"))]
445    {
446        let _ = placement;
447        false
448    }
449}
450
451#[cfg(test)]
452mod tests;
hrx/execution/mod.rs

hrx/execution/
mod.rs
