diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Tensor.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Tensor.cs index 29e3cd48f60fa2..c7df0e00a445c7 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Tensor.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Tensor.cs @@ -34,6 +34,8 @@ public sealed class Tensor internal readonly nint[] _strides; /// If the backing memory is permanently pinned (so not just using a fixed statement). internal readonly bool _isPinned; + /// The offset of the first element in the backing memory. + internal readonly int _memoryOffset; /// /// Creates a new empty Tensor. @@ -44,13 +46,14 @@ internal Tensor() _values = []; _lengths = []; _strides = []; + _memoryOffset = 0; } [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal Tensor(T[]? values, ReadOnlySpan lengths, bool isPinned = false) : this(values, lengths, Array.Empty(), isPinned) { } + internal Tensor(T[]? values, ReadOnlySpan lengths, bool isPinned = false, int memoryOffset = 0) : this(values, lengths, Array.Empty(), isPinned, memoryOffset) { } [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal Tensor(T[]? values, ReadOnlySpan lengths, ReadOnlySpan strides, bool isPinned = false) + internal Tensor(T[]? values, ReadOnlySpan lengths, ReadOnlySpan strides, bool isPinned = false, int memoryOffset = 0) { if (values == null) { @@ -60,10 +63,12 @@ internal Tensor(T[]? values, ReadOnlySpan lengths, ReadOnlySpan stri _values = []; _lengths = []; _strides = []; + _memoryOffset = memoryOffset; return; // returns default } _lengths = lengths.IsEmpty ? [values.Length] : lengths.ToArray(); + _memoryOffset = memoryOffset; _flattenedLength = TensorSpanHelpers.CalculateTotalLength(_lengths); _strides = strides.IsEmpty ? TensorSpanHelpers.CalculateStrides(_lengths, _flattenedLength) : strides.ToArray(); @@ -386,7 +391,7 @@ public Tensor this[Tensor filter] /// Converts this to a pointing to the same backing memory."/> /// /// - public TensorSpan AsTensorSpan() => new TensorSpan(ref MemoryMarshal.GetArrayDataReference(_values), _lengths, _strides, _flattenedLength); + public TensorSpan AsTensorSpan() => new TensorSpan(ref Unsafe.Add(ref MemoryMarshal.GetArrayDataReference(_values), _memoryOffset), _lengths, _strides, _values.Length - _memoryOffset); /// /// Converts this to a pointing to the same backing memory based on the provided ranges."/> @@ -454,26 +459,71 @@ public Tensor this[Tensor filter] /// Forms a slice out of the given tensor /// /// The ranges for the slice - /// as a copy of the provided ranges. - // REVIEW: CURRENTLY DOES A COPY. + /// without copying the provided ranges. public Tensor Slice(params ReadOnlySpan start) { if (start.Length != Lengths.Length) throw new ArgumentOutOfRangeException(nameof(start), "Number of dimensions to slice does not equal the number of dimensions in the span"); - TensorSpan s = AsTensorSpan(start); - T[] values = _isPinned ? GC.AllocateArray(checked((int)s.FlattenedLength), _isPinned) : (new T[s.FlattenedLength]); - var outTensor = new Tensor(values, s.Lengths.ToArray(), _isPinned); - s.CopyTo(outTensor); - return outTensor; + scoped Span lengths; + scoped Span offsets; + nint[]? lengthsArray; + nint[]? offsetsArray; + if (Rank > TensorShape.MaxInlineRank) + { + lengthsArray = ArrayPool.Shared.Rent(Rank); + lengths = lengthsArray.AsSpan(0, Rank); + + offsetsArray = ArrayPool.Shared.Rent(Rank); + offsets = offsetsArray.AsSpan(0, Rank); + } + else + { + lengths = stackalloc nint[Rank]; + offsets = stackalloc nint[Rank]; + + lengthsArray = null; + offsetsArray = null; + } + lengths.Clear(); + offsets.Clear(); + + for (int i = 0; i < start.Length; i++) + { + (offsets[i], lengths[i]) = start[i].GetOffsetAndLength(Lengths[i]); + } + + // When we have an empty Tensor and someone wants to slice all of it, we should return an empty Tensor. + // FlattenedLength is computed everytime so using a local to cache the value. + nint flattenedLength = FlattenedLength; + int memoryOffset = 0; + + if (flattenedLength != 0) + { + for (int i = 0; i < offsets.Length; i++) + { + memoryOffset += (int)(Strides[i] * offsets[i]); + } + } + + if ((memoryOffset >= _values.Length || memoryOffset < 0) && flattenedLength != 0) + ThrowHelper.ThrowIndexOutOfRangeException(); + + Tensor toReturn = new Tensor(_values, lengths, Strides, _isPinned, memoryOffset); + + if (offsetsArray != null) + ArrayPool.Shared.Return(offsetsArray); + if (lengthsArray != null) + ArrayPool.Shared.Return(lengthsArray); + + return toReturn; } /// /// Forms a slice out of the given tensor /// /// The start indexes for the slice - /// as a copy of the provided ranges. - // REVIEW: CURRENTLY DOES A COPY. + /// without copying the provided ranges. public Tensor Slice(params ReadOnlySpan start) { NRange[] ranges = new NRange[start.Length]; @@ -488,8 +538,7 @@ public Tensor Slice(params ReadOnlySpan start) /// Forms a slice out of the given tensor /// /// The start indexes for the slice - /// as a copy of the provided ranges. - // REVIEW: CURRENTLY DOES A COPY. + /// without copying the provided ranges. public Tensor Slice(params ReadOnlySpan startIndex) { NRange[] ranges = new NRange[startIndex.Length]; diff --git a/src/libraries/System.Numerics.Tensors/tests/TensorSpanTests.cs b/src/libraries/System.Numerics.Tensors/tests/TensorSpanTests.cs index b57472302f98aa..0d0fc167b6bfcd 100644 --- a/src/libraries/System.Numerics.Tensors/tests/TensorSpanTests.cs +++ b/src/libraries/System.Numerics.Tensors/tests/TensorSpanTests.cs @@ -239,43 +239,31 @@ public void TensorExtensionsSpanInTOut(TensorPrimitivesSpanInTOut tensorPr T[] data = new T[length]; FillTensor(data); - Tensor x = Tensor.Create(data, tensorLength, []); + Tensor tensor = Tensor.Create(data, tensorLength, []); T expectedOutput = tensorPrimitivesOperation((ReadOnlySpan)data); - T results = tensorOperation(x); + T results = tensorOperation(tensor); Assert.Equal(expectedOutput, results); - float[] testData = [49.788437f, 32.736755f, -0.25761032f, -46.402596f, 4.5581512f, 21.813591f, 44.976646f, 12.691814f, -44.188023f, 40.35988f, -6.999405f, 4.713642f, 5.274975f, 21.312515f, -12.536407f, -34.888573f, -1.90839f, 28.734451f, -38.64155f, -28.840702f, 7.373543f, 18.600182f, 26.007828f, 0.71430206f, -6.8293495f, -13.327972f, -25.149017f, 9.331852f, 40.87751f, 28.321632f, 42.918175f, 25.213333f, -41.392017f, 36.727768f, 26.49012f, 3.8807983f, 24.933182f, -43.050568f, -42.6283f, 18.01947f, -47.62874f, -49.94487f, -1.036602f, -37.086433f, 32.77098f, -12.903477f, -45.100212f, -20.596504f, 33.67714f, 46.864395f, 44.437485f, -44.092155f, 37.122124f, 25.220505f, 41.994873f, -13.3394165f, -28.193134f, -21.329712f, -36.623306f, 3.3981133f, -26.475079f, 16.339478f, -44.07065f, 36.321762f, -24.63433f, 28.652397f, 4.096817f, 33.29615f, -2.3503838f, -7.509815f, 42.943604f, -32.52115f, -0.20326233f, 29.554626f, 18.044052f]; - nint[] testLengths = [5, 3, 5]; - Tensor testTensor = Tensor.Create(testData, testLengths, []); - float[] testSliceData = new float[75]; - testTensor.FlattenTo(testSliceData); - float testExpectedOutput = TensorPrimitives.Sum((ReadOnlySpan)testSliceData); - float testResults = Tensor.Sum(testTensor); - - // Now test if the source is sliced to be non contiguous that it still gives expected result. NRange[] sliceLengths = Helpers.TensorSliceShapes[index].Select(i => new NRange(0, i)).ToArray(); nint sliceFlattenedLength = CalculateTotalLength(Helpers.TensorSliceShapes[index]); - x = x.Slice(sliceLengths); + tensor = tensor.Slice(sliceLengths); T[] sliceData = new T[sliceFlattenedLength]; - x.FlattenTo(sliceData); + tensor.FlattenTo(sliceData); - IEnumerator enumerator = x.GetEnumerator(); + IEnumerator enumerator = tensor.GetEnumerator(); bool cont = enumerator.MoveNext(); - ReadOnlySpan span = MemoryMarshal.CreateSpan(ref x.AsReadOnlyTensorSpan()._reference, (int)x.FlattenedLength); int i = 0; - Assert.True(span.SequenceEqual(sliceData)); + Assert.True(tensor.SequenceEqual(sliceData)); while (cont) { - Assert.Equal(sliceData[i], enumerator.Current); - Assert.Equal(span[i], enumerator.Current); - Assert.Equal(span[i], sliceData[i++]); + Assert.Equal(sliceData[i++], enumerator.Current); cont = enumerator.MoveNext(); } expectedOutput = tensorPrimitivesOperation((ReadOnlySpan)sliceData); - results = tensorOperation(x); + results = tensorOperation(tensor); Assert.Equal(expectedOutput, results); }); diff --git a/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs b/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs index 607d09788db1ce..4982553c55760d 100644 --- a/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs +++ b/src/libraries/System.Numerics.Tensors/tests/TensorTests.cs @@ -1582,12 +1582,12 @@ public static void TensorClearTest() Assert.Equal(0, slice[1, 0]); Assert.Equal(0, slice[1, 1]); - // Since Tensor.Slice does a copy the original tensor shouldn't be modified - Assert.Equal(1, tensor[0, 0]); - Assert.Equal(2, tensor[0, 1]); + // Since Tensor.Slice does do a copy the original tensor should be modified but only in the slice we took. + Assert.Equal(0, tensor[0, 0]); + Assert.Equal(0, tensor[0, 1]); Assert.Equal(3, tensor[0, 2]); - Assert.Equal(4, tensor[1, 0]); - Assert.Equal(5, tensor[1, 1]); + Assert.Equal(0, tensor[1, 0]); + Assert.Equal(0, tensor[1, 1]); Assert.Equal(6, tensor[1, 2]); Assert.Equal(7, tensor[2, 0]); Assert.Equal(8, tensor[2, 1]); @@ -1609,8 +1609,8 @@ public static void TensorClearTest() slice.Clear(); Assert.Equal(0, slice[0]); - // Since Tensor.Slice does a copy the original tensor shouldn't be modified - Assert.Equal(1, tensor[0]); + // Since Tensor.Slice does do a copy the original tensor should be modified but only in the slice we took. + Assert.Equal(0, tensor[0]); Assert.Equal(2, tensor[1]); Assert.Equal(3, tensor[2]); Assert.Equal(4, tensor[3]);