sgemm_sve_arm64_test.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build goexperiment.simd && arm64
     6  
     7  // SGEMM (C += A*B on float32) written with the scalable SVE API, checked
     8  // against a scalar reference and benchmarked against the hand-written
     9  // assembly kernel in sgemm_sve_arm64.s.
    10  
    11  package simd
    12  
    13  import (
    14  	"fmt"
    15  	"simd/archsimd"
    16  	"testing"
    17  )
    18  
    19  // sgemmSIMD computes C += A(m×k) * B(k×n), for any m, k, n.
    20  // A, B, C are row-major matrixes, e.g:
    21  //
    22  //	A (2×3)		B (3×6)					=>	C (2×6)
    23  //	[ 1 2 3 ]	[ 10 11 12 13 14 15 ]		[ 140 146 152 158 164 170 ]
    24  //	[ 4 5 6 ]	[ 20 21 22 23 24 25 ]		[ 320 335 350 365 380 395 ]
    25  //				[ 30 31 32 33 34 35 ]
    26  //
    27  // Matmul is paralleized at the C-row level, the final C is completed in
    28  // in mxn/Float32s.len() steps (the well-aligned case).
    29  // The code documentation in the function body are written assuming 2D slice,
    30  // i.e. a[i][j] is a[i * m + j].
    31  func sgemmSIMD(c, a, b []float32, m, k, n int) {
    32  	if len(c) < m*n || len(a) < m*k || len(b) < k*n {
    33  		panic("sgemmSIMD: insufficient input slice length")
    34  	}
    35  	vl := archsimd.Float32s{}.Len()
    36  	allTrue := archsimd.Mask32sAllTrue()
    37  	// Computes the the i-th row of c
    38  	for i := 0; i < m; i++ {
    39  		// Take a[i]
    40  		ai := a[i*k : i*k+k]
    41  		// c[i] will be completed, take its reference here
    42  		ci := c[i*n : i*n+n]
    43  		j := 0
    44  		// SIMD loop, parallel level vl, parallelize on c[i]
    45  		for ; j+vl <= n; j += vl {
    46  			// acc is the accumulator for c[i][j:j+vl]
    47  			var acc archsimd.Float32s
    48  			off := j // offset in b
    49  			// Run over each column of a[i]
    50  			for p := range k {
    51  				// av is [a[i][p], a[i][p], ..., a[i][p]]
    52  				av := archsimd.BroadcastFloat32s(ai[p])
    53  				// bv is [b[p][j], b[p][j+1], ..., b[p][j+vl]]
    54  				bv := archsimd.LoadFloat32s(b[off : off+vl])
    55  				// Update c[i]:
    56  				// c[i][j] += a[i][p] * b[p][j]
    57  				// c[i][j+1] += a[i][p] * b[p][j+1]
    58  				// ...
    59  				// c[i][j+vl] += a[i][p] * b[p][j+vl]
    60  				//
    61  				// Note: MulAdd is a predicated-only instruction,
    62  				// to provide the API unpredicated the compiler
    63  				// will insert a synthesized all-true predicate
    64  				// for the instruction. However LICM doesn't work
    65  				// on SIMD operations yet, this all-true predicate
    66  				// will be re-generated for every iteration of this
    67  				// loop. To avoid this, we explicitly create
    68  				// an all-true predicate for the instruction, bypassing
    69  				// the synthesized all-true predicate, i.e. essentially
    70  				// doing LICM manually.
    71  				acc = av.MulAdd(bv, acc).IfElse(allTrue, acc)
    72  				// Move b to next row
    73  				off += n
    74  			}
    75  			acc.Store(ci[j : j+vl])
    76  		}
    77  		sgemmTail(ci, ai, b, j, n, allTrue)
    78  	}
    79  }
    80  
    81  // sgemmTail is the masked tail of the inner-most loop of
    82  // sgemmSIMD and sgemmSIMD2, replacing the panicing unmasked loads
    83  // with the masked partial loads.
    84  func sgemmTail(ci, ai, b []float32, j, n int, allTrue archsimd.Mask32s) {
    85  	if j >= n {
    86  		return
    87  	}
    88  	var acc archsimd.Float32s
    89  	off := j
    90  	for p := range ai {
    91  		av := archsimd.BroadcastFloat32s(ai[p])
    92  		bv, _ := archsimd.LoadFloat32sPart(b[off : off+n-j])
    93  		acc = av.MulAdd(bv, acc).IfElse(allTrue, acc)
    94  		off += n
    95  	}
    96  	acc.StorePart(ci[j:])
    97  }
    98  
    99  // sgemmSIMD2 manually unrolled the inner-most loop of sgemmSIMD by 2 times,
   100  // which we can save one broadcast out of it.
   101  func sgemmSIMD2(c, a, b []float32, m, k, n int) {
   102  	if len(c) < m*n || len(a) < m*k || len(b) < k*n {
   103  		panic("sgemmSIMD2: insufficient input slice length")
   104  	}
   105  	vl := archsimd.Float32s{}.Len()
   106  	allTrue := archsimd.Mask32sAllTrue()
   107  	for i := 0; i < m; i++ {
   108  		ai := a[i*k : i*k+k]
   109  		ci := c[i*n : i*n+n]
   110  		j := 0
   111  		for ; j+2*vl <= n; j += 2 * vl {
   112  			var acc0, acc1 archsimd.Float32s
   113  			off := j
   114  			for p := range k {
   115  				av := archsimd.BroadcastFloat32s(ai[p])
   116  				bv0 := archsimd.LoadFloat32s(b[off : off+vl])
   117  				acc0 = av.MulAdd(bv0, acc0).IfElse(allTrue, acc0)
   118  				bv1 := archsimd.LoadFloat32s(b[off+vl : off+2*vl])
   119  				acc1 = av.MulAdd(bv1, acc1).IfElse(allTrue, acc1)
   120  				off += n
   121  			}
   122  			acc0.Store(ci[j : j+vl])
   123  			acc1.Store(ci[j+vl : j+2*vl])
   124  		}
   125  		for ; j+vl <= n; j += vl {
   126  			var acc archsimd.Float32s
   127  			off := j
   128  			for p := range k {
   129  				av := archsimd.BroadcastFloat32s(ai[p])
   130  				bv := archsimd.LoadFloat32s(b[off : off+vl])
   131  				acc = av.MulAdd(bv, acc).IfElse(allTrue, acc)
   132  				off += n
   133  			}
   134  			acc.Store(ci[j : j+vl])
   135  		}
   136  		sgemmTail(ci, ai, b, j, n, allTrue)
   137  	}
   138  }
   139  
   140  // sgemmAsm is asm-equivalent of sgemmSIMD.
   141  // Unlike sgemmSIMD it has no masked tail, so n must be a
   142  // multiple of the vector length.
   143  func sgemmAsm(c, a, b []float32, m, k, n int) {
   144  	if len(c) < m*n || len(a) < m*k || len(b) < k*n {
   145  		panic("sgemmAsm: insufficient input slice length")
   146  	}
   147  	vl := archsimd.Float32s{}.Len()
   148  	if n%vl != 0 {
   149  		panic("sgemmAsm: n must be a multiple of the vector length")
   150  	}
   151  	for i := 0; i < m; i++ {
   152  		for j := 0; j <= n-vl; j += vl {
   153  			sgemmKernelSVE(&c[i*n+j], &a[i*k], &b[j], int64(k), int64(n)*4)
   154  		}
   155  	}
   156  }
   157  
   158  // sgemmScalar is the scalar reference implementation of sgemm.
   159  func sgemmScalar(c, a, b []float32, m, k, n int) {
   160  	if len(c) < m*n || len(a) < m*k || len(b) < k*n {
   161  		panic("sgemmScalar: insufficient input slice length")
   162  	}
   163  	for i := 0; i < m; i++ {
   164  		for j := 0; j < n; j++ {
   165  			s := float32(0)
   166  			for p := 0; p < k; p++ {
   167  				s += a[i*k+p] * b[p*n+j]
   168  			}
   169  			c[i*n+j] = s
   170  		}
   171  	}
   172  }
   173  
   174  // sgemmData returns integer-valued inputs for a(mxk), b(kxn) matrixes.
   175  //
   176  // The test data is integer-valued and small enough that every product and sum is
   177  // exact in float32, so the fused and unfused accumulations agree bit for bit in the scalar impl.
   178  func sgemmData(m, k, n int) (a, b []float32) {
   179  	a = make([]float32, m*k)
   180  	b = make([]float32, k*n)
   181  	for i := range a {
   182  		a[i] = float32(i%9) - 4
   183  	}
   184  	for i := range b {
   185  		b[i] = float32(i%7) - 3
   186  	}
   187  	return
   188  }
   189  
   190  var sgemmImpls = []struct {
   191  	name string
   192  	f    func(c, a, b []float32, m, k, n int)
   193  }{
   194  	{"simd", sgemmSIMD},
   195  	{"simd2", sgemmSIMD2},
   196  	{"asm", sgemmAsm},
   197  }
   198  
   199  func TestSGEMMSVE(t *testing.T) {
   200  	if !archsimd.ARM64.SVE() {
   201  		t.Skip("no SVE")
   202  	}
   203  	vl := archsimd.Float32s{}.Len()
   204  	for _, sz := range []struct{ m, k, n int }{
   205  		{7, 37, 64}, // no tail, full alignment to VL
   206  		{7, 37, 61}, // aligned + unaligned, triggers tail
   207  		{3, 5, 2},   // tail only
   208  	} {
   209  		m, k, n := sz.m, sz.k, sz.n
   210  		a, b := sgemmData(m, k, n)
   211  		want := make([]float32, m*n)
   212  		sgemmScalar(want, a, b, m, k, n)
   213  		for _, impl := range sgemmImpls {
   214  			if impl.name == "asm" && n%vl != 0 {
   215  				continue // the assembly kernel has no masked tail
   216  			}
   217  			got := make([]float32, m*n)
   218  			impl.f(got, a, b, m, k, n)
   219  			for i := range got {
   220  				if got[i] != want[i] {
   221  					t.Errorf("%s %dx%dx%d: c[%d,%d] = %v, want %v", impl.name, m, k, n, i/n, i%n, got[i], want[i])
   222  				}
   223  			}
   224  		}
   225  	}
   226  }
   227  
   228  func BenchmarkSGEMMSVE(b *testing.B) {
   229  	if !archsimd.ARM64.SVE() {
   230  		b.Skip("no SVE")
   231  	}
   232  	for _, s := range []int{64, 256} {
   233  		m, k, n := s, s, s
   234  		am, bm := sgemmData(m, k, n)
   235  		c := make([]float32, m*n)
   236  		impls := append([]struct {
   237  			name string
   238  			f    func(c, a, b []float32, m, k, n int)
   239  		}{{"scalar", sgemmScalar}}, sgemmImpls...)
   240  		for _, impl := range impls {
   241  			b.Run(fmt.Sprintf("%s/%d", impl.name, s), func(b *testing.B) {
   242  				for range b.N {
   243  					impl.f(c, am, bm, m, k, n)
   244  				}
   245  				flops := 2 * float64(m) * float64(k) * float64(n)
   246  				b.ReportMetric(flops*float64(b.N)/b.Elapsed().Seconds()/1e9, "GFLOPS")
   247  			})
   248  		}
   249  	}
   250  }
   251
View as plain text