@@ -11,10 +11,11 @@ namespace AiDotNet.Memory;
1111/// and GC pressure during the forward pass. Tensors are pre-created during warmup
1212/// and recycled across calls.
1313/// </summary>
14- public sealed class ForwardArena < T >
14+ internal sealed class ForwardArena < T >
1515{
1616 private readonly Dictionary < ShapeKey , Tensor < T > [ ] > _slabs = new ( ) ;
1717 private readonly Dictionary < ShapeKey , int > _cursors = new ( ) ;
18+ private ShapeKey [ ] ? _cursorKeysCache ;
1819 private const int DefaultSlabSize = 4 ;
1920 private const int GrowthFactor = 2 ;
2021
@@ -71,9 +72,12 @@ public Tensor<T> RentUninitialized(int[] shape)
7172 [ MethodImpl ( MethodImplOptions . AggressiveInlining ) ]
7273 public void Reset ( )
7374 {
74- var keys = new List < ShapeKey > ( _cursors . Keys ) ;
75- foreach ( var key in keys )
76- _cursors [ key ] = 0 ;
75+ // Use cached keys array to avoid allocation during reset
76+ if ( _cursorKeysCache is null || _cursorKeysCache . Length != _cursors . Count )
77+ _cursorKeysCache = new ShapeKey [ _cursors . Count ] ;
78+ _cursors . Keys . CopyTo ( _cursorKeysCache , 0 ) ;
79+ for ( int i = 0 ; i < _cursorKeysCache . Length ; i ++ )
80+ _cursors [ _cursorKeysCache [ i ] ] = 0 ;
7781 }
7882
7983 /// <summary>
@@ -136,7 +140,7 @@ private Tensor<T> GrowAndRent(ShapeKey key, int[] shape, bool clear = true)
136140
137141 public ShapeKey ( int [ ] shape )
138142 {
139- _dims = shape ;
143+ _dims = ( int [ ] ) shape . Clone ( ) ;
140144 unchecked
141145 {
142146 int hash = ( int ) 2166136261 ;
0 commit comments