diff --git a/internal/kernels/asmgen/go.mod b/internal/kernels/asmgen/go.mod index 1bcce24..8e2266d 100644 --- a/internal/kernels/asmgen/go.mod +++ b/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 diff --git a/internal/kernels/asmgen/go.sum b/internal/kernels/asmgen/go.sum index 725c35b..32bd48d 100644 --- a/internal/kernels/asmgen/go.sum +++ b/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= diff --git a/internal/kernels/asmgen/loong64/gen.go b/internal/kernels/asmgen/loong64/gen.go index 6d39ca4..5392bfb 100644 --- a/internal/kernels/asmgen/loong64/gen.go +++ b/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 diff --git a/internal/kernels/gemm32_loong64.go b/internal/kernels/gemm32_loong64.go new file mode 100644 index 0000000..0701962 --- /dev/null +++ b/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) diff --git a/internal/kernels/gemm32_loong64_test.go b/internal/kernels/gemm32_loong64_test.go new file mode 100644 index 0000000..c193745 --- /dev/null +++ b/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) + } +} diff --git a/internal/kernels/gemm32_other.go b/internal/kernels/gemm32_other.go index 22d2440..227d635 100644 --- a/internal/kernels/gemm32_other.go +++ b/internal/kernels/gemm32_other.go @@ -1,4 +1,4 @@ -//go:build !amd64 && !arm64 && !ppc64le && !s390x +//go:build !amd64 && !arm64 && !ppc64le && !s390x && !loong64 package kernels diff --git a/internal/kernels/sum_loong64.s b/internal/kernels/sum_loong64.s index 1c7257a..5bda738 100644 --- a/internal/kernels/sum_loong64.s +++ b/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