luwen-if 0.4.4

Generic interface to Tenstorrent ai accelerators, abstracting the details of communication
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0

use std::sync::Arc;

use luwen_core::Arch;

use crate::{
    arc_msg::ArcMsgAddr,
    error::{BtWrapper, PlatformError},
    CallbackStorage,
};

use super::{
    communication::chip_comms::load_axi_table, ArcIf, Chip, ChipComms, ChipInterface, Grayskull,
    Wormhole,
};

impl Chip {
    pub fn gs_open<T: Clone + Send + Sync + 'static>(
        arch: Arch,
        backend: CallbackStorage<T>,
    ) -> Result<Grayskull, PlatformError> {
        if let Arch::Grayskull = arch {
            let version = [0u8; 4];
            let version = u32::from_le_bytes(version);

            let arc_if = Arc::new(ArcIf {
                axi_data: load_axi_table("grayskull-axi-pci.bin", version),
            });

            Ok(Grayskull::create(
                Arc::new(backend),
                arc_if.clone(),
                ArcMsgAddr::try_from(arc_if.as_ref() as &dyn ChipComms)?,
            ))
        } else {
            Err(PlatformError::WrongChipArch {
                actual: arch,
                expected: Arch::Grayskull,
                backtrace: BtWrapper::capture(),
            })
        }
    }

    pub fn wh_open<T: Clone + Send + Sync + 'static>(
        arch: Arch,
        backend: CallbackStorage<T>,
    ) -> Result<Wormhole, PlatformError> {
        if let Arch::Wormhole = arch {
            let version = [0u8; 4];
            let version = u32::from_le_bytes(version);

            let arc_if = ArcIf {
                axi_data: load_axi_table("wormhole-axi-pci.bin", version),
            };

            Ok(Wormhole::init(
                false,
                true,
                arc_if,
                Arc::new(backend) as Arc<dyn ChipInterface + Sync + Send>,
            )?)
        } else {
            Err(PlatformError::WrongChipArch {
                actual: arch,
                expected: Arch::Wormhole,
                backtrace: BtWrapper::capture(),
            })
        }
    }

    pub fn open<T: Clone + Send + Sync + 'static>(
        arch: Arch,
        backend: CallbackStorage<T>,
    ) -> Result<Chip, PlatformError> {
        Ok(Chip {
            inner: match arch {
                Arch::Grayskull => Box::new(Self::gs_open(arch, backend)?),
                Arch::Wormhole => Box::new(Self::wh_open(arch, backend)?),
                _ => panic!("Unsupported chip"),
            },
        })
    }
}