luwen_api/chip/communication/
chip_comms.rs

1// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use rust_embed::RustEmbed;
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9
10use super::chip_interface::ChipInterface;
11
12#[derive(Error, Debug)]
13pub enum AxiError {
14    #[error("Invalid path: {key} specifically was not able to find {path}")]
15    InvalidPath { key: String, path: String },
16
17    #[error("Invalid path: {key} specifically was not able to parse {path} as an array def.")]
18    InvalidArrayPath { key: String, path: String },
19
20    #[error("No AXI data table loaded")]
21    NoAxiData,
22
23    #[error("The readbuffer is not large enough to hold the requested data")]
24    ReadBufferTooSmall,
25
26    #[error("The writebuffer is not the same size as the requested field")]
27    WriteBufferMismatch,
28}
29
30#[derive(Clone, Debug, Serialize, Deserialize)]
31pub struct AxiData {
32    pub addr: u64,
33    pub size: u64,
34    pub bits: Option<(u32, u32)>,
35}
36
37#[derive(Debug, Deserialize, Serialize)]
38pub struct MemorySlice {
39    pub name: String,
40    pub offset: u64,
41    pub size: u64,
42    pub array_count: Option<u64>,
43    pub bit_mask: Option<(u32, u32)>,
44    pub children: std::collections::HashMap<String, MemorySlice>,
45}
46
47#[derive(Serialize, Deserialize)]
48pub enum MemorySlices {
49    Flat(std::collections::HashMap<String, AxiData>),
50    Tree(std::collections::HashMap<String, MemorySlice>),
51}
52
53#[derive(RustEmbed)]
54#[folder = "axi-data"]
55struct WHPciData;
56
57pub fn load_axi_table(file: &str, _version: u32) -> MemorySlices {
58    let data = WHPciData::get(file).unwrap();
59    bincode::deserialize(&data.data).unwrap()
60}
61
62/// This is a generic trait which defines the high level chip communication primitives.
63/// It's functions allow for the reading and writing of data to arbitrary noc endpoints on any chip
64/// with the details of how the endpoint is accessed abstracted away.
65///
66/// For the ARC endpoint special functions are defined because unlike most noc endpoints the ARC addresses
67/// are mapped into the pci BAR address space.
68pub trait ChipComms {
69    /// Translate a String path into the corresponding AXI address.
70    fn axi_translate(&self, addr: &str) -> Result<AxiData, AxiError>;
71    /// Read and write to the NOC using AXI address gotten from `axi_translate`.
72    fn axi_read(
73        &self,
74        chip_if: &dyn ChipInterface,
75        addr: u64,
76        data: &mut [u8],
77    ) -> Result<(), Box<dyn std::error::Error>>;
78    fn axi_write(
79        &self,
80        chip_if: &dyn ChipInterface,
81        addr: u64,
82        data: &[u8],
83    ) -> Result<(), Box<dyn std::error::Error>>;
84
85    /// Read and write to a noc endpoint, this could be a local or remote chip.
86    fn noc_read(
87        &self,
88        chip_if: &dyn ChipInterface,
89        noc_id: u8,
90        x: u8,
91        y: u8,
92        addr: u64,
93        data: &mut [u8],
94    ) -> Result<(), Box<dyn std::error::Error>>;
95    fn noc_write(
96        &self,
97        chip_if: &dyn ChipInterface,
98        noc_id: u8,
99        x: u8,
100        y: u8,
101        addr: u64,
102        data: &[u8],
103    ) -> Result<(), Box<dyn std::error::Error>>;
104    fn noc_multicast(
105        &self,
106        chip_if: &dyn ChipInterface,
107        noc_id: u8,
108        start: (u8, u8),
109        end: (u8, u8),
110        addr: u64,
111        data: &[u8],
112    ) -> Result<(), Box<dyn std::error::Error>>;
113    fn noc_broadcast(
114        &self,
115        chip_if: &dyn ChipInterface,
116        noc_id: u8,
117        addr: u64,
118        data: &[u8],
119    ) -> Result<(), Box<dyn std::error::Error>>;
120
121    /// Convenience functions for reading and writing 32 bit values.
122    fn noc_read32(
123        &self,
124        chip_if: &dyn ChipInterface,
125        noc_id: u8,
126        x: u8,
127        y: u8,
128        addr: u64,
129    ) -> Result<u32, Box<dyn std::error::Error>> {
130        let mut value = [0; 4];
131        self.noc_read(chip_if, noc_id, x, y, addr, &mut value)?;
132        Ok(u32::from_le_bytes(value))
133    }
134
135    fn noc_write32(
136        &self,
137        chip_if: &dyn ChipInterface,
138        noc_id: u8,
139        x: u8,
140        y: u8,
141        addr: u64,
142        value: u32,
143    ) -> Result<(), Box<dyn std::error::Error>> {
144        self.noc_write(chip_if, noc_id, x, y, addr, value.to_le_bytes().as_slice())
145    }
146
147    fn noc_multicast32(
148        &self,
149        chip_if: &dyn ChipInterface,
150        noc_id: u8,
151        start: (u8, u8),
152        end: (u8, u8),
153        addr: u64,
154        value: u32,
155    ) -> Result<(), Box<dyn std::error::Error>> {
156        self.noc_multicast(
157            chip_if,
158            noc_id,
159            start,
160            end,
161            addr,
162            value.to_le_bytes().as_slice(),
163        )
164    }
165
166    fn noc_broadcast32(
167        &self,
168        chip_if: &dyn ChipInterface,
169        noc_id: u8,
170        addr: u64,
171        value: u32,
172    ) -> Result<(), Box<dyn std::error::Error>> {
173        self.noc_broadcast(chip_if, noc_id, addr, value.to_le_bytes().as_slice())
174    }
175
176    fn axi_read32(
177        &self,
178        chip_if: &dyn ChipInterface,
179        addr: u64,
180    ) -> Result<u32, Box<dyn std::error::Error>> {
181        let mut value = [0; 4];
182        self.axi_read(chip_if, addr, &mut value)?;
183        Ok(u32::from_le_bytes(value))
184    }
185
186    fn axi_write32(
187        &self,
188        chip_if: &dyn ChipInterface,
189        addr: u64,
190        value: u32,
191    ) -> Result<(), Box<dyn std::error::Error>> {
192        self.axi_write(chip_if, addr, value.to_le_bytes().as_slice())
193    }
194
195    fn axi_sread32(
196        &self,
197        chip_if: &dyn ChipInterface,
198        addr: &str,
199    ) -> Result<u32, Box<dyn std::error::Error>> {
200        let addr = self.axi_translate(addr.as_ref())?.addr;
201
202        let mut value = [0; 4];
203        self.axi_read(chip_if, addr, &mut value)?;
204        Ok(u32::from_le_bytes(value))
205    }
206
207    fn axi_swrite32(
208        &self,
209        chip_if: &dyn ChipInterface,
210        addr: &str,
211        value: u32,
212    ) -> Result<(), Box<dyn std::error::Error>> {
213        let addr = self.axi_translate(addr.as_ref())?.addr;
214
215        self.axi_write(chip_if, addr, &value.to_le_bytes())?;
216        Ok(())
217    }
218}
219
220pub fn axi_translate_tree(
221    data: &std::collections::HashMap<String, MemorySlice>,
222    addr: &str,
223) -> Result<AxiData, AxiError> {
224    let mut it = data;
225
226    let mut offset = 0;
227    let mut size = 0;
228    let mut bits = None;
229    for key in addr.split('.') {
230        let (slice, index) = if !it.contains_key(key) && key.contains('[') {
231            let mut parts = key.split('[');
232            let key = parts.next().ok_or_else(|| AxiError::InvalidArrayPath {
233                key: key.to_string(),
234                path: addr.to_string(),
235            })?;
236            let index = parts
237                .next()
238                .ok_or_else(|| AxiError::InvalidArrayPath {
239                    key: key.to_string(),
240                    path: addr.to_string(),
241                })?
242                .trim_end_matches(']')
243                .parse::<u64>()
244                .map_err(|_| AxiError::InvalidArrayPath {
245                    key: key.to_string(),
246                    path: addr.to_string(),
247                })?;
248
249            (it.get(key), Some(index))
250        } else {
251            (it.get(key), None)
252        };
253
254        if let Some(slice) = slice {
255            it = &slice.children;
256            if let (Some(count), Some(index)) = (slice.array_count, index) {
257                assert!(index < count);
258            } else if index.is_some() ^ slice.array_count.is_some() {
259                dbg!(slice);
260                panic!("Tried to index a non-array or no index for array with {key}")
261            }
262
263            size = slice.size;
264            bits = slice.bit_mask;
265            offset += slice.offset + slice.size * index.unwrap_or(0);
266        } else {
267            return Err(AxiError::InvalidPath {
268                path: addr.to_string(),
269                key: key.to_string(),
270            });
271        }
272    }
273
274    Ok(AxiData {
275        addr: offset,
276        size,
277        bits,
278    })
279}
280
281pub fn axi_translate(data: Option<&MemorySlices>, addr: &str) -> Result<AxiData, AxiError> {
282    match data.ok_or(AxiError::NoAxiData)? {
283        MemorySlices::Flat(data) => {
284            if let Some(data) = data.get(addr) {
285                Ok(data.clone())
286            } else {
287                Err(AxiError::InvalidPath {
288                    path: addr.to_string(),
289                    key: addr.to_string(),
290                })
291            }
292        }
293        MemorySlices::Tree(data) => axi_translate_tree(data, addr),
294    }
295}
296
297pub struct ArcIf {
298    pub axi_data: MemorySlices,
299}
300
301impl ChipComms for ArcIf {
302    fn axi_translate(&self, addr: &str) -> Result<AxiData, AxiError> {
303        axi_translate(Some(&self.axi_data), addr)
304    }
305
306    fn axi_read(
307        &self,
308        chip_if: &dyn ChipInterface,
309        addr: u64,
310        data: &mut [u8],
311    ) -> Result<(), Box<dyn std::error::Error>> {
312        chip_if.axi_read(addr as u32, data)
313    }
314
315    fn axi_write(
316        &self,
317        chip_if: &dyn ChipInterface,
318        addr: u64,
319        data: &[u8],
320    ) -> Result<(), Box<dyn std::error::Error>> {
321        chip_if.axi_write(addr as u32, data)
322    }
323
324    fn noc_read(
325        &self,
326        chip_if: &dyn ChipInterface,
327        noc_id: u8,
328        x: u8,
329        y: u8,
330        addr: u64,
331        data: &mut [u8],
332    ) -> Result<(), Box<dyn std::error::Error>> {
333        chip_if.noc_read(noc_id, x, y, addr, data)
334    }
335
336    fn noc_write(
337        &self,
338        chip_if: &dyn ChipInterface,
339        noc_id: u8,
340        x: u8,
341        y: u8,
342        addr: u64,
343        data: &[u8],
344    ) -> Result<(), Box<dyn std::error::Error>> {
345        chip_if.noc_write(noc_id, x, y, addr, data)
346    }
347
348    fn noc_multicast(
349        &self,
350        chip_if: &dyn ChipInterface,
351        noc_id: u8,
352        start: (u8, u8),
353        end: (u8, u8),
354        addr: u64,
355        data: &[u8],
356    ) -> Result<(), Box<dyn std::error::Error>> {
357        chip_if.noc_multicast(noc_id, start, end, addr, data)
358    }
359
360    fn noc_broadcast(
361        &self,
362        chip_if: &dyn ChipInterface,
363        noc_id: u8,
364        addr: u64,
365        data: &[u8],
366    ) -> Result<(), Box<dyn std::error::Error>> {
367        chip_if.noc_broadcast(noc_id, addr, data)
368    }
369}
370
371pub struct NocIf {
372    pub axi_data: MemorySlices,
373    pub noc_id: u8,
374    pub x: u8,
375    pub y: u8,
376}
377
378impl ChipComms for NocIf {
379    fn axi_translate(&self, addr: &str) -> Result<AxiData, AxiError> {
380        axi_translate(Some(&self.axi_data), addr)
381    }
382
383    fn axi_read(
384        &self,
385        chip_if: &dyn ChipInterface,
386        addr: u64,
387        data: &mut [u8],
388    ) -> Result<(), Box<dyn std::error::Error>> {
389        chip_if.noc_read(self.noc_id, self.x, self.y, addr, data)
390    }
391
392    fn axi_write(
393        &self,
394        chip_if: &dyn ChipInterface,
395        addr: u64,
396        data: &[u8],
397    ) -> Result<(), Box<dyn std::error::Error>> {
398        chip_if.noc_write(self.noc_id, self.x, self.y, addr, data)
399    }
400
401    fn noc_read(
402        &self,
403        chip_if: &dyn ChipInterface,
404        noc_id: u8,
405        x: u8,
406        y: u8,
407        addr: u64,
408        data: &mut [u8],
409    ) -> Result<(), Box<dyn std::error::Error>> {
410        chip_if.noc_read(noc_id, x, y, addr, data)
411    }
412
413    fn noc_write(
414        &self,
415        chip_if: &dyn ChipInterface,
416        noc_id: u8,
417        x: u8,
418        y: u8,
419        addr: u64,
420        data: &[u8],
421    ) -> Result<(), Box<dyn std::error::Error>> {
422        chip_if.noc_write(noc_id, x, y, addr, data)
423    }
424
425    fn noc_multicast(
426        &self,
427        chip_if: &dyn ChipInterface,
428        noc_id: u8,
429        start: (u8, u8),
430        end: (u8, u8),
431        addr: u64,
432        data: &[u8],
433    ) -> Result<(), Box<dyn std::error::Error>> {
434        chip_if.noc_multicast(noc_id, start, end, addr, data)
435    }
436
437    fn noc_broadcast(
438        &self,
439        chip_if: &dyn ChipInterface,
440        noc_id: u8,
441        addr: u64,
442        data: &[u8],
443    ) -> Result<(), Box<dyn std::error::Error>> {
444        chip_if.noc_broadcast(noc_id, addr, data)
445    }
446}
447
448impl ChipComms for Arc<dyn ChipComms> {
449    fn axi_translate(&self, addr: &str) -> Result<AxiData, AxiError> {
450        self.as_ref().axi_translate(addr)
451    }
452
453    fn axi_read(
454        &self,
455        chip_if: &dyn ChipInterface,
456        addr: u64,
457        data: &mut [u8],
458    ) -> Result<(), Box<dyn std::error::Error>> {
459        self.as_ref().axi_read(chip_if, addr, data)
460    }
461
462    fn axi_write(
463        &self,
464        chip_if: &dyn ChipInterface,
465        addr: u64,
466        data: &[u8],
467    ) -> Result<(), Box<dyn std::error::Error>> {
468        self.as_ref().axi_write(chip_if, addr, data)
469    }
470
471    fn noc_read(
472        &self,
473        chip_if: &dyn ChipInterface,
474        noc_id: u8,
475        x: u8,
476        y: u8,
477        addr: u64,
478        data: &mut [u8],
479    ) -> Result<(), Box<dyn std::error::Error>> {
480        self.as_ref().noc_read(chip_if, noc_id, x, y, addr, data)
481    }
482
483    fn noc_write(
484        &self,
485        chip_if: &dyn ChipInterface,
486        noc_id: u8,
487        x: u8,
488        y: u8,
489        addr: u64,
490        data: &[u8],
491    ) -> Result<(), Box<dyn std::error::Error>> {
492        self.as_ref().noc_write(chip_if, noc_id, x, y, addr, data)
493    }
494
495    fn noc_multicast(
496        &self,
497        chip_if: &dyn ChipInterface,
498        noc_id: u8,
499        start: (u8, u8),
500        end: (u8, u8),
501        addr: u64,
502        data: &[u8],
503    ) -> Result<(), Box<dyn std::error::Error>> {
504        self.as_ref()
505            .noc_multicast(chip_if, noc_id, start, end, addr, data)
506    }
507
508    fn noc_broadcast(
509        &self,
510        chip_if: &dyn ChipInterface,
511        noc_id: u8,
512        addr: u64,
513        data: &[u8],
514    ) -> Result<(), Box<dyn std::error::Error>> {
515        self.as_ref().noc_broadcast(chip_if, noc_id, addr, data)
516    }
517}
518
519impl ChipComms for Arc<dyn ChipComms + Send + Sync> {
520    fn axi_translate(&self, addr: &str) -> Result<AxiData, AxiError> {
521        self.as_ref().axi_translate(addr)
522    }
523
524    fn axi_read(
525        &self,
526        chip_if: &dyn ChipInterface,
527        addr: u64,
528        data: &mut [u8],
529    ) -> Result<(), Box<dyn std::error::Error>> {
530        self.as_ref().axi_read(chip_if, addr, data)
531    }
532
533    fn axi_write(
534        &self,
535        chip_if: &dyn ChipInterface,
536        addr: u64,
537        data: &[u8],
538    ) -> Result<(), Box<dyn std::error::Error>> {
539        self.as_ref().axi_write(chip_if, addr, data)
540    }
541
542    fn noc_read(
543        &self,
544        chip_if: &dyn ChipInterface,
545        noc_id: u8,
546        x: u8,
547        y: u8,
548        addr: u64,
549        data: &mut [u8],
550    ) -> Result<(), Box<dyn std::error::Error>> {
551        self.as_ref().noc_read(chip_if, noc_id, x, y, addr, data)
552    }
553
554    fn noc_write(
555        &self,
556        chip_if: &dyn ChipInterface,
557        noc_id: u8,
558        x: u8,
559        y: u8,
560        addr: u64,
561        data: &[u8],
562    ) -> Result<(), Box<dyn std::error::Error>> {
563        self.as_ref().noc_write(chip_if, noc_id, x, y, addr, data)
564    }
565
566    fn noc_multicast(
567        &self,
568        chip_if: &dyn ChipInterface,
569        noc_id: u8,
570        start: (u8, u8),
571        end: (u8, u8),
572        addr: u64,
573        data: &[u8],
574    ) -> Result<(), Box<dyn std::error::Error>> {
575        self.as_ref()
576            .noc_multicast(chip_if, noc_id, start, end, addr, data)
577    }
578
579    fn noc_broadcast(
580        &self,
581        chip_if: &dyn ChipInterface,
582        noc_id: u8,
583        addr: u64,
584        data: &[u8],
585    ) -> Result<(), Box<dyn std::error::Error>> {
586        self.as_ref().noc_broadcast(chip_if, noc_id, addr, data)
587    }
588}
luwen_api/chip/communication/chip_comms.rs

luwen_api/chip/communication/
chip_comms.rs
