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

const simdssaTmpl = `// 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 amd64

import (
	"cmd/compile/internal/ssa"
	"cmd/compile/internal/ssagen"
	"cmd/internal/obj"
	"cmd/internal/obj/x86"
)

func ssaGenSIMDValue(s *ssagen.State, v *ssa.Value) bool {
	p := s.Prog(v.Op.Asm())
	// First arg
	switch v.Op {{"{"}}{{if gt (len .Imms) 0}}
	// Immediates
	case {{.Imms}}:
		imm := v.AuxInt
		if imm < 0 || imm > 255 {
			v.Fatalf("Invalid source selection immediate")
		}
		p.From.Offset = imm
		p.From.Type = obj.TYPE_CONST
{{end}}{{if gt (len .Reg0) 0}}
	// Registers
	case {{.Reg0}}:
		p.From.Type = obj.TYPE_REG
		p.From.Reg = simdReg(v.Args[0])
{{end}}
	default:
		// At least one arg is required.
		return false
	}

	// Second arg
	switch v.Op {{"{"}}{{if gt (len .Reg1) 0}}
	// Registers
	case {{.Reg1}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(simdReg(v.Args[0]))
		} else {
			p.AddRestSourceReg(simdReg(v.Args[1]))
		}{{end}}
	}

	// Third arg
	switch v.Op {{"{"}}{{if gt (len .Reg2) 0}}
	// Registers
	case {{.Reg2}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(simdReg(v.Args[1]))
		} else {
			p.AddRestSourceReg(simdReg(v.Args[2]))
		}{{end}}
	}

	// Fourth arg
	switch v.Op {{"{"}}{{if gt (len .Reg3) 0}}
	case {{.Reg3}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(simdReg(v.Args[2]))
		} else {
			p.AddRestSourceReg(simdReg(v.Args[3]))
		}{{end}}
	}

	// Output
	switch v.Op {{"{"}}{{if gt (len .All) 0}}
	case {{.All}}:
		p.To.Type = obj.TYPE_REG
		p.To.Reg = simdReg(v)
{{end}}
	default:
		// One result is required.
		return false
	}
{{if gt (len .ZeroingMask) 0}}
	// Masked operation are always compiled with zeroing.
	switch v.Op {
	case {{.ZeroingMask}}:
		x86.ParseSuffix(p, "Z")
	}
{{end}}
	return true
}
`

// writeSIMDSSA generates the ssa to prog lowering codes and writes it to simdssa.go
// within the specified directory.
func writeSIMDSSA(directory string, ops []Operation) error {
	var Imms []string
	var All []string
	var ZeroingMask []string
	Regs := map[int][]string{}

	seen := map[string]struct{}{}
	for _, op := range ops {
		asm := op.Asm
		shapeIn, _, maskType, _, _, gOp, err := op.shape()
		if err != nil {
			return err
		}
		if maskType == 2 {
			asm += "Masked"
		}
		asm = fmt.Sprintf("%s%d", asm, *gOp.Out[0].Bits)
		if _, ok := seen[asm]; ok {
			continue
		}
		seen[asm] = struct{}{}
		caseStr := fmt.Sprintf("ssa.OpAMD64%s", asm)
		if shapeIn == OneKmaskIn || shapeIn == OneKmaskConstImmIn {
			if gOp.Zeroing == nil {
				ZeroingMask = append(ZeroingMask, caseStr)
			}
		}
		immCount := 0
		if shapeIn == OneConstImmIn || shapeIn == OneKmaskConstImmIn {
			immCount++
			Imms = append(Imms, caseStr)
		}
		for i := range len(gOp.In) {
			if i > 2 {
				return fmt.Errorf("simdgen does not recognize more than 3 registers: %s", gOp)
			}
			Regs[i+immCount] = append(Regs[i+immCount], caseStr)
		}
		All = append(All, caseStr)
	}

	data := struct {
		Imms        string
		Reg0        string
		Reg1        string
		Reg2        string
		Reg3        string
		All         string
		ZeroingMask string
	}{
		strings.Join(Imms, ",\n\t\t"),
		strings.Join(Regs[0], ",\n\t\t"),
		strings.Join(Regs[1], ",\n\t\t"),
		strings.Join(Regs[2], ",\n\t\t"),
		strings.Join(Regs[3], ",\n\t\t"),
		strings.Join(All, ",\n\t\t"),
		strings.Join(ZeroingMask, ",\n\t\t"),
	}

	file, t, err := openFileAndPrepareTemplate(directory, "src/cmd/compile/internal/amd64/simdssa.go", simdssaTmpl)
	if err != nil {
		return err
	}
	defer file.Close()

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

	return nil
}