et_kernel/simd.rs
1//! Packed-single (PS) SIMD intrinsics for the ET-SoC-1 Minion FP register file.
2//!
3//! The ET-SoC-1 PS extension operates on the 256-bit FP registers (f0..f31),
4//! treating each as a vector of eight single-precision (f32) lanes. PS
5//! instructions share the standard RISC-V FP register file and therefore
6//! require the `f` target feature to be enabled at compile time.
7//!
8//! # Enabling this module
9//!
10//! The module is gated on `cfg(target_feature = "f")`. When building with
11//! `.cargo/config.toml` targeting `riscv64gc-unknown-none-elf`, the F
12//! extension is included in the target triple and the module is available
13//! without any additional flags. On a different target, add `+f` explicitly:
14//!
15//! ```toml
16//! [target.riscv64gc-unknown-none-elf]
17//! rustflags = ["-C", "target-feature=+f"]
18//! ```
19//!
20//! Without `+f` the module is empty; callers must be similarly gated.
21//!
22//! # Register assignment for a 16-row C tile
23//!
24//! A 16-row GEMM C tile occupies the full FP register file (f0..f31) in
25//! register pairs: row N maps to (f[2N], f[2N+1]). The broadcast scratch
26//! register for [`broadcast_ps`] and [`fmul_ps_row`] must not overlap with
27//! the pair being scaled.
28//!
29//! For tiles of at most 14 rows, pass [`PS_SCRATCH_DEFAULT`] (f28) as
30//! `scratch`. For a full 16-row tile, rows 14 (f28/f29) and 15 (f30/f31)
31//! conflict with f28, so the caller must choose a different scratch register:
32//!
33//! 1. Spill the chosen scratch register to the stack (one FP register via
34//! `fsw`/`fld` with a suitable frame slot).
35//! 2. Call [`broadcast_ps`]`(alpha, scratch)`.
36//! 3. Call [`fmul_ps_row`] for the conflicting row.
37//! 4. Restore the scratch register from the stack.
38//!
39//! # Encodings
40//!
41//! Source: `gdb/include/opcode/esperanto-opc.h` in the ET-SoC-1 binutils fork.
42//! All PS arithmetic instructions use the RISC-V custom-3 opcode space (0x7b).
43//!
44//! | Instruction | MATCH | Format | Notes |
45//! |---|---|---|---|
46//! | `fmul.ps` | `0x1000_007b` | R-type | fd = fs1 .* fs2 (element-wise). funct7=0x08, funct3=0. |
47//! | `fbc.ps` | `0x0000_000b` | I-type | Load 4 B from `rs1+imm`, broadcast to all 8 lanes of fd. opcode=0x0b, funct3=0. |
48//! | `fbcx.ps` | `0x0000_300b` | I-type | Broadcast GPR `rs1` to all 8 lanes of fd; imm=0 fixed. opcode=0x0b, funct3=3. |
49//! | `fmvs.x.ps` | `0xe000_207b` | R-type | Extract lane `rs2[2:0]` from PS register `rs1` to GPR `rd`. funct7=0x70, funct3=2. |
50
51#[cfg(target_arch = "riscv64")]
52#[cfg(target_feature = "f")]
53mod inner {
54 use core::arch::asm;
55
56 /// Default PS broadcast scratch register (f28 / ft8).
57 ///
58 /// Safe for C tiles with at most 14 rows (rows 0..=13). Rows 14 and 15
59 /// place C-tile data in f28..f31, conflicting with this value; see the
60 /// module doc for the spill/restore pattern required for full 16-row tiles.
61 pub const PS_SCRATCH_DEFAULT: u8 = 28;
62
63 // Emit two FMUL.PS instructions for a register pair with literal operands.
64 // R-type: opcode=0x7b, funct3=0, funct7=8.
65 // fd = fs1 = f{lo} or f{hi}; fs2 = f{s}.
66 #[rustfmt::skip] // hand-laid .insn operands; keep on one line per instruction
67 macro_rules! fmul2 {
68 ($lo:literal, $hi:literal, $s:literal) => {
69 asm!(
70 concat!(
71 ".insn r 0x7b, 0, 8, f", $lo, ", f", $lo, ", f", $s, "\n",
72 ".insn r 0x7b, 0, 8, f", $hi, ", f", $hi, ", f", $s
73 ),
74 options(nostack, preserves_flags)
75 )
76 };
77 }
78
79 // Dispatch FMUL.PS for register pair ($lo, $hi) over all 32 scratch registers.
80 #[rustfmt::skip] // 32-way scratch dispatch; keep one compact arm per register
81 macro_rules! scale_row {
82 (($lo:literal, $hi:literal), $s:expr) => {
83 match $s {
84 0 => { fmul2!($lo, $hi, 0) },
85 1 => { fmul2!($lo, $hi, 1) },
86 2 => { fmul2!($lo, $hi, 2) },
87 3 => { fmul2!($lo, $hi, 3) },
88 4 => { fmul2!($lo, $hi, 4) },
89 5 => { fmul2!($lo, $hi, 5) },
90 6 => { fmul2!($lo, $hi, 6) },
91 7 => { fmul2!($lo, $hi, 7) },
92 8 => { fmul2!($lo, $hi, 8) },
93 9 => { fmul2!($lo, $hi, 9) },
94 10 => { fmul2!($lo, $hi, 10) },
95 11 => { fmul2!($lo, $hi, 11) },
96 12 => { fmul2!($lo, $hi, 12) },
97 13 => { fmul2!($lo, $hi, 13) },
98 14 => { fmul2!($lo, $hi, 14) },
99 15 => { fmul2!($lo, $hi, 15) },
100 16 => { fmul2!($lo, $hi, 16) },
101 17 => { fmul2!($lo, $hi, 17) },
102 18 => { fmul2!($lo, $hi, 18) },
103 19 => { fmul2!($lo, $hi, 19) },
104 20 => { fmul2!($lo, $hi, 20) },
105 21 => { fmul2!($lo, $hi, 21) },
106 22 => { fmul2!($lo, $hi, 22) },
107 23 => { fmul2!($lo, $hi, 23) },
108 24 => { fmul2!($lo, $hi, 24) },
109 25 => { fmul2!($lo, $hi, 25) },
110 26 => { fmul2!($lo, $hi, 26) },
111 27 => { fmul2!($lo, $hi, 27) },
112 28 => { fmul2!($lo, $hi, 28) },
113 29 => { fmul2!($lo, $hi, 29) },
114 30 => { fmul2!($lo, $hi, 30) },
115 31 => { fmul2!($lo, $hi, 31) },
116 _ => {}
117 }
118 };
119 }
120
121 /// Broadcast `scalar` to all eight lanes of PS register `dest` (0..=31).
122 ///
123 /// Issues `fmv.x.w` to move the scalar's IEEE-754 bit pattern into a
124 /// temporary integer register, then `FBCX.PS f{dest}, tmp`
125 /// (`MATCH_FBCX_PS = 0x0000_300b`; I-type, opcode=0x0b, funct3=3, imm=0).
126 ///
127 /// After this call, `f{dest}` holds `[scalar; 8]` in PS interpretation,
128 /// ready to be used as the `scratch` argument of [`fmul_ps_row`] or
129 /// [`scale_c_row`].
130 ///
131 /// # Safety
132 /// Requires the `f` target feature (guaranteed by the module gate).
133 /// `dest` must not hold live C-tile data that must be preserved; if it
134 /// does, spill and restore it around this call (see module doc).
135 #[inline(always)]
136 #[rustfmt::skip] // tabular FBCX.PS dispatch; keep one aligned arm per register
137 pub unsafe fn broadcast_ps(scalar: f32, dest: u8) {
138 let tmp: u64;
139 asm!(
140 "fmv.x.w {tmp}, {x}",
141 x = in(freg) scalar,
142 tmp = out(reg) tmp,
143 options(nostack, preserves_flags),
144 );
145 // FBCX.PS f{dest}, {tmp}: broadcasts bit pattern of tmp into all 8 PS lanes.
146 match dest {
147 0 => asm!(".insn i 0x0b, 3, f0, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
148 1 => asm!(".insn i 0x0b, 3, f1, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
149 2 => asm!(".insn i 0x0b, 3, f2, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
150 3 => asm!(".insn i 0x0b, 3, f3, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
151 4 => asm!(".insn i 0x0b, 3, f4, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
152 5 => asm!(".insn i 0x0b, 3, f5, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
153 6 => asm!(".insn i 0x0b, 3, f6, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
154 7 => asm!(".insn i 0x0b, 3, f7, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
155 8 => asm!(".insn i 0x0b, 3, f8, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
156 9 => asm!(".insn i 0x0b, 3, f9, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
157 10 => asm!(".insn i 0x0b, 3, f10, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
158 11 => asm!(".insn i 0x0b, 3, f11, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
159 12 => asm!(".insn i 0x0b, 3, f12, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
160 13 => asm!(".insn i 0x0b, 3, f13, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
161 14 => asm!(".insn i 0x0b, 3, f14, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
162 15 => asm!(".insn i 0x0b, 3, f15, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
163 16 => asm!(".insn i 0x0b, 3, f16, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
164 17 => asm!(".insn i 0x0b, 3, f17, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
165 18 => asm!(".insn i 0x0b, 3, f18, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
166 19 => asm!(".insn i 0x0b, 3, f19, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
167 20 => asm!(".insn i 0x0b, 3, f20, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
168 21 => asm!(".insn i 0x0b, 3, f21, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
169 22 => asm!(".insn i 0x0b, 3, f22, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
170 23 => asm!(".insn i 0x0b, 3, f23, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
171 24 => asm!(".insn i 0x0b, 3, f24, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
172 25 => asm!(".insn i 0x0b, 3, f25, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
173 26 => asm!(".insn i 0x0b, 3, f26, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
174 27 => asm!(".insn i 0x0b, 3, f27, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
175 28 => asm!(".insn i 0x0b, 3, f28, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
176 29 => asm!(".insn i 0x0b, 3, f29, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
177 30 => asm!(".insn i 0x0b, 3, f30, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
178 31 => asm!(".insn i 0x0b, 3, f31, {t}, 0", t = in(reg) tmp, options(nostack, preserves_flags)),
179 _ => {}
180 }
181 }
182
183 /// Scale the PS register pair for `row` by the pre-broadcast PS register `scratch`.
184 ///
185 /// Issues two `FMUL.PS` instructions (`MATCH_FMUL_PS = 0x1000_007b`;
186 /// R-type, opcode=0x7b, funct7=0x08, funct3=0):
187 ///
188 /// ```text
189 /// f[2*row] = f[2*row] .* f[scratch]
190 /// f[2*row+1] = f[2*row+1] .* f[scratch]
191 /// ```
192 ///
193 /// The caller must ensure `f[scratch]` holds `[alpha; 8]` before calling
194 /// this function -- typically via a preceding [`broadcast_ps`]`(alpha, scratch)`.
195 /// When scaling multiple rows with the same `alpha`, call [`broadcast_ps`]
196 /// once, then call [`fmul_ps_row`] for each row.
197 ///
198 /// # Panics (debug builds)
199 /// Fires a `debug_assert` if `scratch` equals `2*row` or `2*row+1`
200 /// (the scratch register would overwrite the row data being scaled).
201 ///
202 /// # Safety
203 /// Call `tensor_wait(TensorEvent::Fma)` before this function; the tensor
204 /// co-processor must have finished writing the FP register file before
205 /// any PS operations read it.
206 #[inline(always)]
207 #[rustfmt::skip] // tabular row-pair dispatch; keep one aligned arm per row
208 pub unsafe fn fmul_ps_row(row: u32, scratch: u8) {
209 debug_assert!(
210 scratch != 2 * row as u8 && scratch != 2 * row as u8 + 1,
211 "fmul_ps_row: scratch f{} conflicts with C-tile row {} (f{} and f{})",
212 scratch,
213 row,
214 2 * row,
215 2 * row + 1,
216 );
217 match row {
218 0 => scale_row!((0, 1), scratch),
219 1 => scale_row!((2, 3), scratch),
220 2 => scale_row!((4, 5), scratch),
221 3 => scale_row!((6, 7), scratch),
222 4 => scale_row!((8, 9), scratch),
223 5 => scale_row!((10, 11), scratch),
224 6 => scale_row!((12, 13), scratch),
225 7 => scale_row!((14, 15), scratch),
226 8 => scale_row!((16, 17), scratch),
227 9 => scale_row!((18, 19), scratch),
228 10 => scale_row!((20, 21), scratch),
229 11 => scale_row!((22, 23), scratch),
230 12 => scale_row!((24, 25), scratch),
231 13 => scale_row!((26, 27), scratch),
232 14 => scale_row!((28, 29), scratch),
233 15 => scale_row!((30, 31), scratch),
234 _ => {}
235 }
236 }
237
238 /// Broadcast `alpha` into `f[scratch]`, then scale the PS register pair for `row`.
239 ///
240 /// Convenience wrapper: calls [`broadcast_ps`]`(alpha, scratch)` then
241 /// [`fmul_ps_row`]`(row, scratch)`. The broadcast is repeated on every call;
242 /// when scaling multiple rows with the same `alpha`, call [`broadcast_ps`]
243 /// once and [`fmul_ps_row`] for each row instead.
244 ///
245 /// Pass [`PS_SCRATCH_DEFAULT`] (28) for `scratch` when the C tile has at
246 /// most 14 rows. For full 16-row tiles, see the module-level doc for the
247 /// spill/restore pattern.
248 ///
249 /// # Safety
250 /// Call `tensor_wait(TensorEvent::Fma)` before this function; the tensor
251 /// co-processor must have finished writing the FP register file.
252 #[inline(always)]
253 pub unsafe fn scale_c_row(row: u32, alpha: f32, scratch: u8) {
254 broadcast_ps(alpha, scratch);
255 fmul_ps_row(row, scratch);
256 }
257}
258
259// Re-export the inner items when the feature is present.
260#[cfg(target_arch = "riscv64")]
261#[cfg(target_feature = "f")]
262pub use inner::*;