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"
	"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
}

type simdTypeMap map[int][]simdType

type simdTypePair struct {
	Tsrc simdType
	Tdst simdType
}

const simdTypesTmpl = `// 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 $size, $ts := .TypeMap }}

// v{{$size}} is a tag type that tells the compiler that this is really {{$size}}-bit SIMD
type v{{$size}} struct {
	_{{$size}} struct{}
}

{{- range $i, $tsrc := $ts }}

// {{$tsrc.Name}} is a {{$size}}-bit SIMD vector of {{$tsrc.Lanes}} {{$tsrc.Base}}
type {{$tsrc.Name}} struct {
{{$tsrc.Fields}}
}

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

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

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

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

{{- end}}
{{- 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})
				}
			}
		}
	}
	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)
			}
		}
	}
	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", simdTypesTmpl)
	if err != nil {
		return err
	}
	defer file.Close()

	type templateData struct {
		TypeMap simdTypeMap
	}

	err = t.Execute(file, templateData{typeMap})
	if err != nil {
		return fmt.Errorf("failed to execute template: %w", 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
}