Skip to content

Commit 4dac142

Browse files
authored
Merge pull request #160 from martindevans/GetState_fix
`GetState()` fix
2 parents b8d3fa1 + 832bf7d commit 4dac142

File tree

1 file changed

+12
-20
lines changed

1 file changed

+12
-20
lines changed

LLama/LLamaContext.cs

Lines changed: 12 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -195,32 +195,24 @@ public State GetState()
195195

196196
unsafe
197197
{
198-
var bigMemory = Marshal.AllocHGlobal((nint)stateSize);
199-
var smallMemory = IntPtr.Zero;
198+
// Allocate a chunk of memory large enough to hold the entire state
199+
var memory = Marshal.AllocHGlobal((nint)stateSize);
200200
try
201201
{
202-
// Copy the state data into "big memory", discover the actual size required
203-
var actualSize = _ctx.GetState(bigMemory, stateSize);
202+
// Copy the state data into memory, discover the actual size required
203+
var actualSize = _ctx.GetState(memory, stateSize);
204204

205-
// if big memory is nearly completely full (within 1MB) early exit and skip the extra copying
206-
if (actualSize >= stateSize - 1_000_000)
207-
return new State(bigMemory);
205+
// Shrink to size
206+
memory = Marshal.ReAllocHGlobal(memory, (nint)actualSize);
208207

209-
// Allocate a smaller buffer which is exactly the right size
210-
smallMemory = Marshal.AllocHGlobal((nint)actualSize);
211-
212-
// Copy into the smaller buffer and free the large one to save excess memory usage
213-
Buffer.MemoryCopy(bigMemory.ToPointer(), smallMemory.ToPointer(), actualSize, actualSize);
214-
215-
return new State(smallMemory);
208+
// Wrap memory in a state and return it
209+
memory = IntPtr.Zero;
210+
return new State(memory);
216211
}
217-
catch
212+
finally
218213
{
219-
if (bigMemory != IntPtr.Zero)
220-
Marshal.FreeHGlobal(bigMemory);
221-
if (smallMemory != IntPtr.Zero)
222-
Marshal.FreeHGlobal(smallMemory);
223-
throw;
214+
if (memory != IntPtr.Zero)
215+
Marshal.FreeHGlobal(memory);
224216
}
225217
}
226218
}

0 commit comments

Comments
 (0)