keccakf_amd64: drop NOSPLIT so the runtime can preemp assembly call

gen_keccakf_bmi2.go

@@ -79,8 +79,14 @@ func main() {
 	// Single function: keccakF1600BMI2(a *[200]byte, buf *byte)
 	// When buf != nil, XORs rate bytes into state before permuting.
 	// When buf == nil, just permutes.
+	// Do NOT add NOSPLIT here, and do not shrink the frame below 128 bytes.
+	// The frame holds the round temp state, and its stack check is the only
+	// preemption point in the callers' block loops: assembly is not
+	// async-preemptible, so without it a large Sum256 or Write holds every P
+	// in stop-the-world until it finishes. The arm64 kernel omits NOSPLIT for
+	// the same reason.
 	p("// func keccakF1600BMI2(a *[200]byte, buf *byte)")
-	p("TEXT ·keccakF1600BMI2(SB), NOSPLIT, $%d-16", fsize)
+	p("TEXT ·keccakF1600BMI2(SB), $%d-16", fsize)
 	p("\tMOVQ a+0(FP), DI")
 	p("\tMOVQ buf+8(FP), BX")
 	p("\tTESTQ BX, BX")

keccak_asm_test.go

@@ -4,7 +4,14 @@ package keccak
 
 import (
 	"math/rand"
+	"os"
+	"path/filepath"
+	"regexp"
+	"runtime"
+	"strconv"
+	"strings"
 	"testing"
+	"time"
 )
 
 // TestXorAndPermute verifies the fused XOR+permute assembly entry point
@@ -34,3 +41,111 @@ func TestXorAndPermute(t *testing.T) {
 		}
 	}
 }
+
+// The permutation kernel's stack-growth check is the only preemption point in
+// the absorb and squeeze loops: those loops call nothing else, and assembly is
+// never async-preemptible (runtime/preempt.go, isAsyncSafePoint rejects any
+// FuncFlagAsm frame). Two ways to lose it:
+//
+//	NOSPLIT                  drops the check whatever the frame size
+//	a frame under 128 bytes  makes the assembler mark a leaf NOSPLIT itself
+//	                         (abi.StackSmall, obj6.go/obj7.go)
+//
+// Either one lets a single large Sum256 or Write hold every P in
+// stop-the-world for the length of the hash.
+var kernelTEXT = regexp.MustCompile(`^TEXT\s+·(\w+)\(SB\),\s*(?:([A-Z0-9|+]+),\s*)?\$(\d+)-\d+`)
+
+func TestKernelKeepsStackCheck(t *testing.T) {
+	files, err := filepath.Glob("*.s")
+	if err != nil || len(files) == 0 {
+		t.Fatalf("no assembly files found: %v", err)
+	}
+	for _, f := range files {
+		src, err := os.ReadFile(f)
+		if err != nil {
+			t.Fatal(err)
+		}
+		found := 0
+		for i, line := range strings.Split(string(src), "\n") {
+			m := kernelTEXT.FindStringSubmatch(line)
+			if m == nil {
+				continue
+			}
+			found++
+			name, flags, frame := m[1], m[2], m[3]
+			where := f + ":" + strconv.Itoa(i+1)
+			// Reject the whole flag field, not the spelling: textflag.h
+			// defines NOSPLIT as 4 and the assembler takes the number, so
+			// matching the name alone lets the regression back in.
+			if flags != "" {
+				t.Errorf("%s: %s carries assembler flags %q. NOSPLIT (spelled, or as its "+
+					"numeric value 4) drops the stack-growth check, which is the only "+
+					"preemption point in the block loops calling it; without it a large "+
+					"Sum256 blocks every stop-the-world for the whole hash.",
+					where, name, flags)
+			}
+			if n, _ := strconv.Atoi(frame); n < 128 {
+				t.Errorf("%s: %s has a %d-byte frame, under abi.StackSmall (128). The "+
+					"assembler marks a leaf that small NOSPLIT itself, which drops the "+
+					"stack-growth check. Keep the frame at 128 bytes or more.", where, name, n)
+			}
+		}
+		if found == 0 {
+			t.Errorf("%s: no TEXT directive matched, so nothing in this file is "+
+				"guarded; was the frame size turned into a #define?", f)
+		}
+	}
+}
+
+// gcWait reports the longest a GC had to wait while work ran in another
+// goroutine. Work the runtime cannot preempt shows up here directly.
+func gcWait(work func()) time.Duration {
+	done := make(chan struct{})
+	go func() { defer close(done); work() }()
+	var worst time.Duration
+	for {
+		select {
+		case <-done:
+			return worst
+		default:
+		}
+		start := time.Now()
+		runtime.GC()
+		worst = max(worst, time.Since(start))
+	}
+}
+
+// TestSTWPauseWhileHashing measures the property the TEXT flags only imply: how
+// long the runtime waits for a hashing goroutine to yield.
+//
+// The budget is relative to the same hash computed through the pure-Go fallback,
+// which always has a preemption point, so a slow or loaded runner moves both
+// numbers together and only a real regression separates them.
+func TestSTWPauseWhileHashing(t *testing.T) {
+	if testing.Short() || runtime.GOMAXPROCS(0) < 2 {
+		t.Skip("hashes 32 MiB, and needs GOMAXPROCS >= 2 to see a stalled GC")
+	}
+	if !useASM {
+		t.Skip("hardware acceleration unavailable on this CPU")
+	}
+	buf := make([]byte, 32*1024*1024)
+
+	useASM = false // same call, preemptible implementation
+	floor := gcWait(func() {
+		for range 4 {
+			Sum256(buf)
+		}
+	})
+	useASM = true
+	got := gcWait(func() {
+		for range 4 {
+			Sum256(buf)
+		}
+	})
+
+	t.Logf("worst GC wait: assembly %v, pure Go %v", got, floor)
+	if got > 5*time.Millisecond && got > 4*floor {
+		t.Fatalf("a GC waited %v on the assembly path against %v on the pure-Go path: "+
+			"the block loop has no preemption point", got, floor)
+	}
+}

keccakf_amd64_bmi2.s

@@ -5,7 +5,7 @@
 #include "textflag.h"
 
 // func keccakF1600BMI2(a *[200]byte, buf *byte)
-TEXT ·keccakF1600BMI2(SB), NOSPLIT, $200-16
+TEXT ·keccakF1600BMI2(SB), $200-16
 	MOVQ  a+0(FP), DI
 	MOVQ  buf+8(FP), BX
 	TESTQ BX, BX