From cad3660c27d3bacaf24c9ffaea28a5776b0816f1 Mon Sep 17 00:00:00 2001 From: tannevaled Date: Thu, 8 Oct 2026 10:27:10 +0200 Subject: [PATCH] simd(riscv64): 6x8 RVV GEMM micro-kernel (LMUL=4, vl=8 on any VLEN) Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- docs/perf.md | 22 +++++++++- docs/plan-ndarray.md | 3 +- internal/kernels/asmgen/riscv64/gen.go | 43 +++++++++++++++++++ internal/kernels/gemm_generic.go | 9 ++-- internal/kernels/gemm_riscv64.go | 47 +++++++++++++++++++++ internal/kernels/gemm_riscv64_test.go | 1 + internal/kernels/reduce_riscv64.go | 3 +- internal/kernels/sum_riscv64.s | 57 ++++++++++++++++++++++++++ 9 files changed, 177 insertions(+), 10 deletions(-) create mode 100644 internal/kernels/gemm_riscv64.go diff --git a/README.md b/README.md index 427020b..f278cf8 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ arm64, VSX on ppc64le, LASX on loong64 and RVV on riscv64 when the CPU has them), max/min too on amd64 and arm64, and a **panel-packed, cache-blocked GEMM** with an SIMD-FMA micro-kernel (NEON 4×8; AVX2/FMA 6×8 chosen at run time, SSE2 fallback; VSX -8×8; LASX 8×8), the OpenBLAS/BLIS structure. s390x and the 32-bit targets run +8×8; LASX 8×8; RVV 6×8), the OpenBLAS/BLIS structure. s390x and the 32-bit targets run the same pure-Go code those kernels are tested against. `Exp` and `Log` are ports of Arm's optimized-routines (≤ 0.51 ULP measured, against up to 0.88 and 0.72 for Go's `math.Exp`/`math.Log` on arm64), diff --git a/docs/perf.md b/docs/perf.md index 35022d1..d59a7b1 100644 --- a/docs/perf.md +++ b/docs/perf.md @@ -843,8 +843,8 @@ stay dynamic, so its slower efficiency cores are handled as before. vector-length-agnostic loops (`VSETVLI` strip-mining, LMUL=8, no scalar tail), every instruction a Go mnemonic. max/min stay scalar: RVV's `vfmax`/`vfmin` implement IEEE 754-2019 maximumNumber, which drops a NaN - that NumPy propagates. The GEMM micro-kernel is still the scalar 4×4. On a - SpacemiT X60 (cfarm95, VLEN = 256), one core pinned to an idle CPU, 6 + that NumPy propagates. Since v0.7.1 the GEMM micro-kernel is a 6×8 RVV tile + (see below). On a SpacemiT X60 (cfarm95, VLEN = 256), one core pinned to an idle CPU, 6 interleaved rounds against the scalar code (median): | | 1 Ki | 16 Ki | 256 Ki | @@ -857,6 +857,24 @@ stay dynamic, so its slower efficiency cores are handled as before. `MatVec` 1024² is 2.02× faster and `Dot` 2^20 1.29×. On a riscv64 CPU without V (cfarm94, SiFive U74), the same build runs the scalar path and passes the same tests. + + **GEMM micro-kernel (v0.7.1).** NR is fixed by the packing, while VLEN varies + between RISC-V cores. The V extension guarantees VLEN ≥ 128, so a group of + four registers (LMUL=4) holds at least eight float64, and `VSETIVLI $8` + gives vl = 8 on every core: a 6×8 tile, six accumulator groups, the B row + loaded once per step and each A value fused in with `vfmacc.vf`. Without V, + the same 6×8 tile runs in Go. Measured with `hostload -pick 1` to choose an + idle CPU on a host otherwise 64% busy, and `benchstat` over 8 interleaved runs + each against v0.7.0's scalar 4×4 (all p = 0.000): + + | one core, X60 | 64² | 128² | 256² | 512² | 1024² | + |---|---|---|---|---|---| + | `MatMul` time | −31% | −32% | −45% | −43% | −44% (1.79 s → 1.00 s) | + + `Inner` is −44%, and the small rectangular products are −22% to −38%. At about + 2.1 GFLOP/s on one core this is still far from the hardware: VLEN = 256 with + LMUL=4 uses half of each register group, which keeps one tile portable to + every VLEN at that cost. - The remaining 64-bit Go target — **s390x** — keeps the validated scalar oracles, using the same four-accumulator max/min, direct sqrt loop, and a **scalar 4×4 GEMM micro-kernel** over the packed panels, and diff --git a/docs/plan-ndarray.md b/docs/plan-ndarray.md index 308faa7..86af618 100644 --- a/docs/plan-ndarray.md +++ b/docs/plan-ndarray.md @@ -125,7 +125,8 @@ suite (plus the per-arch CI jobs already wired in `.github/workflows/ci.yml`). dot, sqrt, elementwise, an 8×8 GEMM micro-kernel), through go-asmgen's VSX encoders, and since v0.2.5 loong64 has LASX kernels (used when the CPU reports LASX), and since v0.7.0 riscv64 has RVV kernels (sum, dot, sqrt, - elementwise; used when the CPU reports V). The other targets (s390x, 32-bit) keep the + elementwise, and since v0.7.1 a 6×8 GEMM micro-kernel; used when the CPU + reports V). The other targets (s390x, 32-bit) keep the validated scalar oracles plus a scalar 4×4 GEMM micro-kernel over the packed panels, with packing + cache-blocking + multicore; they have not been measured against NumPy. Split-CI: the pure-Go diff --git a/internal/kernels/asmgen/riscv64/gen.go b/internal/kernels/asmgen/riscv64/gen.go index a31ff39..cbc6fa1 100644 --- a/internal/kernels/asmgen/riscv64/gen.go +++ b/internal/kernels/asmgen/riscv64/gen.go @@ -35,6 +35,7 @@ func main() { } { f.Add(binKernel(k.name, k.op)) } + f.Add(gemmKernel()) if err := os.WriteFile("sum_riscv64.s", []byte(f.String()), 0o644); err != nil { fmt.Fprintln(os.Stderr, err) os.Exit(1) @@ -125,3 +126,45 @@ func binKernel(name, op string) *emit.Function { b.Raw("SUB X14, X13, X13").Raw("BNEZ X13, loop").Ret() return b.Func() } + +// gemmKernel builds gemmMicro6x8RVV(kc int, pa, pb, c *float64, ldc int) for +// kc >= 1: C[r][0:8] += sum_p pa[p*6+r] * pb[p*8:p*8+8] for r < 6, the rows of C +// ldc float64s apart. +// +// NR is fixed by the packing, while VLEN varies between RISC-V cores. The V +// extension guarantees VLEN >= 128, so a group of four registers (LMUL=4) holds +// at least eight float64: vl = 8 on every core. That leaves eight groups, six +// of them accumulators, one for the B row and one for C. Each step loads the B +// row once and fuses each A value in with vfmacc.vf (Vacc += f*Vb). +func gemmKernel() *emit.Function { + sig := riscv64.Layout([]string{"kc", "pa", "pb", "c", "ldc"}, + []riscv64.Type{riscv64.Int64, riscv64.Ptr, riscv64.Ptr, riscv64.Ptr, riscv64.Int64}, nil, nil) + b := riscv64.NewFunc("gemmMicro6x8RVV", sig, 0) + b.LoadArg("kc", "X10").LoadArg("pa", "X11").LoadArg("pb", "X12"). + LoadArg("c", "X13").LoadArg("ldc", "X14") + b.Raw("VSETIVLI $8, E64, M4, TA, MA, X15") + for r := 0; r < 6; r++ { + b.Raw("VMVVI $0, V%d", 4*r) + } + b.Raw("loop:").Raw("VLE64V (X12), V24") + for r := 0; r < 6; r++ { + b.Raw("MOVD %d(X11), F%d", 8*r, r) + } + for r := 0; r < 6; r++ { + b.Raw("VFMACCVF V24, F%d, V%d", r, 4*r) // acc_r += a_r * b + } + b.Raw("ADD $48, X11, X11"). + Raw("ADD $64, X12, X12"). + Raw("ADD $-1, X10, X10"). + Raw("BNEZ X10, loop"). + Raw("SLLI $3, X14, X14") // row stride in bytes + for r := 0; r < 6; r++ { + b.Raw("VLE64V (X13), V28"). + Raw("VFADDVV V%d, V28, V28", 4*r). // C row += acc_r + Raw("VSE64V V28, (X13)") + if r < 5 { + b.Raw("ADD X14, X13, X13") + } + } + return b.Ret().Func() +} diff --git a/internal/kernels/gemm_generic.go b/internal/kernels/gemm_generic.go index deb7869..b2b8b07 100644 --- a/internal/kernels/gemm_generic.go +++ b/internal/kernels/gemm_generic.go @@ -1,12 +1,11 @@ -//go:build !arm64 && !amd64 && !ppc64le && !loong64 +//go:build !arm64 && !amd64 && !ppc64le && !loong64 && !riscv64 package kernels // Generic GEMM micro-kernel for the targets without a vector-double SIMD -// micro-kernel (riscv64, s390x and the 32-bit ones): a portable scalar 4x4 -// register tile over the packed panels. There is no hand-vectorized tile here -// yet (riscv64 has RVV kernels for the other loops; Go's assembler has -// vector-double arithmetic for s390x too, see docs/perf.md) — but the kernel still gets the two structural wins of the packed +// micro-kernel (s390x and the 32-bit ones): a portable scalar 4x4 register +// tile over the packed panels. There is no hand-vectorized tile here yet (Go's +// assembler has vector-double arithmetic for s390x, see docs/perf.md) — but the kernel still gets the two structural wins of the packed // GEMM: the panels are contiguous unit-stride (so the loads are conflict-free, // the L1-collision fix), and the cache-blocking loop nest plus the multicore // fan-out (MatMulP) keep it well ahead of an unblocked scalar loop and beat diff --git a/internal/kernels/gemm_riscv64.go b/internal/kernels/gemm_riscv64.go new file mode 100644 index 0000000..a83dcbf --- /dev/null +++ b/internal/kernels/gemm_riscv64.go @@ -0,0 +1,47 @@ +package kernels + +// riscv64 GEMM micro-kernel: a 6x8 RVV tile (gemmMicro6x8RVV in +// sum_riscv64.s, generated by go-asmgen): six accumulator groups of eight +// float64 (LMUL=4, eight lanes on any VLEN >= 128), the B row loaded once per +// step and each A value fused in with vfmacc.vf. Without V, gemmMicroGo +// computes the same tile in scalar code. + +// MR, NR are the micro-kernel's register-block dimensions. +const ( + MR = 6 + NR = 8 +) + +func gemmMicro(kc int, pa, pb, dst []float64, ldc int) { + if useRVV { + _, _, _ = pa[kc*MR-1], pb[kc*NR-1], dst[(MR-1)*ldc+NR-1] // the kernel's whole reach + gemmMicro6x8RVV(kc, &pa[0], &pb[0], &dst[0], ldc) + return + } + gemmMicroGo(kc, pa, pb, dst, ldc) +} + +// gemmMicroGo is the scalar MR x NR tile: dst[r*ldc+c] += sum_p +// pa[p*MR+r]*pb[p*NR+c], in p order per element. +func gemmMicroGo(kc int, pa, pb, dst []float64, ldc int) { + var acc [MR * NR]float64 + for p := 0; p < kc; p++ { + a := pa[p*MR : p*MR+MR : p*MR+MR] + b := pb[p*NR : p*NR+NR : p*NR+NR] + for r, av := range a { + row := acc[r*NR : r*NR+NR : r*NR+NR] + for c, bv := range b { + row[c] += av * bv + } + } + } + for r := 0; r < MR; r++ { + d := dst[r*ldc : r*ldc+NR : r*ldc+NR] + for c := range d { + d[c] += acc[r*NR+c] + } + } +} + +//go:noescape +func gemmMicro6x8RVV(kc int, pa, pb, c *float64, ldc int) diff --git a/internal/kernels/gemm_riscv64_test.go b/internal/kernels/gemm_riscv64_test.go index 004e3e8..85b0a82 100644 --- a/internal/kernels/gemm_riscv64_test.go +++ b/internal/kernels/gemm_riscv64_test.go @@ -12,5 +12,6 @@ func TestRVVDetected(t *testing.T) { t.Run("scalar", func(t *testing.T) { TestBinSIMD(t) TestSumSIMD(t) + TestGemmParallelYielding(t) // gemmMicroGo against the oracle }) } diff --git a/internal/kernels/reduce_riscv64.go b/internal/kernels/reduce_riscv64.go index 7d44b78..c456ad9 100644 --- a/internal/kernels/reduce_riscv64.go +++ b/internal/kernels/reduce_riscv64.go @@ -2,7 +2,8 @@ package kernels import "os" -// riscv64 RVV kernels generated by go-asmgen (sum_riscv64.s). The vector +// riscv64 RVV kernels generated by go-asmgen (sum_riscv64.s); the GEMM +// micro-kernel is in gemm_riscv64.go. The vector // extension is optional on RISC-V, so each wrapper uses the kernel only when // the kernel reports V in AT_HWCAP, and runs the scalar oracle otherwise. The // loops are vector-length agnostic: they run on any VLEN. diff --git a/internal/kernels/sum_riscv64.s b/internal/kernels/sum_riscv64.s index 057afff..920fe6b 100644 --- a/internal/kernels/sum_riscv64.s +++ b/internal/kernels/sum_riscv64.s @@ -141,3 +141,60 @@ TEXT ·divRVV(SB), NOSPLIT, $0-32 BNEZ X13, loop RET +TEXT ·gemmMicro6x8RVV(SB), NOSPLIT, $0-40 + MOV kc+0(FP), X10 + MOV pa+8(FP), X11 + MOV pb+16(FP), X12 + MOV c+24(FP), X13 + MOV ldc+32(FP), X14 + VSETIVLI $8, E64, M4, TA, MA, X15 + VMVVI $0, V0 + VMVVI $0, V4 + VMVVI $0, V8 + VMVVI $0, V12 + VMVVI $0, V16 + VMVVI $0, V20 + loop: + VLE64V (X12), V24 + MOVD 0(X11), F0 + MOVD 8(X11), F1 + MOVD 16(X11), F2 + MOVD 24(X11), F3 + MOVD 32(X11), F4 + MOVD 40(X11), F5 + VFMACCVF V24, F0, V0 + VFMACCVF V24, F1, V4 + VFMACCVF V24, F2, V8 + VFMACCVF V24, F3, V12 + VFMACCVF V24, F4, V16 + VFMACCVF V24, F5, V20 + ADD $48, X11, X11 + ADD $64, X12, X12 + ADD $-1, X10, X10 + BNEZ X10, loop + SLLI $3, X14, X14 + VLE64V (X13), V28 + VFADDVV V0, V28, V28 + VSE64V V28, (X13) + ADD X14, X13, X13 + VLE64V (X13), V28 + VFADDVV V4, V28, V28 + VSE64V V28, (X13) + ADD X14, X13, X13 + VLE64V (X13), V28 + VFADDVV V8, V28, V28 + VSE64V V28, (X13) + ADD X14, X13, X13 + VLE64V (X13), V28 + VFADDVV V12, V28, V28 + VSE64V V28, (X13) + ADD X14, X13, X13 + VLE64V (X13), V28 + VFADDVV V16, V28, V28 + VSE64V V28, (X13) + ADD X14, X13, X13 + VLE64V (X13), V28 + VFADDVV V20, V28, V28 + VSE64V V28, (X13) + RET +