@ttsc/win32-arm64 0.33.1 → 0.34.0

package/bin/go/src/simd/archsimd/_gen/simdgen/gen_simdrules.go

@@ -10,10 +10,13 @@ import (
 	"slices"
 	"strings"
 	"text/template"
+	"unicode"
+
+	"simd/archsimd/_gen/sgutil"
 )
 
 type tplRuleData struct {
-	tplName        string // e.g. "sftimm"
+	TplName        string // e.g. "sftimm"
 	GoOp           string // e.g. "ShiftAllLeft"
 	GoType         string // e.g. "Uint32x8"
 	Args           string // e.g. "x y"
@@ -26,27 +29,46 @@ type tplRuleData struct {
 	ArgsLoadAddr   string // [Args] with its last vreg arg being a concrete "(VMOVDQUload* ptr mem)", and might contain mask.
 	ArgsAddr       string // [Args] with its last vreg arg being replaced by "ptr", and might contain mask, and with a "mem" at the end.
 	FeatCheck      string // e.g. "v.Block.CPUfeatures.hasFeature(CPUavx512)" -- for a ssa/_gen rules file.
+	RuleCond       string // e.g. "a==0" -- condition for asmRule or argsMatchRule.
+	RuleOut        string // e.g. "y" -- output of an asmRule or argsMatchRule.
+	RuleArgs       string // custom args pattern for argsMatchRule.
+	Rule           string // the whole rule
+}
+
+// Helper type to make template map initialization less repetitive
+// (and also remove a chance for errors.)
+type ruleTemplateMap struct {
+	sgutil.InsertMap[string, *template.Template]
+}
+
+// Add creates a template named "name" after appending "// {{.TmplName}}\n" to the
+// template, and returns the input so that additions may be chained.
+// This helps make template initialization easy to order and easy to read.
+func (rtm *ruleTemplateMap) Add(name string, templ string) *ruleTemplateMap {
+	// Append debugging comment AND end of line
+	templ += " // {{.TplName}}\n"
+	ct := sgutil.TemplateNamed(name, templ)
+	rtm.InsertMap.Put(name, ct)
+	return rtm
 }
 
 var (
-	ruleTemplates = template.Must(template.New("simdRules").Parse(`
-{{define "pureVreg"}}({{.GoOp}}{{.GoType}} {{.Args}}) => ({{.Asm}} {{.ArgsOut}})
-{{end}}
-{{define "maskIn"}}({{.GoOp}}{{.GoType}} {{.Args}} mask) => ({{.Asm}} {{.ArgsOut}} ({{.MaskInConvert}} <types.TypeMask> mask))
-{{end}}
-{{define "maskOut"}}({{.GoOp}}{{.GoType}} {{.Args}}) => ({{.MaskOutConvert}} ({{.Asm}} {{.ArgsOut}}))
-{{end}}
-{{define "maskInMaskOut"}}({{.GoOp}}{{.GoType}} {{.Args}} mask) => ({{.MaskOutConvert}} ({{.Asm}} {{.ArgsOut}} ({{.MaskInConvert}} <types.TypeMask> mask)))
-{{end}}
-{{define "sftimm"}}({{.Asm}} x (MOVQconst [c])) => ({{.Asm}}const [uint8(c)] x)
-{{end}}
-{{define "masksftimm"}}({{.Asm}} x (MOVQconst [c]) mask) => ({{.Asm}}const [uint8(c)] x mask)
-{{end}}
-{{define "vregMem"}}({{.Asm}} {{.ArgsLoadAddr}}) && canMergeLoad(v, l) && clobber(l) => ({{.Asm}}load {{.ArgsAddr}})
-{{end}}
-{{define "vregMemFeatCheck"}}({{.Asm}} {{.ArgsLoadAddr}}) && {{.FeatCheck}} && canMergeLoad(v, l) && clobber(l)=> ({{.Asm}}load {{.ArgsAddr}})
-{{end}}
-`))
+	// ORDER MATTERS.  These should appear in most-to-least-specific order
+	// TODO: vregMemFeatCheck is not necessarily in the right place; there was an order,
+	// this was copied from it, but vregMemFeatCheck was not in it.  It's not clear
+	// it was ever used.
+	// It's also not clear that this is strictly most-to-least-specific?
+	ruleTemplates = new(ruleTemplateMap).
+		Add("masksftimm", `({{.Asm}} x (MOVQconst [c]) mask) => ({{.Asm}}const [amd64CapAVXShift(c)] x mask)`).
+		Add("sftimm", `({{.Asm}} x (MOVQconst [c])) => ({{.Asm}}const [amd64CapAVXShift(c)] x)`).
+		Add("maskInMaskOut", `({{.GoOp}}{{.GoType}} {{.Args}} mask) => ({{.MaskOutConvert}} ({{.Asm}} {{.ArgsOut}} ({{.MaskInConvert}} <types.TypeMask> mask)))`).
+		Add("maskOut", `({{.GoOp}}{{.GoType}} {{.Args}}) => ({{.MaskOutConvert}} ({{.Asm}} {{.ArgsOut}}))`).
+		Add("maskIn", `({{.GoOp}}{{.GoType}} {{.Args}} mask) => ({{.Asm}} {{.ArgsOut}} ({{.MaskInConvert}} <types.TypeMask> mask))`).
+		Add("pureVreg", `({{.GoOp}}{{.GoType}} {{.Args}}) => ({{.Asm}} {{.ArgsOut}})`).
+		Add("vregMem", `({{.Asm}} {{.ArgsLoadAddr}}) && canMergeLoad(v, l) && clobber(l) => ({{.Asm}}load {{.ArgsAddr}})`).
+		Add("vregMemFeatCheck", `({{.Asm}} {{.ArgsLoadAddr}}) && {{.FeatCheck}} && canMergeLoad(v, l) && clobber(l) => ({{.Asm}}load {{.ArgsAddr}})`).
+		Add("asmRule", `({{.Asm}} {{.Args}}) {{.RuleCond}} => {{.RuleOut}}`).
+		Add("specialLower", `{{.Rule}}`)
 )
 
 func (d tplRuleData) MaskOptimization(asmCheck map[string]bool) string {
@@ -85,17 +107,6 @@ func (d tplRuleData) MaskOptimization(asmCheck map[string]bool) string {
 	return fmt.Sprintf("(VMOVDQU%dMasked%s (%s %s) mask) => (%s %s mask)\n", d.ElementSize, size, asmNoMask, d.Args, d.Asm, d.Args)
 }
 
-// SSA rewrite rules need to appear in a most-to-least-specific order.  This works for that.
-var tmplOrder = map[string]int{
-	"masksftimm":    0,
-	"sftimm":        1,
-	"maskInMaskOut": 2,
-	"maskOut":       3,
-	"maskIn":        4,
-	"pureVreg":      5,
-	"vregMem":       6,
-}
-
 func compareTplRuleData(x, y tplRuleData) int {
 	if c := compareNatural(x.GoOp, y.GoOp); c != 0 {
 		return c
@@ -106,37 +117,100 @@ func compareTplRuleData(x, y tplRuleData) int {
 	if c := compareNatural(x.Args, y.Args); c != 0 {
 		return c
 	}
-	if x.tplName == y.tplName {
+	if x.TplName == y.TplName {
 		return 0
 	}
-	xo, xok := tmplOrder[x.tplName]
-	yo, yok := tmplOrder[y.tplName]
-	if !xok {
-		panic(fmt.Errorf("Unexpected template name %s, please add to tmplOrder", x.tplName))
+	return ruleTemplates.Compare(x.TplName, y.TplName)
+}
+
+// parseAsmRule tries to parse given string as it would be asmRule:
+// if <cond> => <out>
+// Return false, "", "" if not matched, otherwise true, condition, rule output.
+// For example:
+// rule:"if a==0 => (VADD4S x y)" can be used to provide addional
+// lowering for an instruction which doesn't support zero immediate encoding:
+// (VUSRA4S [a] x y) && a==0 => (VADD4S x y)
+func parseAsmRule(rule string) (bool, string, string) {
+	arrowIndex := strings.Index(rule, "=>")
+	if arrowIndex == -1 {
+		return false, "", ""
+	}
+
+	condPart := rule[:arrowIndex]
+	outPart := rule[arrowIndex+len("=>"):]
+
+	// Check if condPart starts with "if" followed by at least one space.
+	cond := strings.TrimPrefix(condPart, "if")
+	if cond == condPart || len(cond) == 0 || !unicode.IsSpace(rune(cond[0])) {
+		return false, "", ""
 	}
-	if !yok {
-		panic(fmt.Errorf("Unexpected template name %s, please add to tmplOrder", y.tplName))
+
+	// Trim any spaces around <cond> and <out>.
+	cond = strings.TrimSpace(cond)
+	out := strings.TrimSpace(outPart)
+	if cond == "" || out == "" {
+		return false, "", ""
+	}
+
+	return true, cond, out
+}
+
+// parseArgsMatchRule tries to parse given string as it would be asmRule with custom arguments to match:
+// match <args> [&& <cond>] => <out>
+// earlymatch <args> [&& <cond>] => <out>
+// Return false, "", "", "", false if not matched, otherwise true, args, cond, rule output, isEarly.
+// For example:
+// rule:"match [0] (VMOV%sins [0] _ (MOVDconst [c])) && uint64(c)<= 255 => (VMOVI%a [uint8(c)])" can be used to provide addional
+// lowering for a broadcast to use immediate source:
+// (VDUPSbcast [0] (VMOVSins [0] _ (MOVDconst [c]))) && uint64(c)<= 255 => (VMOVI4S [uint8(c)])
+// The specifiers currently supported are only arm64-specific, it may be generalized in the expandFormatSpecifiers function in future.
+// The "earlymatch" variant uses GoOp instead of Asm on the left-hand side, replacing the default lowering rule.
+func parseArgsMatchRule(rule string) (bool, bool, string) {
+	if strings.HasPrefix(rule, "earlymatch ") {
+		return true, true, rule[len("earlymatch "):]
+	} else if strings.HasPrefix(rule, "match ") {
+		return true, false, rule[len("match "):]
+	} else {
+		return false, false, rule
 	}
-	return xo - yo
+}
+
+// expandFormatSpecifiers replaces format specifiers in s with concrete values
+// derived from elemBits (the element size in bits for the current vector type).
+//
+// Supported specifiers:
+//
+//	%s - lane size letter (B, H, S, D)
+//	%a - arrangement suffix (16B, 8H, 4S, 2D) for neon instructions
+//	%b - bits per lane (8, 16, 32, 64)
+func expandFormatSpecifiers(s string, elemBits int) string {
+	elemLetters := map[int]string{8: "B", 16: "H", 32: "S", 64: "D"}
+	arrangements := map[int]string{8: "16B", 16: "8H", 32: "4S", 64: "2D"}
+	s = strings.ReplaceAll(s, "%s", elemLetters[elemBits])
+	s = strings.ReplaceAll(s, "%a", arrangements[elemBits])
+	s = strings.ReplaceAll(s, "%b", fmt.Sprintf("%d", elemBits))
+	return s
 }
 
 // writeSIMDRules generates the lowering and rewrite rules for ssa and writes it to simdAMD64.rules
 // within the specified directory.
 func writeSIMDRules(ops []Operation) *bytes.Buffer {
 	buffer := new(bytes.Buffer)
-	buffer.WriteString(generatedHeader + "\n")
+	buffer.WriteString(generatedHeader() + "\n")
 
 	// asm -> masked merging rules
 	maskedMergeOpts := make(map[string]string)
 	s2n := map[int]string{8: "B", 16: "W", 32: "D", 64: "Q"}
-	asmCheck := map[string]bool{}
+	asmCheck := map[string]bool{}    // for masked merge optimizations.
+	sftimmCheck := map[string]bool{} // deduplicate sftimm rules
 	var allData []tplRuleData
 	var optData []tplRuleData    // for mask peephole optimizations, and other misc
 	var memOptData []tplRuleData // for memory peephole optimizations
 	memOpSeen := make(map[string]bool)
+	ruleDone := make(map[string]struct{})
 
 	for _, opr := range ops {
-		opInShape, opOutShape, maskType, immType, gOp := opr.shape()
+		opInShape, opOutShape, maskType, immType, gOp, _ := opr.shape()
 		asm := machineOpName(maskType, gOp)
 		vregInCnt := len(gOp.In)
 		if maskType == OneMask {
@@ -162,7 +236,7 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 		}
 		if immType == ConstImm {
 			data.ArgsOut = fmt.Sprintf("[%s] %s", *opr.In[0].Const, data.ArgsOut)
-		} else if immType == VarImm {
+		} else if immType == VarImm || immType == VarImmLim {
 			data.Args = fmt.Sprintf("[a] %s", data.Args)
 			data.ArgsOut = fmt.Sprintf("[a] %s", data.ArgsOut)
 		} else if immType == ConstVarImm {
@@ -182,12 +256,12 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 		}
 		var tplName string
 		// If class overwrite is happening, that's not really a mask but a vreg.
-		if opOutShape == OneVregOut || opOutShape == OneVregOutAtIn || gOp.Out[0].OverwriteClass != nil {
+		if opOutShape == OneVregOut || opOutShape == OneVregOutAtIn || opOutShape == OneVregOutScalar || gOp.Out[0].OverwriteClass != nil {
 			switch opInShape {
 			case OneImmIn:
 				tplName = "pureVreg"
 				data.GoType = goType(gOp)
-			case PureVregIn:
+			case PureVregIn, VlistIn:
 				tplName = "pureVreg"
 				data.GoType = goType(gOp)
 			case OneKmaskImmIn:
@@ -229,19 +303,63 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 
 		if gOp.SpecialLower != nil {
 			if *gOp.SpecialLower == "sftimm" {
-				if data.GoType[0] == 'I' {
-					// only do these for signed types, it is a duplicate rewrite for unsigned
+				if !sftimmCheck[data.Asm] {
+					sftimmCheck[data.Asm] = true
 					sftImmData := data
 					if tplName == "maskIn" {
-						sftImmData.tplName = "masksftimm"
+						sftImmData.TplName = "masksftimm"
 					} else {
-						sftImmData.tplName = "sftimm"
+						sftImmData.TplName = "sftimm"
 					}
 					allData = append(allData, sftImmData)
 					asmCheck[sftImmData.Asm+"const"] = true
 				}
+			} else if ok, cond, out := parseAsmRule(*gOp.SpecialLower); ok {
+				if _, done := ruleDone[data.Asm]; !done {
+					ruleDone[data.Asm] = struct{}{}
+					optData := data
+					optData.TplName = "asmRule"
+					optData.RuleCond = cond
+					if cond != "" {
+						optData.RuleCond = "&& " + cond
+					}
+					optData.RuleOut = out
+					if maskType == OneMask {
+						optData.Args += " mask"
+					}
+					allData = append(allData, optData)
+				}
+			} else if ok, isEarly, rest := parseArgsMatchRule(*gOp.SpecialLower); ok {
+				key := data.Asm
+				if isEarly {
+					key = data.GoOp + data.GoType
+				}
+				if _, done := ruleDone[key]; !done {
+					ruleDone[key] = struct{}{}
+					// Get elemBits from the operation's inputs.
+					elemBits := 0
+					for _, in := range gOp.In {
+						if in.ElemBits != nil {
+							elemBits = *in.ElemBits
+							break
+						}
+					}
+					optData := data
+					optData.Rule = rest
+					optData.Rule = expandFormatSpecifiers(optData.Rule, elemBits)
+					optData.TplName = "specialLower"
+					// %g is the generic operation name
+					optData.Rule = strings.ReplaceAll(optData.Rule, "%g", optData.GoOp+optData.GoType)
+					// %h is the hardware operation name
+					optData.Rule = strings.ReplaceAll(optData.Rule, "%h", optData.Asm)
+
+					allData = append(allData, optData)
+				}
+				if isEarly {
+					continue
+				}
 			} else {
-				panic("simdgen sees unknwon special lower " + *gOp.SpecialLower + ", maybe implement it?")
+				panic("simdgen sees unknown special lower " + *gOp.SpecialLower + ", maybe implement it?")
 			}
 		}
 		if gOp.MemFeatures != nil && *gOp.MemFeatures == "vbcst" {
@@ -293,9 +411,9 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 						"AVX512": "v.Block.CPUfeatures.hasFeature(CPUavx512)",
 					}
 					memOpData.FeatCheck = knownFeatChecks[feat2]
-					memOpData.tplName = "vregMemFeatCheck"
+					memOpData.TplName = "vregMemFeatCheck"
 				} else {
-					memOpData.tplName = "vregMem"
+					memOpData.TplName = "vregMem"
 				}
 				memOptData = append(memOptData, memOpData)
 				asmCheck[memOpData.Asm+"load"] = true
@@ -344,7 +462,7 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 			data.Args = "..."
 			data.ArgsOut = "..."
 		}
-		data.tplName = tplName
+		data.TplName = tplName
 		if opr.NoGenericOps != nil && *opr.NoGenericOps == "true" ||
 			opr.SkipMaskedMethod() {
 			optData = append(optData, data)
@@ -356,16 +474,25 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 
 	slices.SortFunc(allData, compareTplRuleData)
 
+	hiHalfRules := generateHiHalfFoldingRules(ops)
+	for _, rule := range hiHalfRules {
+		buffer.WriteString(rule)
+	}
+
 	for _, data := range allData {
-		if err := ruleTemplates.ExecuteTemplate(buffer, data.tplName, data); err != nil {
-			panic(fmt.Errorf("failed to execute template %s for %s: %w", data.tplName, data.GoOp+data.GoType, err))
+		tpl := ruleTemplates.Get(data.TplName)
+		if tpl == nil {
+			panic(fmt.Errorf("template %s not found", data.TplName))
+		}
+		if err := tpl.Execute(buffer, data); err != nil {
+			panic(fmt.Errorf("failed to execute template %s for %s: %w", data.TplName, data.GoOp+data.GoType, err))
 		}
 	}
 
 	seen := make(map[string]bool)
 
 	for _, data := range optData {
-		if data.tplName == "maskIn" {
+		if data.TplName == "maskIn" {
 			rule := data.MaskOptimization(asmCheck)
 			if seen[rule] {
 				continue
@@ -388,10 +515,104 @@ func writeSIMDRules(ops []Operation) *bytes.Buffer {
 	}
 
 	for _, data := range memOptData {
-		if err := ruleTemplates.ExecuteTemplate(buffer, data.tplName, data); err != nil {
-			panic(fmt.Errorf("failed to execute template %s for %s: %w", data.tplName, data.Asm, err))
+		tpl := ruleTemplates.Get(data.TplName)
+		if tpl == nil {
+			panic(fmt.Errorf("template %s not found", data.TplName))
+		}
+		if err := tpl.Execute(buffer, data); err != nil {
+			panic(fmt.Errorf("failed to execute template %s for %s: %w", data.TplName, data.Asm, err))
 		}
 	}
 
 	return buffer
 }
+
+// Note: SetHi was removed (temporarily?) from 1.27, but may return in some form
+// and these optimizations may apply to the equivalent idioms.
+//
+// generateHiHalfFoldingRules generates folding rules that combine SetHi/HiToLo patterns
+// with narrow/long base operations into the hi-half "2" variant instruction.
+//
+// x.HiToLo() is lowered as (VDUPDextr[1] x)
+//
+// x.SetHi(lo) is lowered as (VMOVDins0 [1] x (VDUPDextr [0] lo))
+//
+// Narrow (e.g., SHRN): SetHi wraps the narrow result.
+//
+//	(VMOVDins0 [1] dst y:(VSHRN4S [a] x)) => (VSHRN2_4S dst [a] x)
+//
+// Unary Long (e.g., USHLL) getting high half (HiToLo) as input:
+//
+//	(VUSHLL4H [a] (VDUPDextr [1] x)) => (VUSHLL2_4H [a] x)
+//
+// Binary Long (e.g., UMULL, SMULL), both inputs from HiToLo:
+//
+//	(VUMULL4H (VDUPDextr [1] x) (VDUPDextr [1] y)) => (VUMULL2_4H x y)
+func generateHiHalfFoldingRules(ops []Operation) []string {
+	seen := make(map[string]bool)
+	var rules []string
+
+	for _, opr := range ops {
+		if opr.HiHalfAsm == nil {
+			continue
+		}
+		kind := opr.hiHalfKind()
+		if kind == "" {
+			continue
+		}
+		_, _, maskType, immType, gOp, _ := opr.shape()
+		asm := machineOpName(maskType, gOp)
+		asm2 := hiHalfOpName(*gOp.HiHalfAsm, gOp)
+
+		if seen[asm] {
+			continue
+		}
+		seen[asm] = true
+
+		vregInCnt := 0
+		for _, in := range gOp.In {
+			if in.Class == "vreg" {
+				vregInCnt++
+			}
+		}
+		hasImm := immType == VarImm || immType == VarImmLim || immType == ConstVarImm
+
+		switch kind {
+		case "narrow":
+			switch vregInCnt {
+			case 1:
+				// Note, at least for 1.27, SetHi is not supported, because it is an emulated instruction,
+				// not part of the scalable set, and though it is intended to mimic the similar instruction
+				// on amd64, it is actually slightly different.
+				//
+				// Narrow: the value used by SetHi may be produced with asm2 (hiHalf) variant.
+				// TODO: Maybe have a separate rule: (VMOVDins0 [c] dst (VDUPDextr [0] src)) => (VMOVDins0 [c] dst src)
+				if hasImm {
+					rules = append(rules, fmt.Sprintf("(VMOVDins0 [1] dst (VDUPDextr [0] (%s [c] y))) => (%s dst [c] y)\n", asm, asm2))
+				} else {
+					rules = append(rules, fmt.Sprintf("(VMOVDins0 [1] dst (VDUPDextr [0] (%s y))) => (%s dst y)\n", asm, asm2))
+				}
+			default:
+				panic("unsupported yet folding narrow ops cases")
+			}
+		case "long":
+			switch vregInCnt {
+			case 1:
+				// Long: input from HiToLo (VDUPDextr [1]), prefer hiHalf variant.
+				if hasImm {
+					rules = append(rules, fmt.Sprintf("(%s [a] (VDUPDextr [1] x)) => (%s [a] x)\n", asm, asm2))
+				} else {
+					rules = append(rules, fmt.Sprintf("(%s (VDUPDextr [1] x)) => (%s x)\n", asm, asm2))
+				}
+			case 2:
+				// Binary Long: both inputs are from HiToLo, fold into hiHalf variant.
+				rules = append(rules, fmt.Sprintf("(%s (VDUPDextr [1] x) (VDUPDextr [1] y)) => (%s x y)\n", asm, asm2))
+			default:
+				panic("unsupported yet folding long ops cases")
+			}
+		}
+	}
+
+	slices.Sort(rules)
+	return rules
+}