internal/simdgen: add generated declarations for AVX2 masked load/store

generates Go dev.simd CL 689335 (which also includes one basic test)

Change-Id: Icd948396a3ca265b307747437efbc0e6f4548c76
Reviewed-on: https://go-review.googlesource.com/c/arch/+/689276
Reviewed-by: Junyang Shao <shaojunyang@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/internal/simdgen/gen_simdIntrinsics.go b/internal/simdgen/gen_simdIntrinsics.go
index 7140eda..d114b4c 100644
--- a/internal/simdgen/gen_simdIntrinsics.go
+++ b/internal/simdgen/gen_simdIntrinsics.go
@@ -69,6 +69,10 @@
 	addF(simdPackage, "{{.Name}}.Store", simdStore(), sys.AMD64)
 {{end}}
 
+{{define "maskedLoadStore"}}	addF(simdPackage, "LoadMasked{{.Name}}", simdMaskedLoad(ssa.OpLoadMasked{{.ElemBits}}), sys.AMD64)
+	addF(simdPackage, "{{.Name}}.StoreMasked", simdMaskedStore(ssa.OpStoreMasked{{.ElemBits}}), sys.AMD64)
+{{end}}
+
 {{define "mask"}}	addF(simdPackage, "{{.Name}}.As{{.VectorCounterpart}}", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, sys.AMD64)
 	addF(simdPackage, "{{.VectorCounterpart}}.As{{.Name}}", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, sys.AMD64)
 	addF(simdPackage, "{{.Name}}.And", opLen2(ssa.OpAnd{{.ReshapedVectorWithAndOr}}, types.TypeVec{{.Size}}), sys.AMD64)
@@ -118,6 +122,14 @@
 		}
 	}
 
+	for _, typ := range typesFromTypeMap(typeMap) {
+		if typ.MaskedLoadStoreFilter() {
+			if err := t.ExecuteTemplate(buffer, "maskedLoadStore", typ); err != nil {
+				panic(fmt.Errorf("failed to execute maskedLoadStore template: %w", err))
+			}
+		}
+	}
+
 	for _, mask := range masksFromTypeMap(typeMap) {
 		if err := t.ExecuteTemplate(buffer, "mask", mask); err != nil {
 			panic(fmt.Errorf("failed to execute mask template: %w", err))
diff --git a/internal/simdgen/gen_simdTypes.go b/internal/simdgen/gen_simdTypes.go
index 98d2b5a..6739b9f 100644
--- a/internal/simdgen/gen_simdTypes.go
+++ b/internal/simdgen/gen_simdTypes.go
@@ -20,8 +20,18 @@
 	Type                    string // Either "mask" or "vreg"
 	VectorCounterpart       string // For mask use only: just replacing the "Mask" in [simdType.Name] with "Int"
 	ReshapedVectorWithAndOr string // For mask use only: vector AND and OR are only available in some shape with element width 32.
-	Size                    int    // The size of the type
-	ElemBits                int    // Size / Lanes
+	Size                    int    // The size of the vector type
+}
+
+func (x simdType) ElemBits() int {
+	return x.Size / x.Lanes
+}
+
+// MaskedLoadStoreFilter encodes which simd type type currently
+// get masked loads/stores generated, it is used in two places,
+// this forces coordination.
+func (x simdType) MaskedLoadStoreFilter() bool {
+	return x.Size < 512 && x.ElemBits() >= 32 && x.Type != "mask"
 }
 
 func compareSimdTypes(x, y simdType) int {
@@ -36,7 +46,7 @@
 		return c
 	}
 	// base type size, 8 < 16 < 32 < 64
-	if c := x.Size/x.Lanes - y.Size/y.Lanes; c != 0 {
+	if c := x.ElemBits() - y.ElemBits(); c != 0 {
 		return c
 	}
 	// vector size last
@@ -78,8 +88,10 @@
 {{.Fields}}
 }
 
-{{- if ne .Type "mask"}}
+{{end}}
+`
 
+const simdLoadStoreTemplate = `
 // Len returns the number of elements in a {{.Name}}
 func (x {{.Name}}) Len() int { return {{.Lanes}} }
 
@@ -92,9 +104,9 @@
 //
 //go:noescape
 func (x {{.Name}}) Store(y *[{{.Lanes}}]{{.Base}})
+`
 
-{{- else}}
-
+const simdMaskFromBitsTemplate = `
 // Load{{.Name}}FromBits constructs a {{.Name}} from a bitmap, where 1 means set for the indexed element, 0 means unset.
 // Only the lower {{.Lanes}} bits of y are used.
 //
@@ -108,9 +120,20 @@
 // CPU Features: AVX512
 //go:noescape
 func (x {{.Name}}) StoreToBits(y *uint64)
+`
 
-{{end}}
-{{end}}
+const simdMaskedLoadStoreTemplate = `
+// LoadMasked{{.Name}} loads a {{.Name}} from an array,
+// at those elements enabled by mask
+//
+//go:noescape
+func LoadMasked{{.Name}}(y *[{{.Lanes}}]{{.Base}}, mask Mask{{.ElemBits}}x{{.Lanes}}) {{.Name}}
+
+// StoreMasked stores a {{.Name}} to an array,
+// at those elements enabled by mask
+//
+//go:noescape
+func (x {{.Name}}) StoreMasked(y *[{{.Lanes}}]{{.Base}}, mask Mask{{.ElemBits}}x{{.Lanes}})
 `
 
 const simdStubsTmpl = `{{define "fileHeader"}}// Code generated by x/arch/internal/simdgen using 'go run . -xedPath $XED_PATH -o godefs -goroot $GOROOT go.yaml types.yaml categories.yaml'; DO NOT EDIT.
@@ -317,14 +340,14 @@
 		if arg.Class == "mask" {
 			vectorCounterpart := strings.ReplaceAll(*arg.Go, "Mask", "Int")
 			reshapedVectorWithAndOr := fmt.Sprintf("Int32x%d", *arg.Bits/32)
-			ret[*arg.Bits] = append(ret[*arg.Bits], simdType{*arg.Go, lanes, base, fields, arg.Class, vectorCounterpart, reshapedVectorWithAndOr, *arg.Bits, *arg.Bits / lanes})
+			ret[*arg.Bits] = append(ret[*arg.Bits], simdType{*arg.Go, lanes, base, fields, arg.Class, vectorCounterpart, reshapedVectorWithAndOr, *arg.Bits})
 			// In case the vector counterpart of a mask is not present, put its vector counterpart typedef into the map as well.
 			if _, ok := seen[vectorCounterpart]; !ok {
 				seen[vectorCounterpart] = struct{}{}
-				ret[*arg.Bits] = append(ret[*arg.Bits], simdType{vectorCounterpart, lanes, base, fields, "vreg", "", "", *arg.Bits, *arg.Bits / lanes})
+				ret[*arg.Bits] = append(ret[*arg.Bits], simdType{vectorCounterpart, lanes, base, fields, "vreg", "", "", *arg.Bits})
 			}
 		} else {
-			ret[*arg.Bits] = append(ret[*arg.Bits], simdType{*arg.Go, lanes, base, fields, arg.Class, "", "", *arg.Bits, *arg.Bits / lanes})
+			ret[*arg.Bits] = append(ret[*arg.Bits], simdType{*arg.Go, lanes, base, fields, arg.Class, "", "", *arg.Bits})
 		}
 	}
 	for _, op := range ops {
@@ -383,6 +406,10 @@
 // writeSIMDTypes generates the simd vector types into a bytes.Buffer
 func writeSIMDTypes(typeMap simdTypeMap) *bytes.Buffer {
 	t := templateOf(simdTypesTemplates, "types_amd64")
+	loadStore := templateOf(simdLoadStoreTemplate, "loadstore_amd64")
+	maskedLoadStore := templateOf(simdMaskedLoadStoreTemplate, "maskedloadstore_amd64")
+	maskFromBits := templateOf(simdMaskFromBitsTemplate, "maskFromBits_amd64")
+
 	buffer := new(bytes.Buffer)
 
 	if err := t.ExecuteTemplate(buffer, "fileHeader", nil); err != nil {
@@ -411,6 +438,21 @@
 			if err := t.ExecuteTemplate(buffer, "typeTmpl", typeDef); err != nil {
 				panic(fmt.Errorf("failed to execute type template for type %s: %w", typeDef.Name, err))
 			}
+			if typeDef.Type != "mask" {
+				if err := loadStore.ExecuteTemplate(buffer, "loadstore_amd64", typeDef); err != nil {
+					panic(fmt.Errorf("failed to execute loadstore template for type %s: %w", typeDef.Name, err))
+				}
+				// restrict to AVX2 masked loads/stores first.
+				if typeDef.MaskedLoadStoreFilter() {
+					if err := maskedLoadStore.ExecuteTemplate(buffer, "maskedloadstore_amd64", typeDef); err != nil {
+						panic(fmt.Errorf("failed to execute maskedloadstore template for type %s: %w", typeDef.Name, err))
+					}
+				}
+			} else {
+				if err := maskFromBits.ExecuteTemplate(buffer, "maskFromBits_amd64", typeDef); err != nil {
+					panic(fmt.Errorf("failed to execute maskFromBits template for type %s: %w", typeDef.Name, err))
+				}
+			}
 		}
 	}