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() - *c = amd64Compiler{ - labels: map[wazeroir.LabelID]*amd64LabelInfo{}, - ir: ir, - cpuFeatures: c.cpuFeatures, - withListener: withListener, + locationStack.reset() + for i := range c.labels { + c.labels[i] = c.labels[i][:0] + } + *c = amd64Compiler{ + 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 + 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 } - c.labels[labelID] = &amd64LabelInfo{} - return c.labels[labelID] + 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 { 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()) + } +}