simd: architecture and vector-size agnostic SIMD intrinsics under a GOEXPERIMENT

[dr2chase]

Proposal Details

This is the first version of the "second level" SIMD api described in #73787; the intent is to provide a portable interface to SIMD vectors that is not tied to a particular architecture or vector length. On architectures where the instructions are supported, the API provides high performance, on those where it is not, it at least provides portability and a competent emulation. The API will also automatically pick the "best" (largest vector length, latest technology) when run on different variants of a different architecture (that is, it will auto-select between AVX, AVX2, and AVX512, depending on what is available).

In this version, the supported vector methods are those in the intersection of the wasm SIMD API and the current amd64 SIMD API; this is a small but still useful set of operations. Starting small also simplifies the emulation for those platforms that do not yet support SIMD, and also increases the chance that newly implemented SIMD support for some other architecture will easily be compatible.

API overview

The vector type names are all the signed, unsigned, and float types, with first letter capitalized, and a plural suffix, for example Int8s or Float64s. Within a given execution of a program all the vectors have the same length in bits, so they may be punned with As<OtherTypeSize>s, for example x.AsFloat64s().

Masks are supported as an abstraction (because all the architectures implement them differently), with a bit width to indicate the size of compatible vector elements, and all non-mask types support the corresponding .Masked(mask) and .Merged(valueIfFalse, mask) methods. Masks can also be manipulated with simple booleans, currently And and Or. On architectures/variants supporting single operation+mask instructions, it is intended that x.Operation(y).Masked(m) pattern-match into that single instruction.

It is possible to convert an architecture-agnostic simd type into a specific architecture's simd type, and back again. This isn't in the prototype yet, the proposed methods names for this are ToArch() any and ToSimd() any followed by a type case that includes the desired types.
For example:

//go:build amd64
...
import (
	"simd"
	"simd/archsimd"
)
func sum(x simd.Float32s) float32 {
	switch a := x.ToArch().(type) {
	case archsimd.Float32x8:
		a = a.AddPairsGrouped(a)
		a = a.AddPairsGrouped(a)
		return a.GetLo().GetElem(0) + a.GetHi().GetElem(0)
	case archsimd.Float32x16:
	...
	case archsimd.Float32x4:
	...
	}

There is a candidate CL; the actual API appears in simd/mocks.go. It is generated automatically by simd/archsimd/_gen/midway/ from an intersection of the SIMD APIs for wasm and amd64 (and other architectures, in the future). The portable API will be expanded over time by various architecture-specific APIs with emulations to fill in the intersection.

As an example, "Int8s" supports the following functions and methods

LoadInt8Slice, LoadInt8SlicePart, Abs, Add, AddSaturated, And, AndNot,
AsFloat32s, AsFloat64s, AsInt16s, AsInt32s, AsInt64s, AsUint16s, AsUint32s,
AsUint64s, AsUint8s, Equal, Greater, Len, Masked, Max, Merge, Min, Neg, Not,
Or, StoreSlice, StoreSlicePart, String, Sub, SubSaturated, ToMask, Xor

and Float32s supports these functions and methods

LoadFloat32Slice, LoadFloat32SlicePart, Add, AsFloat64s, AsInt16s, AsInt32s,
AsInt64s, AsInt8s, AsUint16s, AsUint32s, AsUint64s, AsUint8s, ConvertToInt32,
Div, Equal, Greater, GreaterEqual, Len, Less, LessEqual, Masked, Max, Merge,
Min, Mul, MulAdd, NotEqual, Sqrt, StoreSlice, StoreSlicePart, String, Sub, 

Implementation technique

Code that depends on "simd" is rewritten into multiple versions, and functions and methods that internally depend on simd but do not export that dependence in their signatures will be turned into dispatch functions/methods. Within code compiled for a particular variant, references to other simd-dependent code will be rewritten to reference the same variant; the intent is to both avoid internal dispatch and minimize code cloning. This rewrite also occurs early in the compiler to maximize the chance of inlining (for performance, of course). The rewriting generates names with "@" characters in them, so they cannot be referred to from source code (this may have some effect on tools processing export data or symbols in object files, e.g., debuggers).

[ianlancetaylor]

Since the length is not known when writing the program, I think it would help to show a couple of examples of how to write code that actually uses these new types and functions/methods. Thanks.

[balasanjay]

(copying over my comment from dr2chase/midway#1 and closing that)

Let's say I wanted to write an algorithm specifically using simd.Uint32x8 (for cryptographic reasons, or hash-table sizing reasons (e.g. I'm storing 8 uint32 hash codes in a single vector)).

I would want that to compile to archsimd.Uint32x8 natively for platforms supporting 256-bit operations, but I would want it to also generate an implementation using a pair of archsimd.Uint32x4 values representing lo and hi elements in my logical 8-vector.

(Actually, for that matter, I may even want it to compile it into a variant that is implemented using 8 uint32s that work without any archsimd support at all).

Is the intention for midway/simd package to support use-cases like that (fixed lane size but generic over underlying hardware lane size), or is the use-case limited to code that is "generic" over the number of lanes.
I believe highway supports this via FixedTag<T, kCount> (and what is currently implemented is equivalent to ScalableTag).

[dr2chase]

@ianlancetaylor
The canonical example is inner product:

func ip(x, y []float32) float32 {
	var a simd.Float32s
	var i int
	for i = 0; i < len(x)-a.Len()+1; i += a.Len() {
		u := simd.LoadFloat32Slice(x[i : i+a.Len()])
		v := simd.LoadFloat32Slice(y[i : i+a.Len()])
		a = a.Add(u.Mul(v))
	}
	if i < len(x) {
		a = a.Add(simd.LoadFloat32SlicePart(x[i:]).
			Mul(simd.LoadFloat32SlicePart(y[i:])))
	}
	return sum(a)
}

"Sum" is implemented two different ways, chosen by build tag:

//go:build !amd64
package main
import (
	"simd"
)

func sum(x simd.Float32s) float32 {
	s := make([]float32, x.Len())
	x.StoreSlice(s)
	var r float32
	for _, e := range s {
		r += e
	}
	return r
}

or

//go:build amd64
package main
import (
	"simd"
	"simd/archsimd"
)

func sum(x simd.Float32s) float32 {
	switch a := x.ToArch().(type) { 
	case archsimd.Float32x8:
		a = a.AddPairsGrouped(a)
		a = a.AddPairsGrouped(a)
		return a.GetLo().GetElem(0) + a.GetHi().GetElem(0)
	case archsimd.Float32x16:
		s := make([]float32, a.Len())
		a.StoreSlice(s)
		var r float32
		for _, e := range s {
			r += e
		}
		return r
	case archsimd.Float32x4:
		s := make([]float32, a.Len())
		a.StoreSlice(s)
		var r float32
		for _, e := range s {
			r += e
		}
		return r
	}
	panic("not a known type")
}

[dr2chase]

@balasanjay The use case is code that is not sensitive to the vector width. You can however inquire about the length of a simd.Uint32s and chose different code based on that, and also use build tags (in the case of sum above) to specialize by architecture. The one thing that won't work is converting archsimd.Uint32x8 directly into simd.Uint32s if you happen to be running on AVX512, since in that case simd.Uint32s is actually archsimd.Uint32x16.

This gets more annoying given the scalable SIMD architectures (SVE, RVV) that in theory and sometimes in practice can have more than 256 bits -- code generation to 256-bit vector operations in that case potentially requires a mask applied to all otherwise-larger operations. (Last I checked, real-world SVE was 128, 256, and a 512-bit Japanese supercomputer. There are several RISCV implementations with 512-bit RVV.)

[balasanjay]

@dr2chase Hmm, I see.

I understand the usecase for simd.Uint32s. As you say, it allows you to write portable code that is not sensitive to the vector width. It will compile variations for architectures with archsimd.Uint32x4 support, archsimd.Uint32x8 support, etc.

I think there is also a usecase for simd.Uint32x8. It would allow you to write portable code for dealing with specifically 8 uint32s. It would of course compile a variation for an architecture with archsimd.Uint32x8 support (trivially), but it would also compile a variation for an architecture with just archsimd.Uint32x4 support (by emulating a simd.Uint32x8 as two archsimd.Uint32x4). It would even compile a variation for a scalar architecture (by emulating a simd.Uint32x8 as eight uint32s).

Consider the program I wrote to test out archsimd many months ago: link. Note: this predated the package rename, so simd in that code should be read as archsimd (confusing, I admit). That function is written using archsimd.Uint8x32 because that expresses the computation required by the spec for this encryption algorithm, but I don't want to write separate AVX512, AVX2, SSE, and scalar version of that algorithm. I think using build tags and/or inquiring about the length of a simd.Uint32s are strategies for me writing all those variants. I'd like the compiler to just generate those variants for me. Maybe you're saying that that is out-of-scope for the initial version of the portable simd package.

(ignoring mildly tricky AES situation later in that file)

That file does have 1 tricky aspect which really does seem to require AVX: [the Wide AES](https://github.com/balasanjay/aegis/blob/ae41f298cff5b5a615127c11ab2ac4979384466d/internal/impl/aegis128x2.go#L247). But even that could be emulated with the non-vector AES instructions. Or in the worst case, with a fallback to a software AES implementation. I could buy that a vector AES function specifically should be left out of the portable `simd` package and put elsewhere. But certainly all the simple functions (`Xor`, `And`, `LoadFromSlice`, etc) should be available in the portable package, without resorting to writing `archsimd` code directly.

[aclements]

@balasanjay , thanks for the example. Certainly for now this is out of scope. In particular, emulating fixed widths that are wider than the hardware supports would be quite a lot of work (doable, certainly, but not something we're going to tackle at first).

[aclements]

Pasting our discussion on possible approaches for going between simd and archsimd:

// Option 1: Force a type switch. Totally portable, no "unspeakable types",
// but the caller is going to need build tags anyway to speak the type switch.

func (x Int8s) ToArch() any

// What if we pass Int8x16 on an AVX-512 machine? Panic? Extend? Broadcast?
func Int8sFromArch(any) Int8s


// Option 2: Build-tag methods. Concrete, but annoying for docs. Lots of API
// that portable users don't care about.

//go:build amd64 || wasm
func (x Int8s) ToInt8x16() (archsimd.Int8x16, bool)
//go:build amd64
func (x Int8s) ToInt8x32() (archsimd.Int8x32, bool)
//go:build arm64 || riscv
func (x Int8s) ToInt8s() (archsimd.Int8s, bool) // A little weird!

// Option 2b: Use generics to fold a bunch of the above together.
func Int8sFromArch[T archsimd.Int8x16|archsimd.Int8x32|...](T) Int8s
func (x Int8s) ToArch[T archsimd.Int8x16|archsimd.Int8x32|...]() (T, bool)

[aclements]

To clarify @dr2chase 's sum examples, these show a reducing sum operations, not a vector-wise sum. The current prototype doesn't directly support any reduction operations because no hardware directly supports this. Reduction operations are useful enough we would probably want to add them, and we'd do that by adding emulations to archsimd so they wind up in the intersection of operations exposed in the simd package.

[aclements]

@dr2chase , could you post a complete API? It can be a go doc or a pointer to a CL if there's already a good summary somewhere (maybe mocks.go is that?)

[dr2chase]

Here is most of the API, at today's version: https://go.dev/cl/c/go/+/763100/16/src/simd/mocks.go

This is automatically generated from existing architecture-specific APIS. There may be some operations that we wish were there, and if so, say so, but the proposed mechanism for obtaining more operations is to add emulations to the architecture-specific APIs, which then increases the size of their intersection.

The missing part of the API is the conversions between simd and archsimd. I have a minor preference for a combination of option 1 and 2b; I haven't completely explored the FromArch method/function, but adapted the "sum" code to use the two different options. The answer to

// What if we pass Int8x16 on an AVX-512 machine? Panic? Extend? Broadcast?
func Int8sFromArch(any) Int8s

is "panic", if we don't like panic then it needs to return (Int8s, bool) and it might be slower, too.

As noted above, sum exists because we don't have a ReduceAdd method in the existing archsimd APIs, so it is also not in their intersection. I wrote it like this because it was an easy example of dropping from architecture-agnostic down to architecture-specific code, and I could compile it and run it and compare its behavior. All the type-wrangling vanishes after optimization; because the code is specifically recompiled for each particular vector width, all the type predicates become constants.

Option 1:

func sum(x simd.Float32s) float32 {
	switch a := x.ToArch().(type) {
	case archsimd.Float32x8:
		a = a.AddPairsGrouped(a)
		a = a.AddPairsGrouped(a)
		return a.GetLo().GetElem(0) + a.GetHi().GetElem(0)
	case archsimd.Float32x16:
		b := a.GetLo().Add(a.GetHi())
		b = b.AddPairsGrouped(b)
		b = b.AddPairsGrouped(b)
		return b.GetLo().GetElem(0) + b.GetHi().GetElem(0)
	}
	s := make([]float32, x.Len())
	x.StoreSlice(s)
	var r float32
	for _, e := range s {
		r += e
	}
	return r
}

Option 2b (which required enabling generic methods and then provoked a crash in type-checking, so not currently my favorite):

func sumMaybe(x simd.Float32s) float32 {
	if a, ok := x.ToArchMaybe[archsimd.Float32x8](); ok {
		a = a.AddPairsGrouped(a)
		a = a.AddPairsGrouped(a)
		return a.GetLo().GetElem(0) + a.GetHi().GetElem(0)
	}
	if b, ok := x.ToArchMaybe[archsimd.Float32x16](); ok {
		a := b.GetLo().Add(b.GetHi())
		a = a.AddPairsGrouped(a)
		a = a.AddPairsGrouped(a)
		return a.GetLo().GetElem(0) + a.GetHi().GetElem(0)
	}
	s := make([]float32, x.Len())
	x.StoreSlice(s)
	var r float32
	for _, e := range s {
		r += e
	}
	return r
}