Skip to content

Commit

Permalink
Merge pull request #14 from kelnishi/kelvin-perf
Browse files Browse the repository at this point in the history 
Stack memory optimization
  • Loading branch information
@kelnishi
kelnishi authored Nov 18, 2024
2 parents 77e389e + 882f539 commit 779670a
Show file tree
Hide file tree
Showing 13 changed files with 148 additions and 75 deletions.
4 changes: 2 additions & 2 deletions Wacs.Core/Instructions/Control.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -723,8 +723,8 @@ public override void Execute(ExecContext context)
$"Instruction return failed. Operand stack underflow");
//We're managing separate stacks, so we won't need to shift the operands
// var vals = context.OpStack.PopResults(context.Frame.Type.ResultType);
var frame = context.PopFrame();
context.ResumeSequence(frame.ContinuationAddress);
var address = context.PopFrame();
context.ResumeSequence(address);
}
}

Expand Down
16 changes: 8 additions & 8 deletions Wacs.Core/Instructions/Variable.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ public override string RenderText(ExecContext? context)
if (!context.Frame.Locals.Contains(Index))
return $"{base.RenderText(context)} {Index.Value}";

var value = context.Frame.Locals[Index];
var value = context.Frame.Locals.Get(Index);
string valStr = $" (;>{value}<;)";
return $"{base.RenderText(context)} {Index.Value}{valStr}";
}
Expand All @@ -70,7 +70,7 @@ private static void ValidateLocalGet(IWasmValidationContext context, LocalIdx lo
{
context.Assert(context.Locals.Contains(localIndex),
$"Instruction local.get was invalid. Context Locals did not contain {localIndex}");
var value = context.Locals[localIndex];
var value = context.Locals.Get(localIndex);
context.OpStack.PushType(value.Type);
}

Expand All @@ -81,7 +81,7 @@ private static void ExecuteLocalGet(ExecContext context, LocalIdx localIndex)
context.Assert( context.Frame.Locals.Contains(localIndex),
$"Instruction local.get could not get Local {localIndex}");
//3.
var value = context.Frame.Locals[localIndex];
var value = context.Frame.Locals.Get(localIndex);
//4.
context.OpStack.PushValue(value);
}
Expand All @@ -91,7 +91,7 @@ private static void ValidateLocalSet(IWasmValidationContext context, LocalIdx lo
{
context.Assert(context.Locals.Contains(localIndex),
$"Instruction local.set was invalid. Context Locals did not contain {localIndex}");
var value = context.Locals[localIndex];
var value = context.Locals.Get(localIndex);
context.OpStack.PopType(value.Type);
}

Expand All @@ -104,12 +104,12 @@ private static void ExecuteLocalSet(ExecContext context, LocalIdx localIndex)
//3.
context.Assert( context.OpStack.HasValue,
$"Operand Stack underflow in instruction local.set");
var localValue = context.Frame.Locals[localIndex];
var localValue = context.Frame.Locals.Get(localIndex);
var type = localValue.Type;
//4.
var value = context.OpStack.PopType(type);
//5.
context.Frame.Locals[localIndex] = value;
context.Frame.Locals.Set(localIndex, value);
}

//0x22
Expand All @@ -118,7 +118,7 @@ private static void ValidateLocalTee(IWasmValidationContext context, LocalIdx lo
{
context.Assert(context.Locals.Contains(localIndex),
$"Instruction local.tee was invalid. Context Locals did not contain {localIndex}");
var value = context.Locals[localIndex];
var value = context.Locals.Get(localIndex);
context.OpStack.PopType(value.Type);
context.OpStack.PushType(value.Type);
context.OpStack.PushType(value.Type);
Expand All @@ -131,7 +131,7 @@ private static void ExecuteLocalTee(ExecContext context, LocalIdx localIndex)
//1.
context.Assert( context.OpStack.HasValue,
$"Operand Stack underflow in instruction local.tee");
var localValue = context.Frame.Locals[localIndex];
var localValue = context.Frame.Locals.Get(localIndex);
//2.
var value = context.OpStack.PopType(localValue.Type);
//3.
Expand Down
55 changes: 29 additions & 26 deletions Wacs.Core/Runtime/ExecContext.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
// */

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
Expand Down Expand Up @@ -51,13 +52,15 @@ public ExecContext(Store store, RuntimeAttributes? attributes = default)
{
Store = store;
Attributes = attributes ?? new RuntimeAttributes();
FrameStack = new(Attributes.InitialCallStack, Attributes.GrowCallStack);
LabelStack = new(Attributes.InitialCallStack * 2, Attributes.GrowCallStack);
CallStack = FrameStack.GetSubStack();

var poolPolicy = new LocalsSpacePooledObjectPolicy();
LocalsPool = new DefaultObjectPool<LocalsSpace>(poolPolicy, Attributes.InitialCallStack);
FramePool = new DefaultObjectPool<Frame>(new StackPoolPolicy<Frame>(), Attributes.MaxCallStack)
.Prime(Attributes.InitialCallStack);

LocalsDataPool = ArrayPool<Value>.Create(Attributes.MaxFunctionLocals, Attributes.LocalPoolSize);

CallStack = new (Attributes.InitialCallStack);
LabelStack = new(Attributes.InitialLabelsStack, Attributes.GrowLabelsStack);

_hostReturnSequence = new InstructionSequence(
InstructionFactory.CreateInstruction<InstFuncReturn>(OpCode.Func)
);
Expand All @@ -75,11 +78,10 @@ public ExecContext(Store store, RuntimeAttributes? attributes = default)
public Store Store { get; }
public OpStack OpStack { get; }

private ReusableStack<Frame> FrameStack { get; }
private SubStack<Frame> CallStack { get; }
private ReusableStack<Label> LabelStack { get; }

private DefaultObjectPool<LocalsSpace> LocalsPool { get; }
private ObjectPool<Frame> FramePool { get; }
private ArrayPool<Value> LocalsDataPool { get; }
private Stack<Frame> CallStack { get; }

public Frame Frame => CallStack.Peek();

Expand Down Expand Up @@ -107,14 +109,16 @@ public Frame ReserveFrame(
{
locals ??= Array.Empty<ValType>();

var frame = CallStack.Reserve();
var frame = FramePool.Get();
frame.Module = module;
frame.Type = type;
frame.Labels = LabelStack.GetSubStack();
frame.Index = index;
frame.ContinuationAddress = GetPointer();
frame.Locals = LocalsPool.Get();
frame.Locals.Enscribe(type.ParameterTypes.Types, locals);

int capacity = type.ParameterTypes.Types.Length + locals.Length;
var localData = LocalsDataPool.Rent(capacity);
frame.Locals = new(localData, type.ParameterTypes.Types, locals);

return frame;
}
Expand All @@ -125,18 +129,18 @@ public void PushFrame(Frame frame)
throw new WasmRuntimeException($"Runtime call stack exhausted {CallStack.Count}");

CallStack.Push(frame);
frame.Labels = LabelStack.GetSubStack();
}

public Frame PopFrame()
public InstructionPointer PopFrame()
{
var frame = CallStack.Pop();
frame.Labels.Drop();

LocalsPool.Return(frame.Locals);
frame.Locals = null!;

return frame;
if (frame.Locals.Data != null)
frame.ReturnLocals(LocalsDataPool);

var address = frame.ContinuationAddress;
FramePool.Return(frame);
return address;
}

public void ResetStack(Label label)
Expand All @@ -147,7 +151,6 @@ public void ResetStack(Label label)
}
}


public void FlushCallStack()
{
while (CallStack.Count > 0)
Expand Down Expand Up @@ -193,10 +196,10 @@ public void EnterBlock(Block block, ResultType resultType, ByteCode inst, Stack<
// @Spec 4.4.9.2. Exit Block
public void ExitBlock()
{
var label = Frame.Labels.Pop();
var addr = Frame.PopLabel();
// We manage separate stacks, so we don't need to relocate the operands
// var vals = OpStack.PopResults(label.Type);
ResumeSequence(label.ContinuationAddress);
ResumeSequence(addr);
}

// @Spec 4.4.10.1 Function Invocation
Expand Down Expand Up @@ -243,13 +246,13 @@ private void Invoke(FunctionInstance wasmFunc, FuncIdx idx)
//Load parameters
while (_asideVals.Count > 0)
{
frame.Locals[(LocalIdx)li] = _asideVals.Pop();
frame.Locals.Set((LocalIdx)li, _asideVals.Pop());
li += 1;
}
//Set the Locals to default
for (int ti = 0; li < localCount; ++li, ++ti)
{
frame.Locals[(LocalIdx)li] = new Value(t[ti]);
frame.Locals.Set((LocalIdx)li, new Value(t[ti]));
}

//9.
Expand Down Expand Up @@ -286,10 +289,10 @@ public void FunctionReturn()
// var vals = OpStack.PopResults(Frame.Type.ResultType);
//5.
//6.
var frame = PopFrame();
var address = PopFrame();
//7. split stack, values left in place
//8.
ResumeSequence(frame.ContinuationAddress);
ResumeSequence(address);
}

public IInstruction? Next()
Expand Down
29 changes: 25 additions & 4 deletions Wacs.Core/Runtime/Frame.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,25 @@
// * limitations under the License.
// */

using System.Buffers;
using System.Collections.Generic;
using Microsoft.Extensions.ObjectPool;
using Wacs.Core.OpCodes;
using Wacs.Core.Runtime.Types;
using Wacs.Core.Types;
using Wacs.Core.Utilities;

namespace Wacs.Core.Runtime
{
public class Frame : IReusable<Frame>
public class Frame : IPoolable
{
// public ObjectPool<Label>? LabelPool = null;

public SubStack<Label> Labels;

// public readonly Stack<Label> Labels = new();
public ModuleInstance Module { get; set; } = null!;
public LocalsSpace Locals { get; set; } = null!;
public LocalsSpace Locals { get; set; }
public InstructionPointer ContinuationAddress { get; set; } = InstructionPointer.Nil;
public FunctionType Type { get; set; } = null!;
public FuncIdx Index { get; set; }
Expand All @@ -38,13 +45,19 @@ public void Clear()
{
Labels = default;
Module = default!;
Locals = default!;
Locals = default;
ContinuationAddress = default;
Type = default!;
Index = default!;
FuncId = string.Empty;
}

public void ReturnLocals(ArrayPool<Value> dataPool)
{
dataPool.Return(Locals.Data);
Locals = default!;
}

public bool Equals(Frame other)
{
return ReferenceEquals(this, other);
Expand All @@ -66,8 +79,16 @@ public void ForceLabels(int depth)

while (Labels.Count > depth)
{
Labels.Pop();
PopLabel();
}
}

public InstructionPointer PopLabel()
{
var label = Labels.Pop();
var addr = label.ContinuationAddress;
// LabelPool?.Return(label);
return addr;
}
}
}
2 changes: 1 addition & 1 deletion Wacs.Core/Runtime/Label.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@

namespace Wacs.Core.Runtime
{
public class Label : IReusable<Label>
public class Label : IPoolable
{
public int Arity;

Expand Down
4 changes: 4 additions & 0 deletions Wacs.Core/Runtime/RuntimeAttributes.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,11 @@ public class RuntimeAttributes
public bool Live = true;
public int MaxCallStack = 2048;

public int GrowLabelsStack = 512;
public int InitialLabelsStack = 2048;

public int MaxFunctionLocals = 2048;
public int LocalPoolSize = 64;

public int MaxOpStack = 1024;
public IInstructionFactory InstructionFactory { get; set; } = SpecFactory.Factory;
Expand Down
5 changes: 2 additions & 3 deletions Wacs.Core/Runtime/WasmRuntime.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -675,9 +675,8 @@ private ModuleInstance AllocateModule(Module module)
}

Context.PopFrame();
initFrame.Index = FuncIdx.ElementInitializers;
initFrame.ForceLabels(0);
Context.PushFrame(initFrame);
Frame initFrame2 = Context.ReserveFrame(moduleInstance, FunctionType.Empty, FuncIdx.ElementInitializers);
Context.PushFrame(initFrame2);

//6. Allocate Elements
//12. index ordered element addresses
Expand Down
45 changes: 26 additions & 19 deletions Wacs.Core/Types/IndexSpace.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
// */

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using Wacs.Core.Runtime;
Expand Down Expand Up @@ -179,40 +180,46 @@ public override bool Contains(GlobalIdx idx) =>
idx.Value < _imports.Count + _globals.Count;
}

public class LocalsSpace : AbstractIndexSpace<LocalIdx, Value>
public struct LocalsSpace
{
private readonly List<Value> _data = new();
public Value[]? Data;
private int _capacity;

public LocalsSpace() { }
public int Capacity => _capacity;

public Value Get(LocalIdx idx)
{
if (Data == null)
throw new InvalidOperationException("LocalSpace was used uninitialized.");

return Data[(Index)idx];
}

public override Value this[LocalIdx idx]
public void Set(LocalIdx idx, Value value)
{
get => _data[(Index)idx];
set => _data[(Index)idx] = value;
if (Data == null)
throw new InvalidOperationException("LocalSpace was used uninitialized.");

Data[(Index)idx] = value;
}

public void Enscribe(ValType[] parameters, ValType[] locals)
public LocalsSpace(Value[] data, ValType[] parameters, ValType[] locals)
{
_data.Clear();
// _data.Capacity = parameters.Length + locals.Length;

_capacity = parameters.Length + locals.Length;
Data = data;
int idx = 0;
foreach (var t in parameters)
{
_data.Add(new Value(t));
Data[idx++] = new Value(t);
}
foreach (var t in locals)
{
_data.Add(new Value(t));
Data[idx++] = new Value(t);
}
}

public void Reset()
{
_data.Clear();
}

public override bool Contains(LocalIdx idx) =>
idx.Value < _data.Count;
public bool Contains(LocalIdx idx) =>
idx.Value < _capacity;
}

public class ElementsSpace : AbstractIndexSpace<ElemIdx, Module.ElementSegment>
Expand Down
Loading

0 comments on commit 779670a

Please sign in to comment.