et_abi/lib.rs
1//! Shared host/device ABI for the ET-SoC-1: the kernel-launch argument structs,
2//! defined **once** and used by both the host launcher and the device kernel.
3//!
4//! Kernel arguments are passed by pointer: the host stages an argument struct in
5//! device memory and the firmware delivers its address to the kernel (in `a0`).
6//! Because both the host (x86-64) and the device (RV64) are little-endian, the
7//! in-memory `#[repr(C)]` layout *is* the wire layout, so no explicit
8//! serialisation is needed -- the host takes the struct's bytes and the kernel
9//! reinterprets the pointer. Defining each struct here keeps the two sides from
10//! drifting (mismatched field order, sizes, or padding).
11
12#![no_std]
13
14/// ET-SoC-1 cache-line size, in bytes.
15///
16/// Per-hart outputs are laid out at this stride on both sides: the host strides
17/// its padded arrays by it and the device writes each hart's cell at
18/// `base + hart * CACHE_LINE`. Defining it once here keeps the two from drifting,
19/// which on this software-coherent part would cause silent false-sharing
20/// corruption.
21pub const CACHE_LINE: usize = 64;
22
23/// A wrapper that aligns `T` to a cache-line boundary.
24///
25/// On the ET-SoC-1 (a software-coherent architecture), two values sharing a
26/// cache line that are written by distinct harts without explicit cache
27/// operations cause false-sharing corruption. Wrapping per-hart output data
28/// in `CachePadded` ensures each instance occupies a distinct 64-byte line,
29/// making cross-hart false sharing structurally impossible regardless of the
30/// surrounding allocation layout.
31///
32/// The inner value is accessed directly via the public tuple field `0`.
33///
34/// # Example
35///
36/// ```
37/// use et_abi::CachePadded;
38/// let cell: CachePadded<u64> = CachePadded(0);
39/// assert_eq!(core::mem::align_of::<CachePadded<u64>>(), 64);
40/// ```
41#[repr(align(64))]
42#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
43pub struct CachePadded<T>(pub T);
44
45/// Harts per compute shire on the ET-SoC-1 (architectural constant).
46pub const HARTS_PER_SHIRE: u32 = 64;
47
48/// Harts per neighbourhood on the ET-SoC-1 (architectural constant).
49pub const HARTS_PER_NEIGHBOURHOOD: u32 = 16;
50
51/// A plain-old-data kernel-argument struct exchanged between host and device.
52///
53/// # Safety
54/// Implementors must be `#[repr(C)]`, contain only integer fields with no
55/// padding, and be valid for any bit pattern. Then [`DeviceArgs::as_bytes`] and
56/// [`DeviceArgs::from_ptr`] are a faithful round-trip on little-endian hosts and
57/// devices.
58pub unsafe trait DeviceArgs: Sized + Copy {
59 /// Borrow the struct as its on-wire bytes (host side: stage these in device
60 /// memory as the launch arguments).
61 fn as_bytes(&self) -> &[u8] {
62 // SAFETY: `Self` is repr(C) POD (trait contract), so its bytes are a
63 // valid representation of length `size_of::<Self>()`.
64 unsafe {
65 core::slice::from_raw_parts(
66 self as *const Self as *const u8,
67 core::mem::size_of::<Self>(),
68 )
69 }
70 }
71
72 /// Reinterpret a device-memory pointer as these arguments (device side).
73 ///
74 /// # Safety
75 /// `ptr` must point to at least `size_of::<Self>()` bytes of a valid,
76 /// suitably aligned instance -- e.g. the launch-args region the firmware
77 /// passed in `a0`.
78 unsafe fn from_ptr<'a>(ptr: *const u8) -> &'a Self {
79 // SAFETY: forwarded to the caller's contract on `ptr`.
80 unsafe { &*(ptr as *const Self) }
81 }
82}
83
84// ---------------------------------------------------------------------------
85// Tensor-extension constants
86// ---------------------------------------------------------------------------
87
88/// Required alignment for all matrix pointers and row strides used with the
89/// ET-SoC-1 tensor-load/store instructions. TensorLoad and TensorStore each
90/// require the source or destination address to be 64-byte aligned.
91pub const TENSOR_ALIGN: usize = 64;
92
93/// Number of addressable cache lines in each Minion's L1 scratchpad.
94/// TensorLoad START field is 6 bits, spanning lines 0..47 inclusive.
95pub const SCP_LINES: usize = 48;
96
97/// Bytes per L1 scratchpad line (one cache line).
98pub const SCP_LINE_BYTES: usize = 64;
99
100/// Minion cores per compute shire on the ET-SoC-1.
101/// Each shire has 32 dual-threaded Minion cores (64 harts total).
102pub const MINIONS_PER_SHIRE: u32 = 32;
103
104// ---------------------------------------------------------------------------
105// GEMM tile dimensions
106// ---------------------------------------------------------------------------
107
108/// Number of C output rows computed per tile by TensorFMA32.
109/// Equals the maximum AROWS+1 value (4-bit field, max 15 -> 16 rows).
110pub const GEMM_TILE_M: usize = 16;
111
112/// Inner-dimension (K) slice processed per TensorFMA32 call.
113/// Limited to 16 f32 values per A-matrix row fitting in one 64-byte
114/// scratchpad line (ACOLS field is 4-bit, max 15 -> 16 columns).
115pub const GEMM_TILE_K: usize = 16;
116
117/// Number of f32 output columns produced per TensorFMA32 call (BCOLS=3 gives
118/// 4*(3+1) = 16 columns). Each tile row occupies exactly 64 bytes in the FP
119/// register file. N need not be a multiple of this value; the last tile column
120/// may be partial, with the hardware writing 64 bytes per row regardless --
121/// the caller reads only the N valid columns from the 64-byte-aligned allocation.
122pub const GEMM_TILE_N: usize = 16;
123
124// ---------------------------------------------------------------------------
125// GemmArgs
126// ---------------------------------------------------------------------------
127
128/// Arguments for the single-precision general matrix multiplication (sGEMM)
129/// kernel (`sgemm-rs`), implementing C = alpha*A*B + beta*C.
130///
131/// # Layout invariants (v0.1 restrictions)
132/// - `alpha` must be `1.0` and `beta` must be `0.0`.
133/// - `n` may be any positive integer; partial last-column tiles are handled
134/// transparently via 64-byte-aligned row padding.
135/// - `a`, `b`, `c` must be [`TENSOR_ALIGN`]-byte aligned device addresses.
136/// - `lda`, `ldb`, `ldc` must be multiples of [`TENSOR_ALIGN`] (64 bytes).
137///
138/// All dimensions are in elements; leading dimensions are in bytes.
139///
140/// # ABI layout
141/// The four 8-byte fields (`a`, `b`, `c`, `n_shires`) are grouped first to
142/// give the struct 8-byte alignment with no internal or trailing padding:
143/// `4*8 + 8*4 = 64 bytes` total.
144#[repr(C)]
145#[derive(Clone, Copy, Debug, PartialEq)]
146pub struct GemmArgs {
147 /// Device address of A [M x K], row-major, 64-byte aligned.
148 pub a: u64,
149 /// Device address of B [K x N], row-major, 64-byte aligned.
150 pub b: u64,
151 /// Device address of C [M x N], row-major, 64-byte aligned.
152 pub c: u64,
153 /// Number of participating compute shires. Stored as `u64` to keep
154 /// all 8-byte fields contiguous and the total struct size a multiple
155 /// of the struct's 8-byte alignment. Effective range: 1..=34.
156 pub n_shires: u64,
157 /// Number of rows of A and C (M dimension).
158 pub m: u32,
159 /// Number of columns of B and C (N dimension). May be any positive integer;
160 /// the last output tile column is partial when N is not a multiple of 16.
161 pub n: u32,
162 /// Shared inner dimension (K): columns of A and rows of B.
163 pub k: u32,
164 /// Row stride of A in bytes (multiple of 64).
165 pub lda: u32,
166 /// Row stride of B in bytes (multiple of 64).
167 pub ldb: u32,
168 /// Row stride of C in bytes (multiple of 64).
169 pub ldc: u32,
170 /// A*B scaling factor. Must be `1.0` in v0.1.
171 pub alpha: f32,
172 /// C scaling factor. Must be `0.0` in v0.1.
173 pub beta: f32,
174}
175
176// SAFETY: repr(C); 4 u64 fields followed by 8 u32/f32 fields, ordered by
177// decreasing size -> no padding. 4*8 + 8*4 = 64 bytes, a multiple of the
178// struct's 8-byte alignment.
179unsafe impl DeviceArgs for GemmArgs {}
180const _: () = assert!(core::mem::size_of::<GemmArgs>() == 64);
181
182/// Arguments for the data-parallel reduction kernel (`reduce-rs`).
183#[repr(C)]
184#[derive(Clone, Copy, Debug, PartialEq, Eq)]
185pub struct ReduceArgs {
186 /// Device address of the input array (`n` × `u32`).
187 pub input: u64,
188 /// Device address of the output array (`n_harts` × one `u64` per cache line).
189 pub out: u64,
190 /// Number of input elements.
191 pub n: u32,
192 /// Number of participating harts.
193 pub n_harts: u32,
194}
195
196// SAFETY: repr(C), only u64/u32 fields ordered by decreasing size -> no padding.
197unsafe impl DeviceArgs for ReduceArgs {}
198const _: () = assert!(core::mem::size_of::<ReduceArgs>() == 24);
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203
204 #[test]
205 fn gemm_args_size() {
206 // 4 u64 + 8 u32/f32 = 32 + 32 = 64 bytes, a multiple of 8.
207 assert_eq!(core::mem::size_of::<GemmArgs>(), 64);
208 }
209
210 #[test]
211 fn gemm_args_roundtrip() {
212 let a = GemmArgs {
213 a: 0x0080_0100_0000,
214 b: 0x0080_0200_0000,
215 c: 0x0080_0300_0000,
216 n_shires: 4,
217 m: 128,
218 n: 64,
219 k: 256,
220 lda: 1024, // 256 * 4 bytes, 64-byte aligned
221 ldb: 256, // 64 * 4 bytes, 64-byte aligned
222 ldc: 256, // 64 * 4 bytes, 64-byte aligned
223 alpha: 1.0,
224 beta: 0.0,
225 };
226 let bytes = a.as_bytes();
227 assert_eq!(bytes.len(), 64);
228 let b = unsafe { GemmArgs::from_ptr(bytes.as_ptr()) };
229 assert_eq!(*b, a);
230 // Verify leading dimensions are 64-byte aligned as the kernel requires.
231 assert_eq!(a.lda as usize % TENSOR_ALIGN, 0);
232 assert_eq!(a.ldb as usize % TENSOR_ALIGN, 0);
233 assert_eq!(a.ldc as usize % TENSOR_ALIGN, 0);
234 // Verify N is a multiple of GEMM_TILE_N.
235 assert_eq!(a.n as usize % GEMM_TILE_N, 0);
236 }
237
238 #[test]
239 fn reduce_args_roundtrip() {
240 let a = ReduceArgs {
241 input: 0x0080_0580_1000,
242 out: 0x0080_0590_0000,
243 n: 262_144,
244 n_harts: 64,
245 };
246 let bytes = a.as_bytes();
247 assert_eq!(bytes.len(), 24);
248 // The device would do exactly this from the args pointer.
249 let b = unsafe { ReduceArgs::from_ptr(bytes.as_ptr()) };
250 assert_eq!(*b, a);
251 }
252}