gen_simdIntrinsics.go

     1  // Copyright 2025 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"simd/archsimd/_gen/sgutil"
    11  	"slices"
    12  	"text/template"
    13  )
    14  
    15  // Helper type to make template map initialization less repetitive
    16  // (and also remove a chance for errors.)
    17  type intrinsicTemplateMap struct {
    18  	sgutil.InsertMap[string, *template.Template]
    19  }
    20  
    21  func templateNamed(name string, templ string) *template.Template {
    22  	// Append  end of line
    23  	templ += "\n"
    24  
    25  	t := template.New(name)
    26  
    27  	archInfo := CurrentArch()
    28  	sysArch := "sys." + archInfo.ArchUpper
    29  
    30  	t.Funcs(template.FuncMap{
    31  		"GetSysArch": func() string {
    32  			return sysArch
    33  		},
    34  		"GetArchUpper": func() string {
    35  			return archInfo.ArchUpper
    36  		},
    37  		"Hasmask": func() bool {
    38  			return archInfo.Arch == "amd64"
    39  		},
    40  	})
    41  
    42  	return template.Must(t.Parse(templ))
    43  }
    44  
    45  // Add creates a template named "name" after appending "\n" to the
    46  // template, and returns the input so that additions may be chained.
    47  // This helps make template initialization easy to order and easy to read.
    48  func (rtm *intrinsicTemplateMap) Add(name string, templ string) *intrinsicTemplateMap {
    49  
    50  	rtm.InsertMap.Put(name, templateNamed(name, templ))
    51  	return rtm
    52  }
    53  
    54  // writeSIMDIntrinsics generates the intrinsic mappings and writes it to simdintrinsics.go
    55  // within the specified directory.
    56  func writeSIMDIntrinsics(ops []Operation, typeMap simdTypeMap) *bytes.Buffer {
    57  
    58  	// These are defined here to avoid init-order problems with GetSysArch GetArchUpper etc which depend on flag values
    59  
    60  	var header = templateNamed("header", `package ssagen
    61  
    62  import (
    63  	"cmd/compile/internal/ir"
    64  	"cmd/compile/internal/ssa"
    65  	"cmd/compile/internal/types"
    66  	"cmd/internal/sys"
    67  )
    68  
    69  func simd{{GetArchUpper}}Intrinsics(addF func(pkg, fn string, b intrinsicBuilder, archFamilies ...sys.ArchFamily)) {
    70  `)
    71  
    72  	var intrinsicTemplates = new(intrinsicTemplateMap).
    73  		Add("op1", `		addF(simdPackage, "{{(index .In 0).Go}}.{{.Go}}", opLen1(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    74  		Add("op2", `		addF(simdPackage, "{{(index .In 0).Go}}.{{.Go}}", opLen2(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    75  		Add("op2_21", `		addF(simdPackage, "{{(index .In 0).Go}}.{{.Go}}", opLen2_21(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    76  		Add("op2_21Type1", `addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2_21(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    77  		Add("op3", `		addF(simdPackage, "{{(index .In 0).Go}}.{{.Go}}", opLen3(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    78  		Add("op3_21", `		addF(simdPackage, "{{(index .In 0).Go}}.{{.Go}}", opLen3_21(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    79  		Add("op3_21Type1", `addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen3_21(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    80  		Add("op3_231Type1", `addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen3_231(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    81  		Add("op3_31Zero3", `addF(simdPackage, "{{(index .In 2).Go}}.{{.Go}}", opLen3_31Zero3(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    82  		Add("op4", `		addF(simdPackage, "{{(index .In 0).Go}}.{{.Go}}", opLen4(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    83  		Add("op4_231Type1", `addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen4_231(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    84  		Add("op4_31", `		addF(simdPackage, "{{(index .In 2).Go}}.{{.Go}}", opLen4_31(ssa.Op{{.GenericName}}, {{.SSAType}}), {{GetSysArch}})`).
    85  		Add("op1Imm", `		addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen1Imm(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}, {{(index .In 0).ImmMax}}), {{GetSysArch}})`).
    86  		Add("op1Imm8", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen1Imm8(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    87  		Add("op2Imm", `		addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2Imm(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}, {{(index .In 0).ImmMax}}), {{GetSysArch}})`).
    88  		Add("op2Imm8", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2Imm8(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    89  		Add("op2Imm8_2I", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2Imm8_2I(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    90  		Add("op2Imm_2I", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2Imm_2I(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}, {{(index .In 0).ImmMax}}), {{GetSysArch}})`).
    91  		Add("op2Imm8_II", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2Imm8_II(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    92  		Add("op2Imm8_SHA1RNDS4", `addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen2Imm8_SHA1RNDS4(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    93  		Add("op2ImmVecAsScalar", `addF(simdPackage, "{{(index .In 2).Go}}.{{.Go}}", opLen2Imm(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}, {{(index .In 0).ImmMax}}), {{GetSysArch}})`).
    94  		Add("op3Imm8", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen3Imm8(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    95  		Add("op3Imm8_2I", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen3Imm8_2I(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`).
    96  		Add("op4Imm8", `	addF(simdPackage, "{{(index .In 1).Go}}.{{.Go}}", opLen4Imm8(ssa.Op{{.GenericName}}, {{.SSAType}}, {{(index .In 0).ImmOffset}}), {{GetSysArch}})`)
    97  
    98  	var loadStore = templateNamed("loadStore", `	addF(simdPackage, "Load{{.Name}}Array", simdLoad(), {{GetSysArch}})
    99  	addF(simdPackage, "{{.Name}}.StoreArray", simdStore(), {{GetSysArch}})`)
   100  
   101  	var mask = templateNamed("mask", `	addF(simdPackage, "{{.Name}}.To{{.VectorCounterpart}}", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, {{GetSysArch}})
   102  	addF(simdPackage, "{{.VectorCounterpart}}.asMask", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, {{GetSysArch}})
   103  	addF(simdPackage, "{{.Name}}.And", opLen2(ssa.OpAnd{{.ReshapedVectorWithAndOr}}, types.TypeVec{{.Size}}), {{GetSysArch}})
   104  	addF(simdPackage, "{{.Name}}.Or", opLen2(ssa.OpOr{{.ReshapedVectorWithAndOr}}, types.TypeVec{{.Size}}), {{GetSysArch}})
   105  {{- if eq GetSysArch "sys.ARM64"}}
   106  	addF(simdPackage, "{{.Name}}.Not", opLen1(ssa.OpNot{{.ReshapedVectorWithAndOr}}, types.TypeVec{{.Size}}), {{GetSysArch}})
   107  {{- else}}
   108  	addF(simdPackage, "{{.Name}}FromBits", simdCvtVToMask({{.ElemBits}}, {{.Lanes}}), {{GetSysArch}})
   109  	addF(simdPackage, "{{.Name}}.ToBits", simdCvtMaskToV({{.ElemBits}}, {{.Lanes}}), {{GetSysArch}})
   110  {{- end}}`)
   111  
   112  	var maskedLoadStore = templateNamed("maskedLoadStore", `	addF(simdPackage, "{{.Name}}.StoreArrayMasked", simdMaskedStore(ssa.OpStoreMasked{{.ElemBits}}), sys.AMD64)`)
   113  
   114  	var vectorConversion = templateNamed("vectorConversion", `	addF(simdPackage, "{{.Tsrc.Name}}.As{{.Tdst.Name}}", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, {{GetSysArch}})`)
   115  
   116  	var footer = `}`
   117  
   118  	slices.SortFunc(ops, compareOperations)
   119  
   120  	buffer := new(bytes.Buffer)
   121  	buffer.WriteString(generatedHeader())
   122  
   123  	doTemplate := func(tpl *template.Template, data any) {
   124  		if err := tpl.Execute(buffer, data); err != nil {
   125  			panic(fmt.Errorf("failed to execute template %s: %w", tpl.Name(), err))
   126  		}
   127  	}
   128  
   129  	doTemplate(header, nil)
   130  
   131  	doIntrinsic := func(name string, data any) {
   132  		tpl := intrinsicTemplates.Get(name)
   133  		if tpl == nil {
   134  			panic(fmt.Errorf("template %s not found", name))
   135  		}
   136  		doTemplate(tpl, data)
   137  	}
   138  
   139  	for _, op := range ops {
   140  		if op.NoTypes != nil && *op.NoTypes == "true" {
   141  			continue
   142  		}
   143  		if op.SkipMaskedMethod() {
   144  			continue
   145  		}
   146  		// Cannot have an intrinsic w/o generics, at least for now.
   147  		if op.NoGenericOps != nil && *op.NoGenericOps == "true" {
   148  			continue
   149  		}
   150  		if s, op, err := classifyOp(op); err == nil {
   151  			if s == "op2Imm" {
   152  				idxVecAsScalar, err := checkVecAsScalar(op)
   153  				if err != nil {
   154  					panic(err)
   155  				}
   156  				if idxVecAsScalar >= 0 {
   157  					s += "VecAsScalar"
   158  				}
   159  			}
   160  			doIntrinsic(s, op)
   161  		} else {
   162  			panic(fmt.Errorf("failed to classify op %v: %w", op.Go, err))
   163  		}
   164  	}
   165  
   166  	var TypeDotMethodIntrinsicAMD64 = templateOf(`addF(simdPackage, "{{.TypeDotMethod}}", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, sys.AMD64)
   167  	`, "amd64 type dot method intrinsics")
   168  
   169  	var TypeDotMethodIntrinsicARM64 = templateOf(`addF(simdPackage, "{{.TypeDotMethod}}", func(s *state, n *ir.CallExpr, args []*ssa.Value) *ssa.Value { return args[0] }, sys.ARM64)
   170  	`, "arm64 type dot method intrinsics")
   171  
   172  	for _, conv := range vConvertFromTypeMap(typeMap) {
   173  		// Old As intrinsic
   174  		from, to := &conv.Tsrc, &conv.Tdst
   175  		doTemplate(vectorConversion, conv)
   176  
   177  		// New style factored conversion intrinsics always involve at least one unsigned type
   178  		if from.Name[0] != 'U' && to.Name[0] != 'U' {
   179  			continue
   180  		}
   181  		// Only emit the intrinsic if lanes are equal OR both are unsigned
   182  		if from.Lanes != to.Lanes && (from.Name[0] != 'U' || to.Name[0] != 'U') {
   183  			continue
   184  		}
   185  		var typeDotMethodIntrinsic *template.Template
   186  		switch CurrentArch().Arch {
   187  		case "amd64":
   188  			typeDotMethodIntrinsic = TypeDotMethodIntrinsicAMD64
   189  		case "arm64":
   190  			typeDotMethodIntrinsic = TypeDotMethodIntrinsicARM64
   191  		default:
   192  			panic(fmt.Errorf("unsupported arch %q for type dot method intrinsics", CurrentArch().Arch))
   193  		}
   194  		sgutil.Conversion(from, to).ExecuteIntrinsicTemplateOfTypeDotMethod(buffer, typeDotMethodIntrinsic)
   195  	}
   196  
   197  	for _, typ := range typesFromTypeMap(typeMap) {
   198  		if typ.Type != "mask" {
   199  			loadStore.Execute(buffer, typ)
   200  		}
   201  	}
   202  
   203  	// Masked loads/stores are AVX2/AVX512 only (not available on ARM64 NEON).
   204  	// TODO: Reconsider for ARM64 SVE which supports predicated loads/stores.
   205  	if CurrentArch().Arch == "amd64" {
   206  		for _, typ := range typesFromTypeMap(typeMap) {
   207  			if typ.MaskedLoadStoreFilter() {
   208  				doTemplate(maskedLoadStore, typ)
   209  			}
   210  		}
   211  	}
   212  
   213  	for _, m := range masksFromTypeMap(typeMap) {
   214  		doTemplate(mask, m)
   215  	}
   216  
   217  	buffer.WriteString(footer)
   218  
   219  	return buffer
   220  }
   221
View as plain text