aria_graph/
lib.rs

1//! Zero-copy compute graph: Layer → Op → TensorView + BufferPool.
2
3use aria_kernel::{hadamard_blocked_vec, linear, matmul_dispatch, EngineError, SimdMode};
4use std::sync::Arc;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum DType {
8    F32,
9    F16,
10    U8,
11}
12
13/// Borrowed or owned tensor bytes (mmap / external / pool).
14#[derive(Clone)]
15pub enum TensorBuf {
16    External(Arc<[u8]>),
17    Owned(Vec<u8>),
18}
19
20impl TensorBuf {
21    pub fn as_slice(&self) -> &[u8] {
22        match self {
23            Self::External(a) => a.as_ref(),
24            Self::Owned(v) => v.as_slice(),
25        }
26    }
27}
28
29#[derive(Clone)]
30pub struct TensorView {
31    pub dtype: DType,
32    pub shape: Vec<usize>,
33    pub strides: Vec<usize>,
34    pub buf: TensorBuf,
35    pub offset: usize,
36    pub len: usize,
37}
38
39impl TensorView {
40    pub fn from_f32(data: Vec<f32>, shape: Vec<usize>) -> Self {
41        let nbytes = data.len() * 4;
42        let mut bytes = Vec::with_capacity(nbytes);
43        for v in data {
44            bytes.extend_from_slice(&v.to_le_bytes());
45        }
46        let strides = row_major_strides(&shape, 4);
47        Self {
48            dtype: DType::F32,
49            shape,
50            strides,
51            buf: TensorBuf::Owned(bytes),
52            offset: 0,
53            len: nbytes,
54        }
55    }
56
57    pub fn from_external(bytes: Arc<[u8]>, dtype: DType, shape: Vec<usize>, offset: usize, len: usize) -> Self {
58        let elem = match dtype {
59            DType::F32 => 4,
60            DType::F16 => 2,
61            DType::U8 => 1,
62        };
63        let strides = row_major_strides(&shape, elem);
64        Self {
65            dtype,
66            shape,
67            strides,
68            buf: TensorBuf::External(bytes),
69            offset,
70            len,
71        }
72    }
73
74    pub fn as_f32_slice(&self) -> Result<&[f32], EngineError> {
75        if self.dtype != DType::F32 {
76            return Err(EngineError::ShapeMismatch("expected f32 tensor".into()));
77        }
78        let bytes = &self.buf.as_slice()[self.offset..self.offset + self.len];
79        if !bytes.len().is_multiple_of(4) {
80            return Err(EngineError::Format("f32 byte length not aligned".into()));
81        }
82        let ptr = bytes.as_ptr() as *const f32;
83        Ok(unsafe { std::slice::from_raw_parts(ptr, bytes.len() / 4) })
84    }
85
86    pub fn to_f32_vec(&self) -> Result<Vec<f32>, EngineError> {
87        Ok(self.as_f32_slice()?.to_vec())
88    }
89}
90
91fn row_major_strides(shape: &[usize], elem: usize) -> Vec<usize> {
92    let mut strides = vec![0; shape.len()];
93    let mut acc = elem;
94    for i in (0..shape.len()).rev() {
95        strides[i] = acc;
96        acc *= shape[i].max(1);
97    }
98    strides
99}
100
101#[derive(Default)]
102pub struct BufferPool {
103    buffers: Vec<Vec<u8>>,
104}
105
106impl BufferPool {
107    pub fn new() -> Self {
108        Self::default()
109    }
110
111    pub fn alloc(&mut self, nbytes: usize) -> &mut [u8] {
112        self.buffers.push(vec![0u8; nbytes]);
113        self.buffers.last_mut().unwrap().as_mut_slice()
114    }
115
116    pub fn reuse_count(&self) -> usize {
117        self.buffers.len()
118    }
119}
120
121#[derive(Debug, Clone)]
122pub enum Op {
123    /// y = x @ W^T  (W: [out, in])
124    Linear { out_f: usize, in_f: usize },
125    /// Generic matmul
126    MatMul {
127        a_rows: usize,
128        a_cols: usize,
129        b_rows: usize,
130        b_cols: usize,
131    },
132    /// Fused path: matmul with rotated W then blocked unrotate on output features.
133    HdmLinear {
134        out_f: usize,
135        in_f: usize,
136        hadamard_seed: Option<i64>,
137    },
138}
139
140pub struct Node {
141    pub op: Op,
142    pub inputs: Vec<usize>,
143    pub output: usize,
144}
145
146pub struct Graph {
147    pub tensors: Vec<Option<TensorView>>,
148    pub nodes: Vec<Node>,
149    pub mode: SimdMode,
150}
151
152impl Graph {
153    pub fn new(mode: SimdMode) -> Self {
154        Self {
155            tensors: Vec::new(),
156            nodes: Vec::new(),
157            mode,
158        }
159    }
160
161    pub fn push_tensor(&mut self, t: TensorView) -> usize {
162        let id = self.tensors.len();
163        self.tensors.push(Some(t));
164        id
165    }
166
167    pub fn add_node(&mut self, node: Node) {
168        self.nodes.push(node);
169    }
170
171    pub fn execute(&mut self, _pool: &mut BufferPool) -> Result<(), EngineError> {
172        for node in &self.nodes {
173            match &node.op {
174                Op::Linear { out_f, in_f } => {
175                    let x = self.tensors[node.inputs[0]]
176                        .as_ref()
177                        .ok_or_else(|| EngineError::Format("missing input tensor".into()))?
178                        .to_f32_vec()?;
179                    let w = self.tensors[node.inputs[1]]
180                        .as_ref()
181                        .ok_or_else(|| EngineError::Format("missing weight tensor".into()))?
182                        .to_f32_vec()?;
183                    let y = linear(&x, &w, *out_f, *in_f)?;
184                    let shape = if x.len() == *in_f {
185                        vec![*out_f]
186                    } else {
187                        vec![x.len() / *in_f, *out_f]
188                    };
189                    self.tensors[node.output] = Some(TensorView::from_f32(y, shape));
190                }
191                Op::HdmLinear {
192                    out_f,
193                    in_f,
194                    hadamard_seed,
195                } => {
196                    let x = self.tensors[node.inputs[0]]
197                        .as_ref()
198                        .ok_or_else(|| EngineError::Format("missing input tensor".into()))?
199                        .to_f32_vec()?;
200                    let w = self.tensors[node.inputs[1]]
201                        .as_ref()
202                        .ok_or_else(|| EngineError::Format("missing weight tensor".into()))?
203                        .to_f32_vec()?;
204                    // y = W_rot @ x  (or batched), then blocked unrotate on out features.
205                    let mut y = linear(&x, &w, *out_f, *in_f)?;
206                    let batch = y.len() / *out_f;
207                    for b in 0..batch {
208                        let sl = b * *out_f..(b + 1) * *out_f;
209                        hadamard_blocked_vec(&mut y[sl], *hadamard_seed, true)?;
210                    }
211                    let shape = if x.len() == *in_f {
212                        vec![*out_f]
213                    } else {
214                        vec![batch, *out_f]
215                    };
216                    self.tensors[node.output] = Some(TensorView::from_f32(y, shape));
217                }
218                Op::MatMul {
219                    a_rows,
220                    a_cols,
221                    b_rows,
222                    b_cols,
223                } => {
224                    let a = self.tensors[node.inputs[0]]
225                        .as_ref()
226                        .ok_or_else(|| EngineError::Format("missing A".into()))?
227                        .to_f32_vec()?;
228                    let b = self.tensors[node.inputs[1]]
229                        .as_ref()
230                        .ok_or_else(|| EngineError::Format("missing B".into()))?
231                        .to_f32_vec()?;
232                    let c = matmul_dispatch(
233                        &a, *a_rows, *a_cols, &b, *b_rows, *b_cols, self.mode,
234                    )?;
235                    self.tensors[node.output] =
236                        Some(TensorView::from_f32(c, vec![*a_rows, *b_cols]));
237                }
238            }
239        }
240        Ok(())
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247
248    fn run_hdm(
249        w_rot: Vec<f32>,
250        x: Vec<f32>,
251        out_f: usize,
252        in_f: usize,
253        seed: Option<i64>,
254    ) -> Vec<f32> {
255        let batch = x.len() / in_f;
256        let mut g = Graph::new(SimdMode::Scalar);
257        let xi = g.push_tensor(TensorView::from_f32(
258            x,
259            if batch == 1 {
260                vec![in_f]
261            } else {
262                vec![batch, in_f]
263            },
264        ));
265        let wi = g.push_tensor(TensorView::from_f32(w_rot, vec![out_f, in_f]));
266        let yi = g.push_tensor(TensorView::from_f32(
267            vec![0.0; batch * out_f],
268            if batch == 1 {
269                vec![out_f]
270            } else {
271                vec![batch, out_f]
272            },
273        ));
274        g.add_node(Node {
275            op: Op::HdmLinear {
276                out_f,
277                in_f,
278                hadamard_seed: seed,
279            },
280            inputs: vec![xi, wi],
281            output: yi,
282        });
283        g.execute(&mut BufferPool::new()).unwrap();
284        g.tensors[yi].as_ref().unwrap().to_f32_vec().unwrap()
285    }
286
287    #[test]
288    fn hdm_linear_matches_orig_weight() {
289        use aria_kernel::hadamard_blocked_rows;
290        let out_f = 10usize;
291        let in_f = 4usize;
292        let seed = Some(3i64);
293        let mut w_orig: Vec<f32> = (0..out_f * in_f)
294            .map(|i| (i as f32) * 0.07 - 0.3)
295            .collect();
296        let x = vec![0.5f32, -0.2, 0.1, 0.3];
297        let y_ref = linear(&x, &w_orig, out_f, in_f).unwrap();
298
299        // Rotate weight rows for HDM path.
300        hadamard_blocked_rows(&mut w_orig, out_f, in_f, seed, false).unwrap();
301        let y = run_hdm(w_orig, x, out_f, in_f, seed);
302        assert_eq!(y.len(), y_ref.len());
303        for (a, b) in y.iter().zip(y_ref.iter()) {
304            assert!((a - b).abs() < 1e-4, "{a} vs {b}");
305        }
306    }
307
308    #[test]
309    fn hdm_linear_batch_and_unsigned() {
310        use aria_kernel::hadamard_blocked_rows;
311        let out_f = 10usize;
312        let in_f = 3usize;
313        let mut w: Vec<f32> = (0..out_f * in_f)
314            .map(|i| (i as f32) * 0.05 - 0.2)
315            .collect();
316        let x = vec![
317            0.2f32, -0.1, 0.4, // batch 0
318            -0.3, 0.5, 0.1, // batch 1
319        ];
320        let y_ref = linear(&x, &w, out_f, in_f).unwrap();
321        hadamard_blocked_rows(&mut w, out_f, in_f, None, false).unwrap();
322        let y = run_hdm(w, x, out_f, in_f, None);
323        assert_eq!(y.len(), 2 * out_f);
324        for (a, b) in y.iter().zip(y_ref.iter()) {
325            assert!((a - b).abs() < 1e-4, "{a} vs {b}");
326        }
327    }
328
329    #[test]
330    fn hdm_linear_shape_err() {
331        let mut g = Graph::new(SimdMode::Scalar);
332        let x = g.push_tensor(TensorView::from_f32(vec![1.0, 2.0], vec![2]));
333        let w = g.push_tensor(TensorView::from_f32(vec![1.0; 6], vec![3, 2]));
334        let y = g.push_tensor(TensorView::from_f32(vec![0.0; 3], vec![3]));
335        g.add_node(Node {
336            op: Op::HdmLinear {
337                out_f: 3,
338                in_f: 3, // mismatch vs weight / x
339                hadamard_seed: Some(0),
340            },
341            inputs: vec![x, w],
342            output: y,
343        });
344        let err = g.execute(&mut BufferPool::new()).unwrap_err();
345        assert!(matches!(err, EngineError::ShapeMismatch(_)));
346    }
347
348    #[test]
349    fn linear_dispatch() {
350        let mut g = Graph::new(SimdMode::Scalar);
351        let x = g.push_tensor(TensorView::from_f32(vec![1.0, 0.0], vec![2]));
352        // W 2x2 identity rows
353        let w = g.push_tensor(TensorView::from_f32(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]));
354        let y = g.push_tensor(TensorView::from_f32(vec![0.0, 0.0], vec![2]));
355        g.add_node(Node {
356            op: Op::Linear { out_f: 2, in_f: 2 },
357            inputs: vec![x, w],
358            output: y,
359        });
360        let mut pool = BufferPool::new();
361        g.execute(&mut pool).unwrap();
362        let out = g.tensors[y].as_ref().unwrap().to_f32_vec().unwrap();
363        assert_eq!(out, vec![1.0, 0.0]);
364    }
365
366    #[test]
367    fn external_zero_copy() {
368        let bytes: Arc<[u8]> = Arc::from([0u8, 0, 0x80, 0x3f].as_slice()); // 1.0f32 LE
369        let t = TensorView::from_external(bytes, DType::F32, vec![1], 0, 4);
370        assert_eq!(t.to_f32_vec().unwrap(), vec![1.0]);
371    }
372
373    #[test]
374    fn matmul_shape_err() {
375        let mut g = Graph::new(SimdMode::Scalar);
376        let a = g.push_tensor(TensorView::from_f32(vec![1.0], vec![1, 1]));
377        let b = g.push_tensor(TensorView::from_f32(vec![1.0, 2.0], vec![2, 1]));
378        let c = g.push_tensor(TensorView::from_f32(vec![0.0], vec![1, 1]));
379        g.add_node(Node {
380            op: Op::MatMul {
381                a_rows: 1,
382                a_cols: 1,
383                b_rows: 2,
384                b_cols: 1,
385            },
386            inputs: vec![a, b],
387            output: c,
388        });
389        let err = g.execute(&mut BufferPool::new()).unwrap_err();
390        assert!(matches!(err, EngineError::ShapeMismatch(_)));
391    }
392}
aria_graph/lib.rs

aria_graph/
lib.rs
