Repository: tetratelabs/wazero
Pull request #1327 (2023-03-31 05:58:37)

compiler: stop using map for label infos
_arm64_
```
goos: darwin
goarch: arm64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
                                    │   old.txt   │              new.txt              │
                                    │   sec/op    │   sec/op     vs base              │
Compilation/with_extern_cache-10      133.9µ ± 1%   134.4µ ± 1%       ~ (p=0.456 n=7)
Compilation/without_extern_cache-10   2.125m ± 1%   2.058m ± 1%  -3.14% (p=0.001 n=7)
geomean                               533.3µ        526.0µ       -1.37%

                                    │   old.txt    │              new.txt               │
                                    │     B/op     │     B/op      vs base              │
Compilation/with_extern_cache-10      53.39Ki ± 0%   53.38Ki ± 0%       ~ (p=0.779 n=7)
Compilation/without_extern_cache-10   1.115Mi ± 0%   1.082Mi ± 0%  -2.99% (p=0.001 n=7)
geomean                               246.9Ki        243.2Ki       -1.52%

                                    │   old.txt   │               new.txt               │
                                    │  allocs/op  │  allocs/op   vs base                │
Compilation/with_extern_cache-10       979.0 ± 0%    979.0 ± 0%       ~ (p=1.000 n=7) ¹
Compilation/without_extern_cache-10   8.440k ± 0%   7.956k ± 0%  -5.73% (p=0.001 n=7)
geomean                               2.875k        2.791k       -2.91%
¹ all samples are equal
```


_amd64_
```
goos: darwin
goarch: amd64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
cpu: VirtualApple @ 2.50GHz
                                    │ old_amd64.txt │           new_amd64.txt           │
                                    │    sec/op     │   sec/op     vs base              │
Compilation/with_extern_cache-10        328.8µ ± 1%   332.2µ ± 2%       ~ (p=0.097 n=7)
Compilation/without_extern_cache-10     4.679m ± 0%   4.550m ± 1%  -2.76% (p=0.001 n=7)
geomean                                 1.240m        1.229m       -0.89%

                                    │ old_amd64.txt │           new_amd64.txt            │
                                    │     B/op      │     B/op      vs base              │
Compilation/with_extern_cache-10       53.63Ki ± 0%   53.64Ki ± 0%       ~ (p=0.607 n=7)
Compilation/without_extern_cache-10    1.124Mi ± 0%   1.090Mi ± 0%  -3.00% (p=0.001 n=7)
geomean                                248.5Ki        244.7Ki       -1.51%

                                    │ old_amd64.txt │            new_amd64.txt            │
                                    │   allocs/op   │  allocs/op   vs base                │
Compilation/with_extern_cache-10         981.0 ± 0%    981.0 ± 0%       ~ (p=1.000 n=7) ¹
Compilation/without_extern_cache-10     9.391k ± 0%   8.906k ± 0%  -5.16% (p=0.001 n=7)
geomean                                 3.035k        2.956k       -2.62%

```


Patch:
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index ae48cb77..eb3e8dfb 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -89,7 +89,7 @@ type amd64Compiler struct {
 	// and each item is either placed in register or the actual memory stack.
 	locationStack runtimeValueLocationStack
 	// labels hold per wazeroir label specific information in this function.
-	labels map[wazeroir.LabelID]*amd64LabelInfo
+	labels [wazeroir.LabelKindNum][]amd64LabelInfo
 	// stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation.
 	stackPointerCeil uint64
 	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the max stack pointer is determined BEFORE generating native code.
@@ -106,17 +106,22 @@ func newAmd64Compiler() compiler {
 	return c
 }
 
+// Init implements compiler.Init.
 func (c *amd64Compiler) Init(ir *wazeroir.CompilationResult, withListener bool) {
-	assembler, vstack := c.assembler, c.locationStack
+	assembler, locationStack := c.assembler, c.locationStack
 	assembler.Reset()
-	vstack.reset()
+	locationStack.reset()
+	for i := range c.labels {
+		c.labels[i] = c.labels[i][:0]
+	}
 	*c = amd64Compiler{
-		labels:       map[wazeroir.LabelID]*amd64LabelInfo{},
-		ir:           ir,
-		cpuFeatures:  c.cpuFeatures,
-		withListener: withListener,
+		ir:            ir,
+		assembler:     assembler,
+		locationStack: locationStack,
+		cpuFeatures:   c.cpuFeatures,
+		withListener:  withListener,
+		labels:        c.labels,
 	}
-	c.assembler, c.locationStack = assembler, vstack
 }
 
 // runtimeValueLocationStack implements compilerImpl.runtimeValueLocationStack for the amd64 architecture.
@@ -157,12 +162,16 @@ type amd64LabelInfo struct {
 }
 
 func (c *amd64Compiler) label(labelID wazeroir.LabelID) *amd64LabelInfo {
-	ret, ok := c.labels[labelID]
-	if ok {
-		return ret
-	}
-	c.labels[labelID] = &amd64LabelInfo{}
-	return c.labels[labelID]
+	kind := labelID.Kind()
+	frames := c.labels[kind]
+	frameID := labelID.FrameID()
+	// If the frameID is not allocated yet, expand the slice by twice of the diff,
+	// so that we could reduce the allocation in the subsequent compilation.
+	if diff := frameID - len(frames) + 1; diff > 0 {
+		frames = append(frames, make([]amd64LabelInfo, diff*2)...)
+		c.labels[kind] = frames
+	}
+	return &frames[frameID]
 }
 
 // compileBuiltinFunctionCheckExitCode implements compiler.compileBuiltinFunctionCheckExitCode for the amd64 architecture.
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index f5dd550f..1329a93d 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -23,7 +23,7 @@ type arm64Compiler struct {
 	// and each item is either placed in register or the actual memory stack.
 	locationStack runtimeValueLocationStack
 	// labels maps a label (e.g. ".L1_then") to *arm64LabelInfo.
-	labels map[wazeroir.LabelID]*arm64LabelInfo
+	labels [wazeroir.LabelKindNum][]arm64LabelInfo
 	// stackPointerCeil is the greatest stack pointer value (from runtimeValueLocationStack) seen during compilation.
 	stackPointerCeil uint64
 	// onStackPointerCeilDeterminedCallBack hold a callback which are called when the ceil of stack pointer is determined before generating native code.
@@ -38,12 +38,18 @@ func newArm64Compiler() compiler {
 	}
 }
 
+// Init implements compiler.Init.
 func (c *arm64Compiler) Init(ir *wazeroir.CompilationResult, withListener bool) {
-	assembler, vstack := c.assembler, c.locationStack
+	assembler, locationStack := c.assembler, c.locationStack
 	assembler.Reset()
-	vstack.reset()
-	*c = arm64Compiler{labels: map[wazeroir.LabelID]*arm64LabelInfo{}, ir: ir, withListener: withListener}
-	c.assembler, c.locationStack = assembler, vstack
+	locationStack.reset()
+	for i := range c.labels {
+		c.labels[i] = c.labels[i][:0]
+	}
+	*c = arm64Compiler{
+		assembler: assembler, locationStack: locationStack,
+		ir: ir, withListener: withListener, labels: c.labels,
+	}
 }
 
 var (
@@ -135,13 +141,17 @@ type arm64LabelInfo struct {
 	initialStack runtimeValueLocationStack
 }
 
-func (c *arm64Compiler) label(labelKey wazeroir.LabelID) *arm64LabelInfo {
-	ret, ok := c.labels[labelKey]
-	if ok {
-		return ret
+func (c *arm64Compiler) label(labelID wazeroir.LabelID) *arm64LabelInfo {
+	kind := labelID.Kind()
+	frames := c.labels[kind]
+	frameID := labelID.FrameID()
+	// If the frameID is not allocated yet, expand the slice by twice of the diff,
+	// so that we could reduce the allocation in the subsequent compilation.
+	if diff := frameID - len(frames) + 1; diff > 0 {
+		frames = append(frames, make([]arm64LabelInfo, diff*2)...)
+		c.labels[kind] = frames
 	}
-	c.labels[labelKey] = &arm64LabelInfo{}
-	return c.labels[labelKey]
+	return &frames[frameID]
 }
 
 // runtimeValueLocationStack implements compilerImpl.runtimeValueLocationStack for the amd64 architecture.
diff --git a/internal/wazeroir/operations.go b/internal/wazeroir/operations.go
index a8be4bd8..fe7c03d3 100644
--- a/internal/wazeroir/operations.go
+++ b/internal/wazeroir/operations.go
@@ -819,6 +819,16 @@ type Label struct {
 // LabelID is the unique identifiers for blocks in a single function.
 type LabelID uint64
 
+// Kind returns the LabelKind encoded in this LabelID.
+func (l LabelID) Kind() LabelKind {
+	return LabelKind(uint32(l))
+}
+
+// FrameID returns the frame id encoded in this LabelID.
+func (l LabelID) FrameID() int {
+	return int(uint32(l >> 32))
+}
+
 // ID returns the LabelID for this Label.
 func (l Label) ID() (id LabelID) {
 	id = LabelID(l.Kind) | LabelID(l.FrameID)<<32
@@ -863,6 +873,7 @@ const (
 	// we have the continuation block (of if-block) corresponding to "return" opcode.
 	LabelKindContinuation
 	LabelKindReturn
+	LabelKindNum
 )
 
 func (l Label) asBranchTargetDrop() BranchTargetDrop {


Test patch:
diff --git a/internal/wazeroir/operations_test.go b/internal/wazeroir/operations_test.go
index 42c47250..1e5ae8f4 100644
--- a/internal/wazeroir/operations_test.go
+++ b/internal/wazeroir/operations_test.go
@@ -65,3 +65,12 @@ func TestUnionOperation_String(t *testing.T) {
 		require.NotEqual(t, "", op.String())
 	}
 }
+
+func TestLabelID(t *testing.T) {
+	for k := LabelKind(0); k < LabelKindNum; k++ {
+		l := Label{Kind: k, FrameID: 12345}
+		id := l.ID()
+		require.Equal(t, k, id.Kind())
+		require.Equal(t, int(l.FrameID), id.FrameID())
+	}
+}