Skip to content

Commit 3a8c982

Browse files
michaelgsharpjtschuster
authored andcommitted
Empty slice fix (dotnet#107316)
* empty slice fix * PR comments and more testing * fixing check * fixing array rent init * fixes from PR comments
1 parent d3569ed commit 3a8c982

File tree

6 files changed

+285
-141
lines changed

6 files changed

+285
-141
lines changed

src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/ReadOnlyTensorSpan.cs

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
using System.Linq;
88
using System.Runtime.CompilerServices;
99
using System.Runtime.InteropServices;
10-
using static System.Runtime.InteropServices.JavaScript.JSType;
1110
using EditorBrowsableAttribute = System.ComponentModel.EditorBrowsableAttribute;
1211
using EditorBrowsableState = System.ComponentModel.EditorBrowsableState;
1312

@@ -528,14 +527,14 @@ public void CopyTo(scoped TensorSpan<T> destination)
528527
if (Rank > TensorShape.MaxInlineRank)
529528
{
530529
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
531-
curIndexes = curIndexesArray;
532-
curIndexes = curIndexes.Slice(0, Rank);
530+
curIndexes = curIndexesArray.AsSpan(0, Rank);
533531
}
534532
else
535533
{
536534
curIndexesArray = null;
537535
curIndexes = stackalloc nint[Rank];
538536
}
537+
curIndexes.Clear();
539538

540539
nint copiedValues = 0;
541540
TensorSpan<T> slice = destination.Slice(_shape.Lengths);
@@ -576,14 +575,14 @@ public bool TryCopyTo(scoped TensorSpan<T> destination)
576575
if (Rank > TensorShape.MaxInlineRank)
577576
{
578577
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
579-
curIndexes = curIndexesArray;
580-
curIndexes = curIndexes.Slice(0, Rank);
578+
curIndexes = curIndexesArray.AsSpan(0, Rank);
581579
}
582580
else
583581
{
584582
curIndexesArray = null;
585583
curIndexes = stackalloc nint[Rank];
586584
}
585+
curIndexes.Clear();
587586

588587
nint copiedValues = 0;
589588
TensorSpan<T> slice = destination.Slice(_shape.Lengths);
@@ -656,34 +655,58 @@ public ReadOnlyTensorSpan<T> Slice(params scoped ReadOnlySpan<NRange> ranges)
656655
if (ranges.Length != Lengths.Length)
657656
throw new ArgumentOutOfRangeException(nameof(ranges), "Number of dimensions to slice does not equal the number of dimensions in the span");
658657

658+
ReadOnlyTensorSpan<T> toReturn;
659659
scoped Span<nint> lengths;
660660
scoped Span<nint> offsets;
661+
nint[]? lengthsArray;
662+
nint[]? offsetsArray;
661663
if (Rank > TensorShape.MaxInlineRank)
662664
{
663-
lengths = stackalloc nint[Rank];
664-
offsets = stackalloc nint[Rank];
665+
lengthsArray = ArrayPool<nint>.Shared.Rent(Rank);
666+
lengths = lengthsArray.AsSpan(0, Rank);
667+
668+
offsetsArray = ArrayPool<nint>.Shared.Rent(Rank);
669+
offsets = offsetsArray.AsSpan(0, Rank);
665670
}
666671
else
667672
{
668-
lengths = new nint[Rank];
669-
offsets = new nint[Rank];
673+
lengths = stackalloc nint[Rank];
674+
offsets = stackalloc nint[Rank];
675+
676+
lengthsArray = null;
677+
offsetsArray = null;
670678
}
679+
lengths.Clear();
680+
offsets.Clear();
671681

672682
for (int i = 0; i < ranges.Length; i++)
673683
{
674684
(offsets[i], lengths[i]) = ranges[i].GetOffsetAndLength(Lengths[i]);
675685
}
676686

687+
// FlattenedLength is computed everytime so using a local to cache the value.
688+
nint flattenedLength = FlattenedLength;
677689
nint index = 0;
678-
for (int i = 0; i < offsets.Length; i++)
690+
691+
if (flattenedLength != 0)
679692
{
680-
index += Strides[i] * (offsets[i]);
693+
for (int i = 0; i < offsets.Length; i++)
694+
{
695+
index += Strides[i] * (offsets[i]);
696+
}
681697
}
682698

683-
if (index >= _shape._memoryLength || index < 0)
699+
if ((index >= _shape._memoryLength || index < 0) && flattenedLength != 0)
684700
ThrowHelper.ThrowIndexOutOfRangeException();
685701

686-
return new ReadOnlyTensorSpan<T>(ref Unsafe.Add(ref _reference, index), lengths, _shape.Strides, _shape._memoryLength - index);
702+
toReturn = new ReadOnlyTensorSpan<T>(ref Unsafe.Add(ref _reference, index), lengths, _shape.Strides, _shape._memoryLength - index);
703+
704+
if (offsetsArray != null)
705+
ArrayPool<nint>.Shared.Return(offsetsArray);
706+
if (lengthsArray != null)
707+
ArrayPool<nint>.Shared.Return(lengthsArray);
708+
709+
return toReturn;
687710
}
688711

689712
/// <summary>
@@ -700,14 +723,14 @@ public bool TryFlattenTo(scoped Span<T> destination)
700723
if (Rank > TensorShape.MaxInlineRank)
701724
{
702725
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
703-
curIndexes = curIndexesArray;
704-
curIndexes = curIndexes.Slice(0, Rank);
726+
curIndexes = curIndexesArray.AsSpan(0, Rank);
705727
}
706728
else
707729
{
708730
curIndexesArray = null;
709731
curIndexes = stackalloc nint[Rank];
710732
}
733+
curIndexes.Clear();
711734

712735
nint copiedValues = 0;
713736
while (copiedValues < _shape.FlattenedLength)
@@ -741,14 +764,14 @@ public void FlattenTo(scoped Span<T> destination)
741764
if (Rank > TensorShape.MaxInlineRank)
742765
{
743766
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
744-
curIndexes = curIndexesArray;
745-
curIndexes = curIndexes.Slice(0, Rank);
767+
curIndexes = curIndexesArray.AsSpan(0, Rank);
746768
}
747769
else
748770
{
749771
curIndexesArray = null;
750772
curIndexes = stackalloc nint[Rank];
751773
}
774+
curIndexes.Clear();
752775

753776
nint copiedValues = 0;
754777
while (copiedValues < _shape.FlattenedLength)

0 commit comments

Comments
 (0)