internal/simdgen/gen_simdMachineOps.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"
	"sort"
)

const simdMachineOpsTmpl = `// 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.
package main

func simdAMD64Ops(fp1fp1, fp2fp1, fp2m1, fp2m1fp1, fp2m1m1 regInfo) []opData {
	return []opData{
{{- range .OpsData }}
		{name: "{{.OpName}}", argLength: {{.OpInLen}}, reg: {{.RegInfo}}, asm: "{{.Asm}}", commutative: {{.Comm}}, typ: "{{.Type}}"},
{{- end }}
{{- range .OpsDataImm }}
		{name: "{{.OpName}}", argLength: {{.OpInLen}}, reg: {{.RegInfo}}, asm: "{{.Asm}}", aux: "Int8", commutative: {{.Comm}}, typ: "{{.Type}}"},
{{- end }}
	}
}
`

// writeSIMDMachineOps generates the machine ops and writes it to simdAMD64ops.go
// within the specified directory.
func writeSIMDMachineOps(directory string, ops []Operation) error {
	file, t, err := openFileAndPrepareTemplate(directory, "src/cmd/compile/internal/ssa/_gen/simdAMD64ops.go", simdMachineOpsTmpl)
	if err != nil {
		return err
	}
	defer file.Close()
	type opData struct {
		sortKey string
		OpName  string
		Asm     string
		OpInLen int
		RegInfo string
		Comm    string
		Type    string
	}
	type machineOpsData struct {
		OpsData    []opData
		OpsDataImm []opData
	}
	seen := map[string]struct{}{}
	regInfoSet := map[string]bool{"fp1fp1": true, "fp2fp1": true, "fp2m1": true, "fp2m1fp1": true, "fp2m1m1": true}
	opsData := make([]opData, 0)
	opsDataImm := make([]opData, 0)
	for _, op := range ops {
		shapeIn, shapeOut, maskType, _, _, gOp, err := op.shape()
		if err != nil {
			return err
		}
		asm := gOp.Asm
		if maskType == OneMask {
			asm += "Masked"
		}
		asm = fmt.Sprintf("%s%d", asm, *gOp.Out[0].Bits)
		// TODO: all our masked operations are now zeroing, we need to generate machine ops with merging masks, maybe copy
		// one here with a name suffix "Merging". The rewrite rules will need them.
		if _, ok := seen[asm]; ok {
			continue
		}
		seen[asm] = struct{}{}
		var regInfo string
		// Process input reg shapes.
		var vRegInCnt, kMaskInCnt, vRegOutCnt, kMaskOutCnt int
		for _, in := range gOp.In {
			if in.Class == "vreg" {
				vRegInCnt++
			} else if in.Class == "mask" {
				kMaskInCnt++
			}
		}
		for _, out := range gOp.Out {
			// If class overwrite is happening, that's not really a mask but a vreg.
			if out.Class == "vreg" || out.OverwriteClass != nil {
				vRegOutCnt++
			} else if out.Class == "mask" {
				kMaskOutCnt++
			}
		}
		var vRegInS, kMaskInS, vRegOutS, kMaskOutS string
		if vRegInCnt > 0 {
			vRegInS = fmt.Sprintf("fp%d", vRegInCnt)
		}
		if kMaskInCnt > 0 {
			kMaskInS = fmt.Sprintf("m%d", kMaskInCnt)
		}
		if vRegOutCnt > 0 {
			vRegOutS = fmt.Sprintf("fp%d", vRegOutCnt)
		}
		if kMaskOutCnt > 0 {
			kMaskOutS = fmt.Sprintf("m%d", kMaskOutCnt)
		}
		regInfo = fmt.Sprintf("%s%s%s%s", vRegInS, kMaskInS, vRegOutS, kMaskOutS)
		if _, ok := regInfoSet[regInfo]; !ok {
			return fmt.Errorf("unsupported register constraint, please update the template and AMD64Ops.go: %s", regInfo)
		}
		var outType string
		if shapeOut == OneVregOut || gOp.Out[0].OverwriteClass != nil {
			// If class overwrite is happening, that's not really a mask but a vreg.
			outType = fmt.Sprintf("Vec%d", *gOp.Out[0].Bits)
		} else if shapeOut == OneKmaskOut {
			outType = "Mask"
		} else {
			return fmt.Errorf("simdgen does not recognize this output shape: %+v", shapeOut)
		}
		if shapeIn == OneConstImmIn || shapeIn == OneKmaskConstImmIn {
			opsDataImm = append(opsDataImm, opData{*gOp.In[0].Go + gOp.Go, asm, gOp.Asm, len(gOp.In), regInfo, gOp.Commutative, outType})
		} else {
			opsData = append(opsData, opData{*gOp.In[0].Go + gOp.Go, asm, gOp.Asm, len(gOp.In), regInfo, gOp.Commutative, outType})
		}
	}
	sort.Slice(opsData, func(i, j int) bool {
		return opsData[i].sortKey < opsData[j].sortKey
	})
	sort.Slice(opsDataImm, func(i, j int) bool {
		return opsDataImm[i].sortKey < opsDataImm[j].sortKey
	})
	err = t.Execute(file, machineOpsData{opsData, opsDataImm})
	if err != nil {
		return fmt.Errorf("failed to execute template: %w", err)
	}

	return nil
}