virtio_accel_xdna/
artifact.rs

1//! The crate-local precompiled XDNA artifact format.
2//!
3//! The backend accepts a self-describing container carrying a prebuilt `final.xclbin` and its
4//! unfolded `insts.bin` (the XAie transaction stream), the entry point name, and the per-slot
5//! binding plan: the exact byte size of every runtime input and output. The sizes are load-bearing:
6//! the compiled transaction stream DMAs fixed tensor extents, so `submit` must reject a binding
7//! whose length differs from what the kernel will actually transfer (mirroring the OpenVINO
8//! backend's per-slot `byte_len` check). This format is portable and `unsafe`-free: a host or tool
9//! builds it with [`encode`] and the backend parses it with [`PrecompiledArtifact::parse`].
10//!
11//! Layout (little-endian), header then payloads:
12//!
13//! | field | bytes |
14//! |---|---|
15//! | magic `b"XDNP"` | 4 |
16//! | version (`= 2`) | 4 |
17//! | input count | 4 |
18//! | output count | 4 |
19//! | entry-name length | 4 |
20//! | xclbin length | 8 |
21//! | insts length | 8 |
22//! | per-slot byte sizes (inputs, then outputs) | 8 × (inputs + outputs) |
23//! | entry-name bytes | entry-name length |
24//! | xclbin bytes | xclbin length |
25//! | insts bytes | insts length |
26
27use virtio_accel_core::{ArtifactFormat, BackendError};
28
29/// Artifact-format tag for a precompiled XDNA container ("XDNP" in ASCII).
30pub const XDNA_PRECOMPILED_FORMAT: ArtifactFormat = match ArtifactFormat::new(0x5844_4e50) {
31    Some(format) => format,
32    None => unreachable!(),
33};
34
35const MAGIC: [u8; 4] = *b"XDNP";
36const VERSION: u32 = 2;
37const HEADER_LEN: usize = 4 + 4 + 4 + 4 + 4 + 8 + 8;
38
39/// A parsed precompiled artifact: borrowed payload views plus the owned per-slot binding plan.
40#[derive(Clone, Debug)]
41pub struct PrecompiledArtifact<'a> {
42    pub xclbin: &'a [u8],
43    pub insts: &'a [u8],
44    pub entry: &'a str,
45    pub inputs: usize,
46    pub outputs: usize,
47    /// Exact byte size of each binding slot: inputs occupy `0..inputs`, outputs follow.
48    pub slot_bytes: Vec<u64>,
49}
50
51impl<'a> PrecompiledArtifact<'a> {
52    /// Parse and validate a container. Malformed input is `InvalidArgument`; an unknown version is
53    /// `Incompatible`.
54    pub fn parse(bytes: &'a [u8]) -> Result<Self, BackendError> {
55        if bytes.len() < HEADER_LEN || bytes[0..4] != MAGIC {
56            return Err(BackendError::InvalidArgument);
57        }
58        let u32_at = |offset: usize| {
59            u32::from_le_bytes(bytes[offset..offset + 4].try_into().expect("4 bytes"))
60        };
61        let u64_at = |offset: usize| {
62            u64::from_le_bytes(bytes[offset..offset + 8].try_into().expect("8 bytes"))
63        };
64        if u32_at(4) != VERSION {
65            return Err(BackendError::Incompatible);
66        }
67        let inputs = u32_at(8) as usize;
68        let outputs = u32_at(12) as usize;
69        let entry_len = u32_at(16) as usize;
70        let xclbin_len = usize::try_from(u64_at(20)).map_err(|_| BackendError::InvalidArgument)?;
71        let insts_len = usize::try_from(u64_at(28)).map_err(|_| BackendError::InvalidArgument)?;
72
73        // Every section must lie within the container, and the counts must be nonzero and bounded:
74        // a program needs at least one input to read and one output to invalidate after dispatch.
75        if inputs == 0 || outputs == 0 {
76            return Err(BackendError::InvalidArgument);
77        }
78        let total = match inputs.checked_add(outputs) {
79            Some(total) if total <= 256 => total,
80            _ => return Err(BackendError::InvalidArgument),
81        };
82        // HRX consumes the TXN stream as little-endian u32 words.
83        if insts_len % 4 != 0 || insts_len == 0 || xclbin_len == 0 {
84            return Err(BackendError::InvalidArgument);
85        }
86        let sizes_end = HEADER_LEN
87            .checked_add(total * 8)
88            .ok_or(BackendError::InvalidArgument)?;
89        let entry_end = sizes_end
90            .checked_add(entry_len)
91            .ok_or(BackendError::InvalidArgument)?;
92        let xclbin_end = entry_end
93            .checked_add(xclbin_len)
94            .ok_or(BackendError::InvalidArgument)?;
95        let insts_end = xclbin_end
96            .checked_add(insts_len)
97            .ok_or(BackendError::InvalidArgument)?;
98        if insts_end != bytes.len() {
99            return Err(BackendError::InvalidArgument);
100        }
101        let mut slot_bytes = Vec::with_capacity(total);
102        for index in 0..total {
103            let size = u64_at(HEADER_LEN + index * 8);
104            // A zero-sized slot cannot describe a real tensor transfer.
105            if size == 0 {
106                return Err(BackendError::InvalidArgument);
107            }
108            slot_bytes.push(size);
109        }
110        let entry = core::str::from_utf8(&bytes[sizes_end..entry_end])
111            .map_err(|_| BackendError::InvalidArgument)?;
112        if entry.is_empty() || entry.contains('\0') {
113            return Err(BackendError::InvalidArgument);
114        }
115        Ok(Self {
116            xclbin: &bytes[entry_end..xclbin_end],
117            insts: &bytes[xclbin_end..insts_end],
118            entry,
119            inputs,
120            outputs,
121            slot_bytes,
122        })
123    }
124}
125
126/// Build a container from its parts (for hosts and tools producing precompiled artifacts). The
127/// slot sizes are the exact per-binding byte counts: inputs first, then outputs; the counts are
128/// derived from the two slices' lengths.
129pub fn encode(
130    entry: &str,
131    input_bytes: &[u64],
132    output_bytes: &[u64],
133    xclbin: &[u8],
134    insts: &[u8],
135) -> Vec<u8> {
136    let total = input_bytes.len() + output_bytes.len();
137    let mut out =
138        Vec::with_capacity(HEADER_LEN + total * 8 + entry.len() + xclbin.len() + insts.len());
139    out.extend_from_slice(&MAGIC);
140    out.extend_from_slice(&VERSION.to_le_bytes());
141    out.extend_from_slice(&(input_bytes.len() as u32).to_le_bytes());
142    out.extend_from_slice(&(output_bytes.len() as u32).to_le_bytes());
143    out.extend_from_slice(&(entry.len() as u32).to_le_bytes());
144    out.extend_from_slice(&(xclbin.len() as u64).to_le_bytes());
145    out.extend_from_slice(&(insts.len() as u64).to_le_bytes());
146    for size in input_bytes.iter().chain(output_bytes) {
147        out.extend_from_slice(&size.to_le_bytes());
148    }
149    out.extend_from_slice(entry.as_bytes());
150    out.extend_from_slice(xclbin);
151    out.extend_from_slice(insts);
152    out
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn round_trips_through_encode_and_parse() {
161        let bytes = encode("MLIR_AIE", &[128, 64], &[256], &[0xAB; 40], &[0x11; 8]);
162        let parsed = PrecompiledArtifact::parse(&bytes).expect("valid container");
163        assert_eq!(parsed.entry, "MLIR_AIE");
164        assert_eq!((parsed.inputs, parsed.outputs), (2, 1));
165        assert_eq!(parsed.slot_bytes, [128, 64, 256]);
166        assert_eq!(parsed.xclbin, &[0xAB; 40]);
167        assert_eq!(parsed.insts, &[0x11; 8]);
168    }
169
170    #[test]
171    fn rejects_bad_magic_version_and_alignment() {
172        assert!(matches!(
173            PrecompiledArtifact::parse(b"nope"),
174            Err(BackendError::InvalidArgument)
175        ));
176        let mut bytes = encode("MLIR_AIE", &[4], &[4], &[0u8; 4], &[0u8; 4]);
177        bytes[4] = 1; // an old (or unknown) version
178        assert!(matches!(
179            PrecompiledArtifact::parse(&bytes),
180            Err(BackendError::Incompatible)
181        ));
182        // insts length not a multiple of 4.
183        let bad = encode("MLIR_AIE", &[4], &[4], &[0u8; 4], &[0u8; 6]);
184        assert!(matches!(
185            PrecompiledArtifact::parse(&bad),
186            Err(BackendError::InvalidArgument)
187        ));
188    }
189
190    #[test]
191    fn rejects_missing_outputs_and_zero_sized_slots() {
192        // No outputs: the program could never publish a result the host can read back.
193        let no_outputs = encode("MLIR_AIE", &[4], &[], &[0u8; 4], &[0u8; 4]);
194        assert!(matches!(
195            PrecompiledArtifact::parse(&no_outputs),
196            Err(BackendError::InvalidArgument)
197        ));
198        let no_inputs = encode("MLIR_AIE", &[], &[4], &[0u8; 4], &[0u8; 4]);
199        assert!(matches!(
200            PrecompiledArtifact::parse(&no_inputs),
201            Err(BackendError::InvalidArgument)
202        ));
203        let zero_slot = encode("MLIR_AIE", &[0], &[4], &[0u8; 4], &[0u8; 4]);
204        assert!(matches!(
205            PrecompiledArtifact::parse(&zero_slot),
206            Err(BackendError::InvalidArgument)
207        ));
208    }
209
210    #[test]
211    fn rejects_truncated_container() {
212        let bytes = encode("MLIR_AIE", &[40], &[8], &[0xAB; 40], &[0x11; 8]);
213        assert!(matches!(
214            PrecompiledArtifact::parse(&bytes[..bytes.len() - 4]),
215            Err(BackendError::InvalidArgument)
216        ));
217    }
218}
virtio_accel_xdna/artifact.rs

virtio_accel_xdna/
artifact.rs
