Skip to content

Commit f43f8ee

Browse files
Copilottarekgh
andcommitted
Remove stackalloc for State array; always use ArrayPool since method is only called for >128 bytes
Co-authored-by: tarekgh <10833894+tarekgh@users.noreply.github.com>
1 parent 8984a3d commit f43f8ee

File tree

1 file changed

+3
-9
lines changed

1 file changed

+3
-9
lines changed

src/Microsoft.ML.Tokenizers/Utils/BytePairEncoder.cs

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -155,12 +155,9 @@ public int CompareTo(MergeEntry other)
155155

156156
private static (int Id, int TokenIndex, int TokenLength)[] BytePairEncodeLarge(ReadOnlyMemory<byte> mergingBytes, IReadOnlyDictionary<ReadOnlyMemory<byte>, int> ranks, ReadOnlySpan<int> indexMappingSpan)
157157
{
158-
State[]? statePoolArray = null;
159158
int stateLength = mergingBytes.Length;
160-
Span<State> state = stateLength <= 256 ?
161-
stackalloc State[256] :
162-
(statePoolArray = ArrayPool<State>.Shared.Rent(stateLength));
163-
state = state.Slice(0, stateLength);
159+
State[] statePoolArray = ArrayPool<State>.Shared.Rent(stateLength);
160+
Span<State> state = statePoolArray.AsSpan(0, stateLength);
164161

165162
state[0] = new State
166163
{
@@ -280,10 +277,7 @@ void PotentialMerge(Span<State> stateSpan, PriorityQueue<MergeEntry> heapQueue,
280277
currentIndex = state[currentIndex].End;
281278
}
282279

283-
if (statePoolArray is not null)
284-
{
285-
ArrayPool<State>.Shared.Return(statePoolArray);
286-
}
280+
ArrayPool<State>.Shared.Return(statePoolArray);
287281

288282
return resultList.ToArray();
289283
}

0 commit comments

Comments
 (0)