mask_asm: Note implementation may not be perfect

nhooyr · nhooyr · commit fee373961a05 · 2024-02-22T03:58:17.000-08:00
diff --git a/go.mod b/go.mod
@@ -1,3 +1,5 @@
 module nhooyr.io/websocket
 
 go 1.19
+
+require golang.org/x/sys v0.17.0 // indirect
diff --git a/go.sum b/go.sum
@@ -0,0 +1,2 @@
+golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
+golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
diff --git a/mask_arm64.s b/mask_arm64.s
@@ -15,8 +15,6 @@ TEXT ·maskAsm(SB), NOSPLIT, $0-28
 	CMP   $64, R1
 	BLT   less_than_64
 
-// TODO: align memory like amd64
-
 loop_64:
 	VLD1   (R0), [V1.B16, V2.B16, V3.B16, V4.B16]
 	VEOR   V1.B16, V0.B16, V1.B16
@@ -29,6 +27,7 @@ loop_64:
 	BGE    loop_64
 
 less_than_64:
+	CBZ    R1, end
 	TBZ    $5, R1, less_than_32
 	VLD1   (R0), [V1.B16, V2.B16]
 	VEOR   V1.B16, V0.B16, V1.B16
diff --git a/mask_asm.go b/mask_asm.go
@@ -2,14 +2,21 @@
 
 package websocket
 
+import "golang.org/x/sys/cpu"
+
 func mask(b []byte, key uint32) uint32 {
 	if len(b) > 0 {
 		return maskAsm(&b[0], len(b), key)
 	}
 	return key
 }
 
-var useAVX2 = true
+var useAVX2 = cpu.X86.HasAVX2
 
+// @nhooyr: I am not confident that the amd64 or the arm64 implementations of this
+// function are perfect. There are almost certainly missing optimizations or
+// opportunities for // simplification. I'm confident there are no bugs though.
+// For example, the arm64 implementation doesn't align memory like the amd64.
+// Or the amd64 implementation could use AVX512 instead of just AVX2.
 //go:noescape
 func maskAsm(b *byte, len int, key uint32) uint32

-Original file line number
+Diff line change
@@ @@ -1,3 +1,5 @@ @@
 module nhooyr.io/websocket
 go 1.19
++
 +require golang.org/x/sys v0.17.0 // indirect
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=`
	`2`	`+golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=`