From e83228525a714a077c78830b8ce561d231aaabf3 Mon Sep 17 00:00:00 2001 From: ckristian Date: Fri, 7 Aug 2026 09:26:15 +0200 Subject: [PATCH 1/5] riscv64: RVV 1.0 f32 matmul kernels tract had no RISC-V backend, so rv64gc ran the generic Rust kernels for every matmul. Add an RVV 1.0 f32 mmm tier, detecting V from the AT_HWCAP bit that Linux never sets for the incompatible 0.7.1 draft and VLEN from the vlenb CSR. Because VLEN is a runtime property while MR must be a const generic, each kernel fixes (MR, NR, LMUL) and pins vl to MR, which is correct wherever VLMAX >= MR and short below it, so dispatch is gated on the hart's VLMAX reaching MR and the kernel re-checks the granted vl before running. An assembler probe keeps toolchains predating RVV 1.0 on the generic fallback. --- linalg/Cargo.toml | 4 + linalg/build.rs | 68 ++++++++++ linalg/riscv64/rvv/dispatcher.j2 | 42 +++++++ linalg/riscv64/rvv/dummy_rvv.S | 24 ++++ linalg/riscv64/rvv/rvv_mmm.S.j2 | 206 +++++++++++++++++++++++++++++++ linalg/src/lib.rs | 5 + linalg/src/riscv64.rs | 113 +++++++++++++++++ linalg/src/riscv64/rvv.rs | 95 ++++++++++++++ 8 files changed, 557 insertions(+) create mode 100644 linalg/riscv64/rvv/dispatcher.j2 create mode 100644 linalg/riscv64/rvv/dummy_rvv.S create mode 100644 linalg/riscv64/rvv/rvv_mmm.S.j2 create mode 100644 linalg/src/riscv64.rs create mode 100644 linalg/src/riscv64/rvv.rs diff --git a/linalg/Cargo.toml b/linalg/Cargo.toml index 7dfaf9b1c6..597ec48b50 100644 --- a/linalg/Cargo.toml +++ b/linalg/Cargo.toml @@ -28,6 +28,10 @@ rayon = { workspace = true, optional = true } scan_fmt.workspace = true tract-data.workspace = true +# RVV detection reads the V bit out of AT_HWCAP via getauxval(3). +[target.'cfg(target_arch = "riscv64")'.dependencies] +libc.workspace = true + [build-dependencies] cc.workspace = true half.workspace = true diff --git a/linalg/build.rs b/linalg/build.rs index e769f46731..00860bb4c9 100644 --- a/linalg/build.rs +++ b/linalg/build.rs @@ -53,6 +53,21 @@ fn assembler_supports_dotprod() -> bool { .is_ok() } +// Probe whether the target assembler can encode ratified RVV 1.0. `.option +// arch` arrived in binutils 2.36 and the 1.0 encodings in 2.38; anything older +// either rejects the directive outright or knows only the incompatible 0.7.1 +// draft. When the probe fails we skip the RVV kernels and the `tract_rvv` cfg, +// and dispatch falls back to the generic Rust kernels. +fn assembler_supports_rvv() -> bool { + cc::Build::new() + .file("riscv64/rvv/dummy_rvv.S") + .cargo_metadata(false) + .cargo_warnings(false) + .warnings(false) + .try_compile("tract_rvv_probe") + .is_ok() +} + // Probe whether the target assembler can encode `vpdpbusd ymm` (AVX-512 VNNI // with AVX-512 VL, i.e. the 256-bit form). binutils gained this in ~2.30 // (2018); the Debian stretch toolchain ships 2.28 and rejects the mnemonic. @@ -256,6 +271,8 @@ fn main() { // Set below only when the assembler accepts the `{vex}` prefix on // VPDPBUSD (binutils >= 2.36) -- needed for the AVX-VNNI ymm kernel. println!("cargo:rustc-check-cfg=cfg(tract_avxvnni)"); + // Set below only when the riscv64 assembler probe for RVV 1.0 passes. + println!("cargo:rustc-check-cfg=cfg(tract_rvv)"); match arch.as_ref() { "x86_64" => { @@ -502,10 +519,61 @@ fn main() { config.cc().files(files).compile("arm64fp16") } } + "riscv64" if assembler_supports_rvv() => { + let files = render_rvv_kernels("f32", "4", "+v", RVV_F32_KERNELS, &suffix); + println!("cargo:rustc-cfg=tract_rvv"); + cc::Build::new().files(files).compile("rvv"); + } _ => {} } } +/// `(geometry, MR, NR, LMUL)`. The geometry string ends up in the exported +/// symbol, and the Rust side derives each kernel's dispatch predicate from the +/// same MR and LMUL, so a hart with `VLMAX < MR` never sees it. +/// +/// LMUL is the smallest that both reaches MR on the narrowest hart the kernel +/// targets and leaves room for NR accumulator groups plus one for A. +/// +/// 8x8 m2 VLEN >= 128 universal GEMM tile +/// 16x8 m2 VLEN >= 256 SpacemiT K1 / X100, twice the tile for free +/// 32x1 m8 VLEN >= 128 universal GEMV +/// 64x1 m8 VLEN >= 256 wider GEMV where the registers allow it +const RVV_F32_KERNELS: &[(&str, &str, &str, &str)] = &[ + ("8x8", "8", "8", "2"), + ("16x8", "16", "8", "2"), + ("32x1", "32", "1", "8"), + ("64x1", "64", "1", "8"), +]; + +fn render_rvv_kernels( + dt: &'static str, + esize: &'static str, + arch: &'static str, + kernels: &[(&'static str, &'static str, &'static str, &'static str)], + suffix: &str, +) -> Vec { + let out_dir = path::PathBuf::from(var("OUT_DIR")); + kernels + .iter() + .map(|(geo, mr, nr, lmul)| { + let tmpl = path::Path::new("riscv64/rvv/rvv_mmm.S.j2"); + let out = out_dir.join(format!("rvv_mmm_{dt}_{geo}_{suffix}.S")); + let globals = [ + ("dt", dt), + ("esize", esize), + ("arch", arch), + ("geo", *geo), + ("mr", *mr), + ("nr", *nr), + ("lmul", *lmul), + ]; + preprocess_file(tmpl, &out, &globals, suffix, false); + out + }) + .collect() +} + type Variant = (&'static str, Vec<&'static str>); fn preprocess_files( diff --git a/linalg/riscv64/rvv/dispatcher.j2 b/linalg/riscv64/rvv/dispatcher.j2 new file mode 100644 index 0000000000..e03a5f232a --- /dev/null +++ b/linalg/riscv64/rvv/dispatcher.j2 @@ -0,0 +1,42 @@ +// vim: ft=asm + +// Walks the FusedKerSpec array in a0 and vectors to the handler for each +// entry's discriminant. Layout is #[repr(C, usize)]: an 8-byte discriminant at +// +0 then the payload at +8/+16/+24/+32, 40 bytes per entry. +// +// The table is `j` instructions rather than addresses so it stays PC-relative +// (the kernels are linked into PIE binaries, and a table of absolute addresses +// in .text would need runtime relocations). `.option norvc` keeps the assembler +// from compressing any of them to 2 bytes, which would break the fixed stride +// the index arithmetic below assumes. + +.non_linear: + addi a0, a0, -40 + +.non_linear_loop: + addi a0, a0, 40 + ld t0, 0(a0) + + li t1, {{ jump_table | length }} + bgeu t0, t1, .unsupported + + lla t1, .jmp_table + slli t0, t0, 2 + add t1, t1, t0 + jr t1 + +.option push +.option norvc +.jmp_table: +{% for j in jump_table %} + j .{{j}} +{% endfor %} +.option pop + +.unsupported: + li a0, 1 + ret + +.done: + li a0, 0 + ret diff --git a/linalg/riscv64/rvv/dummy_rvv.S b/linalg/riscv64/rvv/dummy_rvv.S new file mode 100644 index 0000000000..aaf57a4cc6 --- /dev/null +++ b/linalg/riscv64/rvv/dummy_rvv.S @@ -0,0 +1,24 @@ +// Build-time capability probe for the assembler, used by build.rs +// (assembler_supports_rvv). `.option arch` landed in binutils 2.36 and the +// ratified RVV 1.0 encodings in 2.38; toolchains older than that either reject +// the directive or silently know only the incompatible 0.7.1 draft. If this +// file fails to assemble, build.rs skips the RVV kernels and the `tract_rvv` +// cfg, and the runtime falls back to the generic Rust kernels. Not linked into +// anything. +// +// The instruction mix is deliberately the one the kernels actually need: the +// vector-scalar FMA that carries the inner loop, and the strided load/store +// pair the tile store and add_unicast are built on. +.option arch, +v +.text +.globl tract_rvv_probe +tract_rvv_probe: + vsetivli t0, 8, e32, m2, ta, ma + vle32.v v8, (a0) + vfmacc.vf v16, ft0, v8 + vfrsub.vf v16, v16, ft0 + vmfgt.vf v0, v16, ft0 + vmerge.vvm v16, v18, v16, v0 + vlse32.v v4, (a0), a1 + vsse32.v v16, (a0), a1 + ret diff --git a/linalg/riscv64/rvv/rvv_mmm.S.j2 b/linalg/riscv64/rvv/rvv_mmm.S.j2 new file mode 100644 index 0000000000..0e1a54519c --- /dev/null +++ b/linalg/riscv64/rvv/rvv_mmm.S.j2 @@ -0,0 +1,206 @@ +// vim: ft=asm +{% set mr = mr | int %}{% set nr = nr | int %}{% set lmul = lmul | int %}{% set esize = esize | int %} +{% set ew = esize * 8 %} +{% set fl = "flw" if esize == 4 else "flh" %} +{% set fzero = "fmv.w.x" if esize == 4 else "fmv.h.x" %} +{% set va = 8 %} +{% set vs = 24 if lmul == 8 else 4 %} + +// {{mr}}x{{nr}} {{dt}} matmul tile, RVV 1.0, SEW={{ew}}, LMUL={{lmul}}. +// +// v{{va}} packed A column, {{mr}} {{dt}} +// v{{vs}} scratch: per-row operand, unicast, leaky relu +// v0 leaky relu mask +{% for j in range(nr) %} +// v{{ 16 + j * lmul }} accumulator, C tile column {{j}} +{% endfor %} +// +// Accumulators start at v16 so the LMUL=8 case stays 8-aligned. Only +// caller-saved registers are touched, so there is no frame. +// +// `vl` is pinned to {{mr}} for the whole call, which is what makes this kernel +// VLEN-specific: vsetvli grants min(AVL, VLMAX), so a hart with VLMAX < {{mr}} +// would compute a short tile rather than fail. Dispatch is gated on +// vlmax_{{dt}}({{lmul}}) >= {{mr}}; the check below is the backstop for the +// two disagreeing, and cannot fire on a hart with no vector unit at all -- +// there the guard instruction is itself illegal. + +.option arch, {{arch}} +.text +.align 2 + +.global {{G}}rvv_mmm_{{dt}}_{{geo}}_{{suffix}} +{{G}}rvv_mmm_{{dt}}_{{geo}}_{{suffix}}: + + li t0, {{mr}} + vsetvli t1, t0, e{{ew}}, m{{lmul}}, ta, ma + bne t1, t0, .unsupported + +{% include "dispatcher.j2" %} + +// Packed A is column-major panels of {{mr}}, so a K step reads one contiguous +// column; packed B is row-major panels of {{nr}}, so it reads {{nr}} scalars. +// vfmacc.vf sources its scalar straight from an f register, so no broadcast is +// needed. +.add_mat_mul: + ld t1, 8(a0) // k + ld a1, 16(a0) // pa + ld a2, 24(a0) // pb + + beqz t1, .non_linear_loop + +.packed_packed_loop_1: + vle{{ew}}.v v{{va}}, (a1) + addi a1, a1, {{ mr * esize }} +{% for j in range(nr) %} + {{fl}} ft{{j}}, {{ j * esize }}(a2) +{% endfor %} + addi a2, a2, {{ nr * esize }} +{% for j in range(nr) %} + vfmacc.vf v{{ 16 + j * lmul }}, ft{{j}}, v{{va}} +{% endfor %} + addi t1, t1, -1 + bnez t1, .packed_packed_loop_1 + + j .non_linear_loop + +.clear: +{% for j in range(nr) %} + vmv.v.i v{{ 16 + j * lmul }}, 0 +{% endfor %} + j .non_linear_loop + +.load_tile: + ld t1, 8(a0) +{% for j in range(nr) %} + vle{{ew}}.v v{{ 16 + j * lmul }}, (t1) +{% if not loop.last %} + addi t1, t1, {{ mr * esize }} +{% endif %} +{% endfor %} + j .non_linear_loop + +{# ScalarSub is `scalar - acc` and ScalarSubF is `acc - scalar`, so the plain + form maps to the reversed vector op. #} +{% for label, op in [ + ("scalar_min", "vfmin.vf"), + ("scalar_max", "vfmax.vf"), + ("scalar_mul", "vfmul.vf"), + ("scalar_add", "vfadd.vf"), + ("scalar_sub", "vfrsub.vf"), + ("scalar_sub_flipped", "vfsub.vf"), +] %} +.{{label}}: + {{fl}} ft0, 8(a0) +{% for j in range(nr) %} + {{op}} v{{ 16 + j * lmul }}, v{{ 16 + j * lmul }}, ft0 +{% endfor %} + j .non_linear_loop +{% endfor %} + +.leaky_relu: + {{fl}} ft0, 8(a0) + {{fzero}} ft1, zero +{% for j in range(nr) %} + vfmul.vf v{{vs}}, v{{ 16 + j * lmul }}, ft0 + vmfgt.vf v0, v{{ 16 + j * lmul }}, ft1 + vmerge.vvm v{{ 16 + j * lmul }}, v{{vs}}, v{{ 16 + j * lmul }}, v0 +{% endfor %} + j .non_linear_loop + +{# Per-row operand at +8: {{mr}} values, one per row, so one lane each and a + single unit-stride load serves every accumulator. #} +{% for label, op, flipped in [ + ("per_row_min", "vfmin.vv", false), + ("per_row_max", "vfmax.vv", false), + ("per_row_mul", "vfmul.vv", false), + ("per_row_add", "vfadd.vv", false), + ("per_row_sub", "vfsub.vv", false), + ("per_row_sub_flipped", "vfsub.vv", true), +] %} +.{{label}}: + ld t1, 8(a0) + vle{{ew}}.v v{{vs}}, (t1) +{% for j in range(nr) %} +{% if flipped %} + {{op}} v{{ 16 + j * lmul }}, v{{ 16 + j * lmul }}, v{{vs}} +{% else %} + {{op}} v{{ 16 + j * lmul }}, v{{vs}}, v{{ 16 + j * lmul }} +{% endif %} +{% endfor %} + j .non_linear_loop +{% endfor %} + +{# Per-col operand at +8: {{nr}} values, one per accumulator. #} +{% for label, op in [ + ("per_col_min", "vfmin.vf"), + ("per_col_max", "vfmax.vf"), + ("per_col_mul", "vfmul.vf"), + ("per_col_add", "vfadd.vf"), + ("per_col_sub", "vfrsub.vf"), + ("per_col_sub_flipped", "vfsub.vf"), +] %} +.{{label}}: + ld t1, 8(a0) +{% for j in range(nr) %} + {{fl}} ft0, {{ j * esize }}(t1) + {{op}} v{{ 16 + j * lmul }}, v{{ 16 + j * lmul }}, ft0 +{% endfor %} + j .non_linear_loop +{% endfor %} + +.add_row_col_products: + ld t1, 8(a0) // rows, {{mr}} values + ld t2, 16(a0) // cols, {{nr}} values + vle{{ew}}.v v{{vs}}, (t1) +{% for j in range(nr) %} + {{fl}} ft0, {{ j * esize }}(t2) + vfmacc.vf v{{ 16 + j * lmul }}, ft0, v{{vs}} +{% endfor %} + j .non_linear_loop + +{# OutputStoreKer at +8: ptr, row_byte_stride, col_byte_stride, item_size. + C[i][j] lives at ptr + i*row + j*col, and accumulator j holds column j one + row per lane, so a column is exactly one strided access -- no lane-by-lane + peeling for the non-contiguous case, unlike the NEON and AVX ports. #} +.add_unicast: + ld t1, 8(a0) + ld t2, 16(a0) + ld t3, 24(a0) +{% for j in range(nr) %} + vlse{{ew}}.v v{{vs}}, (t1), t2 + vfadd.vv v{{ 16 + j * lmul }}, v{{ 16 + j * lmul }}, v{{vs}} +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.store: + ld t1, 8(a0) + ld t2, 16(a0) + ld t3, 24(a0) + + li t4, {{esize}} + bne t2, t4, .store_strided +{% for j in range(nr) %} + vse{{ew}}.v v{{ 16 + j * lmul }}, (t1) +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.store_strided: +{% for j in range(nr) %} + vsse{{ew}}.v v{{ 16 + j * lmul }}, (t1), t2 +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.q_scale: +.q_shl: +.q_shr: + j .unsupported diff --git a/linalg/src/lib.rs b/linalg/src/lib.rs index feba41da18..ca4b34e295 100644 --- a/linalg/src/lib.rs +++ b/linalg/src/lib.rs @@ -49,6 +49,9 @@ pub fn has_fp16() -> bool { #[cfg(any(target_arch = "arm", target_arch = "armv7", target_arch = "arm"))] pub mod arm32; +#[cfg(target_arch = "riscv64")] +pub mod riscv64; + #[cfg(all(target_family = "wasm", target_feature = "simd128"))] pub mod wasm; @@ -268,6 +271,8 @@ pub fn best() -> Ops { arm32::plug(&mut ops); #[cfg(target_arch = "aarch64")] arm64::plug(&mut ops); + #[cfg(target_arch = "riscv64")] + riscv64::plug(&mut ops); #[cfg(all(target_family = "wasm", target_feature = "simd128"))] wasm::plug(&mut ops); diff --git a/linalg/src/riscv64.rs b/linalg/src/riscv64.rs new file mode 100644 index 0000000000..7be9ea2b74 --- /dev/null +++ b/linalg/src/riscv64.rs @@ -0,0 +1,113 @@ +//! RISC-V (rv64) backend for the ratified Vector extension, RVV 1.0. +//! +//! Kernels are assembly rendered from jinja, as on x86_64 and arm64, because +//! Rust exposes no stable RVV intrinsics and `-C target-feature=+v` is itself +//! unstable. +//! +//! RVV is vector-length agnostic -- `VLEN` is a runtime property of the hart -- +//! while `MR` and `NR` must be const generics, since they select the packing +//! format. Kernels reconcile the two by fixing `(MR, NR, LMUL)` and pinning +//! `vl` to `MR`. `vsetvli` clamps to `VLMAX`, so such a kernel is correct +//! wherever `VLMAX >= MR`, merely leaves lanes idle when `VLMAX > MR`, and +//! computes a short tile when `VLMAX < MR`. Every kernel is therefore gated on +//! [`vlmax_f32`] reaching its `MR`. + +use crate::Ops; + +// `tract_rvv` is set by build.rs only when the assembler could encode RVV 1.0; +// without it the kernel symbols do not exist and dispatch stays generic. +#[cfg(tract_rvv)] +mod rvv; +#[cfg(tract_rvv)] +pub use rvv::*; + +/// `AT_HWCAP` -- see `getauxval(3)`. +const AT_HWCAP: libc::c_ulong = 16; + +/// Bit `'V' - 'A'` of the single-letter extension bitmap Linux puts in +/// `AT_HWCAP` (`COMPAT_HWCAP_ISA_V`). +/// +/// This bit alone is a sufficient RVV 1.0 gate: Linux sets it only for the +/// ratified extension, never for the incompatible 0.7.1 draft implemented by +/// the Allwinner D1 and Sophgo SG2042. +const HWCAP_ISA_V: libc::c_ulong = 1 << (b'V' - b'A'); + +fn hwcap() -> libc::c_ulong { + // SAFETY: getauxval is thread-safe and takes a scalar; it returns 0 for an + // unknown type, which reads here as "no vector unit". + unsafe { libc::getauxval(AT_HWCAP) } +} + +/// Reads `vlenb`, the read-only CSR holding `VLEN / 8`. +/// +/// Callers must establish [`has_rvv`] first: without vector state the CSR read +/// raises an illegal instruction, which arrives as SIGILL and cannot be +/// recovered from. +fn read_vlenb() -> usize { + let vlenb: usize; + // SAFETY: guarded by has_rvv(). `csrr` on a read-only CSR has no side + // effects. The CSR is named numerically so the assembler does not need to + // know about vector extensions to encode it. + unsafe { + std::arch::asm!("csrr {out}, 0xC22", out = out(reg) vlenb, options(nomem, nostack, preserves_flags)); + } + vlenb +} + +lazy_static::lazy_static! { + static ref HAS_RVV: bool = hwcap() & HWCAP_ISA_V != 0; + + static ref VLENB: usize = if *HAS_RVV { read_vlenb() } else { 0 }; + +} + +/// Whether the hart implements the ratified RVV 1.0 vector extension. +pub fn has_rvv() -> bool { + *HAS_RVV +} + +/// Vector register width in bytes (`VLEN / 8`); 0 without RVV. +pub fn vlenb() -> usize { + *VLENB +} + +/// `VLMAX = LMUL * VLEN / SEW` for 32-bit elements: the largest `vl` this hart +/// can grant. A kernel of tile height `MR` is dispatchable only where this +/// reaches `MR`. +pub fn vlmax_f32(lmul: usize) -> usize { + vlenb() * lmul / std::mem::size_of::() +} + +pub fn plug(_ops: &mut Ops) { + if has_rvv() { + #[cfg(tract_rvv)] + rvv::plug(_ops); + } +} + +#[cfg(test)] +mod test { + use super::*; + + /// Detection must be internally consistent and must not trap. Written to + /// pass on a hart without V as well, so it stays meaningful under + /// qemu-riscv64 with and without `v=true`. + #[test] + fn detection_is_coherent() { + eprintln!( + "rvv={} VLEN={} vlmax_f32(lmul=1,2,4)={:?}", + has_rvv(), + vlenb() * 8, + [vlmax_f32(1), vlmax_f32(2), vlmax_f32(4)], + ); + if has_rvv() { + let vlenb = vlenb(); + assert!(vlenb >= 16, "RVV 1.0 mandates VLEN >= 128, got VLEN={}", vlenb * 8); + assert!(vlenb.is_power_of_two(), "VLEN must be a power of two, got {}", vlenb * 8); + assert_eq!(vlmax_f32(1), vlenb / 4); + assert_eq!(vlmax_f32(4), vlenb); + } else { + assert_eq!(vlenb(), 0); + } + } +} diff --git a/linalg/src/riscv64/rvv.rs b/linalg/src/riscv64/rvv.rs new file mode 100644 index 0000000000..df3ea17f9e --- /dev/null +++ b/linalg/src/riscv64/rvv.rs @@ -0,0 +1,95 @@ +//! RVV 1.0 matmul kernels. +//! +//! Each kernel pins `vl` to its own `MR`, so it is correct only where +//! `VLMAX >= MR`. The `where(...)` predicates are that constraint, and they +//! read the same `(MR, LMUL)` the assembly was rendered from -- see +//! `RVV_F32_KERNELS` in build.rs, which must stay in step with them. +//! +//! `vlmax_f32` returns 0 without RVV, so the predicates also subsume the +//! `has_rvv` check. + +use crate::Ops; +use crate::frame::mmm::ImplementationQuality::ManuallyOptimized; + +use super::vlmax_f32; + +const VLMAX_M2_GE_8: fn() -> bool = || vlmax_f32(2) >= 8; +const VLMAX_M2_GE_16: fn() -> bool = || vlmax_f32(2) >= 16; +const VLMAX_M8_GE_32: fn() -> bool = || vlmax_f32(8) >= 32; +const VLMAX_M8_GE_64: fn() -> bool = || vlmax_f32(8) >= 64; + +MMMExternKernel!(rvv_mmm_f32_8x8 ( 8, 8)@(16, 16) where(VLMAX_M2_GE_8) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_16x8 (16, 8)@(16, 16) where(VLMAX_M2_GE_16) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_32x1 (32, 1)@(16, 16) where(VLMAX_M8_GE_32) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_64x1 (64, 1)@(16, 16) where(VLMAX_M8_GE_64) quality(ManuallyOptimized)); + +/// `(name, MR, LMUL)` mirroring the build.rs kernel table. +#[cfg(test)] +const GEOMETRIES: &[(&str, usize, usize)] = + &[("8x8", 8, 2), ("16x8", 16, 2), ("32x1", 32, 8), ("64x1", 64, 8)]; + +pub fn plug(ops: &mut Ops) { + ops.mmm_impls.extend_from_slice(&[ + rvv_mmm_f32_8x8.mmm(), + rvv_mmm_f32_16x8.mmm(), + rvv_mmm_f32_32x1.mmm(), + rvv_mmm_f32_64x1.mmm(), + ]); +} + +#[cfg(test)] +mod test { + use super::*; + use crate::frame::mmm::{FusedKerSpec, MatMatMulKer}; + + fn supported() -> [bool; 4] { + [ + rvv_mmm_f32_8x8.is_supported_here(), + rvv_mmm_f32_16x8.is_supported_here(), + rvv_mmm_f32_32x1.is_supported_here(), + rvv_mmm_f32_64x1.is_supported_here(), + ] + } + + /// The generated kernel suites early-return on an unsupported kernel and + /// count as passes, so a green run says nothing about whether the VLEN + /// predicates are right. This asserts the dispatch set directly. + /// + /// The permissive direction is the dangerous one: a kernel whose MR exceeds + /// VLMAX computes a short tile, and short is not the same as failing. + #[test] + fn dispatch_matches_vlen() { + let vlenb = super::super::vlenb(); + for ((name, mr, lmul), got) in GEOMETRIES.iter().zip(supported()) { + let want = vlenb * lmul / std::mem::size_of::() >= *mr; + eprintln!("VLEN={} {name}: {got} (want {want})", vlenb * 8); + assert_eq!(got, want, "{name} dispatch disagrees with VLEN={}", vlenb * 8); + } + } + + /// The `vsetvli` guard heading every kernel backstops the predicates above. + /// Vacuous on a hart wide enough for all of them, hence no assertion on + /// finding a candidate. + /// + /// Meaningful only where V is present: the guard is itself a vector + /// instruction, so it covers "unit too narrow for this tile" but not "no + /// unit", where calling at all is a SIGILL rather than a return code. + #[test] + fn oversized_tile_refuses_to_run() { + if !super::super::has_rvv() { + return; + } + let runners: [&dyn Fn() -> isize; 4] = [ + &|| rvv_mmm_f32_8x8.kernel(&[FusedKerSpec::Done]), + &|| rvv_mmm_f32_16x8.kernel(&[FusedKerSpec::Done]), + &|| rvv_mmm_f32_32x1.kernel(&[FusedKerSpec::Done]), + &|| rvv_mmm_f32_64x1.kernel(&[FusedKerSpec::Done]), + ]; + for (((name, ..), ok), run) in GEOMETRIES.iter().zip(supported()).zip(runners) { + if !ok { + assert_eq!(run(), 1, "{name} ran on a hart whose VLMAX is below its MR"); + eprintln!("{name}: correctly refused"); + } + } + } +} From c61dde7926540908135dfee5ab11ce7e7c479108 Mon Sep 17 00:00:00 2001 From: ckristian Date: Fri, 7 Aug 2026 11:54:54 +0200 Subject: [PATCH 2/5] ci: cross-test riscv64 under qemu The RVV kernels are gated on the hart's vector length, so a single emulated width would leave half the kernel set untested. Add riscv64gc to the qemu cross-test platforms twice, at VLEN 256 and 128, which select disjoint halves. -cpu max rather than a profile model because the generic rv64 model cannot run Debian's riscv64 glibc at all. --- .github/workflows/cross-platform.yml | 2 ++ .travis/cross.sh | 28 +++++++++++++++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cross-platform.yml b/.github/workflows/cross-platform.yml index cb6a754a15..66d36c8df3 100644 --- a/.github/workflows/cross-platform.yml +++ b/.github/workflows/cross-platform.yml @@ -71,6 +71,8 @@ jobs: - aarch64-unknown-linux-musl - cortexa53-unknown-linux-musl - armv7-unknown-linux-musl + - riscv64gc-unknown-linux-gnu + - rvv128-unknown-linux-gnu - aarch64-linux-android - armv7-linux-androideabi - i686-linux-android diff --git a/.travis/cross.sh b/.travis/cross.sh index f69a2cbe93..ef8dc1f478 100755 --- a/.travis/cross.sh +++ b/.travis/cross.sh @@ -138,7 +138,8 @@ case "$PLATFORM" in ;; "aarch64-unknown-linux-gnu" | "armv6vfp-unknown-linux-gnueabihf" | "armv7-unknown-linux-gnueabihf" | \ - "aarch64-unknown-linux-musl" | "armv7-unknown-linux-musl" | "cortexa53-unknown-linux-musl" ) + "aarch64-unknown-linux-musl" | "armv7-unknown-linux-musl" | "cortexa53-unknown-linux-musl" | \ + "riscv64gc-unknown-linux-gnu" | "rvv128-unknown-linux-gnu" ) ensure_cargo_dinghy case "$PLATFORM" in @@ -202,6 +203,31 @@ case "$PLATFORM" in [ -d "$CUSTOM_TC" ] || curl -s https://tract-test-assets.tract.rs/toolchains/armv7l-linux-musleabihf-cross.tgz | tar zx export TARGET_CFLAGS="-mfpu=neon" ;; + # RVV is vector-length agnostic and the mmm kernels are gated on the + # hart's VLEN, so the two entries below differ only in vlen: 256 is + # the SpacemiT K1/X100 shape, 128 the Sophgo SG2044 one, and they + # select disjoint halves of the kernel set. + # + # -cpu max rather than a profile model: rva23u64 would describe real + # silicon more closely but predates neither the CI image's qemu nor + # its glibc safely, and the generic rv64 model cannot run Debian's + # riscv64 glibc at all (it SIGILLs on a trivial static binary). + "riscv64gc-unknown-linux-gnu") + export ARCH=riscv64 + export QEMU_ARCH=riscv64 + export LIBC_ARCH=riscv64 + export QEMU_OPTS="-cpu max,vlen=256" + export RUSTC_TRIPLE=riscv64gc-unknown-linux-gnu + export DEBIAN_TRIPLE=riscv64-linux-gnu + ;; + "rvv128-unknown-linux-gnu") + export ARCH=riscv64 + export QEMU_ARCH=riscv64 + export LIBC_ARCH=riscv64 + export QEMU_OPTS="-cpu max,vlen=128" + export RUSTC_TRIPLE=riscv64gc-unknown-linux-gnu + export DEBIAN_TRIPLE=riscv64-linux-gnu + ;; *) echo "unsupported platform $PLATFORM" exit 1 From 80a800f5aaf51ba0cfb48c037ad31e6309d90c68 Mon Sep 17 00:00:00 2001 From: ckristian Date: Fri, 7 Aug 2026 09:29:52 +0200 Subject: [PATCH 3/5] riscv64: f16 RVV matmul kernels behind Zvfh The riscv64 tier covered f32 only, so f16 matmul fell back to the generic kernels even on parts with native half-precision vectors, which includes the SpacemiT X60 in the K1. Add an f16 mmm tier from the same template at SEW=16, where VLMAX doubles and so does every tile height for a given LMUL and VLEN. Zvfh is read from the /proc/cpuinfo isa line, since RVA23 mandates only Zvfhmin and that cannot hold an f16 accumulator, and a second assembler probe keeps toolchains predating Zvfh on the f32 tier alone. --- linalg/build.rs | 37 ++++++++- linalg/riscv64/rvv/dummy_rvv_zvfh.S | 17 ++++ linalg/src/riscv64.rs | 40 +++++++++- linalg/src/riscv64/rvv.rs | 119 +++++++++++++++++++++------- 4 files changed, 183 insertions(+), 30 deletions(-) create mode 100644 linalg/riscv64/rvv/dummy_rvv_zvfh.S diff --git a/linalg/build.rs b/linalg/build.rs index 00860bb4c9..e4632bfbbd 100644 --- a/linalg/build.rs +++ b/linalg/build.rs @@ -68,6 +68,20 @@ fn assembler_supports_rvv() -> bool { .is_ok() } +// Probe whether the target assembler can encode Zvfh (f16 vector arithmetic). +// Zvfh reached binutils later than base RVV 1.0, so a toolchain can assemble +// the f32 kernels and still reject these. When the probe fails we skip the f16 +// kernels and the `tract_rvv_zvfh` cfg, and f16 matmul stays generic. +fn assembler_supports_zvfh() -> bool { + cc::Build::new() + .file("riscv64/rvv/dummy_rvv_zvfh.S") + .cargo_metadata(false) + .cargo_warnings(false) + .warnings(false) + .try_compile("tract_rvv_zvfh_probe") + .is_ok() +} + // Probe whether the target assembler can encode `vpdpbusd ymm` (AVX-512 VNNI // with AVX-512 VL, i.e. the 256-bit form). binutils gained this in ~2.30 // (2018); the Debian stretch toolchain ships 2.28 and rejects the mnemonic. @@ -273,6 +287,8 @@ fn main() { println!("cargo:rustc-check-cfg=cfg(tract_avxvnni)"); // Set below only when the riscv64 assembler probe for RVV 1.0 passes. println!("cargo:rustc-check-cfg=cfg(tract_rvv)"); + // Set below only when the riscv64 assembler probe for Zvfh also passes. + println!("cargo:rustc-check-cfg=cfg(tract_rvv_zvfh)"); match arch.as_ref() { "x86_64" => { @@ -520,8 +536,18 @@ fn main() { } } "riscv64" if assembler_supports_rvv() => { - let files = render_rvv_kernels("f32", "4", "+v", RVV_F32_KERNELS, &suffix); + let mut files = render_rvv_kernels("f32", "4", "+v", RVV_F32_KERNELS, &suffix); println!("cargo:rustc-cfg=tract_rvv"); + if assembler_supports_zvfh() { + files.extend(render_rvv_kernels( + "f16", + "2", + "+v, +zvfh, +zfhmin", + RVV_F16_KERNELS, + &suffix, + )); + println!("cargo:rustc-cfg=tract_rvv_zvfh"); + } cc::Build::new().files(files).compile("rvv"); } _ => {} @@ -546,6 +572,15 @@ const RVV_F32_KERNELS: &[(&str, &str, &str, &str)] = &[ ("64x1", "64", "1", "8"), ]; +/// As [`RVV_F32_KERNELS`]; SEW=16 doubles VLMAX, so every tile is twice as +/// tall for the same LMUL and VLEN. +const RVV_F16_KERNELS: &[(&str, &str, &str, &str)] = &[ + ("16x8", "16", "8", "2"), + ("32x8", "32", "8", "2"), + ("64x1", "64", "1", "8"), + ("128x1", "128", "1", "8"), +]; + fn render_rvv_kernels( dt: &'static str, esize: &'static str, diff --git a/linalg/riscv64/rvv/dummy_rvv_zvfh.S b/linalg/riscv64/rvv/dummy_rvv_zvfh.S new file mode 100644 index 0000000000..db4af4a46c --- /dev/null +++ b/linalg/riscv64/rvv/dummy_rvv_zvfh.S @@ -0,0 +1,17 @@ +// Build-time capability probe for the assembler, used by build.rs +// (assembler_supports_zvfh). Zvfh reached binutils later than base RVV 1.0, so +// a toolchain can assemble the f32 kernels and still reject these. When the +// probe fails we skip the f16 kernels and the `tract_rvv_zvfh` cfg, and f16 +// matmul falls back to the generic Rust kernels. Not linked into anything. +.option arch, +v, +zvfh, +zfhmin +.text +.globl tract_rvv_zvfh_probe +tract_rvv_zvfh_probe: + vsetivli t0, 16, e16, m2, ta, ma + flh ft0, 0(a0) + fmv.h.x ft1, zero + vle16.v v8, (a0) + vfmacc.vf v16, ft0, v8 + vmfgt.vf v0, v16, ft1 + vsse16.v v16, (a0), a1 + ret diff --git a/linalg/src/riscv64.rs b/linalg/src/riscv64.rs index 7be9ea2b74..5cc6fae89f 100644 --- a/linalg/src/riscv64.rs +++ b/linalg/src/riscv64.rs @@ -10,7 +10,7 @@ //! `vl` to `MR`. `vsetvli` clamps to `VLMAX`, so such a kernel is correct //! wherever `VLMAX >= MR`, merely leaves lanes idle when `VLMAX > MR`, and //! computes a short tile when `VLMAX < MR`. Every kernel is therefore gated on -//! [`vlmax_f32`] reaching its `MR`. +//! [`vlmax_f32`] or [`vlmax_f16`] reaching its `MR`. use crate::Ops; @@ -54,11 +54,35 @@ fn read_vlenb() -> usize { vlenb } +/// Splits the kernel-canonicalised `isa` line of /proc/cpuinfo, e.g. +/// `rv64imafdcv_zicsr_zvfh_zvl256b`, into lowercase extension tokens. +/// +/// This is the only source for multi-letter extensions; `AT_HWCAP` carries +/// bits for the single-letter ones alone. +fn isa_extensions() -> Vec { + #[cfg(test)] + crate::setup_test_logger(); + let Ok(cpu_info) = std::fs::read_to_string("/proc/cpuinfo") else { + log::warn!("Could not read /proc/cpuinfo. CPU feature detection may be impaired."); + return vec![]; + }; + let Some(line) = cpu_info.lines().find(|line| line.trim_start().starts_with("isa")) else { + log::warn!("No \"isa :\" line in /proc/cpuinfo. CPU feature detection may be impaired."); + return vec![]; + }; + let Some((_, isa)) = line.split_once(':') else { return vec![] }; + isa.trim().split('_').map(|s| s.to_lowercase()).collect() +} + lazy_static::lazy_static! { static ref HAS_RVV: bool = hwcap() & HWCAP_ISA_V != 0; static ref VLENB: usize = if *HAS_RVV { read_vlenb() } else { 0 }; + /// Zvfh, not Zvfhmin: the latter offers only f16<->f32 conversion and so + /// cannot carry an f16 accumulator. RVA23 mandates Zvfhmin and leaves Zvfh + /// optional, so the profile is not enough to infer it. + static ref HAS_ZVFH: bool = *HAS_RVV && isa_extensions().iter().any(|e| e == "zvfh"); } /// Whether the hart implements the ratified RVV 1.0 vector extension. @@ -66,6 +90,11 @@ pub fn has_rvv() -> bool { *HAS_RVV } +/// Whether the hart implements Zvfh (native f16 vector arithmetic). +pub fn has_zvfh() -> bool { + *HAS_ZVFH +} + /// Vector register width in bytes (`VLEN / 8`); 0 without RVV. pub fn vlenb() -> usize { *VLENB @@ -78,6 +107,11 @@ pub fn vlmax_f32(lmul: usize) -> usize { vlenb() * lmul / std::mem::size_of::() } +/// As [`vlmax_f32`], for 16-bit elements. +pub fn vlmax_f16(lmul: usize) -> usize { + vlenb() * lmul / std::mem::size_of::() +} + pub fn plug(_ops: &mut Ops) { if has_rvv() { #[cfg(tract_rvv)] @@ -95,9 +129,10 @@ mod test { #[test] fn detection_is_coherent() { eprintln!( - "rvv={} VLEN={} vlmax_f32(lmul=1,2,4)={:?}", + "rvv={} VLEN={} zvfh={} vlmax_f32(lmul=1,2,4)={:?}", has_rvv(), vlenb() * 8, + has_zvfh(), [vlmax_f32(1), vlmax_f32(2), vlmax_f32(4)], ); if has_rvv() { @@ -108,6 +143,7 @@ mod test { assert_eq!(vlmax_f32(4), vlenb); } else { assert_eq!(vlenb(), 0); + assert!(!has_zvfh(), "Zvfh cannot be present without V"); } } } diff --git a/linalg/src/riscv64/rvv.rs b/linalg/src/riscv64/rvv.rs index df3ea17f9e..6ca1a135a9 100644 --- a/linalg/src/riscv64/rvv.rs +++ b/linalg/src/riscv64/rvv.rs @@ -3,30 +3,62 @@ //! Each kernel pins `vl` to its own `MR`, so it is correct only where //! `VLMAX >= MR`. The `where(...)` predicates are that constraint, and they //! read the same `(MR, LMUL)` the assembly was rendered from -- see -//! `RVV_F32_KERNELS` in build.rs, which must stay in step with them. +//! `RVV_F32_KERNELS` and `RVV_F16_KERNELS` in build.rs, which must stay in +//! step with them. //! -//! `vlmax_f32` returns 0 without RVV, so the predicates also subsume the -//! `has_rvv` check. +//! `vlmax_f32` and `vlmax_f16` return 0 without RVV, so the predicates also +//! subsume the `has_rvv` check. use crate::Ops; use crate::frame::mmm::ImplementationQuality::ManuallyOptimized; -use super::vlmax_f32; +use super::{vlmax_f16, vlmax_f32}; -const VLMAX_M2_GE_8: fn() -> bool = || vlmax_f32(2) >= 8; -const VLMAX_M2_GE_16: fn() -> bool = || vlmax_f32(2) >= 16; -const VLMAX_M8_GE_32: fn() -> bool = || vlmax_f32(8) >= 32; -const VLMAX_M8_GE_64: fn() -> bool = || vlmax_f32(8) >= 64; +const VLMAX_F32_M2_GE_8: fn() -> bool = || vlmax_f32(2) >= 8; +const VLMAX_F32_M2_GE_16: fn() -> bool = || vlmax_f32(2) >= 16; +const VLMAX_F32_M8_GE_32: fn() -> bool = || vlmax_f32(8) >= 32; +const VLMAX_F32_M8_GE_64: fn() -> bool = || vlmax_f32(8) >= 64; -MMMExternKernel!(rvv_mmm_f32_8x8 ( 8, 8)@(16, 16) where(VLMAX_M2_GE_8) quality(ManuallyOptimized)); -MMMExternKernel!(rvv_mmm_f32_16x8 (16, 8)@(16, 16) where(VLMAX_M2_GE_16) quality(ManuallyOptimized)); -MMMExternKernel!(rvv_mmm_f32_32x1 (32, 1)@(16, 16) where(VLMAX_M8_GE_32) quality(ManuallyOptimized)); -MMMExternKernel!(rvv_mmm_f32_64x1 (64, 1)@(16, 16) where(VLMAX_M8_GE_64) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_8x8 ( 8, 8)@(16, 16) where(VLMAX_F32_M2_GE_8) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_16x8 (16, 8)@(16, 16) where(VLMAX_F32_M2_GE_16) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_32x1 (32, 1)@(16, 16) where(VLMAX_F32_M8_GE_32) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_f32_64x1 (64, 1)@(16, 16) where(VLMAX_F32_M8_GE_64) quality(ManuallyOptimized)); -/// `(name, MR, LMUL)` mirroring the build.rs kernel table. +#[cfg(tract_rvv_zvfh)] +mod zvfh { + use super::*; + use crate::f16; + + /// f16 arithmetic needs Zvfh on top of a wide enough vector unit; the + /// profile does not imply it, and Zvfhmin alone cannot hold an f16 + /// accumulator. + const VLMAX_F16_M2_GE_16: fn() -> bool = || super::super::has_zvfh() && vlmax_f16(2) >= 16; + const VLMAX_F16_M2_GE_32: fn() -> bool = || super::super::has_zvfh() && vlmax_f16(2) >= 32; + const VLMAX_F16_M8_GE_64: fn() -> bool = || super::super::has_zvfh() && vlmax_f16(8) >= 64; + const VLMAX_F16_M8_GE_128: fn() -> bool = || super::super::has_zvfh() && vlmax_f16(8) >= 128; + + MMMExternKernel!(rvv_mmm_f16_16x8 ( 16, 8)@(16, 16) where(VLMAX_F16_M2_GE_16) quality(ManuallyOptimized)); + MMMExternKernel!(rvv_mmm_f16_32x8 ( 32, 8)@(16, 16) where(VLMAX_F16_M2_GE_32) quality(ManuallyOptimized)); + MMMExternKernel!(rvv_mmm_f16_64x1 ( 64, 1)@(16, 16) where(VLMAX_F16_M8_GE_64) quality(ManuallyOptimized)); + MMMExternKernel!(rvv_mmm_f16_128x1(128, 1)@(16, 16) where(VLMAX_F16_M8_GE_128) quality(ManuallyOptimized)); +} + +/// `(name, MR, LMUL, element size)` mirroring the build.rs kernel tables. #[cfg(test)] -const GEOMETRIES: &[(&str, usize, usize)] = - &[("8x8", 8, 2), ("16x8", 16, 2), ("32x1", 32, 8), ("64x1", 64, 8)]; +const GEOMETRIES: &[(&str, usize, usize, usize)] = &[ + ("f32 8x8", 8, 2, 4), + ("f32 16x8", 16, 2, 4), + ("f32 32x1", 32, 8, 4), + ("f32 64x1", 64, 8, 4), + #[cfg(tract_rvv_zvfh)] + ("f16 16x8", 16, 2, 2), + #[cfg(tract_rvv_zvfh)] + ("f16 32x8", 32, 2, 2), + #[cfg(tract_rvv_zvfh)] + ("f16 64x1", 64, 8, 2), + #[cfg(tract_rvv_zvfh)] + ("f16 128x1", 128, 8, 2), +]; pub fn plug(ops: &mut Ops) { ops.mmm_impls.extend_from_slice(&[ @@ -35,6 +67,13 @@ pub fn plug(ops: &mut Ops) { rvv_mmm_f32_32x1.mmm(), rvv_mmm_f32_64x1.mmm(), ]); + #[cfg(tract_rvv_zvfh)] + ops.mmm_impls.extend_from_slice(&[ + zvfh::rvv_mmm_f16_16x8.mmm(), + zvfh::rvv_mmm_f16_32x8.mmm(), + zvfh::rvv_mmm_f16_64x1.mmm(), + zvfh::rvv_mmm_f16_128x1.mmm(), + ]); } #[cfg(test)] @@ -42,13 +81,22 @@ mod test { use super::*; use crate::frame::mmm::{FusedKerSpec, MatMatMulKer}; - fn supported() -> [bool; 4] { - [ + fn supported() -> Vec { + #[allow(unused_mut)] + let mut v = vec![ rvv_mmm_f32_8x8.is_supported_here(), rvv_mmm_f32_16x8.is_supported_here(), rvv_mmm_f32_32x1.is_supported_here(), rvv_mmm_f32_64x1.is_supported_here(), - ] + ]; + #[cfg(tract_rvv_zvfh)] + v.extend([ + zvfh::rvv_mmm_f16_16x8.is_supported_here(), + zvfh::rvv_mmm_f16_32x8.is_supported_here(), + zvfh::rvv_mmm_f16_64x1.is_supported_here(), + zvfh::rvv_mmm_f16_128x1.is_supported_here(), + ]); + v } /// The generated kernel suites early-return on an unsupported kernel and @@ -60,10 +108,17 @@ mod test { #[test] fn dispatch_matches_vlen() { let vlenb = super::super::vlenb(); - for ((name, mr, lmul), got) in GEOMETRIES.iter().zip(supported()) { - let want = vlenb * lmul / std::mem::size_of::() >= *mr; - eprintln!("VLEN={} {name}: {got} (want {want})", vlenb * 8); - assert_eq!(got, want, "{name} dispatch disagrees with VLEN={}", vlenb * 8); + for ((name, mr, lmul, esize), got) in GEOMETRIES.iter().zip(supported()) { + let mut want = vlenb * lmul / esize >= *mr; + if *esize == 2 { + want &= super::super::has_zvfh(); + } + eprintln!( + "VLEN={} zvfh={} {name}: {got} (want {want})", + vlenb * 8, + super::super::has_zvfh() + ); + assert_eq!(got, want, "{name} dispatch disagrees with this hart"); } } @@ -79,12 +134,22 @@ mod test { if !super::super::has_rvv() { return; } - let runners: [&dyn Fn() -> isize; 4] = [ - &|| rvv_mmm_f32_8x8.kernel(&[FusedKerSpec::Done]), - &|| rvv_mmm_f32_16x8.kernel(&[FusedKerSpec::Done]), - &|| rvv_mmm_f32_32x1.kernel(&[FusedKerSpec::Done]), - &|| rvv_mmm_f32_64x1.kernel(&[FusedKerSpec::Done]), + #[allow(unused_mut)] + let mut runners: Vec isize>> = vec![ + Box::new(|| rvv_mmm_f32_8x8.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_f32_16x8.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_f32_32x1.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_f32_64x1.kernel(&[FusedKerSpec::Done])), ]; + #[cfg(tract_rvv_zvfh)] + if super::super::has_zvfh() { + runners.extend:: isize>>>(vec![ + Box::new(|| zvfh::rvv_mmm_f16_16x8.kernel(&[FusedKerSpec::Done])), + Box::new(|| zvfh::rvv_mmm_f16_32x8.kernel(&[FusedKerSpec::Done])), + Box::new(|| zvfh::rvv_mmm_f16_64x1.kernel(&[FusedKerSpec::Done])), + Box::new(|| zvfh::rvv_mmm_f16_128x1.kernel(&[FusedKerSpec::Done])), + ]); + } for (((name, ..), ok), run) in GEOMETRIES.iter().zip(supported()).zip(runners) { if !ok { assert_eq!(run(), 1, "{name} ran on a hart whose VLMAX is below its MR"); From 0b3bce0fab82587a9186616d896abeb0061e9e65 Mon Sep 17 00:00:00 2001 From: ckristian Date: Fri, 7 Aug 2026 10:17:05 +0200 Subject: [PATCH 4/5] riscv64: i8 and i32 RVV matmul kernels The riscv64 tier had no integer kernels, so quantised models fell back to the generic ones for every matmul. Add an i32 accumulator tier handling both packings the frame offers: i32 x i32 at e32, and i8 x i8 through a loop at e16 where vle8.v picks EEW=8 off the instruction and vwmacc.vx widens straight into the e32 accumulators, which keeps VLMAX identical between the two so one vl serves both. QScale, RoundingShiftRight and ShiftLeft follow the sign-magnitude reference in generic/rounding.rs, widening to e64 for the multiply; the negative-value mask taken before the magnitude overwrites it also serves the MinusInf and PlusInf nudges, which differ only where the result is zero anyway. --- linalg/build.rs | 30 +++- linalg/riscv64/rvv/rvv_mmm_i32.S.j2 | 255 ++++++++++++++++++++++++++++ linalg/riscv64/rvv/rvv_mmm_i32_q.j2 | 145 ++++++++++++++++ linalg/src/riscv64/rvv.rs | 48 +++++- 4 files changed, 475 insertions(+), 3 deletions(-) create mode 100644 linalg/riscv64/rvv/rvv_mmm_i32.S.j2 create mode 100644 linalg/riscv64/rvv/rvv_mmm_i32_q.j2 diff --git a/linalg/build.rs b/linalg/build.rs index e4632bfbbd..ac3c92d41b 100644 --- a/linalg/build.rs +++ b/linalg/build.rs @@ -536,10 +536,20 @@ fn main() { } } "riscv64" if assembler_supports_rvv() => { - let mut files = render_rvv_kernels("f32", "4", "+v", RVV_F32_KERNELS, &suffix); + const F32: &str = "riscv64/rvv/rvv_mmm.S.j2"; + let mut files = render_rvv_kernels(F32, "f32", "4", "+v", RVV_F32_KERNELS, &suffix); + files.extend(render_rvv_kernels( + "riscv64/rvv/rvv_mmm_i32.S.j2", + "i32", + "4", + "+v", + RVV_I32_KERNELS, + &suffix, + )); println!("cargo:rustc-cfg=tract_rvv"); if assembler_supports_zvfh() { files.extend(render_rvv_kernels( + F32, "f16", "2", "+v, +zvfh, +zfhmin", @@ -574,6 +584,21 @@ const RVV_F32_KERNELS: &[(&str, &str, &str, &str)] = &[ /// As [`RVV_F32_KERNELS`]; SEW=16 doubles VLMAX, so every tile is twice as /// tall for the same LMUL and VLEN. +/// As [`RVV_F32_KERNELS`], for the i32 accumulator tier. LMUL here is the one +/// the i8 inner loop runs at; the accumulators, and therefore the dispatch +/// predicate, sit at twice it. +/// +/// 8x8 m1 VLEN >= 128 +/// 16x8 m1 VLEN >= 256 +/// 16x1 m2 VLEN >= 128 +/// 32x1 m2 VLEN >= 256 +const RVV_I32_KERNELS: &[(&str, &str, &str, &str)] = &[ + ("8x8", "8", "8", "1"), + ("16x8", "16", "8", "1"), + ("16x1", "16", "1", "2"), + ("32x1", "32", "1", "2"), +]; + const RVV_F16_KERNELS: &[(&str, &str, &str, &str)] = &[ ("16x8", "16", "8", "2"), ("32x8", "32", "8", "2"), @@ -582,6 +607,7 @@ const RVV_F16_KERNELS: &[(&str, &str, &str, &str)] = &[ ]; fn render_rvv_kernels( + tmpl: &str, dt: &'static str, esize: &'static str, arch: &'static str, @@ -592,7 +618,7 @@ fn render_rvv_kernels( kernels .iter() .map(|(geo, mr, nr, lmul)| { - let tmpl = path::Path::new("riscv64/rvv/rvv_mmm.S.j2"); + let tmpl = path::Path::new(tmpl); let out = out_dir.join(format!("rvv_mmm_{dt}_{geo}_{suffix}.S")); let globals = [ ("dt", dt), diff --git a/linalg/riscv64/rvv/rvv_mmm_i32.S.j2 b/linalg/riscv64/rvv/rvv_mmm_i32.S.j2 new file mode 100644 index 0000000000..7e639489be --- /dev/null +++ b/linalg/riscv64/rvv/rvv_mmm_i32.S.j2 @@ -0,0 +1,255 @@ +// vim: ft=asm +{% set mr = mr | int %}{% set nr = nr | int %}{% set lmul = lmul | int %} +{% set al = lmul * 2 %} +{% set vs = 4 %} +{% set va = 8 %} +{% set q0 = 8 %} +{% set q1 = 12 if al == 2 else 24 %} + +// {{mr}}x{{nr}} i32 matmul tile, RVV 1.0. +// +// v{{vs}} scratch: A for the i32 packing, per-row operand, unicast +// v{{va}} packed A column widened to i16, i8 packing only +// v{{q0}}/v{{q1}} i64 scratch for the quantised ops +// v0 mask +{% for j in range(nr) %} +// v{{ 16 + j * al }} accumulator, C tile column {{j}} +{% endfor %} +// +// vtype is e32/m{{al}} everywhere except inside the i8 inner loop and the +// quantised ops, which set and restore it themselves. The i8 loop runs at +// e16/m{{lmul}} so that `vle8.v` picks up EEW=8 from the instruction and +// `vwmacc.vx` widens straight into the e32 accumulators; that also makes +// VLMAX identical to the e32 one, so a single `vl` of {{mr}} serves both. +// +// The tile is correct only where VLMAX >= {{mr}}, which is vlmax_f32({{al}}); +// the check below backstops a disagreement with the dispatch predicate. + +.option arch, +v +.text +.align 2 + +.global {{G}}rvv_mmm_i32_{{geo}}_{{suffix}} +{{G}}rvv_mmm_i32_{{geo}}_{{suffix}}: + + li t0, {{mr}} + vsetvli t1, t0, e32, m{{al}}, ta, ma + bne t1, t0, .unsupported + +{% include "dispatcher.j2" %} + +// AddMatMul { k: +8, pa: +16, pb: +24, packing: +32 }. Packing 0 is i32 x i32, +// packing 1 is i8 x i8 accumulating into i32. +.add_mat_mul: + ld t1, 8(a0) + ld a1, 16(a0) + ld a2, 24(a0) + ld t2, 32(a0) + + beqz t1, .non_linear_loop + + li t3, 1 + beq t2, t3, .packed_packed_loop_1_i8i8 + +.packed_packed_loop_1: + vle32.v v{{vs}}, (a1) + addi a1, a1, {{ mr * 4 }} +{% for j in range(nr) %} + lw t3, {{ j * 4 }}(a2) + vmacc.vx v{{ 16 + j * al }}, t3, v{{vs}} +{% endfor %} + addi a2, a2, {{ nr * 4 }} + addi t1, t1, -1 + bnez t1, .packed_packed_loop_1 + + j .non_linear_loop + +.packed_packed_loop_1_i8i8: + li t0, {{mr}} + vsetvli t0, t0, e16, m{{lmul}}, ta, ma + +.packed_packed_loop_1_i8i8_inner: + vle8.v v{{vs}}, (a1) + addi a1, a1, {{mr}} + vsext.vf2 v{{va}}, v{{vs}} +{% for j in range(nr) %} + lb t3, {{j}}(a2) + vwmacc.vx v{{ 16 + j * al }}, t3, v{{va}} +{% endfor %} + addi a2, a2, {{nr}} + addi t1, t1, -1 + bnez t1, .packed_packed_loop_1_i8i8_inner + + li t0, {{mr}} + vsetvli t0, t0, e32, m{{al}}, ta, ma + j .non_linear_loop + +.clear: +{% for j in range(nr) %} + vmv.v.i v{{ 16 + j * al }}, 0 +{% endfor %} + j .non_linear_loop + +.load_tile: + ld t1, 8(a0) +{% for j in range(nr) %} + vle32.v v{{ 16 + j * al }}, (t1) +{% if not loop.last %} + addi t1, t1, {{ mr * 4 }} +{% endif %} +{% endfor %} + j .non_linear_loop + +{# ScalarSub is `scalar - acc` and ScalarSubF is `acc - scalar`, so the plain + form maps to the reversed vector op. #} +{% for label, op in [ + ("scalar_min", "vmin.vx"), + ("scalar_max", "vmax.vx"), + ("scalar_mul", "vmul.vx"), + ("scalar_add", "vadd.vx"), + ("scalar_sub", "vrsub.vx"), + ("scalar_sub_flipped", "vsub.vx"), +] %} +.{{label}}: + lw t1, 8(a0) +{% for j in range(nr) %} + {{op}} v{{ 16 + j * al }}, v{{ 16 + j * al }}, t1 +{% endfor %} + j .non_linear_loop +{% endfor %} + +.leaky_relu: + lw t1, 8(a0) +{% for j in range(nr) %} + vmul.vx v{{vs}}, v{{ 16 + j * al }}, t1 + vmsgt.vx v0, v{{ 16 + j * al }}, x0 + vmerge.vvm v{{ 16 + j * al }}, v{{vs}}, v{{ 16 + j * al }}, v0 +{% endfor %} + j .non_linear_loop + +{# Per-row operand at +8: {{mr}} values, one per row, so one lane each. #} +{% for label, op, flipped in [ + ("per_row_min", "vmin.vv", false), + ("per_row_max", "vmax.vv", false), + ("per_row_mul", "vmul.vv", false), + ("per_row_add", "vadd.vv", false), + ("per_row_sub", "vsub.vv", false), + ("per_row_sub_flipped", "vsub.vv", true), +] %} +.{{label}}: + ld t1, 8(a0) + vle32.v v{{vs}}, (t1) +{% for j in range(nr) %} +{% if flipped %} + {{op}} v{{ 16 + j * al }}, v{{ 16 + j * al }}, v{{vs}} +{% else %} + {{op}} v{{ 16 + j * al }}, v{{vs}}, v{{ 16 + j * al }} +{% endif %} +{% endfor %} + j .non_linear_loop +{% endfor %} + +{# Per-col operand at +8: {{nr}} values, one per accumulator. #} +{% for label, op in [ + ("per_col_min", "vmin.vx"), + ("per_col_max", "vmax.vx"), + ("per_col_mul", "vmul.vx"), + ("per_col_add", "vadd.vx"), + ("per_col_sub", "vrsub.vx"), + ("per_col_sub_flipped", "vsub.vx"), +] %} +.{{label}}: + ld t1, 8(a0) +{% for j in range(nr) %} + lw t2, {{ j * 4 }}(t1) + {{op}} v{{ 16 + j * al }}, v{{ 16 + j * al }}, t2 +{% endfor %} + j .non_linear_loop +{% endfor %} + +.add_row_col_products: + ld t1, 8(a0) + ld t2, 16(a0) + vle32.v v{{vs}}, (t1) +{% for j in range(nr) %} + lw t3, {{ j * 4 }}(t2) + vmacc.vx v{{ 16 + j * al }}, t3, v{{vs}} +{% endfor %} + j .non_linear_loop + +{# OutputStoreKer at +8: ptr, row_byte_stride, col_byte_stride, item_size. + Accumulator j holds column j one row per lane, so a column is exactly one + strided access whatever the strides. The i8 store narrows twice, e32 to e16 + to e8, since RVV narrows by one power of two at a time. #} +.add_unicast: + ld t1, 8(a0) + ld t2, 16(a0) + ld t3, 24(a0) + ld t5, 32(a0) + + li t4, 4 + bne t5, t4, .add_unicast_i8 +{% for j in range(nr) %} + vlse32.v v{{vs}}, (t1), t2 + vadd.vv v{{ 16 + j * al }}, v{{ 16 + j * al }}, v{{vs}} +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.add_unicast_i8: +{% for j in range(nr) %} + vlse8.v v{{vs}}, (t1), t2 + vsext.vf4 v{{va}}, v{{vs}} + vadd.vv v{{ 16 + j * al }}, v{{ 16 + j * al }}, v{{va}} +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.store: + ld t1, 8(a0) + ld t2, 16(a0) + ld t3, 24(a0) + ld t5, 32(a0) + + li t4, 4 + bne t5, t4, .store_i8 + + bne t2, t4, .store_strided +{% for j in range(nr) %} + vse32.v v{{ 16 + j * al }}, (t1) +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.store_strided: +{% for j in range(nr) %} + vsse32.v v{{ 16 + j * al }}, (t1), t2 +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + j .non_linear_loop + +.store_i8: + li t0, {{mr}} +{% for j in range(nr) %} + vsetvli t4, t0, e16, m{{lmul}}, ta, ma + vnsra.wi v{{vs}}, v{{ 16 + j * al }}, 0 + vsetvli t4, t0, e8, {{ "mf2" if lmul == 1 else "m1" }}, ta, ma + vnsra.wi v{{va}}, v{{vs}}, 0 + vsse8.v v{{va}}, (t1), t2 +{% if not loop.last %} + add t1, t1, t3 +{% endif %} +{% endfor %} + li t0, {{mr}} + vsetvli t4, t0, e32, m{{al}}, ta, ma + j .non_linear_loop + +{% include "rvv_mmm_i32_q.j2" %} diff --git a/linalg/riscv64/rvv/rvv_mmm_i32_q.j2 b/linalg/riscv64/rvv/rvv_mmm_i32_q.j2 new file mode 100644 index 0000000000..f05a926913 --- /dev/null +++ b/linalg/riscv64/rvv/rvv_mmm_i32_q.j2 @@ -0,0 +1,145 @@ +// vim: ft=asm +{# + Quantised fused ops for the i32 tier, shared by every geometry. + + All three follow the reference in generic/rounding.rs, which rounds in + sign-magnitude rather than by arithmetic shift: + + half = 1 << (shift - 1) + r = (|v| + half + nudge) >> shift + out = signum(v) * r + + with nudge selecting the tie behaviour. signum needs no special case at + v == 0: half + nudge never exceeds 1 << (shift - 1), which shifts to zero + for every shift >= 1, so all six policies already yield 0 there. + + MinusInf and PlusInf nudge on the sign of the *original* value, which is + gone once |v| overwrites it -- but the negative mask in v0 was taken before + that, and PlusInf's v <= 0 differs from v < 0 only at v == 0, where the + result is 0 whatever the nudge. So v0 alone serves both. + + Callers set t1 = MR, t2 = shift, t3 = half, and leave vtype at the width + the arithmetic runs in. +#} + +{% macro nudge(policy, acc, q) %} +{% if policy == "zero" %} + vadd.vi v{{acc}}, v{{acc}}, -1 +{% elif policy == "away" %} +{% elif policy == "minus_inf" %} + vmv.v.i v{{q}}, -1 + vmerge.vxm v{{q}}, v{{q}}, x0, v0 + vadd.vv v{{acc}}, v{{acc}}, v{{q}} +{% elif policy == "plus_inf" %} + li t4, -1 + vmv.v.i v{{q}}, 0 + vmerge.vxm v{{q}}, v{{q}}, t4, v0 + vadd.vv v{{acc}}, v{{acc}}, v{{q}} +{% elif policy == "even" %} + vsra.vx v{{q}}, v{{acc}}, t2 + vand.vi v{{q}}, v{{q}}, 1 + vadd.vi v{{q}}, v{{q}}, -1 + vadd.vv v{{acc}}, v{{acc}}, v{{q}} +{% elif policy == "odd" %} + vsra.vx v{{q}}, v{{acc}}, t2 + vand.vi v{{q}}, v{{q}}, 1 + vsub.vv v{{acc}}, v{{acc}}, v{{q}} +{% endif %} +{% endmacro %} + +{# Rounds one register group in place at the current vtype width, using a + single scratch group. The mask is taken before |v| overwrites the sign. #} +{% macro round_in_place(acc, policy, q) %} + vmslt.vx v0, v{{acc}}, x0 + vneg.v v{{q}}, v{{acc}} + vmerge.vvm v{{acc}}, v{{acc}}, v{{q}}, v0 + {{ nudge(policy, acc, q) }} + vadd.vx v{{acc}}, v{{acc}}, t3 + vsra.vx v{{acc}}, v{{acc}}, t2 + vneg.v v{{q}}, v{{acc}} + vmerge.vvm v{{acc}}, v{{acc}}, v{{q}}, v0 +{% endmacro %} + +{# RoundingPolicy has Native=0 first, and Native is the one value no kernel + implements, so it falls through to .unsupported. #} +{% macro dispatch(prefix) %} + li t4, 1 + beq t5, t4, .{{prefix}}_zero + li t4, 2 + beq t5, t4, .{{prefix}}_away + li t4, 3 + beq t5, t4, .{{prefix}}_minus_inf + li t4, 4 + beq t5, t4, .{{prefix}}_plus_inf + li t4, 5 + beq t5, t4, .{{prefix}}_even + li t4, 6 + beq t5, t4, .{{prefix}}_odd + j .unsupported +{% endmacro %} + +// RoundingShiftRight(shift: +8, policy: +16). i32 throughout, so the +// accumulators are rounded where they sit. +.q_shr: + ld t2, 8(a0) + ld t5, 16(a0) + addi t4, t2, -1 + li t3, 1 + sll t3, t3, t4 + {{ dispatch("q_shr") }} + +{% for policy in ["zero", "away", "minus_inf", "plus_inf", "even", "odd"] %} +.q_shr_{{policy}}: +{% for j in range(nr) %} + {{ round_in_place(16 + j * al, policy, q1) }} +{% endfor %} + j .non_linear_loop +{% endfor %} + +.q_shl: + ld t2, 8(a0) +{% for j in range(nr) %} + vsll.vx v{{ 16 + j * al }}, v{{ 16 + j * al }}, t2 +{% endfor %} + j .non_linear_loop + +// QScale(shift: +8, policy: +16, mult: +24). The product is i64, so this +// widens to e64, rounds by shift + 31 there and truncates back, matching the +// reference which computes in i64 and casts. A non-positive effective shift +// means a plain left shift instead. +.q_scale: + li t1, {{mr}} + ld t2, 8(a0) + ld t5, 16(a0) + lw t6, 24(a0) + addi t2, t2, 31 + blez t2, .q_scale_shl + addi t4, t2, -1 + li t3, 1 + sll t3, t3, t4 + {{ dispatch("q_scale") }} + +{% for policy in ["zero", "away", "minus_inf", "plus_inf", "even", "odd"] %} +.q_scale_{{policy}}: +{% for j in range(nr) %} + vsetvli t0, t1, e64, m{{ al * 2 }}, ta, ma + vsext.vf2 v{{q0}}, v{{ 16 + j * al }} + vmul.vx v{{q0}}, v{{q0}}, t6 + {{ round_in_place(q0, policy, q1) }} + vsetvli t0, t1, e32, m{{al}}, ta, ma + vnsra.wi v{{ 16 + j * al }}, v{{q0}}, 0 +{% endfor %} + j .non_linear_loop +{% endfor %} + +.q_scale_shl: + neg t2, t2 +{% for j in range(nr) %} + vsetvli t0, t1, e64, m{{ al * 2 }}, ta, ma + vsext.vf2 v{{q0}}, v{{ 16 + j * al }} + vmul.vx v{{q0}}, v{{q0}}, t6 + vsll.vx v{{q0}}, v{{q0}}, t2 + vsetvli t0, t1, e32, m{{al}}, ta, ma + vnsra.wi v{{ 16 + j * al }}, v{{q0}}, 0 +{% endfor %} + j .non_linear_loop diff --git a/linalg/src/riscv64/rvv.rs b/linalg/src/riscv64/rvv.rs index 6ca1a135a9..77fb9adc56 100644 --- a/linalg/src/riscv64/rvv.rs +++ b/linalg/src/riscv64/rvv.rs @@ -11,11 +11,14 @@ use crate::Ops; use crate::frame::mmm::ImplementationQuality::ManuallyOptimized; +use crate::pack::PackedFormat; use super::{vlmax_f16, vlmax_f32}; const VLMAX_F32_M2_GE_8: fn() -> bool = || vlmax_f32(2) >= 8; const VLMAX_F32_M2_GE_16: fn() -> bool = || vlmax_f32(2) >= 16; +const VLMAX_F32_M4_GE_16: fn() -> bool = || vlmax_f32(4) >= 16; +const VLMAX_F32_M4_GE_32: fn() -> bool = || vlmax_f32(4) >= 32; const VLMAX_F32_M8_GE_32: fn() -> bool = || vlmax_f32(8) >= 32; const VLMAX_F32_M8_GE_64: fn() -> bool = || vlmax_f32(8) >= 64; @@ -24,6 +27,31 @@ MMMExternKernel!(rvv_mmm_f32_16x8 (16, 8)@(16, 16) where(VLMAX_F32_M2_GE_16 MMMExternKernel!(rvv_mmm_f32_32x1 (32, 1)@(16, 16) where(VLMAX_F32_M8_GE_32) quality(ManuallyOptimized)); MMMExternKernel!(rvv_mmm_f32_64x1 (64, 1)@(16, 16) where(VLMAX_F32_M8_GE_64) quality(ManuallyOptimized)); +MMMExternKernel!(rvv_mmm_i32_8x8(8, 8)@(16, 16) + where(VLMAX_F32_M2_GE_8) + packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 8, 16), PackedFormat::new(DatumType::I8, 8, 16)); + quality(ManuallyOptimized) + store(i8) +); +MMMExternKernel!(rvv_mmm_i32_16x8(16, 8)@(16, 16) + where(VLMAX_F32_M2_GE_16) + packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 16, 16), PackedFormat::new(DatumType::I8, 8, 16)); + quality(ManuallyOptimized) + store(i8) +); +MMMExternKernel!(rvv_mmm_i32_16x1(16, 1)@(16, 1) + where(VLMAX_F32_M4_GE_16) + packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 16, 16), PackedFormat::new(DatumType::I8, 1, 1)); + quality(ManuallyOptimized) + store(i8) +); +MMMExternKernel!(rvv_mmm_i32_32x1(32, 1)@(16, 1) + where(VLMAX_F32_M4_GE_32) + packing[1] = i8i8 => |k| k.with_packing(PackedFormat::new(DatumType::I8, 32, 16), PackedFormat::new(DatumType::I8, 1, 1)); + quality(ManuallyOptimized) + store(i8) +); + #[cfg(tract_rvv_zvfh)] mod zvfh { use super::*; @@ -43,13 +71,19 @@ mod zvfh { MMMExternKernel!(rvv_mmm_f16_128x1(128, 1)@(16, 16) where(VLMAX_F16_M8_GE_128) quality(ManuallyOptimized)); } -/// `(name, MR, LMUL, element size)` mirroring the build.rs kernel tables. +/// `(name, MR, LMUL, element size)` mirroring the build.rs kernel tables. The +/// i32 entries carry the accumulator LMUL, twice the one their table lists, +/// because that is what their dispatch predicate is written against. #[cfg(test)] const GEOMETRIES: &[(&str, usize, usize, usize)] = &[ ("f32 8x8", 8, 2, 4), ("f32 16x8", 16, 2, 4), ("f32 32x1", 32, 8, 4), ("f32 64x1", 64, 8, 4), + ("i32 8x8", 8, 2, 4), + ("i32 16x8", 16, 2, 4), + ("i32 16x1", 16, 4, 4), + ("i32 32x1", 32, 4, 4), #[cfg(tract_rvv_zvfh)] ("f16 16x8", 16, 2, 2), #[cfg(tract_rvv_zvfh)] @@ -66,6 +100,10 @@ pub fn plug(ops: &mut Ops) { rvv_mmm_f32_16x8.mmm(), rvv_mmm_f32_32x1.mmm(), rvv_mmm_f32_64x1.mmm(), + rvv_mmm_i32_8x8.mmm(), + rvv_mmm_i32_16x8.mmm(), + rvv_mmm_i32_16x1.mmm(), + rvv_mmm_i32_32x1.mmm(), ]); #[cfg(tract_rvv_zvfh)] ops.mmm_impls.extend_from_slice(&[ @@ -88,6 +126,10 @@ mod test { rvv_mmm_f32_16x8.is_supported_here(), rvv_mmm_f32_32x1.is_supported_here(), rvv_mmm_f32_64x1.is_supported_here(), + rvv_mmm_i32_8x8.is_supported_here(), + rvv_mmm_i32_16x8.is_supported_here(), + rvv_mmm_i32_16x1.is_supported_here(), + rvv_mmm_i32_32x1.is_supported_here(), ]; #[cfg(tract_rvv_zvfh)] v.extend([ @@ -140,6 +182,10 @@ mod test { Box::new(|| rvv_mmm_f32_16x8.kernel(&[FusedKerSpec::Done])), Box::new(|| rvv_mmm_f32_32x1.kernel(&[FusedKerSpec::Done])), Box::new(|| rvv_mmm_f32_64x1.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_i32_8x8.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_i32_16x8.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_i32_16x1.kernel(&[FusedKerSpec::Done])), + Box::new(|| rvv_mmm_i32_32x1.kernel(&[FusedKerSpec::Done])), ]; #[cfg(tract_rvv_zvfh)] if super::super::has_zvfh() { From 64de462dfe70fa62622ad3008334ae10933405da Mon Sep 17 00:00:00 2001 From: ckristian Date: Fri, 7 Aug 2026 10:49:34 +0200 Subject: [PATCH 5/5] riscv64: RVV element-wise and reduction kernels Every element-wise binary op and every f32 reduction still ran the generic kernels on riscv64. Add RVV versions of by-scalar and unicast mul, add, sub, subf, min and max, plus the max, min and sum reductions. These strip-mine on vsetvli rather than fixing a tile, so unlike the matmul kernels they are vector-length agnostic and need no VLEN predicate, and being Rust asm! blocks rather than .S files they need only rustc's own assembler, so they sit outside the tract_rvv cfg that records whether an external one could encode RVV. --- linalg/src/lib.rs | 4 ++ linalg/src/riscv64.rs | 80 +++++++++++++++++------ linalg/src/riscv64/by_scalar.rs | 70 ++++++++++++++++++++ linalg/src/riscv64/reduce.rs | 109 ++++++++++++++++++++++++++++++++ linalg/src/riscv64/unicast.rs | 72 +++++++++++++++++++++ 5 files changed, 316 insertions(+), 19 deletions(-) create mode 100644 linalg/src/riscv64/by_scalar.rs create mode 100644 linalg/src/riscv64/reduce.rs create mode 100644 linalg/src/riscv64/unicast.rs diff --git a/linalg/src/lib.rs b/linalg/src/lib.rs index ca4b34e295..55fbfdaa2c 100644 --- a/linalg/src/lib.rs +++ b/linalg/src/lib.rs @@ -310,12 +310,16 @@ fn register_all_unicast(registry: &mut LinalgRegistry) { generic::register_all_unicast(registry); #[cfg(target_arch = "aarch64")] arm64::register_all_unicast(registry); + #[cfg(target_arch = "riscv64")] + riscv64::register_all_unicast(registry); } fn register_all_by_scalar(registry: &mut LinalgRegistry) { generic::register_all_by_scalar(registry); #[cfg(target_arch = "aarch64")] arm64::register_all_by_scalar(registry); + #[cfg(target_arch = "riscv64")] + riscv64::register_all_by_scalar(registry); } pub type LinalgFn = dyn Fn(&mut TensorView, &TensorView) -> TractResult<()> + Send + Sync; diff --git a/linalg/src/riscv64.rs b/linalg/src/riscv64.rs index 5cc6fae89f..6f498a868f 100644 --- a/linalg/src/riscv64.rs +++ b/linalg/src/riscv64.rs @@ -1,25 +1,38 @@ //! RISC-V (rv64) backend for the ratified Vector extension, RVV 1.0. //! -//! Kernels are assembly rendered from jinja, as on x86_64 and arm64, because -//! Rust exposes no stable RVV intrinsics and `-C target-feature=+v` is itself -//! unstable. +//! Everything here is assembly, because Rust exposes no stable RVV intrinsics +//! and `-C target-feature=+v` is itself unstable. //! -//! RVV is vector-length agnostic -- `VLEN` is a runtime property of the hart -- -//! while `MR` and `NR` must be const generics, since they select the packing -//! format. Kernels reconcile the two by fixing `(MR, NR, LMUL)` and pinning -//! `vl` to `MR`. `vsetvli` clamps to `VLMAX`, so such a kernel is correct -//! wherever `VLMAX >= MR`, merely leaves lanes idle when `VLMAX > MR`, and -//! computes a short tile when `VLMAX < MR`. Every kernel is therefore gated on -//! [`vlmax_f32`] or [`vlmax_f16`] reaching its `MR`. - -use crate::Ops; - -// `tract_rvv` is set by build.rs only when the assembler could encode RVV 1.0; -// without it the kernel symbols do not exist and dispatch stays generic. +//! The element-wise and reduction kernels strip-mine on `vsetvli`, so they are +//! vector-length agnostic and run on any hart with V. The matmul kernels +//! cannot be: `MR` and `NR` are const generics because they select the packing +//! format, while `VLEN` is only known at run time. Those kernels fix +//! `(MR, NR, LMUL)` and pin `vl` to `MR`, which `vsetvli` clamps to `VLMAX` -- +//! correct wherever `VLMAX >= MR`, idle lanes above it, and a short tile below +//! it. Each is therefore gated on [`vlmax_f32`] or [`vlmax_f16`] reaching its +//! `MR`. + +use crate::frame::by_scalar::ByScalarKer; +use crate::frame::element_wise::ElementWiseKer; +use crate::frame::reduce::ReduceKer; +use crate::frame::unicast::UnicastKer; +use crate::{BinOp, DatumType, LinalgRegistry, Ops}; + +// The element-wise and reduction kernels are Rust `asm!` blocks, so they need +// only rustc's own assembler and are always compiled. The matmul kernels are +// `.S` files, and exist only when build.rs found an external assembler able to +// encode RVV 1.0 -- which is what `tract_rvv` records. +mod by_scalar; +mod reduce; #[cfg(tract_rvv)] mod rvv; +mod unicast; + +pub use by_scalar::*; +pub use reduce::*; #[cfg(tract_rvv)] pub use rvv::*; +pub use unicast::*; /// `AT_HWCAP` -- see `getauxval(3)`. const AT_HWCAP: libc::c_ulong = 16; @@ -112,11 +125,40 @@ pub fn vlmax_f16(lmul: usize) -> usize { vlenb() * lmul / std::mem::size_of::() } -pub fn plug(_ops: &mut Ops) { - if has_rvv() { - #[cfg(tract_rvv)] - rvv::plug(_ops); +pub fn plug(ops: &mut Ops) { + if !has_rvv() { + return; + } + ops.mul_by_scalar_f32 = Box::new(|| rvv_mul_by_scalar_f32::ew()); + ops.max_f32 = Box::new(|| rvv_max_f32::red()); + ops.min_f32 = Box::new(|| rvv_min_f32::red()); + ops.sum_f32 = Box::new(|| rvv_sum_f32::red()); + #[cfg(tract_rvv)] + rvv::plug(ops); +} + +pub(crate) fn register_all_by_scalar(registry: &mut LinalgRegistry) { + if !has_rvv() { + return; + } + registry.insert((BinOp::Mul, DatumType::F32), Box::new(|| rvv_mul_by_scalar_f32::bin())); + registry.insert((BinOp::Add, DatumType::F32), Box::new(|| rvv_add_by_scalar_f32::bin())); + registry.insert((BinOp::Sub, DatumType::F32), Box::new(|| rvv_sub_by_scalar_f32::bin())); + registry.insert((BinOp::SubF, DatumType::F32), Box::new(|| rvv_subf_by_scalar_f32::bin())); + registry.insert((BinOp::Min, DatumType::F32), Box::new(|| rvv_min_by_scalar_f32::bin())); + registry.insert((BinOp::Max, DatumType::F32), Box::new(|| rvv_max_by_scalar_f32::bin())); +} + +pub(crate) fn register_all_unicast(registry: &mut LinalgRegistry) { + if !has_rvv() { + return; } + registry.insert((BinOp::Mul, DatumType::F32), Box::new(|| rvv_unicast_mul_f32::bin())); + registry.insert((BinOp::Add, DatumType::F32), Box::new(|| rvv_unicast_add_f32::bin())); + registry.insert((BinOp::Sub, DatumType::F32), Box::new(|| rvv_unicast_sub_f32::bin())); + registry.insert((BinOp::SubF, DatumType::F32), Box::new(|| rvv_unicast_subf_f32::bin())); + registry.insert((BinOp::Min, DatumType::F32), Box::new(|| rvv_unicast_min_f32::bin())); + registry.insert((BinOp::Max, DatumType::F32), Box::new(|| rvv_unicast_max_f32::bin())); } #[cfg(test)] diff --git a/linalg/src/riscv64/by_scalar.rs b/linalg/src/riscv64/by_scalar.rs new file mode 100644 index 0000000000..0633c22a10 --- /dev/null +++ b/linalg/src/riscv64/by_scalar.rs @@ -0,0 +1,70 @@ +//! Element-wise `buf op scalar` over f32, RVV 1.0. +//! +//! The loop is strip-mined on `vsetvli`, so it is vector-length agnostic and +//! carries no VLEN predicate, unlike the matmul kernels: `vl` is whatever the +//! hart grants and the tail falls out of the loop condition. `nr` is therefore +//! only the frame's chunking granularity, not a tile width. + +macro_rules! rvv_by_scalar { + ($func: ident, $op: expr) => { + by_scalar_impl_wrap!( + f32, + $func, + 4, + 4, + f32, + #[inline(never)] + fn run(buf: &mut [f32], s: f32) { + assert!(!buf.is_empty()); + let len = buf.len(); + let ptr = buf.as_mut_ptr(); + // SAFETY: `len` elements are in bounds from `ptr` by + // construction, and the loop advances both together. + unsafe { + std::arch::asm!( + concat!(" + .option push + .option arch, +v + 2: + vsetvli t0, {len}, e32, m8, ta, ma + vle32.v v8, ({ptr}) + ", $op, " + vse32.v v8, ({ptr}) + slli t1, t0, 2 + add {ptr}, {ptr}, t1 + sub {len}, {len}, t0 + bnez {len}, 2b + .option pop + "), + len = inout(reg) len => _, + ptr = inout(reg) ptr => _, + s = in(freg) s, + out("t0") _, + out("t1") _, + options(nostack), + ); + } + } + ); + }; +} + +rvv_by_scalar!(rvv_mul_by_scalar_f32, "vfmul.vf v8, v8, {s}"); +rvv_by_scalar!(rvv_add_by_scalar_f32, "vfadd.vf v8, v8, {s}"); +rvv_by_scalar!(rvv_sub_by_scalar_f32, "vfsub.vf v8, v8, {s}"); +rvv_by_scalar!(rvv_subf_by_scalar_f32, "vfrsub.vf v8, v8, {s}"); +rvv_by_scalar!(rvv_min_by_scalar_f32, "vfmin.vf v8, v8, {s}"); +rvv_by_scalar!(rvv_max_by_scalar_f32, "vfmax.vf v8, v8, {s}"); + +#[cfg(test)] +mod test { + use super::*; + use crate::riscv64::has_rvv; + + crate::by_scalar_frame_tests!(has_rvv(), f32, rvv_mul_by_scalar_f32, |a, b| a * b); + crate::by_scalar_frame_tests!(has_rvv(), f32, rvv_add_by_scalar_f32, |a, b| a + b); + crate::by_scalar_frame_tests!(has_rvv(), f32, rvv_sub_by_scalar_f32, |a, b| a - b); + crate::by_scalar_frame_tests!(has_rvv(), f32, rvv_subf_by_scalar_f32, |a, b| b - a); + crate::by_scalar_frame_tests!(has_rvv(), f32, rvv_min_by_scalar_f32, |a, b| a.min(b)); + crate::by_scalar_frame_tests!(has_rvv(), f32, rvv_max_by_scalar_f32, |a, b| a.max(b)); +} diff --git a/linalg/src/riscv64/reduce.rs b/linalg/src/riscv64/reduce.rs new file mode 100644 index 0000000000..8ea0f00259 --- /dev/null +++ b/linalg/src/riscv64/reduce.rs @@ -0,0 +1,109 @@ +//! Whole-slice f32 reductions, RVV 1.0. +//! +//! The running result lives in element 0 of v1 and is fed back as the +//! reduction's scalar operand each round, so strip-mining needs no separate +//! accumulator vector and no final horizontal step. Reduction operands are +//! LMUL=1 whatever vtype says, which is why v1 stays clear of the v8 group. +//! +//! `vfredusum` is the unordered sum: it may reassociate, as the reference +//! kernels on other targets also do. + +macro_rules! rvv_reduce { + ($func: ident, $neutral: expr, $red: expr, $reduce_two: item) => { + reduce_impl_wrap!( + f32, + $func, + 4, + 4, + (), + $neutral, + #[inline(never)] + fn run(buf: &[f32], _: ()) -> f32 { + assert!(!buf.is_empty()); + let len = buf.len(); + let ptr = buf.as_ptr(); + let out: f32; + // SAFETY: `len` elements are in bounds from `ptr`, and the + // loop advances both together. + unsafe { + std::arch::asm!( + concat!(" + .option push + .option arch, +v + vsetivli t0, 1, e32, m1, ta, ma + vfmv.s.f v1, {neutral} + 2: + vsetvli t0, {len}, e32, m8, ta, ma + vle32.v v8, ({ptr}) + ", $red, " + slli t1, t0, 2 + add {ptr}, {ptr}, t1 + sub {len}, {len}, t0 + bnez {len}, 2b + vfmv.f.s {out}, v1 + .option pop + "), + len = inout(reg) len => _, + ptr = inout(reg) ptr => _, + neutral = in(freg) $neutral, + out = lateout(freg) out, + out("t0") _, + out("t1") _, + options(nostack), + ); + } + out + }, + $reduce_two + ); + }; +} + +rvv_reduce!( + rvv_max_f32, + f32::MIN, + "vfredmax.vs v1, v8, v1", + #[inline(never)] + fn reduce_two(a: f32, b: f32) -> f32 { + a.max(b) + } +); + +rvv_reduce!( + rvv_min_f32, + f32::MAX, + "vfredmin.vs v1, v8, v1", + #[inline(never)] + fn reduce_two(a: f32, b: f32) -> f32 { + a.min(b) + } +); + +rvv_reduce!( + rvv_sum_f32, + 0f32, + "vfredusum.vs v1, v8, v1", + #[inline(never)] + fn reduce_two(a: f32, b: f32) -> f32 { + a + b + } +); + +#[cfg(test)] +mod test { + use super::*; + use crate::riscv64::has_rvv; + + mod max { + use super::*; + crate::max_frame_tests!(has_rvv(), f32, rvv_max_f32); + } + mod min { + use super::*; + crate::min_frame_tests!(has_rvv(), f32, rvv_min_f32); + } + mod sum { + use super::*; + crate::sum_frame_tests!(has_rvv(), f32, rvv_sum_f32); + } +} diff --git a/linalg/src/riscv64/unicast.rs b/linalg/src/riscv64/unicast.rs new file mode 100644 index 0000000000..5b436c77c5 --- /dev/null +++ b/linalg/src/riscv64/unicast.rs @@ -0,0 +1,72 @@ +//! Element-wise `a op b` over two f32 slices, RVV 1.0. +//! +//! Strip-mined on `vsetvli` like [`super::by_scalar`], so vector-length +//! agnostic and unpredicated. + +macro_rules! rvv_unicast { + ($func: ident, $op: expr) => { + unicast_impl_wrap!( + f32, + $func, + 4, + 4, + #[inline(never)] + fn run(a: &mut [f32], b: &[f32]) { + assert!(a.len() == b.len()); + assert!(!a.is_empty()); + let len = a.len(); + let a_ptr = a.as_mut_ptr(); + let b_ptr = b.as_ptr(); + // SAFETY: both slices hold `len` elements and the loop + // advances all three cursors in step. + unsafe { + std::arch::asm!( + concat!(" + .option push + .option arch, +v + 2: + vsetvli t0, {len}, e32, m8, ta, ma + vle32.v v8, ({a}) + vle32.v v16, ({b}) + ", $op, " + vse32.v v8, ({a}) + slli t1, t0, 2 + add {a}, {a}, t1 + add {b}, {b}, t1 + sub {len}, {len}, t0 + bnez {len}, 2b + .option pop + "), + len = inout(reg) len => _, + a = inout(reg) a_ptr => _, + b = inout(reg) b_ptr => _, + out("t0") _, + out("t1") _, + options(nostack), + ); + } + } + ); + }; +} + +rvv_unicast!(rvv_unicast_mul_f32, "vfmul.vv v8, v8, v16"); +rvv_unicast!(rvv_unicast_add_f32, "vfadd.vv v8, v8, v16"); +rvv_unicast!(rvv_unicast_sub_f32, "vfsub.vv v8, v8, v16"); +rvv_unicast!(rvv_unicast_subf_f32, "vfsub.vv v8, v16, v8"); +rvv_unicast!(rvv_unicast_min_f32, "vfmin.vv v8, v8, v16"); +rvv_unicast!(rvv_unicast_max_f32, "vfmax.vv v8, v8, v16"); + +#[cfg(test)] +mod test { + use super::*; + use crate::riscv64::has_rvv; + use proptest::strategy::Strategy; + + crate::unicast_frame_tests!(has_rvv(), f32, rvv_unicast_mul_f32, |a, b| a * b); + crate::unicast_frame_tests!(has_rvv(), f32, rvv_unicast_add_f32, |a, b| a + b); + crate::unicast_frame_tests!(has_rvv(), f32, rvv_unicast_sub_f32, |a, b| a - b); + crate::unicast_frame_tests!(has_rvv(), f32, rvv_unicast_subf_f32, |a, b| b - a); + crate::unicast_frame_tests!(has_rvv(), f32, rvv_unicast_min_f32, |a, b| a.min(b)); + crate::unicast_frame_tests!(has_rvv(), f32, rvv_unicast_max_f32, |a, b| a.max(b)); +}