cache: includes cpu features in cache keys (#2220)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

internal/engine/wazevo/backend/isa/amd64/machine_test.go

@@ -443,18 +443,21 @@ L2:
 	}
 }
 
-// mockCpuFlags implements platform.CpuFeatureFlags
+// mockCpuFlags implements platform.CpuFeatureFlags.
 type mockCpuFlags struct {
 	flags      platform.CpuFeature
 	extraFlags platform.CpuFeature
 }
 
-// Has implements the method of the same name in platform.CpuFeatureFlags
+// Has implements the method of the same name in platform.CpuFeatureFlags.
 func (f *mockCpuFlags) Has(flag platform.CpuFeature) bool {
 	return (f.flags & flag) != 0
 }
 
-// HasExtra implements the method of the same name in platform.CpuFeatureFlags
+// HasExtra implements the method of the same name in platform.CpuFeatureFlags.
 func (f *mockCpuFlags) HasExtra(flag platform.CpuFeature) bool {
 	return (f.extraFlags & flag) != 0
 }
+
+// Raw implements the method of the same name in platform.CpuFeatureFlags.
+func (f *mockCpuFlags) Raw() uint64 { return 0 }

internal/engine/wazevo/engine_cache.go

@@ -31,6 +31,13 @@ func fileCacheKey(m *wasm.Module) (ret filecache.Key) {
 	s := sha256.New()
 	s.Write(m.ID[:])
 	s.Write(magic)
+	// Write the CPU features so that we can cache the compiled module for the same CPU.
+	// This prevents the incompatible CPU features from being used.
+	cpu := platform.CpuFeatures.Raw()
+	// Reuse the `ret` buffer to write the first 8 bytes of the CPU features so that we can avoid the allocation.
+	binary.LittleEndian.PutUint64(ret[:8], cpu)
+	s.Write(ret[:8])
+	// Finally, write the hash to the ret buffer.
 	s.Sum(ret[:0])
 	return
 }

internal/platform/cpuid.go

@@ -6,6 +6,9 @@ type CpuFeatureFlags interface {
 	Has(cpuFeature CpuFeature) bool
 	// HasExtra returns true when the specified extraFlag (represented as uint64) is supported
 	HasExtra(cpuFeature CpuFeature) bool
+	// Raw returns the raw bitset that represents CPU features used by wazero. This can be used for cache keying.
+	// For now, we only use four features, so uint64 is enough.
+	Raw() uint64
 }
 
 type CpuFeature uint64
@@ -17,9 +20,11 @@ const (
 	CpuFeatureAmd64SSE4_1 CpuFeature = 1 << 19
 	// CpuFeatureAmd64SSE4_2 is the flag to query CpuFeatureFlags.Has for SSEv4.2 capabilities on amd64
 	CpuFeatureAmd64SSE4_2 CpuFeature = 1 << 20
+	// Note: when adding new features, ensure that the feature is included in CpuFeatureFlags.Raw.
 )
 
 const (
 	// CpuExtraFeatureAmd64ABM is the flag to query CpuFeatureFlags.HasExtra for Advanced Bit Manipulation capabilities (e.g. LZCNT) on amd64
 	CpuExtraFeatureAmd64ABM CpuFeature = 1 << 5
+	// Note: when adding new features, ensure that the feature is included in CpuFeatureFlags.Raw.
 )

internal/platform/cpuid_amd64.go

@@ -2,10 +2,10 @@
 
 package platform
 
-// CpuFeatures exposes the capabilities for this CPU, queried via the Has, HasExtra methods
-var CpuFeatures CpuFeatureFlags = loadCpuFeatureFlags()
+// CpuFeatures exposes the capabilities for this CPU, queried via the Has, HasExtra methods.
+var CpuFeatures = loadCpuFeatureFlags()
 
-// cpuFeatureFlags implements CpuFeatureFlags interface
+// cpuFeatureFlags implements CpuFeatureFlags interface.
 type cpuFeatureFlags struct {
 	flags      uint64
 	extraFlags uint64
@@ -15,13 +15,13 @@ type cpuFeatureFlags struct {
 // 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
+// cpuidAsBitmap combines the result of invoking cpuid to uint64 bitmap.
 func cpuidAsBitmap(arg1, arg2 uint32) uint64 {
 	_ /* eax */, _ /* ebx */, ecx, edx := cpuid(arg1, arg2)
 	return (uint64(edx) << 32) | uint64(ecx)
 }
 
-// loadStandardRange load flags from the standard range, panics otherwise
+// loadStandardRange load flags from the standard range, panics otherwise.
 func loadStandardRange(id uint32) uint64 {
 	// ensure that the id is in the valid range, returned by cpuid(0,0)
 	maxRange, _, _, _ := cpuid(0, 0)
@@ -31,7 +31,7 @@ func loadStandardRange(id uint32) uint64 {
 	return cpuidAsBitmap(id, 0)
 }
 
-// loadStandardRange load flags from the extended range, panics otherwise
+// loadStandardRange load flags from the extended range, panics otherwise.
 func loadExtendedRange(id uint32) uint64 {
 	// ensure that the id is in the valid range, returned by cpuid(0x80000000,0)
 	maxRange, _, _, _ := cpuid(0x80000000, 0)
@@ -48,12 +48,32 @@ func loadCpuFeatureFlags() CpuFeatureFlags {
 	}
 }
 
-// Has implements the same method on the CpuFeatureFlags interface
+// Has implements the same method on the CpuFeatureFlags interface.
 func (f *cpuFeatureFlags) Has(cpuFeature CpuFeature) bool {
 	return (f.flags & uint64(cpuFeature)) != 0
 }
 
-// HasExtra implements the same method on the CpuFeatureFlags interface
+// HasExtra implements the same method on the CpuFeatureFlags interface.
 func (f *cpuFeatureFlags) HasExtra(cpuFeature CpuFeature) bool {
 	return (f.extraFlags & uint64(cpuFeature)) != 0
 }
+
+// Raw implements the same method on the CpuFeatureFlags interface.
+func (f *cpuFeatureFlags) Raw() uint64 {
+	// Below, we only set the first 4 bits for the features we care about,
+	// instead of setting all the unnecessary bits obtained from the CPUID instruction.
+	var ret uint64
+	if f.Has(CpuFeatureAmd64SSE3) {
+		ret = 1 << 0
+	}
+	if f.Has(CpuFeatureAmd64SSE4_1) {
+		ret |= 1 << 1
+	}
+	if f.Has(CpuFeatureAmd64SSE4_2) {
+		ret |= 1 << 2
+	}
+	if f.HasExtra(CpuExtraFeatureAmd64ABM) {
+		ret |= 1 << 3
+	}
+	return ret
+}

internal/platform/cpuid_amd64_test.go

@@ -16,3 +16,15 @@ func TestAmd64CpuId_cpuHasFeature(t *testing.T) {
 	require.True(t, flags.HasExtra(CpuExtraFeatureAmd64ABM))
 	require.False(t, flags.HasExtra(1<<6)) // some other value
 }
+
+func TestAmd64CpuFeatureFlags_Raw(t *testing.T) {
+	flags := cpuFeatureFlags{
+		flags:      uint64(CpuFeatureAmd64SSE3 | CpuFeatureAmd64SSE4_1 | CpuFeatureAmd64SSE4_2),
+		extraFlags: uint64(CpuExtraFeatureAmd64ABM),
+	}
+	require.Equal(t, uint64(0b1111), flags.Raw())
+	flags.flags = 0
+	require.Equal(t, uint64(0b1000), flags.Raw())
+	flags.extraFlags = 0
+	require.Equal(t, uint64(0), flags.Raw())
+}

internal/platform/cpuid_unsupported.go

@@ -4,11 +4,14 @@ package platform
 
 var CpuFeatures CpuFeatureFlags = &cpuFeatureFlags{}
 
-// cpuFeatureFlags implements CpuFeatureFlags for unsupported platforms
+// cpuFeatureFlags implements CpuFeatureFlags for unsupported platforms.
 type cpuFeatureFlags struct{}
 
-// Has implements the same method on the CpuFeatureFlags interface
+// Has implements the same method on the CpuFeatureFlags interface.
 func (c *cpuFeatureFlags) Has(cpuFeature CpuFeature) bool { return false }
 
-// HasExtra implements the same method on the CpuFeatureFlags interface
+// HasExtra implements the same method on the CpuFeatureFlags interface.
 func (c *cpuFeatureFlags) HasExtra(cpuFeature CpuFeature) bool { return false }
+
+// Raw implements the same method on the CpuFeatureFlags interface.
+func (c *cpuFeatureFlags) Raw() uint64 { return 0 }