Make ARM SVE code vector length agnostic

klauspost · Aug 20, 2024 · 5b12fc2 · 5b12fc2
1 parent 3412d52
commit 5b12fc2
Show file tree

Hide file tree

Showing 5 changed files with 465 additions and 227 deletions.
diff --git a/_gen/gen-arm-sve.go b/_gen/gen-arm-sve.go
@@ -9,6 +9,7 @@ import (
 	"log"
 	"os"
 	"regexp"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -359,3 +360,143 @@ func genArmSve() {
 	fromAvx2ToSve()
 	addEarlyExit("Sve")
 }
+
+func assemble(sve string) string {
+	opcode, err := sve_as.Assemble(sve)
+	if err != nil {
+		return fmt.Sprintf("    WORD $0x00000000 // %s", sve)
+	} else {
+		return fmt.Sprintf("    WORD $0x%08x // %s", opcode, sve)
+	}
+}
+
+func addArmSveVectorLength() (addInits []string) {
+	const filename = "../galois_gen_arm64.s"
+	asmOut := &bytes.Buffer{}
+
+	file, err := os.Open(filename)
+	if err != nil {
+		return
+	}
+	defer file.Close()
+
+	// Create a scanner to read the file line by line
+	scanner := bufio.NewScanner(file)
+
+	routine := ""
+	addInits = make([]string, 0)
+
+	// Iterate over each line
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		if strings.HasPrefix(line, "TEXT ·") {
+			routine = line
+		}
+
+		correctShift := func(shift, vl string) {
+			if strings.Contains(line, " // lsr ") && strings.HasSuffix(strings.TrimSpace(line), ", "+shift) {
+				instr := strings.Split(strings.TrimSpace(line), "// lsr ")[1]
+				args := strings.Split(instr, ", ")
+				if len(args) == 3 && args[0] == args[1] {
+					// keep the original right shift, but reverse the effect (so effectively
+					// clearing out the lower bits so we cannot do eg. "half loops" )
+					line += "\n"
+					line += assemble(fmt.Sprintf("lsl %s, %s, %s", args[0], args[1], shift)) + "\n"
+					line += assemble(fmt.Sprintf("rdvl x16, %s", vl)) + "\n"
+					line += assemble(fmt.Sprintf("udiv %s, %s, x16", args[0], args[1]))
+				}
+			}
+		}
+
+		correctShift("#6", "#2")
+		correctShift("#5", "#1")
+
+		if strings.Contains(line, " // add ") && strings.HasSuffix(strings.TrimSpace(line), "#64") {
+			instr := strings.Split(strings.TrimSpace(line), "// add ")[1]
+			args := strings.Split(instr, ", ")
+			if len(args) == 3 && args[0] == args[1] {
+				line = assemble(fmt.Sprintf("addvl %s, %s, #2", args[0], args[1]))
+			}
+		}
+
+		if strings.Contains(line, " // add ") && strings.HasSuffix(strings.TrimSpace(line), "#32") {
+			instr := strings.Split(strings.TrimSpace(line), "// add ")[1]
+			args := strings.Split(instr, ", ")
+			if len(args) == 3 && args[0] == args[1] {
+				line = assemble(fmt.Sprintf("addvl %s, %s, #1", args[0], args[1]))
+			}
+		}
+
+		if strings.Contains(line, " // add ") && strings.HasSuffix(strings.TrimSpace(line), "#4") {
+			// mark routine as needing initialization of register 17
+			addInits = append(addInits, routine)
+			line = assemble("add x15, x15, x17")
+		}
+
+		asmOut.WriteString(line + "\n")
+	}
+
+	// Check for any errors that occurred during scanning
+	if err = scanner.Err(); err != nil {
+		log.Fatal(err)
+	} else if err = os.WriteFile("../galois_gen_arm64.s", asmOut.Bytes(), 0644); err != nil {
+		log.Fatal(err)
+	}
+
+	return
+}
+
+func addArmSveInitializations(addInits []string) {
+
+	const filename = "../galois_gen_arm64.s"
+	asmOut := &bytes.Buffer{}
+
+	file, err := os.Open(filename)
+	if err != nil {
+		return
+	}
+	defer file.Close()
+
+	// Create a scanner to read the file line by line
+	scanner := bufio.NewScanner(file)
+	routine := ""
+	checkNextLine := false
+
+	// Iterate over each line
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		if strings.HasPrefix(line, "TEXT ·") {
+			routine = line
+		}
+
+		if strings.Contains(line, "// Load number of input shards") {
+			checkNextLine = true
+		} else {
+			if checkNextLine {
+				idx := slices.IndexFunc(addInits, func(s string) bool { return s == routine })
+				if idx != -1 {
+					line += "\n"
+					line += assemble("rdvl x17, #1") + "\n"
+					line += assemble("lsr  x17, x17, #3")
+				}
+				checkNextLine = false
+			}
+		}
+
+		asmOut.WriteString(line + "\n")
+	}
+
+	// Check for any errors that occurred during scanning
+	if err = scanner.Err(); err != nil {
+		log.Fatal(err)
+	} else if err = os.WriteFile("../galois_gen_arm64.s", asmOut.Bytes(), 0644); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func genArmSveAllVl() {
+	addInits := addArmSveVectorLength()
+	addArmSveInitializations(addInits)
+}
diff --git a/_gen/gen.go b/_gen/gen.go
@@ -93,6 +93,7 @@ func main() {
 	if pshufb {
 		genArmSve()
 		genArmNeon()
+		genArmSveAllVl()
 	}
 	Generate()
 }

diff --git a/_gen/go.mod b/_gen/go.mod
@@ -9,7 +9,7 @@ require (
 
 require (
 	github.com/fwessels/avxTwo2sve v0.0.0-20240611172111-6b8528700471 // indirect
-	github.com/fwessels/sve-as v0.0.0-20240611015707-daffc010447f // indirect
+	github.com/fwessels/sve-as v0.0.0-20240817192210-83d5dbff9505 // indirect
 	golang.org/x/mod v0.6.0 // indirect
 	golang.org/x/sys v0.1.0 // indirect
 	golang.org/x/tools v0.2.0 // indirect

diff --git a/_gen/go.sum b/_gen/go.sum
@@ -2,6 +2,8 @@ github.com/fwessels/avxTwo2sve v0.0.0-20240611172111-6b8528700471 h1:omdgAKxePZx
 github.com/fwessels/avxTwo2sve v0.0.0-20240611172111-6b8528700471/go.mod h1:9+ibRsEIs0vLXkalKCGEbZfVS4fafeIvMvM9GvIsdeQ=
 github.com/fwessels/sve-as v0.0.0-20240611015707-daffc010447f h1:HQud3yIU82LdkQzHEYiSJs73wCHjprIqeZE9JvSjKbQ=
 github.com/fwessels/sve-as v0.0.0-20240611015707-daffc010447f/go.mod h1:j3s7EY79XxNMyjx/54Vo6asZafWU4yijB+KIfj4hrh8=
+github.com/fwessels/sve-as v0.0.0-20240817192210-83d5dbff9505 h1:oKLoVXrXDsNNTdNLsSbEu18Vy0Z0b1yeanl5TG4qSyU=
+github.com/fwessels/sve-as v0.0.0-20240817192210-83d5dbff9505/go.mod h1:j3s7EY79XxNMyjx/54Vo6asZafWU4yijB+KIfj4hrh8=
 github.com/klauspost/asmfmt v1.3.1 h1:7xZi1N7s9gTLbqiM8KUv8TLyysavbTRGBT5/ly0bRtw=
 github.com/klauspost/asmfmt v1.3.1/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
 github.com/mmcloughlin/avo v0.5.1-0.20221128045730-bf1d05562091 h1:C2c8ttOBeyhs1SvyCXVPCFd0EqtPiTKGnMWQ+JkM0Lc=