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 00423f5eba..a5d5369f42 100755 --- a/.travis/cross.sh +++ b/.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 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 d76f5d510d..75de3f645a 100644 --- a/linalg/build.rs +++ b/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,106 @@ fn main() { config.cc().files(files).compile("arm64fp16") } } + "riscv64" if assembler_supports_rvv() => { + 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", + 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. +/// 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"), + ("64x1", "64", "1", "8"), + ("128x1", "128", "1", "8"), +]; + +fn render_rvv_kernels( + tmpl: &str, + 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(tmpl); + 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/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/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/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/hwbench/bandwidth.rs b/linalg/src/hwbench/bandwidth.rs index 7806dd51de..57da2ad210 100644 --- a/linalg/src/hwbench/bandwidth.rs +++ b/linalg/src/hwbench/bandwidth.rs @@ -102,6 +102,27 @@ fn load_a_slice(slice: &[u8], loops: usize) { } } +/// Portable fallback for targets with no hand-written streaming load. +/// +/// One volatile read per 64-byte line is enough to pull every line from +/// memory, so the traffic measured is the same; without the wide loads of the +/// specialised versions there are fewer misses in flight, so an out-of-order +/// core may not saturate and the figure reads low rather than wrong. +#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "arm")))] +#[inline(never)] +fn load_a_slice(slice: &[u8], loops: usize) { + for _ in 0..loops { + let mut offset = 0; + while offset < slice.len() { + // SAFETY: `offset` is below `slice.len()`, so the read is in bounds + // and `u64`-sized reads stay within the line they start in because + // the buffer is 1024-byte aligned. + unsafe { std::ptr::read_volatile(slice.as_ptr().add(offset) as *const u64) }; + offset += 64; + } + } +} + fn bandwidth_seq(slice_len: usize, threads: usize) -> f64 { #[cfg(target_arch = "x86_64")] unsafe { 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..5cc6fae89f --- /dev/null +++ b/linalg/src/riscv64.rs @@ -0,0 +1,149 @@ +//! 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`] 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. +#[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 +} + +/// 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. +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 +} + +/// `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::() +} + +/// 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)] + 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={} 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() { + 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); + assert!(!has_zvfh(), "Zvfh cannot be present without V"); + } + } +} diff --git a/linalg/src/riscv64/rvv.rs b/linalg/src/riscv64/rvv.rs new file mode 100644 index 0000000000..77fb9adc56 --- /dev/null +++ b/linalg/src/riscv64/rvv.rs @@ -0,0 +1,206 @@ +//! 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` and `RVV_F16_KERNELS` in build.rs, which must stay in +//! step with them. +//! +//! `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 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; + +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)); + +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::*; + 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. 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)] + ("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(&[ + rvv_mmm_f32_8x8.mmm(), + 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(&[ + 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)] +mod test { + use super::*; + use crate::frame::mmm::{FusedKerSpec, MatMatMulKer}; + + 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(), + 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([ + 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 + /// 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, 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"); + } + } + + /// 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; + } + #[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])), + 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() { + 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"); + eprintln!("{name}: correctly refused"); + } + } + } +}