hrx/npu/
mod.rs

1//! XDNA2 programs and fixed instruction specializations.
2//!
3//! Programs are explicitly trusted native code. Buffers and execution are managed
4//! through [`crate::execution`]; the [`raw`] API requires external synchronization.
5
6pub mod compiler;
7pub mod provision;
8pub mod raw;
9use crate::{Error, Result, execution::KernelContract};
10use std::{path::Path, sync::Arc};
11
12pub(crate) fn load_shim() -> Result<libloading::Library> {
13    let directory = provision::resolve()?;
14    load_shim_from(&directory)
15}
16
17fn load_shim_from(directory: &Path) -> Result<libloading::Library> {
18    if directory.join("component.json").is_file() {
19        provision::Manifest::load(directory.join("component.json"))?.verify(directory)?;
20    }
21    let path = directory.join("libhrx_npu.so.1");
22    let library = unsafe { libloading::Library::new(&path) }.map_err(|e| {
23        Error::from(e).context(format!(
24            "loading {}; run hrx prepare (or check HRX_NPU_RUNTIME_DIR)",
25            path.display()
26        ))
27    })?;
28    let abi = unsafe { library.get::<unsafe extern "C" fn() -> u32>(b"hrx_npu_abi_version\0") }?;
29    if unsafe { abi() } != 1 {
30        return Err(Error::Unsupported("NPU shim ABI must be 1".into()));
31    }
32    Ok(library)
33}
34/// An immutable, resident NPU device image. Cloning reuses the native context.
35#[derive(Clone)]
36pub struct NpuProgram {
37    pub(crate) inner: Arc<ProgramInner>,
38    pub(crate) budget: Option<crate::residency::MemoryBudget>,
39}
40pub(crate) struct ProgramInner {
41    pub context: raw::Context,
42    pub device: i32,
43    pub identity: String,
44}
45impl NpuProgram {
46    /// Load an xclbin on an explicitly selected NPU.
47    /// # Safety
48    /// The image is trusted native code for this device; loading it can start
49    /// tile programs. It must not perform unbound DMA or access host memory
50    /// except through subsequent correctly contracted invocations.
51    pub unsafe fn load(device: i32, path: impl AsRef<Path>) -> Result<Self> {
52        let path = std::fs::canonicalize(path)?;
53        let identity = crate::bundle::file_digest(&path)?;
54        type Programs = std::collections::BTreeMap<(i32, String), std::sync::Weak<ProgramInner>>;
55        static PROGRAMS: std::sync::OnceLock<std::sync::Mutex<Programs>> =
56            std::sync::OnceLock::new();
57        let mut programs = PROGRAMS
58            .get_or_init(Default::default)
59            .lock()
60            .map_err(|_| Error::Message("NPU program registry poisoned".into()))?;
61        if let Some(inner) = programs
62            .get(&(device, identity.clone()))
63            .and_then(std::sync::Weak::upgrade)
64        {
65            return Ok(Self {
66                inner,
67                budget: None,
68            });
69        }
70        programs.retain(|_, program| program.strong_count() != 0);
71        let path_str = path
72            .to_str()
73            .ok_or_else(|| Error::Message("xclbin path is not UTF-8".into()))?;
74        let context = unsafe { raw::Context::new(device, path_str) }.map_err(Error::Message)?;
75        if crate::bundle::file_digest(&path)? != identity {
76            return Err(Error::Message("xclbin changed while loading".into()));
77        }
78        let inner = Arc::new(ProgramInner {
79            context,
80            device,
81            identity: identity.clone(),
82        });
83        programs.insert((device, identity), Arc::downgrade(&inner));
84        Ok(Self {
85            inner,
86            budget: None,
87        })
88    }
89    /// Content identity of the resident xclbin.
90    pub fn identity(&self) -> &str {
91        &self.inner.identity
92    }
93    /// Bind a fixed instruction specialization to this image and contract.
94    /// Instructions inherit the loading runtime's optional allocation budget;
95    /// native runs retain that charge until completion or quarantine teardown.
96    /// # Safety
97    /// Instructions must match this exact image and target. All accesses must
98    /// remain within the contract, with the declared read/write modes, for every
99    /// accepted binding. Argument aliasing is rejected by the safe graph API.
100    pub unsafe fn kernel(
101        &self,
102        instructions: &[u8],
103        contract: KernelContract,
104    ) -> Result<NpuKernel> {
105        contract.validate()?;
106        if !contract.constants.is_empty()
107            || instructions.is_empty()
108            || !instructions.len().is_multiple_of(4)
109            || instructions.len() / 4 > u32::MAX as usize
110        {
111            return Err(Error::Message(
112                "MLIR_AIE requires a nonempty u32 instruction stream and no extra scalar arguments"
113                    .into(),
114            ));
115        }
116        let context = &self.inner.context;
117        let reservation = self
118            .budget
119            .as_ref()
120            .map(|budget| budget.reserve(instructions.len()))
121            .transpose()?;
122        let insts = context
123            .alloc_bo_reserved(
124                instructions.len(),
125                raw::BoKind::Cacheable,
126                context.group_id(1).map_err(Error::Message)?,
127                reservation,
128            )
129            .map_err(Error::Message)?;
130        insts.write(instructions).map_err(Error::Message)?;
131        let groups = (0..contract.bindings.len())
132            .map(|i| context.group_id((i + 3) as i32).map_err(Error::Message))
133            .collect::<Result<Vec<_>>>()?;
134        Ok(NpuKernel {
135            inner: Arc::new(KernelInner {
136                program: self.clone(),
137                instructions: insts,
138                contract,
139                groups,
140            }),
141        })
142    }
143}
144/// A reusable fixed NPU instruction specialization and its trusted contract.
145#[derive(Clone)]
146pub struct NpuKernel {
147    pub(crate) inner: Arc<KernelInner>,
148}
149pub(crate) struct KernelInner {
150    pub program: NpuProgram,
151    pub instructions: raw::Bo,
152    pub contract: KernelContract,
153    pub groups: Vec<i32>,
154}
155impl NpuKernel {
156    /// The contract checked for every graph binding.
157    pub fn contract(&self) -> &KernelContract {
158        &self.inner.contract
159    }
160}
hrx/npu/mod.rs

hrx/npu/
mod.rs
