perf(gemm): native float32 GEMM on riscv64

internal/kernels/asmgen/riscv64/gen.go

@@ -37,6 +37,7 @@ func main() {
 		f.Add(bin32Kernel(k.name[:3]+"32RVV", k.op))
 	}
 	f.Add(gemmKernel())
+	f.Add(gemm32Kernel())
 	f.Add(blocks8Kernel())
 	f.Add(sqrt32Kernel())
 	for _, k := range []struct {
@@ -276,6 +277,45 @@ func gemmKernel() *emit.Function {
 	return b.Ret().Func()
 }
 
+// gemm32Kernel builds gemmMicro6x16RVV(kc int, pa, pb, c *float32, ldc int) for
+// kc >= 1: C[r][0:16] += sum_p pa[p*6+r] * pb[p*16:p*16+16] for r < 6, the rows
+// of C ldc float32s apart. gemmKernel in float32: with E32 a group of four
+// registers (LMUL=4) holds at least sixteen floats on any VLEN >= 128, so
+// vl = 16 everywhere, and the same six accumulator groups, B group and C group
+// fit.
+func gemm32Kernel() *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("gemmMicro6x16RVV", sig, 0)
+	b.LoadArg("kc", "X10").LoadArg("pa", "X11").LoadArg("pb", "X12").
+		LoadArg("c", "X13").LoadArg("ldc", "X14")
+	b.Raw("VSETIVLI $16, E32, M4, TA, MA, X15")
+	for r := 0; r < 6; r++ {
+		b.Raw("VMVVI $0, V%d", 4*r)
+	}
+	b.Raw("loop:").Raw("VLE32V (X12), V24")
+	for r := 0; r < 6; r++ {
+		b.Raw("MOVF %d(X11), F%d", 4*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 $24, X11, X11").
+		Raw("ADD $64, X12, X12").
+		Raw("ADD $-1, X10, X10").
+		Raw("BNEZ X10, loop").
+		Raw("SLLI $2, X14, X14") // row stride in bytes
+	for r := 0; r < 6; r++ {
+		b.Raw("VLE32V (X13), V28").
+			Raw("VFADDVV V%d, V28, V28", 4*r). // C row += acc_r
+			Raw("VSE32V V28, (X13)")
+		if r < 5 {
+			b.Raw("ADD X14, X13, X13")
+		}
+	}
+	return b.Ret().Func()
+}
+
 // blocks8Kernel builds blocks8RVV(r, a *float32, blocks int): for each of the
 // blocks of eight float32 at a, in order, r[j] += a[8i+j] (NumPy's
 // pairwise-sum block). VSETIVLI $8 with E32 and M2 gives vl = 8 on every

internal/kernels/gemm32_other.go

@@ -1,4 +1,4 @@
-//go:build !amd64 && !arm64 && !ppc64le && !s390x
+//go:build !amd64 && !arm64 && !ppc64le && !s390x && !riscv64
 
 package kernels
 

internal/kernels/gemm32_riscv64.go

@@ -0,0 +1,24 @@
+package kernels
+
+// gemm32 is the float32 GEMM on the 6x16 RVV kernel (gemmMicro6x16RVV in
+// sum_riscv64.s, generated by go-asmgen; see asmgen/riscv64/gen.go). Without
+// the V extension there is none, and float32 products are computed in float64.
+var gemm32 = newGemm32(useRVV)
+
+func newGemm32(rvv bool) *gemmKernel[float32] {
+	if !rvv {
+		return nil
+	}
+	return &gemmKernel[float32]{mr: 6, nr: 16, micro: gemmMicro32}
+}
+
+// gemmMicro32 adds the 6 x 16 tile sum_p pa[p*6..]*pb[p*16..] into the C block
+// at dst[0:], whose rows are ldc apart. pa, pb hold at least kc*6 and kc*16
+// packed (zero-padded) floats; dst has at least 5*ldc + 16 elements.
+func gemmMicro32(kc int, pa, pb, dst []float32, ldc int) {
+	_, _, _ = pa[kc*6-1], pb[kc*16-1], dst[5*ldc+15] // the kernel's whole reach
+	gemmMicro6x16RVV(kc, &pa[0], &pb[0], &dst[0], ldc)
+}
+
+//go:noescape
+func gemmMicro6x16RVV(kc int, pa, pb, c *float32, ldc int)

internal/kernels/gemm32_riscv64_test.go

@@ -0,0 +1,14 @@
+package kernels
+
+import "testing"
+
+// TestGemm32NeedsRVV: without the V extension there is no float32 kernel, and
+// HasGemm32 sends float32 products to the float64 GEMM.
+func TestGemm32NeedsRVV(t *testing.T) {
+	if newGemm32(false) != nil {
+		t.Fatal("a float32 GEMM without V")
+	}
+	if g := newGemm32(true); g == nil || g.mr != 6 || g.nr != 16 {
+		t.Fatalf("newGemm32(true) = %+v", g)
+	}
+}

internal/kernels/sum_riscv64.s

@@ -274,6 +274,63 @@ TEXT ·gemmMicro6x8RVV(SB), NOSPLIT, $0-40
 	VSE64V V28, (X13)
 	RET
 
+TEXT ·gemmMicro6x16RVV(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 $16, E32, M4, TA, MA, X15
+	VMVVI $0, V0
+	VMVVI $0, V4
+	VMVVI $0, V8
+	VMVVI $0, V12
+	VMVVI $0, V16
+	VMVVI $0, V20
+	loop:
+	VLE32V (X12), V24
+	MOVF 0(X11), F0
+	MOVF 4(X11), F1
+	MOVF 8(X11), F2
+	MOVF 12(X11), F3
+	MOVF 16(X11), F4
+	MOVF 20(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 $24, X11, X11
+	ADD $64, X12, X12
+	ADD $-1, X10, X10
+	BNEZ X10, loop
+	SLLI $2, X14, X14
+	VLE32V (X13), V28
+	VFADDVV V0, V28, V28
+	VSE32V V28, (X13)
+	ADD X14, X13, X13
+	VLE32V (X13), V28
+	VFADDVV V4, V28, V28
+	VSE32V V28, (X13)
+	ADD X14, X13, X13
+	VLE32V (X13), V28
+	VFADDVV V8, V28, V28
+	VSE32V V28, (X13)
+	ADD X14, X13, X13
+	VLE32V (X13), V28
+	VFADDVV V12, V28, V28
+	VSE32V V28, (X13)
+	ADD X14, X13, X13
+	VLE32V (X13), V28
+	VFADDVV V16, V28, V28
+	VSE32V V28, (X13)
+	ADD X14, X13, X13
+	VLE32V (X13), V28
+	VFADDVV V20, V28, V28
+	VSE32V V28, (X13)
+	RET
+
 TEXT ·blocks8RVV(SB), NOSPLIT, $0-24
 	MOV r+0(FP), X10
 	MOV a+8(FP), X11