tpu-sg2002 0.1.0

TPU driver in Rust for SG2002 SoC.
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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

//! TPU data types and structures.

use crate::{PhysAddr, TpuError, registers::CPU_ENGINE_DESCRIPTOR_NUM};
use crate::platform::TPU_DMABUF_HEADER_M;

/// DMA buffer header used by the TPU firmware.
#[repr(C)]
#[derive(Debug, Clone, Copy, Default)]
pub struct DmaHeader {
    pub dmabuf_magic_m: u16,
    pub dmabuf_magic_s: u16,
    pub dmabuf_size: u32,
    pub cpu_desc_count: u32,
    pub bd_desc_count: u32,
    pub tdma_desc_count: u32,
    pub tpu_clk_rate: u32,
    pub pmubuf_size: u32,
    pub pmubuf_offset: u32,
    pub arraybase_0_l: u32,
    pub arraybase_0_h: u32,
    pub arraybase_1_l: u32,
    pub arraybase_1_h: u32,
    pub arraybase_2_l: u32,
    pub arraybase_2_h: u32,
    pub arraybase_3_l: u32,
    pub arraybase_3_h: u32,
    pub arraybase_4_l: u32,
    pub arraybase_4_h: u32,
    pub arraybase_5_l: u32,
    pub arraybase_5_h: u32,
    pub arraybase_6_l: u32,
    pub arraybase_6_h: u32,
    pub arraybase_7_l: u32,
    pub arraybase_7_h: u32,
    pub reserve: [u32; 8],
}

impl DmaHeader {
    pub const MAGIC: u16 = TPU_DMABUF_HEADER_M;

    pub fn is_valid(&self) -> bool {
        // Keep compatibility with the original C driver, which validates only magic_m.
        self.dmabuf_magic_m == Self::MAGIC
    }

    pub fn arraybase(&self, idx: usize) -> u64 {
        let (lo, hi) = match idx {
            0 => (self.arraybase_0_l, self.arraybase_0_h),
            1 => (self.arraybase_1_l, self.arraybase_1_h),
            2 => (self.arraybase_2_l, self.arraybase_2_h),
            3 => (self.arraybase_3_l, self.arraybase_3_h),
            4 => (self.arraybase_4_l, self.arraybase_4_h),
            5 => (self.arraybase_5_l, self.arraybase_5_h),
            6 => (self.arraybase_6_l, self.arraybase_6_h),
            7 => (self.arraybase_7_l, self.arraybase_7_h),
            _ => (0, 0),
        };
        (lo as u64) | ((hi as u64) << 32)
    }

    pub fn arraybase_l(&self) -> [u32; 8] {
        [
            self.arraybase_0_l, self.arraybase_1_l, self.arraybase_2_l, self.arraybase_3_l,
            self.arraybase_4_l, self.arraybase_5_l, self.arraybase_6_l, self.arraybase_7_l,
        ]
    }
}

/// Borrowed view of a TPU command DMA buffer.
#[derive(Debug, Clone, Copy)]
pub struct DmabufView<'a> {
    pub header: &'a DmaHeader,
    pub descs: &'a [CpuSyncDesc],
}

/// Parse a TPU command DMA buffer from a raw CPU-accessible pointer.
///
/// # Safety
///
/// Caller must ensure `dmabuf_vaddr` points to readable memory of length
/// `dmabuf_len` for the entire returned lifetime.
pub unsafe fn parse_dmabuf_view<'a>(dmabuf_vaddr: *const u8, dmabuf_len: usize) -> Result<DmabufView<'a>, TpuError> {
    if dmabuf_vaddr.is_null() {
        return Err(TpuError::InvalidParameter);
    }

    let header_size = core::mem::size_of::<DmaHeader>();
    if dmabuf_len < header_size {
        return Err(TpuError::InvalidParameter);
    }

    let header = unsafe { &*(dmabuf_vaddr as *const DmaHeader) };
    let desc_count = header.cpu_desc_count as usize;
    let desc_bytes = desc_count
        .checked_mul(core::mem::size_of::<CpuSyncDesc>())
        .ok_or(TpuError::InvalidParameter)?;
    let required = header_size
        .checked_add(desc_bytes)
        .ok_or(TpuError::InvalidParameter)?;
    if required > dmabuf_len {
        return Err(TpuError::InvalidParameter);
    }

    let desc_ptr = unsafe { dmabuf_vaddr.add(header_size) as *const CpuSyncDesc };
    let descs = unsafe { core::slice::from_raw_parts(desc_ptr, desc_count) };

    Ok(DmabufView { header, descs })
}

/// CPU sync descriptor embedded in command buffers.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct CpuSyncDesc {
    pub op_type: u32,
    pub num_bd: u32,
    pub num_gdma: u32,
    pub offset_bd: u32,
    pub offset_gdma: u32,
    pub reserved: [u32; 2],
    pub str_data: [u8; (CPU_ENGINE_DESCRIPTOR_NUM - 7) * core::mem::size_of::<u32>()],
}

impl Default for CpuSyncDesc {
    fn default() -> Self {
        Self {
            op_type: 0,
            num_bd: 0,
            num_gdma: 0,
            offset_bd: 0,
            offset_gdma: 0,
            reserved: [0; 2],
            str_data: [0; (CPU_ENGINE_DESCRIPTOR_NUM - 7) * core::mem::size_of::<u32>()],
        }
    }
}

impl CpuSyncDesc {
    pub fn bd_count(&self) -> u32 { self.num_bd & 0xFFFF }
    pub fn tdma_count(&self) -> u32 { self.num_gdma & 0xFFFF }
}

/// TPU device configuration.
#[derive(Debug, Clone, Copy, Default)]
pub struct TpuConfig {
    pub pmu_buf_size: u32,
    pub pmu_buf_paddr: PhysAddr,
}

/// TDMA PIO submission info.
#[repr(C)]
#[derive(Debug, Clone, Copy, Default)]
pub struct TpuTdmaPioInfo {
    pub paddr_src: u64,
    pub paddr_dst: u64,
    pub h: u32,
    pub w_bytes: u32,
    pub stride_bytes_src: u32,
    pub stride_bytes_dst: u32,
    pub enable_2d: u32,
    pub leng_bytes: u32,
}

impl TpuTdmaPioInfo {
    pub fn is_2d(&self) -> bool { self.enable_2d != 0 }
}