internal/simdgen/gen_simdTypes.go

// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"slices"
	"sort"
	"strings"
)

type simdType struct {
	Name                    string // The go type name of this simd type, for example Int32x4.
	Lanes                   int    // The number of elements in this vector/mask.
	Base                    string // The element's type, like for Int32x4 it will be int32.
	Fields                  string // The struct fields, it should be right formatted.
	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
}

func compareSimdTypes(x, y simdType) int {
	c := strings.Compare(x.Name, y.Name)
	if c != 0 {
		return c
	}
	return strings.Compare(x.Type, y.Type)
}

type simdTypeMap map[int][]simdType

type simdTypePair struct {
	Tsrc simdType
	Tdst simdType
}

func compareSimdTypePairs(x, y simdTypePair) int {
	c := compareSimdTypes(x.Tsrc, y.Tsrc)
	if c != 0 {
		return c
	}
	return compareSimdTypes(x.Tdst, y.Tdst)
}

const simdTypesTemplates = `{{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.

//go:build goexperiment.simd

package simd
{{end}}

{{define "sizeTmpl"}}
// v{{.}} is a tag type that tells the compiler that this is really {{.}}-bit SIMD
type v{{.}} struct {
	_{{.}} struct{}
}
{{end}}

{{define "typeTmpl"}}
// {{.Name}} is a {{.Size}}-bit SIMD vector of {{.Lanes}} {{.Base}}
type {{.Name}} struct {
{{.Fields}}
}

{{- if ne .Type "mask"}}

// Len returns the number of elements in a {{.Name}}
func (x {{.Name}}) Len() int { return {{.Lanes}} }

// Load{{.Name}} loads a {{.Name}} from an array
//
//go:noescape
func Load{{.Name}}(y *[{{.Lanes}}]{{.Base}}) {{.Name}}

// Store stores a {{.Name}} to an array
//
//go:noescape
func (x {{.Name}}) Store(y *[{{.Lanes}}]{{.Base}})

{{- end}}
{{end}}
`

const simdStubsTmpl = `// 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.

//go:build goexperiment.simd

package simd

{{- range .OpsLen1}}

// Asm: {{.Asm}}, Arch: {{.Extension}}{{if .Documentation}}, Doc: {{.Documentation}}{{end}}
func (x {{(index .In 0).Go}}) {{.Go}}() {{(index .Out 0).Go}}

{{- end}}
{{- range .OpsLen2}}

// Asm: {{.Asm}}, Arch: {{.Extension}}{{if .Documentation}}, Doc: {{.Documentation}}{{end}}
func (x {{(index .In 0).Go}}) {{.Go}}(y {{(index .In 1).Go}}) {{(index .Out 0).Go}}

{{- end}}
{{- range .OpsLen3}}

// Asm: {{.Asm}}, Arch: {{.Extension}}{{if .Documentation}}, Doc: {{.Documentation}}{{end}}
func (x {{(index .In 0).Go}}) {{.Go}}(y {{(index .In 1).Go}}, z {{(index .In 2).Go}}) {{(index .Out 0).Go}}

{{- end}}
{{- range .VectorConversions }}

// {{.Tdst.Name}} converts from {{.Tsrc.Name}} to {{.Tdst.Name}}
func (from {{.Tsrc.Name}}) As{{.Tdst.Name}}() (to {{.Tdst.Name}})

{{- end}}
{{- range .Masks }}

// converts from {{.Name}} to {{.VectorCounterpart}}
func (from {{.Name}}) As{{.VectorCounterpart}}() (to {{.VectorCounterpart}})

// converts from {{.VectorCounterpart}} to {{.Name}}
func (from {{.VectorCounterpart}}) As{{.Name}}() (to {{.Name}})

func (x {{.Name}}) And(y {{.Name}}) {{.Name}}

func (x {{.Name}}) Or(y {{.Name}}) {{.Name}}

{{- end}}
`

// parseSIMDTypes groups go simd types by their vector sizes, and
// returns a map whose key is the vector size, value is the simd type.
func parseSIMDTypes(ops []Operation) simdTypeMap {
	// TODO: maybe instead of going over ops, let's try go over types.yaml.
	ret := map[int][]simdType{}
	seen := map[string]struct{}{}
	processArg := func(arg Operand) {
		if arg.Class == "immediate" {
			// Immediates are not encoded as vector types.
			return
		}
		if _, ok := seen[*arg.Go]; ok {
			return
		}
		seen[*arg.Go] = struct{}{}
		lanes := *arg.Bits / *arg.ElemBits
		base := fmt.Sprintf("%s%d", *arg.Base, *arg.ElemBits)
		tagFieldNameS := fmt.Sprintf("%sx%d", base, lanes)
		tagFieldS := fmt.Sprintf("%s v%d", tagFieldNameS, *arg.Bits)
		valFieldS := fmt.Sprintf("vals%s[%d]%s", strings.Repeat(" ", len(tagFieldNameS)-3), lanes, base)
		fields := fmt.Sprintf("\t%s\n\t%s", tagFieldS, valFieldS)
		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})
			// 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})
			}
		} else {
			ret[*arg.Bits] = append(ret[*arg.Bits], simdType{*arg.Go, lanes, base, fields, arg.Class, "", "", *arg.Bits})
		}
	}
	for _, op := range ops {
		for _, arg := range op.In {
			processArg(arg)
		}
		for _, arg := range op.Out {
			processArg(arg)
		}
	}
	return ret
}

func vConvertFromTypeMap(typeMap simdTypeMap) []simdTypePair {
	v := []simdTypePair{}
	for _, ts := range typeMap {
		for i, tsrc := range ts {
			for j, tdst := range ts {
				if i != j && tsrc.Type == tdst.Type && tsrc.Type == "vreg" {
					v = append(v, simdTypePair{tsrc, tdst})
				}
			}
		}
	}
	slices.SortFunc(v, compareSimdTypePairs)
	return v
}

func masksFromTypeMap(typeMap simdTypeMap) []simdType {
	m := []simdType{}
	for _, ts := range typeMap {
		for _, tsrc := range ts {
			if tsrc.Type == "mask" {
				m = append(m, tsrc)
			}
		}
	}
	slices.SortFunc(m, compareSimdTypes)
	return m
}

// writeSIMDTypes generates the simd vector type and writes it to types_amd64.go
// within the specified directory.
func writeSIMDTypes(directory string, typeMap simdTypeMap) error {
	file, t, err := openFileAndPrepareTemplate(directory, "src/"+simdPackage+"/types_amd64.go", simdTypesTemplates)
	if err != nil {
		return err
	}
	defer file.Close()

	if err := t.ExecuteTemplate(file, "fileHeader", nil); err != nil {
		return fmt.Errorf("failed to execute fileHeader template: %w", err)
	}

	sizes := make([]int, 0, len(typeMap))
	for size := range typeMap {
		sizes = append(sizes, size)
	}
	sort.Ints(sizes)

	for _, size := range sizes {
		if err := t.ExecuteTemplate(file, "sizeTmpl", size); err != nil {
			return fmt.Errorf("failed to execute size template for size %d: %w", size, err)
		}
		for _, typeDef := range typeMap[size] {
			if err := t.ExecuteTemplate(file, "typeTmpl", typeDef); err != nil {
				return fmt.Errorf("failed to execute type template for type %s: %w", typeDef.Name, err)
			}
		}
	}

	return nil
}

// writeSIMDStubs generates the simd vector intrinsic stubs and writes it to stubs_amd64.go
// within the specified directory.
func writeSIMDStubs(directory string, ops []Operation, typeMap simdTypeMap) error {
	file, t, err := openFileAndPrepareTemplate(directory, "src/"+simdPackage+"/stubs_amd64.go", simdStubsTmpl)
	if err != nil {
		return err
	}
	defer file.Close()
	opsLen1, opsLen2, opsLen3, err := genericOpsByLen(ops)
	if err != nil {
		return err
	}

	type templateData struct {
		OpsLen1           []Operation
		OpsLen2           []Operation
		OpsLen3           []Operation
		VectorConversions []simdTypePair
		Masks             []simdType
	}

	err = t.Execute(file, templateData{opsLen1, opsLen2, opsLen3, vConvertFromTypeMap(typeMap), masksFromTypeMap(typeMap)})
	if err != nil {
		return fmt.Errorf("failed to execute template : %w", err)
	}

	return nil
}