Skip to content

Commit

Permalink
[release/9.0] Empty slice fix (#107880)
Browse files Browse the repository at this point in the history
* empty slice fix

* PR comments and more testing

* fixing check

* fixing array rent init

* fixes from PR comments

---------

Co-authored-by: Michael Sharp <misharp@microsoft.com>
  • Loading branch information
github-actions[bot] and michaelgsharp committed Sep 16, 2024
1 parent 31cef05 commit c43d0e2
Show file tree
Hide file tree
Showing 6 changed files with 285 additions and 141 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
using System.Linq;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using static System.Runtime.InteropServices.JavaScript.JSType;
using EditorBrowsableAttribute = System.ComponentModel.EditorBrowsableAttribute;
using EditorBrowsableState = System.ComponentModel.EditorBrowsableState;

Expand Down Expand Up @@ -526,14 +525,14 @@ public void CopyTo(scoped TensorSpan<T> destination)
if (Rank > TensorShape.MaxInlineRank)
{
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
curIndexes = curIndexesArray;
curIndexes = curIndexes.Slice(0, Rank);
curIndexes = curIndexesArray.AsSpan(0, Rank);
}
else
{
curIndexesArray = null;
curIndexes = stackalloc nint[Rank];
}
curIndexes.Clear();

nint copiedValues = 0;
TensorSpan<T> slice = destination.Slice(_shape.Lengths);
Expand Down Expand Up @@ -573,14 +572,14 @@ public bool TryCopyTo(scoped TensorSpan<T> destination)
if (Rank > TensorShape.MaxInlineRank)
{
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
curIndexes = curIndexesArray;
curIndexes = curIndexes.Slice(0, Rank);
curIndexes = curIndexesArray.AsSpan(0, Rank);
}
else
{
curIndexesArray = null;
curIndexes = stackalloc nint[Rank];
}
curIndexes.Clear();

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

ReadOnlyTensorSpan<T> toReturn;
scoped Span<nint> lengths;
scoped Span<nint> offsets;
nint[]? lengthsArray;
nint[]? offsetsArray;
if (Rank > TensorShape.MaxInlineRank)
{
lengths = stackalloc nint[Rank];
offsets = stackalloc nint[Rank];
lengthsArray = ArrayPool<nint>.Shared.Rent(Rank);
lengths = lengthsArray.AsSpan(0, Rank);

offsetsArray = ArrayPool<nint>.Shared.Rent(Rank);
offsets = offsetsArray.AsSpan(0, Rank);
}
else
{
lengths = new nint[Rank];
offsets = new nint[Rank];
lengths = stackalloc nint[Rank];
offsets = stackalloc nint[Rank];

lengthsArray = null;
offsetsArray = null;
}
lengths.Clear();
offsets.Clear();

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

// FlattenedLength is computed everytime so using a local to cache the value.
nint flattenedLength = FlattenedLength;
nint index = 0;
for (int i = 0; i < offsets.Length; i++)

if (flattenedLength != 0)
{
index += Strides[i] * (offsets[i]);
for (int i = 0; i < offsets.Length; i++)
{
index += Strides[i] * (offsets[i]);
}
}

if (index >= _shape._memoryLength || index < 0)
if ((index >= _shape._memoryLength || index < 0) && flattenedLength != 0)
ThrowHelper.ThrowIndexOutOfRangeException();

return new ReadOnlyTensorSpan<T>(ref Unsafe.Add(ref _reference, index), lengths, _shape.Strides, _shape._memoryLength - index);
toReturn = new ReadOnlyTensorSpan<T>(ref Unsafe.Add(ref _reference, index), lengths, _shape.Strides, _shape._memoryLength - index);

if (offsetsArray != null)
ArrayPool<nint>.Shared.Return(offsetsArray);
if (lengthsArray != null)
ArrayPool<nint>.Shared.Return(lengthsArray);

return toReturn;
}

/// <summary>
Expand All @@ -698,14 +721,14 @@ public bool TryFlattenTo(scoped Span<T> destination)
if (Rank > TensorShape.MaxInlineRank)
{
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
curIndexes = curIndexesArray;
curIndexes = curIndexes.Slice(0, Rank);
curIndexes = curIndexesArray.AsSpan(0, Rank);
}
else
{
curIndexesArray = null;
curIndexes = stackalloc nint[Rank];
}
curIndexes.Clear();

nint copiedValues = 0;
while (copiedValues < _shape.FlattenedLength)
Expand Down Expand Up @@ -739,14 +762,14 @@ public void FlattenTo(scoped Span<T> destination)
if (Rank > TensorShape.MaxInlineRank)
{
curIndexesArray = ArrayPool<nint>.Shared.Rent(Rank);
curIndexes = curIndexesArray;
curIndexes = curIndexes.Slice(0, Rank);
curIndexes = curIndexesArray.AsSpan(0, Rank);
}
else
{
curIndexesArray = null;
curIndexes = stackalloc nint[Rank];
}
curIndexes.Clear();

nint copiedValues = 0;
while (copiedValues < _shape.FlattenedLength)
Expand Down
Loading

0 comments on commit c43d0e2

Please sign in to comment.