riscv64: 2 of 4: 16 RVV matmul kernels behind Zvfh

.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

.travis/cross.sh

@@ -127,7 +127,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" )
 
         case "$PLATFORM" in
             "aarch64-unknown-linux-gnu")
@@ -190,6 +191,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

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

linalg/build.rs

@@ -53,6 +53,35 @@ 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 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.
@@ -256,6 +285,10 @@ 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)");
+    // 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" => {
@@ -501,10 +534,80 @@ fn main() {
                 config.cc().files(files).compile("arm64fp16")
             }
         }
+        "riscv64" if assembler_supports_rvv() => {
+            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");
+        }
         _ => {}
     }
 }
 
+/// `(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"),
+];
+
+/// 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,
+    arch: &'static str,
+    kernels: &[(&'static str, &'static str, &'static str, &'static str)],
+    suffix: &str,
+) -> Vec<path::PathBuf> {
+    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(

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

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

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