perf(gemm): native float32 GEMM on loong64

internal/kernels/asmgen/go.mod

@@ -6,4 +6,4 @@ module github.com/go-ndarray/ndarray/internal/kernels/asmgen
 
 go 1.27.2
 
-require github.com/go-asmgen/asmgen v0.22.0
+require github.com/go-asmgen/asmgen v0.23.0

internal/kernels/asmgen/go.sum

@@ -1,2 +1,2 @@
-github.com/go-asmgen/asmgen v0.22.0 h1:ypPrlxd1V4u3NpsIDCIhNxtpVJzQuPZsPWGGUyt8GgA=
-github.com/go-asmgen/asmgen v0.22.0/go.mod h1:FHWXZt1MrejOwSShl4tPYjI2GPEfNAJuXJUX6Z+ruyQ=
+github.com/go-asmgen/asmgen v0.23.0 h1:K0L4gQ/CqK7gwuVrCWfrYU8S/OEsUt4qNftkjHIchXU=
+github.com/go-asmgen/asmgen v0.23.0/go.mod h1:FHWXZt1MrejOwSShl4tPYjI2GPEfNAJuXJUX6Z+ruyQ=

internal/kernels/asmgen/loong64/gen.go

@@ -41,6 +41,7 @@ func main() {
 		f.Add(bin32Kernel(k.name[:3]+"32LSX", k.vec[1:len(k.vec)-1]+"F", k.sca[:len(k.sca)-1]+"F"))
 	}
 	f.Add(gemmKernel())
+	f.Add(gemm32Kernel())
 	f.Add(blocks8Kernel())
 	f.Add(sqrt32Kernel())
 	abs64 := f.Data("abs64", lanes(8, 0x7fffffffffffffff))
@@ -364,6 +365,45 @@ func gemmKernel() *emit.Function {
 	return b.Func()
 }
 
+// gemm32Kernel builds gemmMicro8x16LASX(kc int, pa, pb, c *float32, ldc int):
+//
+//	for p in [0,kc):  C[r][col] += pa[p*8+r] * pb[p*16+col]   (r<8, col<16)
+//
+// gemmKernel at twice the lanes: 16 accumulators X0..X15 (row r in X(2r),
+// X(2r+1), eight columns each), the B row in X16, X17, each A value broadcast
+// with xvldrepl.w (XVMOVQ .W8) into X18..X25, and xvfmadd.s, which go-asmgen
+// encodes (XVFMADDF, v0.23.0) because Go's assembler has no vector float FMA.
+func gemm32Kernel() *emit.Function {
+	sig := loong64.Layout([]string{"kc", "pa", "pb", "c", "ldc"},
+		[]loong64.Type{loong64.Int64, loong64.Ptr, loong64.Ptr, loong64.Ptr, loong64.Int64}, nil, nil)
+	b := loong64.NewFunc("gemmMicro8x16LASX", sig, 0)
+	b.LoadArg("kc", "R5").LoadArg("pa", "R6").LoadArg("pb", "R7").LoadArg("c", "R8").LoadArg("ldc", "R9")
+	b.Raw("SLLV $2, R9") // ldc -> bytes
+	zero(b, 16)
+	b.Raw("BEQ R5, sstore")
+	b.Label("sloop")
+	b.Raw("XVMOVQ 0(R7), X16").Raw("XVMOVQ 32(R7), X17")
+	for r := 0; r < 8; r++ {
+		b.Raw("XVMOVQ %d(R6), X%d.W8", 4*r, 18+r) // xvldrepl.w: pa[r] in all eight lanes
+	}
+	for r := 0; r < 8; r++ {
+		b.XVFMADDF(2*r, 18+r, 16, 2*r)
+		b.XVFMADDF(2*r+1, 18+r, 17, 2*r+1)
+	}
+	b.Raw("ADDV $32, R6").Raw("ADDV $64, R7").Raw("ADDV $-1, R5").Raw("BNE R5, sloop")
+	b.Label("sstore")
+	for r := 0; r < 8; r++ {
+		b.Raw("XVMOVQ 0(R8), X16").Raw("XVMOVQ 32(R8), X17")
+		b.Raw("XVADDF X%d, X16, X16", 2*r).Raw("XVADDF X%d, X17, X17", 2*r+1)
+		b.Raw("XVMOVQ X16, 0(R8)").Raw("XVMOVQ X17, 32(R8)")
+		if r < 7 {
+			b.Raw("ADDV R9, R8")
+		}
+	}
+	b.Ret()
+	return b.Func()
+}
+
 // blocks8Kernel builds blocks8LSX(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). The eight accumulators are split over two 128-bit LSX

internal/kernels/gemm32_loong64.go

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

internal/kernels/gemm32_loong64_test.go

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

internal/kernels/gemm32_other.go

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

internal/kernels/sum_loong64.s

@@ -598,6 +598,119 @@ gstore:
 	XVMOVQ X17, 32(R8)
 	RET
 
+TEXT ·gemmMicro8x16LASX(SB), NOSPLIT, $0-40
+	MOVV kc+0(FP), R5
+	MOVV pa+8(FP), R6
+	MOVV pb+16(FP), R7
+	MOVV c+24(FP), R8
+	MOVV ldc+32(FP), R9
+	SLLV $2, R9
+	XVXORV X0, X0, X0
+	XVXORV X1, X1, X1
+	XVXORV X2, X2, X2
+	XVXORV X3, X3, X3
+	XVXORV X4, X4, X4
+	XVXORV X5, X5, X5
+	XVXORV X6, X6, X6
+	XVXORV X7, X7, X7
+	XVXORV X8, X8, X8
+	XVXORV X9, X9, X9
+	XVXORV X10, X10, X10
+	XVXORV X11, X11, X11
+	XVXORV X12, X12, X12
+	XVXORV X13, X13, X13
+	XVXORV X14, X14, X14
+	XVXORV X15, X15, X15
+	BEQ R5, sstore
+sloop:
+	XVMOVQ 0(R7), X16
+	XVMOVQ 32(R7), X17
+	XVMOVQ 0(R6), X18.W8
+	XVMOVQ 4(R6), X19.W8
+	XVMOVQ 8(R6), X20.W8
+	XVMOVQ 12(R6), X21.W8
+	XVMOVQ 16(R6), X22.W8
+	XVMOVQ 20(R6), X23.W8
+	XVMOVQ 24(R6), X24.W8
+	XVMOVQ 28(R6), X25.W8
+	WORD $0x0a104240 // XVFMADDF X0, X16, X18, X0
+	WORD $0x0a10c641 // XVFMADDF X1, X17, X18, X1
+	WORD $0x0a114262 // XVFMADDF X2, X16, X19, X2
+	WORD $0x0a11c663 // XVFMADDF X3, X17, X19, X3
+	WORD $0x0a124284 // XVFMADDF X4, X16, X20, X4
+	WORD $0x0a12c685 // XVFMADDF X5, X17, X20, X5
+	WORD $0x0a1342a6 // XVFMADDF X6, X16, X21, X6
+	WORD $0x0a13c6a7 // XVFMADDF X7, X17, X21, X7
+	WORD $0x0a1442c8 // XVFMADDF X8, X16, X22, X8
+	WORD $0x0a14c6c9 // XVFMADDF X9, X17, X22, X9
+	WORD $0x0a1542ea // XVFMADDF X10, X16, X23, X10
+	WORD $0x0a15c6eb // XVFMADDF X11, X17, X23, X11
+	WORD $0x0a16430c // XVFMADDF X12, X16, X24, X12
+	WORD $0x0a16c70d // XVFMADDF X13, X17, X24, X13
+	WORD $0x0a17432e // XVFMADDF X14, X16, X25, X14
+	WORD $0x0a17c72f // XVFMADDF X15, X17, X25, X15
+	ADDV $32, R6
+	ADDV $64, R7
+	ADDV $-1, R5
+	BNE R5, sloop
+sstore:
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X0, X16, X16
+	XVADDF X1, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X2, X16, X16
+	XVADDF X3, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X4, X16, X16
+	XVADDF X5, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X6, X16, X16
+	XVADDF X7, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X8, X16, X16
+	XVADDF X9, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X10, X16, X16
+	XVADDF X11, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X12, X16, X16
+	XVADDF X13, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	ADDV R9, R8
+	XVMOVQ 0(R8), X16
+	XVMOVQ 32(R8), X17
+	XVADDF X14, X16, X16
+	XVADDF X15, X17, X17
+	XVMOVQ X16, 0(R8)
+	XVMOVQ X17, 32(R8)
+	RET
+
 TEXT ·blocks8LSX(SB), NOSPLIT, $0-24
 	MOVV r+0(FP), R4
 	MOVV a+8(FP), R5