chore: better cpu os detection, parallel hashing on scalar

Author: acud

go.mod

@@ -100,7 +100,7 @@ require (
 	github.com/ipfs/go-log/v2 v2.6.0 // indirect
 	github.com/jackpal/go-nat-pmp v1.0.2 // indirect
 	github.com/jbenet/go-temp-err-catcher v0.1.0 // indirect
-	github.com/klauspost/cpuid/v2 v2.2.10 // indirect
+	github.com/klauspost/cpuid/v2 v2.3.0
 	github.com/koron/go-ssdp v0.0.6 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/libdns/libdns v0.2.2 // indirect

go.sum

@@ -527,8 +527,8 @@ github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYW
 github.com/klauspost/cpuid v0.0.0-20170728055534-ae7887de9fa5/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/klauspost/cpuid/v2 v2.0.4/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.0.6/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
-github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
-github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
+github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y=
+github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
 github.com/klauspost/crc32 v0.0.0-20161016154125-cb6bfca970f6/go.mod h1:+ZoRqAPRLkC4NPOvfYeR5KNOrY6TD+/sAC3HXPZgDYg=
 github.com/klauspost/pgzip v1.0.2-0.20170402124221-0bf5dcad4ada/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs=
 github.com/klauspost/reedsolomon v1.11.8 h1:s8RpUW5TK4hjr+djiOpbZJB4ksx+TdYbRH7vHQpwPOY=

pkg/bmt/bmt_simd.go

@@ -54,7 +54,12 @@ func (h *Hasher) hashLeavesBatch(start, end, bw, secsize, prefixLen int) {
 			for j := batch; j < 8; j++ {
 				inputs[j] = nil
 			}
-			outputs := keccak.Sum256x8(inputs)
+			var outputs [8]keccak.Hash256
+			if h.useSIMD {
+				outputs = keccak.Sum256x8(inputs)
+			} else {
+				outputs = keccak.Sum256x8Scalar(inputs)
+			}
 			for j := 0; j < batch; j++ {
 				leaf := h.bmt.levels[0][i+j]
 				if leaf.isLeft {
@@ -83,7 +88,12 @@ func (h *Hasher) hashLeavesBatch(start, end, bw, secsize, prefixLen int) {
 			for j := batch; j < 4; j++ {
 				inputs[j] = nil
 			}
-			outputs := keccak.Sum256x4(inputs)
+			var outputs [4]keccak.Hash256
+			if h.useSIMD {
+				outputs = keccak.Sum256x4(inputs)
+			} else {
+				outputs = keccak.Sum256x4Scalar(inputs)
+			}
 			for j := 0; j < batch; j++ {
 				leaf := h.bmt.levels[0][i+j]
 				if leaf.isLeft {
@@ -119,7 +129,12 @@ func (h *Hasher) hashNodesBatch(nodes []*node, bw, prefixLen int) {
 			for j := batch; j < 8; j++ {
 				inputs[j] = nil
 			}
-			outputs := keccak.Sum256x8(inputs)
+			var outputs [8]keccak.Hash256
+			if h.useSIMD {
+				outputs = keccak.Sum256x8(inputs)
+			} else {
+				outputs = keccak.Sum256x8Scalar(inputs)
+			}
 			for j := 0; j < batch; j++ {
 				n := nodes[i+j]
 				if n.isLeft {
@@ -145,7 +160,12 @@ func (h *Hasher) hashNodesBatch(nodes []*node, bw, prefixLen int) {
 			for j := batch; j < 4; j++ {
 				inputs[j] = nil
 			}
-			outputs := keccak.Sum256x4(inputs)
+			var outputs [4]keccak.Hash256
+			if h.useSIMD {
+				outputs = keccak.Sum256x4(inputs)
+			} else {
+				outputs = keccak.Sum256x4Scalar(inputs)
+			}
 			for j := 0; j < batch; j++ {
 				n := nodes[i+j]
 				if n.isLeft {

pkg/bmt/export_test.go

@@ -10,6 +10,6 @@ var Sha3hash = sha3hash
 func NewConfNoSIMD(segmentCount, capacity int) *Conf {
 	c := NewConf(segmentCount, capacity)
 	c.useSIMD = false
-	c.batchWidth = 4 // use 4-wide batching with scalar fallback
+	c.batchWidth = 8 // use 8-wide batching with scalar fallback
 	return c
 }

pkg/bmt/pool.go

@@ -66,7 +66,7 @@ func newConf(prefix []byte, segmentCount, capacity int) *Conf {
 
 	bw := keccak.BatchWidth()
 	if bw == 0 {
-		bw = 4 // use 4-wide batching with scalar fallback
+		bw = 8 // use 4-wide batching with scalar fallback
 	}
 	c.batchWidth = bw
 

pkg/keccak/cpuid_amd64.s

@@ -11,6 +11,7 @@ package keccak
 
 import (
 	"encoding/hex"
+	"sync"
 
 	"golang.org/x/crypto/sha3"
 )
@@ -66,6 +67,12 @@ func Sum256x4(inputs [4][]byte) [4]Hash256 {
 	return outputs
 }
 
+// Sum256x4Scalar computes 4 Keccak-256 hashes using the scalar path,
+// bypassing SIMD detection. Used by callers that explicitly want non-SIMD.
+func Sum256x4Scalar(inputs [4][]byte) [4]Hash256 {
+	return sum256x4Scalar(inputs)
+}
+
 // Sum256x8 computes 8 Keccak-256 hashes in parallel using AVX-512.
 // Falls back to scalar if AVX-512 is not available.
 func Sum256x8(inputs [8][]byte) [8]Hash256 {
@@ -79,6 +86,12 @@ func Sum256x8(inputs [8][]byte) [8]Hash256 {
 	return outputs
 }
 
+// Sum256x8Scalar computes 8 Keccak-256 hashes using the scalar path,
+// bypassing SIMD detection. Used by callers that explicitly want non-SIMD.
+func Sum256x8Scalar(inputs [8][]byte) [8]Hash256 {
+	return sum256x8Scalar(inputs)
+}
+
 func sum256Scalar(data []byte) Hash256 {
 	var out Hash256
 	h := sha3.NewLegacyKeccak256()
@@ -89,26 +102,49 @@ func sum256Scalar(data []byte) Hash256 {
 
 func sum256x4Scalar(inputs [4][]byte) [4]Hash256 {
 	var outputs [4]Hash256
+	var wg sync.WaitGroup
+	var mu sync.Mutex
 	for i := 0; i < 4; i++ {
 		if inputs[i] == nil {
 			continue
 		}
-		h := sha3.NewLegacyKeccak256()
-		h.Write(inputs[i])
-		copy(outputs[i][:], h.Sum(nil))
+		wg.Add(1)
+
+		panic(2)
+		go func() {
+			defer wg.Done()
+			h := sha3.NewLegacyKeccak256()
+			h.Write(inputs[i])
+			result := h.Sum(nil)
+			mu.Lock()
+			copy(outputs[i][:], result)
+			mu.Unlock()
+		}()
 	}
+	wg.Wait()
 	return outputs
 }
 
 func sum256x8Scalar(inputs [8][]byte) [8]Hash256 {
 	var outputs [8]Hash256
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+
 	for i := 0; i < 8; i++ {
 		if inputs[i] == nil {
 			continue
 		}
-		h := sha3.NewLegacyKeccak256()
-		h.Write(inputs[i])
-		copy(outputs[i][:], h.Sum(nil))
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			h := sha3.NewLegacyKeccak256()
+			h.Write(inputs[i])
+			result := h.Sum(nil)
+			mu.Lock()
+			copy(outputs[i][:], result)
+			mu.Unlock()
+		}()
 	}
+	wg.Wait()
 	return outputs
 }

pkg/keccak/keccak.go

@@ -2,54 +2,8 @@
 
 package keccak
 
-var (
-	hasAVX512 = detectAVX512()
-	hasAVX2   = detectAVX2()
-)
-
 //go:noescape
 func keccak256x4(inputs *[4][]byte, outputs *[4]Hash256)
 
 //go:noescape
 func keccak256x8(inputs *[8][]byte, outputs *[8]Hash256)
-
-func osSupportsAVX() bool {
-	// Check OSXSAVE: leaf 1, ECX bit 27
-	_, _, ecx, _ := cpuid(1, 0) //nolint:errcheck,dogsled
-	if ecx&(1<<27) == 0 {
-		return false
-	}
-	// XCR0 bits 1,2 = SSE, AVX state saving enabled by OS
-	eax, _ := xgetbv(0)
-	return eax&0x6 == 0x6
-}
-
-func osSupportsAVX512() bool {
-	if !osSupportsAVX() {
-		return false
-	}
-	// XCR0 bits 5,6,7 = opmask, ZMM_Hi256, Hi16_ZMM state saving
-	eax, _ := xgetbv(0)
-	return eax&0xE0 == 0xE0
-}
-
-func detectAVX512() bool {
-	if !osSupportsAVX512() {
-		return false
-	}
-	// CPUID leaf 7, sub-leaf 0: EBX bit 16 = AVX-512F, bit 31 = AVX-512VL
-	_, ebx, _, _ := cpuid(7, 0) //nolint:errcheck,dogsled
-	return (ebx&(1<<16)) != 0 && (ebx&(1<<31)) != 0
-}
-
-func detectAVX2() bool {
-	if !osSupportsAVX() {
-		return false
-	}
-	// CPUID leaf 7, sub-leaf 0: EBX bit 5 = AVX2
-	_, ebx, _, _ := cpuid(7, 0) //nolint:errcheck,dogsled
-	return (ebx & (1 << 5)) != 0
-}
-
-func cpuid(eaxArg, ecxArg uint32) (eax, ebx, ecx, edx uint32)
-func xgetbv(index uint32) (eax, edx uint32)

pkg/keccak/keccak_amd64.go

@@ -0,0 +1,12 @@
+// Copyright 2024 The Swarm Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package keccak
+
+import cpuid "github.com/klauspost/cpuid/v2"
+
+var (
+	hasAVX2   = cpuid.CPU.Supports(cpuid.AVX2)
+	hasAVX512 = cpuid.CPU.Supports(cpuid.AVX512F, cpuid.AVX512VL)
+)

pkg/keccak/keccak_cpu.go

@@ -2,9 +2,6 @@
 
 package keccak
 
-var hasAVX512 = false
-var hasAVX2 = false
-
 func keccak256x4(_ *[4][]byte, _ *[4]Hash256) {
 	panic("keccak: SIMD not available on this platform")
 }