add support for querying cpuid

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

internal/engine/compiler/cpuid_amd64.go

@@ -0,0 +1,43 @@
+package compiler
+
+const (
+	cpuFeatureSSE3   = uint64(1)
+	cpuFeatureSSE4_1 = uint64(1) << 19
+	cpuFeatureSSE4_2 = uint64(1) << 20
+)
+
+const (
+	cpuExtraFeatureABM = uint64(1) << 5 // Advanced Bit Manipulation
+)
+
+type cpuFeatureFlags struct {
+	flags      uint64
+	extraFlags uint64
+}
+
+func loadCpuFeatureFlags() cpuFeatureFlags {
+	return cpuFeatureFlags{
+		flags:      cpuidAsBitmap(1, 0),
+		extraFlags: cpuidAsBitmap(0x80000001, 0),
+	}
+}
+
+// cpuid exposes the CPUID instruction to the Go layer (https://www.amd.com/system/files/TechDocs/25481.pdf)
+// implemented in impl_amd64.s
+func cpuid(arg1, arg2 uint32) (eax, ebx, ecx, edx uint32)
+
+// cpuidAsBitmap combines the result of invoking cpuid to uint64 bitmap
+func cpuidAsBitmap(arg1, arg2 uint32) uint64 {
+	_ /* eax */, _ /* ebx */, ecx, edx := cpuid(1, 0)
+	return (uint64(edx) << 32) | uint64(ecx)
+}
+
+// has invokes cpuid and returns true when the specified flag (represented as uint64) is supported
+func (f *cpuFeatureFlags) has(cpuFeature uint64) bool {
+	return (f.flags & cpuFeature) != 0
+}
+
+// has invokes cpuid and returns true when the specified extraFlag (represented as uint64) is supported
+func (f *cpuFeatureFlags) hasExtra(cpuFeature uint64) bool {
+	return (f.extraFlags & cpuFeature) != 0
+}

internal/engine/compiler/cpuid_amd64.s

@@ -0,0 +1,14 @@
+#include "textflag.h"
+
+// lifted from github.com/intel-go/cpuid and src/internal/cpu/cpu_x86.s
+// func cpuid(arg1, arg2 uint32) (eax, ebx, ecx, edx uint32)
+TEXT ·cpuid(SB), NOSPLIT, $0-24
+	MOVL arg1+0(FP), AX
+	MOVL arg2+4(FP), CX
+	CPUID
+	MOVL AX, eax+8(FP)
+	MOVL BX, ebx+12(FP)
+	MOVL CX, ecx+16(FP)
+	MOVL DX, edx+20(FP)
+	RET
+

internal/engine/compiler/cpuid_amd64_test.go

@@ -0,0 +1,18 @@
+package compiler
+
+import (
+	"testing"
+
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestAmd64CpuId_cpuHasFeature(t *testing.T) {
+	flags := cpuFeatureFlags{
+		flags:      cpuFeatureSSE3,
+		extraFlags: cpuExtraFeatureABM,
+	}
+	require.True(t, flags.has(cpuFeatureSSE3))
+	require.False(t, flags.has(cpuFeatureSSE4_2))
+	require.True(t, flags.hasExtra(cpuExtraFeatureABM))
+	require.False(t, flags.hasExtra(1<<6)) // some other value
+}

internal/engine/compiler/impl_amd64.go

@@ -8,7 +8,6 @@ import (
 	"bytes"
 	"fmt"
 	"math"
-	"runtime"
 
 	"github.com/tetratelabs/wazero/internal/asm"
 	"github.com/tetratelabs/wazero/internal/asm/amd64"
@@ -83,8 +82,9 @@ func (c *amd64Compiler) compileNOP() asm.Node {
 }
 
 type amd64Compiler struct {
-	assembler amd64.Assembler
-	ir        *wazeroir.CompilationResult
+	assembler   amd64.Assembler
+	ir          *wazeroir.CompilationResult
+	cpuFeatures cpuFeatureFlags
 	// locationStack holds the state of wazeroir virtual stack.
 	// and each item is either placed in register or the actual memory stack.
 	locationStack *runtimeValueLocationStack
@@ -103,6 +103,7 @@ func newAmd64Compiler() compiler {
 	c := &amd64Compiler{
 		assembler:     amd64.NewAssembler(),
 		locationStack: newRuntimeValueLocationStack(),
+		cpuFeatures:   loadCpuFeatureFlags(),
 	}
 	return c
 }
@@ -114,6 +115,7 @@ func (c *amd64Compiler) Init(ir *wazeroir.CompilationResult, withListener bool)
 	*c = amd64Compiler{
 		labels:       map[string]*amd64LabelInfo{},
 		ir:           ir,
+		cpuFeatures:  c.cpuFeatures,
 		withListener: withListener,
 		currentLabel: wazeroir.EntrypointLabel,
 	}
@@ -1170,7 +1172,7 @@ func (c *amd64Compiler) compileClz(o *wazeroir.OperationClz) error {
 		return err
 	}
 
-	if runtime.GOOS != "darwin" && runtime.GOOS != "freebsd" {
+	if c.cpuFeatures.hasExtra(cpuExtraFeatureABM) {
 		if o.Type == wazeroir.UnsignedInt32 {
 			c.assembler.CompileRegisterToRegister(amd64.LZCNTL, target.register, target.register)
 		} else {
@@ -1233,7 +1235,7 @@ func (c *amd64Compiler) compileCtz(o *wazeroir.OperationCtz) error {
 		return err
 	}
 
-	if runtime.GOOS != "darwin" && runtime.GOOS != "freebsd" {
+	if c.cpuFeatures.hasExtra(cpuExtraFeatureABM) {
 		if o.Type == wazeroir.UnsignedInt32 {
 			c.assembler.CompileRegisterToRegister(amd64.TZCNTL, target.register, target.register)
 		} else {

internal/engine/compiler/impl_amd64_test.go

@@ -441,6 +441,61 @@ func TestAmd64Compiler_preventCrossedTargetdRegisters(t *testing.T) {
 	}
 }
 
+// Relates to #1111: older AMD64 CPUs do not support the LZCNT instruction
+// CPUID should be used instead. We simulate presence/absence of the feature
+// by overriding the field in the corresponding struct.
+func TestAmd64Compiler_ensureClz_ABM(t *testing.T) {
+	tests := []struct {
+		name        string
+		cpuFeatures cpuFeatureFlags
+		expectedLen int
+	}{
+		{
+			name:        "with ABM",
+			expectedLen: 140,
+			cpuFeatures: cpuFeatureFlags{
+				flags:      0,
+				extraFlags: cpuExtraFeatureABM,
+			},
+		},
+		{
+			name:        "without ABM",
+			expectedLen: 156,
+			cpuFeatures: cpuFeatureFlags{
+				flags:      0,
+				extraFlags: 0, // no flags, thus no ABM, i.e. no LZCNT
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			env := newCompilerEnvironment()
+
+			newCompiler := func() compiler {
+				c := newCompiler().(*amd64Compiler)
+				// override auto-detected CPU features with the test case
+				c.cpuFeatures = tt.cpuFeatures
+				return c
+			}
+
+			compiler := env.requireNewCompiler(t, newCompiler, nil)
+			err := compiler.compilePreamble()
+			require.NoError(t, err)
+
+			err = compiler.compileClz(&wazeroir.OperationClz{Type: wazeroir.UnsignedInt64})
+			require.NoError(t, err)
+			//
+			err = compiler.compileReturnFunction()
+			require.NoError(t, err)
+
+			code, _, err := compiler.compile()
+			require.NoError(t, err)
+
+			require.Equal(t, tt.expectedLen, len(code))
+		})
+	}
+}
+
 // collectRegistersFromRuntimeValues returns the registers occupied by locs.
 func collectRegistersFromRuntimeValues(locs []*runtimeValueLocation) []asm.Register {
 	out := make([]asm.Register, len(locs))