Repository: tetratelabs/wazero
Pull request #2220 (2024-05-28 01:52:08)

cache: includes cpu features in cache keys
there's a discussion about "offline" compilation, and if that happens, 
the CPU feature incompatibility becomes the problem.
This hardens the cache against it by including the cpu id in the file cache keys 
so that the content will be used exactly on the same CPUs.

Note that this impacts only amd64 target.


Patch:
diff --git a/config.go b/config.go
index 819a76df..abc0d4b2 100644
--- a/config.go
+++ b/config.go
@@ -148,7 +148,7 @@ type RuntimeConfig interface {
 	//	customSections := c.CustomSections()
 	WithCustomSections(bool) RuntimeConfig
 
-	// WithCloseOnContextDone ensures the executions of functions to be closed under one of the following circumstances:
+	// WithCloseOnContextDone ensures the executions of functions to be terminated under one of the following circumstances:
 	//
 	// 	- context.Context passed to the Call method of api.Function is canceled during execution. (i.e. ctx by context.WithCancel)
 	// 	- context.Context passed to the Call method of api.Function reaches timeout during execution. (i.e. ctx by context.WithTimeout or context.WithDeadline)
@@ -159,6 +159,8 @@ type RuntimeConfig interface {
 	// entire underlying OS thread which runs the api.Function call. See "Why it's safe to execute runtime-generated
 	// machine codes against async Goroutine preemption" section in RATIONALE.md for detail.
 	//
+	// Upon the termination of the function executions, api.Module is closed.
+	//
 	// Note that this comes with a bit of extra cost when enabled. The reason is that internally this forces
 	// interpreter and compiler runtimes to insert the periodical checks on the conditions above. For that reason,
 	// this is disabled by default.
diff --git a/internal/engine/wazevo/engine_cache.go b/internal/engine/wazevo/engine_cache.go
index f7c0450a..e49353dc 100644
--- a/internal/engine/wazevo/engine_cache.go
+++ b/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
 }
diff --git a/internal/platform/cpuid.go b/internal/platform/cpuid.go
index 25d7d3fd..0dc6ec19 100644
--- a/internal/platform/cpuid.go
+++ b/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.
 )
diff --git a/internal/platform/cpuid_amd64.go b/internal/platform/cpuid_amd64.go
index 8c9f1a9f..fbdb5393 100644
--- a/internal/platform/cpuid_amd64.go
+++ b/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
+}
diff --git a/internal/platform/cpuid_unsupported.go b/internal/platform/cpuid_unsupported.go
index 8ae826d3..291bcea6 100644
--- a/internal/platform/cpuid_unsupported.go
+++ b/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 }


Test patch:
diff --git a/internal/engine/wazevo/backend/isa/amd64/machine_test.go b/internal/engine/wazevo/backend/isa/amd64/machine_test.go
index 5dd05081..b870ae4b 100644
--- a/internal/engine/wazevo/backend/isa/amd64/machine_test.go
+++ b/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 }
diff --git a/internal/platform/cpuid_amd64_test.go b/internal/platform/cpuid_amd64_test.go
index 001914fe..626f63e5 100644
--- a/internal/platform/cpuid_amd64_test.go
+++ b/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())
+}