diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Common/TensorPrimitives.IIndexOfMinMaxOperator.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Common/TensorPrimitives.IIndexOfMinMaxOperator.cs index 12d290ab014f04..c1424050bb62a9 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Common/TensorPrimitives.IIndexOfMinMaxOperator.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/Common/TensorPrimitives.IIndexOfMinMaxOperator.cs @@ -1,25 +1,59 @@ -// Licensed to the .NET Foundation under one or more agreements. +// Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. using System.Diagnostics; using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; using System.Runtime.Intrinsics; namespace System.Numerics.Tensors { public static unsafe partial class TensorPrimitives { + /// The operations an IndexOfMin/Max-style search needs from its ordering. + /// + /// returns whether x should replace y: strictly better, or an equal value the + /// operator orders by sign (-0 before +0 and the like). and + /// must return an element that ranks no worse than any input, and for floating-point types must + /// propagate NaN, so that a block containing a NaN reduces to NaN. + /// private interface IIndexOfMinMaxOperator { static abstract T Aggregate(Vector128 value); static abstract T Aggregate(Vector256 value); static abstract T Aggregate(Vector512 value); + static abstract T Reduce(T x, T y); + static abstract Vector128 Reduce(Vector128 x, Vector128 y); + static abstract Vector256 Reduce(Vector256 x, Vector256 y); + static abstract Vector512 Reduce(Vector512 x, Vector512 y); static abstract bool Compare(T x, T y); static abstract Vector128 Compare(Vector128 x, Vector128 y); static abstract Vector256 Compare(Vector256 x, Vector256 y); static abstract Vector512 Compare(Vector512 x, Vector512 y); } + /// Number of vectors per block in the block-reduction search (256 ints per block at 256 bits). + private const int BlockVectors = 32; + + /// + /// Finds the index of the best element of under , or -1 for an empty span, + /// with a two-pass block reduction: pass 1 reduces every block to its best element with a pure vector loop (one load and one + /// per vector, no index tracking and no blends) and + /// remembers the first block whose best beats the running result; pass 2 scans only that block for the first element the result + /// does not beat. + /// + /// + /// + /// The reduction is what makes this run at memory bandwidth regardless of the input pattern; an index-vector loop needs a compare + /// and two selects per element. Indices never live in vector lanes, so element sizes need no special handling. + /// + /// + /// Ties: within a block, the reduction picks some tied-best element and the scan then returns the earliest element it does not beat, + /// which is the earliest tied-best element; across blocks the strict keeps the + /// earliest block. For floating-point types a block containing a NaN reduces to NaN (the reductions propagate it), and the index of + /// the first NaN of that block is returned, which is the first NaN overall because earlier blocks contained none. + /// + /// private static int IndexOfMinMaxCore(ReadOnlySpan x) where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { @@ -30,26 +64,17 @@ private static int IndexOfMinMaxCore(ReadOnlySpan x) if (Vector512.IsHardwareAccelerated && Vector512.IsSupported && x.Length >= Vector512.Count) { - return sizeof(T) == 8 ? IndexOfMinMaxVectorized512Size4Plus(x) : - sizeof(T) == 4 ? IndexOfMinMaxVectorized512Size4Plus(x) : - sizeof(T) == 2 ? IndexOfMinMaxVectorized512Size2(x) : - IndexOfMinMaxVectorized512Size1(x); + return IndexOfMinMaxBlocks512(x); } if (Vector256.IsHardwareAccelerated && Vector256.IsSupported && x.Length >= Vector256.Count) { - return sizeof(T) == 8 ? IndexOfMinMaxVectorized256Size4Plus(x) : - sizeof(T) == 4 ? IndexOfMinMaxVectorized256Size4Plus(x) : - sizeof(T) == 2 ? IndexOfMinMaxVectorized256Size2(x) : - IndexOfMinMaxVectorized256Size1(x); + return IndexOfMinMaxBlocks256(x); } if (Vector128.IsHardwareAccelerated && Vector128.IsSupported && x.Length >= Vector128.Count) { - return sizeof(T) == 8 ? IndexOfMinMaxVectorized128Size4Plus(x) : - sizeof(T) == 4 ? IndexOfMinMaxVectorized128Size4Plus(x) : - sizeof(T) == 2 ? IndexOfMinMaxVectorized128Size2(x) : - IndexOfMinMaxVectorized128Size1(x); + return IndexOfMinMaxBlocks128(x); } return IndexOfMinMaxFallback(x); @@ -82,697 +107,499 @@ private static int IndexOfMinMaxFallback(ReadOnlySpan x) return resultIndex; } - private static int IndexOfMinMaxVectorized128Size4Plus(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator where TInt : IBinaryInteger + /// + /// Debug-only contract check for : no element among the first + /// elements at beats under , and none is NaN + /// (a NaN would have had to propagate into ). + /// + private static bool NoElementBeats(ref T xRef, int length, T value) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { - Debug.Assert(sizeof(T) == 4 || sizeof(T) == 8); - Debug.Assert(typeof(TInt) == typeof(uint) || typeof(TInt) == typeof(ulong)); - Debug.Assert(sizeof(TInt) == sizeof(T)); - - // Initialize result by reading first vector and quick return if possible. - Vector128 result = Vector128.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + for (int i = 0; i < length; i++) { - Vector128 nanMask = Vector128.IsNaN(result); - if (nanMask != Vector128.Zero) + T element = Unsafe.Add(ref xRef, i); + if (T.IsNaN(element) || TOperator.Compare(element, value)) { - return IndexOfFirstMatch(nanMask); + return false; } } - // Initialize indices. - Vector128 indexIncrement = Vector128.Create(TInt.CreateTruncating(Vector128.Count)); - Vector128 resultIndex = Vector128.Indices; - Vector128 currentIndex = resultIndex + indexIncrement; - ReadOnlySpan span = x.Slice(Vector128.Count); + return true; + } - while (!span.IsEmpty) + /// See . + private static int IndexOfMinMaxBlocks128(ReadOnlySpan x) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + { + Debug.Assert(Vector128.IsHardwareAccelerated && Vector128.IsSupported); + Debug.Assert(x.Length >= Vector128.Count); + Debug.Assert(sizeof(T) is 1 or 2 or 4 or 8); + + int blockSize = BlockVectors * Vector128.Count; + int length = x.Length; + ref T xRef = ref MemoryMarshal.GetReference(x); + + // Pass 1: reduce every block to its best element; the first block whose best beats the running result wins ties. + T result = xRef; + int resultBlock = -1; + for (int i = 0; i < length; i += blockSize) { - Vector128 current; - if (span.Length >= Vector128.Count) - { - current = Vector128.Create(span); - span = span.Slice(Vector128.Count); - } - else - { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector128.Count; - current = Vector128.Create(x.Slice(start)); - currentIndex = Vector128.Create(TInt.CreateTruncating(start)) + Vector128.Indices; - span = ReadOnlySpan.Empty; - } + int blockLength = Math.Min(blockSize, length - i); + T blockResult = BlockReduce128(ref Unsafe.Add(ref xRef, i), blockLength); - // Quick return if possible. if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) { - Vector128 nanMask = Vector128.IsNaN(current); - if (nanMask != Vector128.Zero) + if (T.IsNaN(blockResult)) { - return int.CreateTruncating(currentIndex.ToScalar()) + IndexOfFirstMatch(nanMask); + return i + IndexOfFirstNaN128(ref Unsafe.Add(ref xRef, i), blockLength); } } - // Get mask for which lanes that should have result updated. - Vector128 mask = TOperator.Compare(current, result); + Debug.Assert(NoElementBeats(ref Unsafe.Add(ref xRef, i), blockLength, blockResult), + "Reduce/Aggregate must return an element no other element of the block beats under Compare, and must propagate NaN."); - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex = ElementWiseSelect(mask.As(), currentIndex, resultIndex); - currentIndex += indexIncrement; + if (resultBlock < 0 || TOperator.Compare(blockResult, result)) + { + result = blockResult; + resultBlock = i; + } } - { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector128 aggMask = ~Vector128.Equals(result.As(), Vector128.Create(aggResult).As()); - Vector128 aggIndex = resultIndex | aggMask; - return int.CreateTruncating(HorizontalAggregate>(aggIndex)); - } + Debug.Assert(resultBlock >= 0); + + // Pass 2: the first element of the winning block that the result does not beat, i.e. the first tied-best element. + return resultBlock + IndexOfFirstNotBeaten128(ref Unsafe.Add(ref xRef, resultBlock), Math.Min(blockSize, length - resultBlock), result); } - private static int IndexOfMinMaxVectorized128Size2(ReadOnlySpan x) + /// Reduces elements starting at to their best element (no bounds checks). + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static T BlockReduce128(ref T xRef, int length) where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { - Debug.Assert(sizeof(T) == 2); + Debug.Assert(length >= 1); - // Initialize result by reading first vector and quick return if possible. - Vector128 result = Vector128.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) - { - Vector128 nanMask = Vector128.IsNaN(result); - if (nanMask != Vector128.Zero) - { - return IndexOfFirstMatch(nanMask); - } - } + nuint count = (nuint)Vector128.Count; + nuint end = (nuint)length; + T result; + nuint i; - // Initialize indices. - Vector128 indexIncrement = Vector128.Create((uint)Vector128.Count); - Vector128 resultIndex1 = Vector128.Indices; - Vector128 resultIndex2 = resultIndex1 + indexIncrement; - Vector128 currentIndex = resultIndex2 + indexIncrement; - ReadOnlySpan span = x.Slice(Vector128.Count); - - while (!span.IsEmpty) + if (end >= 2 * count) { - Vector128 current; - if (span.Length >= Vector128.Count) - { - current = Vector128.Create(span); - span = span.Slice(Vector128.Count); - } - else + // Two independent accumulators so that consecutive reductions do not serialize on one register. + Vector128 acc1 = Vector128.LoadUnsafe(ref xRef); + Vector128 acc2 = Vector128.LoadUnsafe(ref xRef, count); + nuint last = end - 2 * count; + for (i = 2 * count; i <= last; i += 2 * count) { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector128.Count; - current = Vector128.Create(x.Slice(start)); - currentIndex = Vector128.Create((uint)start) + Vector128.Indices; - span = ReadOnlySpan.Empty; + acc1 = TOperator.Reduce(acc1, Vector128.LoadUnsafe(ref xRef, i)); + acc2 = TOperator.Reduce(acc2, Vector128.LoadUnsafe(ref xRef, i + count)); } - // Quick return if possible. - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + if (i + count <= end) { - Vector128 nanMask = Vector128.IsNaN(current); - if (nanMask != Vector128.Zero) - { - return (int)currentIndex.ToScalar() + IndexOfFirstMatch(nanMask); - } + acc1 = TOperator.Reduce(acc1, Vector128.LoadUnsafe(ref xRef, i)); + i += count; } - // Get mask for which lanes that should have result updated, also widen it for updating the indices. - Vector128 mask = TOperator.Compare(current, result); - (Vector128 mask1, Vector128 mask2) = Vector128.Widen(mask.AsInt16()); - - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex1 = ElementWiseSelect(mask1.AsUInt32(), currentIndex, resultIndex1); - currentIndex += indexIncrement; - resultIndex2 = ElementWiseSelect(mask2.AsUInt32(), currentIndex, resultIndex2); - currentIndex += indexIncrement; + result = TOperator.Aggregate(TOperator.Reduce(acc1, acc2)); } - + else if (end >= count) { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector128 aggMask = ~Vector128.Equals(result.AsInt16(), Vector128.Create(aggResult).AsInt16()); - - (Vector128 mask1, Vector128 mask2) = Vector128.Widen(aggMask); - Vector128 aggIndex = resultIndex1 | mask1.AsUInt32(); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex2 | mask2.AsUInt32()); + result = TOperator.Aggregate(Vector128.LoadUnsafe(ref xRef)); + i = count; + } + else + { + result = xRef; + i = 1; + } - return (int)HorizontalAggregate>(aggIndex); + for (; i < end; i++) + { + result = TOperator.Reduce(result, Unsafe.Add(ref xRef, i)); } + + return result; } - private static int IndexOfMinMaxVectorized128Size1(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + /// Index of the first NaN among elements starting at ; there must be one. + [MethodImpl(MethodImplOptions.NoInlining)] // cold: called at most once per search; keeps the caller within the inlining budget + private static int IndexOfFirstNaN128(ref T xRef, int length) + where T : INumber { - Debug.Assert(sizeof(T) == 1); + int count = Vector128.Count; + int i = 0; - // Initialize result by reading first vector and quick return if possible. - Vector128 result = Vector128.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + for (; i + count <= length; i += count) { - Vector128 nanMask = Vector128.IsNaN(result); + Vector128 nanMask = Vector128.IsNaN(Vector128.LoadUnsafe(ref xRef, (nuint)i)); if (nanMask != Vector128.Zero) { - return IndexOfFirstMatch(nanMask); + return i + IndexOfFirstMatch(nanMask); } } - // Initialize indices. - Vector128 indexIncrement = Vector128.Create((uint)Vector128.Count); - Vector128 resultIndex1 = Vector128.Indices; - Vector128 resultIndex2 = resultIndex1 + indexIncrement; - Vector128 resultIndex3 = resultIndex2 + indexIncrement; - Vector128 resultIndex4 = resultIndex3 + indexIncrement; - Vector128 currentIndex = resultIndex4 + indexIncrement; - ReadOnlySpan span = x.Slice(Vector128.Count); - - while (!span.IsEmpty) + for (; i < length; i++) { - Vector128 current; - if (span.Length >= Vector128.Count) - { - current = Vector128.Create(span); - span = span.Slice(Vector128.Count); - } - else - { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector128.Count; - current = Vector128.Create(x.Slice(start)); - currentIndex = Vector128.Create((uint)start) + Vector128.Indices; - span = ReadOnlySpan.Empty; - } - - // Quick return if possible. - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + if (T.IsNaN(Unsafe.Add(ref xRef, i))) { - Vector128 nanMask = Vector128.IsNaN(current); - if (nanMask != Vector128.Zero) - { - return (int)currentIndex.ToScalar() + IndexOfFirstMatch(nanMask); - } + return i; } - - // Get mask for which lanes that should have result updated, also widen it for updating the indices. - Vector128 mask = TOperator.Compare(current, result); - (Vector128 lowerMask, Vector128 upperMask) = Vector128.Widen(mask.AsSByte()); - (Vector128 mask1, Vector128 mask2) = Vector128.Widen(lowerMask); - (Vector128 mask3, Vector128 mask4) = Vector128.Widen(upperMask); - - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex1 = ElementWiseSelect(mask1.AsUInt32(), currentIndex, resultIndex1); - currentIndex += indexIncrement; - resultIndex2 = ElementWiseSelect(mask2.AsUInt32(), currentIndex, resultIndex2); - currentIndex += indexIncrement; - resultIndex3 = ElementWiseSelect(mask3.AsUInt32(), currentIndex, resultIndex3); - currentIndex += indexIncrement; - resultIndex4 = ElementWiseSelect(mask4.AsUInt32(), currentIndex, resultIndex4); - currentIndex += indexIncrement; } - { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector128 aggMask = ~Vector128.Equals(result.AsSByte(), Vector128.Create(aggResult).AsSByte()); - - (Vector128 lowerMask, Vector128 upperMask) = Vector128.Widen(aggMask); - (Vector128 mask1, Vector128 mask2) = Vector128.Widen(lowerMask); - (Vector128 mask3, Vector128 mask4) = Vector128.Widen(upperMask); - Vector128 aggIndex = resultIndex1 | mask1.AsUInt32(); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex2 | mask2.AsUInt32()); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex3 | mask3.AsUInt32()); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex4 | mask4.AsUInt32()); - - return (int)HorizontalAggregate>(aggIndex); - } + Debug.Fail("A NaN was expected in the block."); + return -1; } - private static int IndexOfMinMaxVectorized256Size4Plus(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator where TInt : IBinaryInteger + /// + /// Index of the first element among elements starting at that + /// does not beat under ; must be the + /// block's best element, so this is the first element tied with it (equal, or an equal-magnitude tie the operator does not order). + /// + [MethodImpl(MethodImplOptions.NoInlining)] // called once per search; keeps the caller within the inlining budget + private static int IndexOfFirstNotBeaten128(ref T xRef, int length, T value) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { - Debug.Assert(sizeof(T) == 4 || sizeof(T) == 8); - Debug.Assert(typeof(TInt) == typeof(uint) || typeof(TInt) == typeof(ulong)); - Debug.Assert(sizeof(TInt) == sizeof(T)); + int count = Vector128.Count; + Vector128 best = Vector128.Create(value); + int i = 0; - // Initialize result by reading first vector and quick return if possible. - Vector256 result = Vector256.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + for (; i + count <= length; i += count) { - Vector256 nanMask = Vector256.IsNaN(result); - if (nanMask != Vector256.Zero) + var bits = (~TOperator.Compare(best, Vector128.LoadUnsafe(ref xRef, (nuint)i))).ExtractMostSignificantBits(); + if (bits != 0) { - return IndexOfFirstMatch(nanMask); + return i + BitOperations.TrailingZeroCount(bits); } } - // Initialize indices. - Vector256 indexIncrement = Vector256.Create(TInt.CreateTruncating(Vector256.Count)); - Vector256 resultIndex = Vector256.Indices; - Vector256 currentIndex = resultIndex + indexIncrement; - ReadOnlySpan span = x.Slice(Vector256.Count); - - while (!span.IsEmpty) + for (; i < length; i++) { - Vector256 current; - if (span.Length >= Vector256.Count) - { - current = Vector256.Create(span); - span = span.Slice(Vector256.Count); - } - else + if (!TOperator.Compare(value, Unsafe.Add(ref xRef, i))) { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector256.Count; - current = Vector256.Create(x.Slice(start)); - currentIndex = Vector256.Create(TInt.CreateTruncating(start)) + Vector256.Indices; - span = ReadOnlySpan.Empty; + return i; } + } + + Debug.Fail("The block's best element was expected in the block."); + return -1; + } + + /// See . + private static int IndexOfMinMaxBlocks256(ReadOnlySpan x) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + { + Debug.Assert(Vector256.IsHardwareAccelerated && Vector256.IsSupported); + Debug.Assert(x.Length >= Vector256.Count); + Debug.Assert(sizeof(T) is 1 or 2 or 4 or 8); + + int blockSize = BlockVectors * Vector256.Count; + int length = x.Length; + ref T xRef = ref MemoryMarshal.GetReference(x); + + // Pass 1: reduce every block to its best element; the first block whose best beats the running result wins ties. + T result = xRef; + int resultBlock = -1; + for (int i = 0; i < length; i += blockSize) + { + int blockLength = Math.Min(blockSize, length - i); + T blockResult = BlockReduce256(ref Unsafe.Add(ref xRef, i), blockLength); - // Quick return if possible. if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) { - Vector256 nanMask = Vector256.IsNaN(current); - if (nanMask != Vector256.Zero) + if (T.IsNaN(blockResult)) { - return int.CreateTruncating(currentIndex.ToScalar()) + IndexOfFirstMatch(nanMask); + return i + IndexOfFirstNaN256(ref Unsafe.Add(ref xRef, i), blockLength); } } - // Get mask for which lanes that should have result updated. - Vector256 mask = TOperator.Compare(current, result); + Debug.Assert(NoElementBeats(ref Unsafe.Add(ref xRef, i), blockLength, blockResult), + "Reduce/Aggregate must return an element no other element of the block beats under Compare, and must propagate NaN."); - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex = ElementWiseSelect(mask.As(), currentIndex, resultIndex); - currentIndex += indexIncrement; + if (resultBlock < 0 || TOperator.Compare(blockResult, result)) + { + result = blockResult; + resultBlock = i; + } } - { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector256 aggMask = ~Vector256.Equals(result.As(), Vector256.Create(aggResult).As()); - Vector256 aggIndex = resultIndex | aggMask; - return int.CreateTruncating(HorizontalAggregate>(aggIndex)); - } + Debug.Assert(resultBlock >= 0); + + // Pass 2: the first element of the winning block that the result does not beat, i.e. the first tied-best element. + return resultBlock + IndexOfFirstNotBeaten256(ref Unsafe.Add(ref xRef, resultBlock), Math.Min(blockSize, length - resultBlock), result); } - private static int IndexOfMinMaxVectorized256Size2(ReadOnlySpan x) + /// Reduces elements starting at to their best element (no bounds checks). + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static T BlockReduce256(ref T xRef, int length) where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { - Debug.Assert(sizeof(T) == 2); + Debug.Assert(length >= 1); - // Initialize result by reading first vector and quick return if possible. - Vector256 result = Vector256.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) - { - Vector256 nanMask = Vector256.IsNaN(result); - if (nanMask != Vector256.Zero) - { - return IndexOfFirstMatch(nanMask); - } - } - - // Initialize indices. - Vector256 indexIncrement = Vector256.Create((uint)Vector256.Count); - Vector256 resultIndex1 = Vector256.Indices; - Vector256 resultIndex2 = resultIndex1 + indexIncrement; - Vector256 currentIndex = resultIndex2 + indexIncrement; - ReadOnlySpan span = x.Slice(Vector256.Count); + nuint count = (nuint)Vector256.Count; + nuint end = (nuint)length; + T result; + nuint i; - while (!span.IsEmpty) + if (end >= 2 * count) { - Vector256 current; - if (span.Length >= Vector256.Count) - { - current = Vector256.Create(span); - span = span.Slice(Vector256.Count); - } - else + // Two independent accumulators so that consecutive reductions do not serialize on one register. + Vector256 acc1 = Vector256.LoadUnsafe(ref xRef); + Vector256 acc2 = Vector256.LoadUnsafe(ref xRef, count); + nuint last = end - 2 * count; + for (i = 2 * count; i <= last; i += 2 * count) { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector256.Count; - current = Vector256.Create(x.Slice(start)); - currentIndex = Vector256.Create((uint)start) + Vector256.Indices; - span = ReadOnlySpan.Empty; + acc1 = TOperator.Reduce(acc1, Vector256.LoadUnsafe(ref xRef, i)); + acc2 = TOperator.Reduce(acc2, Vector256.LoadUnsafe(ref xRef, i + count)); } - // Quick return if possible. - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + if (i + count <= end) { - Vector256 nanMask = Vector256.IsNaN(current); - if (nanMask != Vector256.Zero) - { - return (int)currentIndex.ToScalar() + IndexOfFirstMatch(nanMask); - } + acc1 = TOperator.Reduce(acc1, Vector256.LoadUnsafe(ref xRef, i)); + i += count; } - // Get mask for which lanes that should have result updated, also widen it for updating the indices. - Vector256 mask = TOperator.Compare(current, result); - (Vector256 mask1, Vector256 mask2) = Vector256.Widen(mask.AsInt16()); - - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex1 = ElementWiseSelect(mask1.AsUInt32(), currentIndex, resultIndex1); - currentIndex += indexIncrement; - resultIndex2 = ElementWiseSelect(mask2.AsUInt32(), currentIndex, resultIndex2); - currentIndex += indexIncrement; + result = TOperator.Aggregate(TOperator.Reduce(acc1, acc2)); } - + else if (end >= count) { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector256 aggMask = ~Vector256.Equals(result.AsInt16(), Vector256.Create(aggResult).AsInt16()); - - (Vector256 mask1, Vector256 mask2) = Vector256.Widen(aggMask); - Vector256 aggIndex = resultIndex1 | mask1.AsUInt32(); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex2 | mask2.AsUInt32()); + result = TOperator.Aggregate(Vector256.LoadUnsafe(ref xRef)); + i = count; + } + else + { + result = xRef; + i = 1; + } - return (int)HorizontalAggregate>(aggIndex); + for (; i < end; i++) + { + result = TOperator.Reduce(result, Unsafe.Add(ref xRef, i)); } + + return result; } - private static int IndexOfMinMaxVectorized256Size1(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + /// Index of the first NaN among elements starting at ; there must be one. + [MethodImpl(MethodImplOptions.NoInlining)] // cold: called at most once per search; keeps the caller within the inlining budget + private static int IndexOfFirstNaN256(ref T xRef, int length) + where T : INumber { - Debug.Assert(sizeof(T) == 1); + int count = Vector256.Count; + int i = 0; - // Initialize result by reading first vector and quick return if possible. - Vector256 result = Vector256.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + for (; i + count <= length; i += count) { - Vector256 nanMask = Vector256.IsNaN(result); + Vector256 nanMask = Vector256.IsNaN(Vector256.LoadUnsafe(ref xRef, (nuint)i)); if (nanMask != Vector256.Zero) { - return IndexOfFirstMatch(nanMask); + return i + IndexOfFirstMatch(nanMask); } } - // Initialize indices. - Vector256 indexIncrement = Vector256.Create((uint)Vector256.Count); - Vector256 resultIndex1 = Vector256.Indices; - Vector256 resultIndex2 = resultIndex1 + indexIncrement; - Vector256 resultIndex3 = resultIndex2 + indexIncrement; - Vector256 resultIndex4 = resultIndex3 + indexIncrement; - Vector256 currentIndex = resultIndex4 + indexIncrement; - ReadOnlySpan span = x.Slice(Vector256.Count); - - while (!span.IsEmpty) + for (; i < length; i++) { - Vector256 current; - if (span.Length >= Vector256.Count) - { - current = Vector256.Create(span); - span = span.Slice(Vector256.Count); - } - else - { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector256.Count; - current = Vector256.Create(x.Slice(start)); - currentIndex = Vector256.Create((uint)start) + Vector256.Indices; - span = ReadOnlySpan.Empty; - } - - // Quick return if possible. - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + if (T.IsNaN(Unsafe.Add(ref xRef, i))) { - Vector256 nanMask = Vector256.IsNaN(current); - if (nanMask != Vector256.Zero) - { - return (int)currentIndex.ToScalar() + IndexOfFirstMatch(nanMask); - } + return i; } - - // Get mask for which lanes that should have result updated, also widen it for updating the indices. - Vector256 mask = TOperator.Compare(current, result); - (Vector256 lowerMask, Vector256 upperMask) = Vector256.Widen(mask.AsSByte()); - (Vector256 mask1, Vector256 mask2) = Vector256.Widen(lowerMask); - (Vector256 mask3, Vector256 mask4) = Vector256.Widen(upperMask); - - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex1 = ElementWiseSelect(mask1.AsUInt32(), currentIndex, resultIndex1); - currentIndex += indexIncrement; - resultIndex2 = ElementWiseSelect(mask2.AsUInt32(), currentIndex, resultIndex2); - currentIndex += indexIncrement; - resultIndex3 = ElementWiseSelect(mask3.AsUInt32(), currentIndex, resultIndex3); - currentIndex += indexIncrement; - resultIndex4 = ElementWiseSelect(mask4.AsUInt32(), currentIndex, resultIndex4); - currentIndex += indexIncrement; } - { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector256 aggMask = ~Vector256.Equals(result.AsSByte(), Vector256.Create(aggResult).AsSByte()); - - (Vector256 lowerMask, Vector256 upperMask) = Vector256.Widen(aggMask); - (Vector256 mask1, Vector256 mask2) = Vector256.Widen(lowerMask); - (Vector256 mask3, Vector256 mask4) = Vector256.Widen(upperMask); - Vector256 aggIndex = resultIndex1 | mask1.AsUInt32(); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex2 | mask2.AsUInt32()); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex3 | mask3.AsUInt32()); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex4 | mask4.AsUInt32()); - - return (int)HorizontalAggregate>(aggIndex); - } + Debug.Fail("A NaN was expected in the block."); + return -1; } - private static int IndexOfMinMaxVectorized512Size4Plus(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator where TInt : IBinaryInteger + /// + /// Index of the first element among elements starting at that + /// does not beat under ; must be the + /// block's best element, so this is the first element tied with it (equal, or an equal-magnitude tie the operator does not order). + /// + [MethodImpl(MethodImplOptions.NoInlining)] // called once per search; keeps the caller within the inlining budget + private static int IndexOfFirstNotBeaten256(ref T xRef, int length, T value) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { - Debug.Assert(sizeof(T) == 4 || sizeof(T) == 8); - Debug.Assert(typeof(TInt) == typeof(uint) || typeof(TInt) == typeof(ulong)); - Debug.Assert(sizeof(TInt) == sizeof(T)); + int count = Vector256.Count; + Vector256 best = Vector256.Create(value); + int i = 0; - // Initialize result by reading first vector and quick return if possible. - Vector512 result = Vector512.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + for (; i + count <= length; i += count) { - Vector512 nanMask = Vector512.IsNaN(result); - if (nanMask != Vector512.Zero) + var bits = (~TOperator.Compare(best, Vector256.LoadUnsafe(ref xRef, (nuint)i))).ExtractMostSignificantBits(); + if (bits != 0) { - return IndexOfFirstMatch(nanMask); + return i + BitOperations.TrailingZeroCount(bits); } } - // Initialize indices. - Vector512 indexIncrement = Vector512.Create(TInt.CreateTruncating(Vector512.Count)); - Vector512 resultIndex = Vector512.Indices; - Vector512 currentIndex = resultIndex + indexIncrement; - ReadOnlySpan span = x.Slice(Vector512.Count); - - while (!span.IsEmpty) + for (; i < length; i++) { - Vector512 current; - if (span.Length >= Vector512.Count) - { - current = Vector512.Create(span); - span = span.Slice(Vector512.Count); - } - else + if (!TOperator.Compare(value, Unsafe.Add(ref xRef, i))) { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector512.Count; - current = Vector512.Create(x.Slice(start)); - currentIndex = Vector512.Create(TInt.CreateTruncating(start)) + Vector512.Indices; - span = ReadOnlySpan.Empty; + return i; } + } + + Debug.Fail("The block's best element was expected in the block."); + return -1; + } + + /// See . + private static int IndexOfMinMaxBlocks512(ReadOnlySpan x) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + { + Debug.Assert(Vector512.IsHardwareAccelerated && Vector512.IsSupported); + Debug.Assert(x.Length >= Vector512.Count); + Debug.Assert(sizeof(T) is 1 or 2 or 4 or 8); + + int blockSize = BlockVectors * Vector512.Count; + int length = x.Length; + ref T xRef = ref MemoryMarshal.GetReference(x); + + // Pass 1: reduce every block to its best element; the first block whose best beats the running result wins ties. + T result = xRef; + int resultBlock = -1; + for (int i = 0; i < length; i += blockSize) + { + int blockLength = Math.Min(blockSize, length - i); + T blockResult = BlockReduce512(ref Unsafe.Add(ref xRef, i), blockLength); - // Quick return if possible. if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) { - Vector512 nanMask = Vector512.IsNaN(current); - if (nanMask != Vector512.Zero) + if (T.IsNaN(blockResult)) { - return int.CreateTruncating(currentIndex.ToScalar()) + IndexOfFirstMatch(nanMask); + return i + IndexOfFirstNaN512(ref Unsafe.Add(ref xRef, i), blockLength); } } - // Get mask for which lanes that should have result updated. - Vector512 mask = TOperator.Compare(current, result); + Debug.Assert(NoElementBeats(ref Unsafe.Add(ref xRef, i), blockLength, blockResult), + "Reduce/Aggregate must return an element no other element of the block beats under Compare, and must propagate NaN."); - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex = ElementWiseSelect(mask.As(), currentIndex, resultIndex); - currentIndex += indexIncrement; + if (resultBlock < 0 || TOperator.Compare(blockResult, result)) + { + result = blockResult; + resultBlock = i; + } } - { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector512 aggMask = ~Vector512.Equals(result.As(), Vector512.Create(aggResult).As()); - Vector512 aggIndex = resultIndex | aggMask; - return int.CreateTruncating(HorizontalAggregate>(aggIndex)); - } + Debug.Assert(resultBlock >= 0); + + // Pass 2: the first element of the winning block that the result does not beat, i.e. the first tied-best element. + return resultBlock + IndexOfFirstNotBeaten512(ref Unsafe.Add(ref xRef, resultBlock), Math.Min(blockSize, length - resultBlock), result); } - private static int IndexOfMinMaxVectorized512Size2(ReadOnlySpan x) + /// Reduces elements starting at to their best element (no bounds checks). + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static T BlockReduce512(ref T xRef, int length) where T : INumber where TOperator : struct, IIndexOfMinMaxOperator { - Debug.Assert(sizeof(T) == 2); + Debug.Assert(length >= 1); - // Initialize result by reading first vector and quick return if possible. - Vector512 result = Vector512.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) - { - Vector512 nanMask = Vector512.IsNaN(result); - if (nanMask != Vector512.Zero) - { - return IndexOfFirstMatch(nanMask); - } - } - - // Initialize indices. - Vector512 indexIncrement = Vector512.Create((uint)Vector512.Count); - Vector512 resultIndex1 = Vector512.Indices; - Vector512 resultIndex2 = resultIndex1 + indexIncrement; - Vector512 currentIndex = resultIndex2 + indexIncrement; - ReadOnlySpan span = x.Slice(Vector512.Count); + nuint count = (nuint)Vector512.Count; + nuint end = (nuint)length; + T result; + nuint i; - while (!span.IsEmpty) + if (end >= 2 * count) { - Vector512 current; - if (span.Length >= Vector512.Count) + // Two independent accumulators so that consecutive reductions do not serialize on one register. + Vector512 acc1 = Vector512.LoadUnsafe(ref xRef); + Vector512 acc2 = Vector512.LoadUnsafe(ref xRef, count); + nuint last = end - 2 * count; + for (i = 2 * count; i <= last; i += 2 * count) { - current = Vector512.Create(span); - span = span.Slice(Vector512.Count); - } - else - { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector512.Count; - current = Vector512.Create(x.Slice(start)); - currentIndex = Vector512.Create((uint)start) + Vector512.Indices; - span = ReadOnlySpan.Empty; + acc1 = TOperator.Reduce(acc1, Vector512.LoadUnsafe(ref xRef, i)); + acc2 = TOperator.Reduce(acc2, Vector512.LoadUnsafe(ref xRef, i + count)); } - // Quick return if possible. - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + if (i + count <= end) { - Vector512 nanMask = Vector512.IsNaN(current); - if (nanMask != Vector512.Zero) - { - return (int)currentIndex.ToScalar() + IndexOfFirstMatch(nanMask); - } + acc1 = TOperator.Reduce(acc1, Vector512.LoadUnsafe(ref xRef, i)); + i += count; } - // Get mask for which lanes that should have result updated, also widen it for updating the indices. - Vector512 mask = TOperator.Compare(current, result); - (Vector512 mask1, Vector512 mask2) = Vector512.Widen(mask.AsInt16()); - - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex1 = ElementWiseSelect(mask1.AsUInt32(), currentIndex, resultIndex1); - currentIndex += indexIncrement; - resultIndex2 = ElementWiseSelect(mask2.AsUInt32(), currentIndex, resultIndex2); - currentIndex += indexIncrement; + result = TOperator.Aggregate(TOperator.Reduce(acc1, acc2)); } - + else if (end >= count) { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector512 aggMask = ~Vector512.Equals(result.AsInt16(), Vector512.Create(aggResult).AsInt16()); - - (Vector512 mask1, Vector512 mask2) = Vector512.Widen(aggMask); - Vector512 aggIndex = resultIndex1 | mask1.AsUInt32(); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex2 | mask2.AsUInt32()); + result = TOperator.Aggregate(Vector512.LoadUnsafe(ref xRef)); + i = count; + } + else + { + result = xRef; + i = 1; + } - return (int)HorizontalAggregate>(aggIndex); + for (; i < end; i++) + { + result = TOperator.Reduce(result, Unsafe.Add(ref xRef, i)); } + + return result; } - private static int IndexOfMinMaxVectorized512Size1(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + /// Index of the first NaN among elements starting at ; there must be one. + [MethodImpl(MethodImplOptions.NoInlining)] // cold: called at most once per search; keeps the caller within the inlining budget + private static int IndexOfFirstNaN512(ref T xRef, int length) + where T : INumber { - Debug.Assert(sizeof(T) == 1); + int count = Vector512.Count; + int i = 0; - // Initialize result by reading first vector and quick return if possible. - Vector512 result = Vector512.Create(x); - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) + for (; i + count <= length; i += count) { - Vector512 nanMask = Vector512.IsNaN(result); + Vector512 nanMask = Vector512.IsNaN(Vector512.LoadUnsafe(ref xRef, (nuint)i)); if (nanMask != Vector512.Zero) { - return IndexOfFirstMatch(nanMask); + return i + IndexOfFirstMatch(nanMask); } } - // Initialize indices. - Vector512 indexIncrement = Vector512.Create((uint)Vector512.Count); - Vector512 resultIndex1 = Vector512.Indices; - Vector512 resultIndex2 = resultIndex1 + indexIncrement; - Vector512 resultIndex3 = resultIndex2 + indexIncrement; - Vector512 resultIndex4 = resultIndex3 + indexIncrement; - Vector512 currentIndex = resultIndex4 + indexIncrement; - ReadOnlySpan span = x.Slice(Vector512.Count); - - while (!span.IsEmpty) + for (; i < length; i++) { - Vector512 current; - if (span.Length >= Vector512.Count) + if (T.IsNaN(Unsafe.Add(ref xRef, i))) { - current = Vector512.Create(span); - span = span.Slice(Vector512.Count); - } - else - { - // Process a final back-shifted to cover remaining elements in x in one vector. - int start = x.Length - Vector512.Count; - current = Vector512.Create(x.Slice(start)); - currentIndex = Vector512.Create((uint)start) + Vector512.Indices; - span = ReadOnlySpan.Empty; + return i; } + } - // Quick return if possible. - if (typeof(T) == typeof(float) || typeof(T) == typeof(double)) - { - Vector512 nanMask = Vector512.IsNaN(current); - if (nanMask != Vector512.Zero) - { - return (int)currentIndex.ToScalar() + IndexOfFirstMatch(nanMask); - } - } + Debug.Fail("A NaN was expected in the block."); + return -1; + } - // Get mask for which lanes that should have result updated, also widen it for updating the indices. - Vector512 mask = TOperator.Compare(current, result); - (Vector512 lowerMask, Vector512 upperMask) = Vector512.Widen(mask.AsSByte()); - (Vector512 mask1, Vector512 mask2) = Vector512.Widen(lowerMask); - (Vector512 mask3, Vector512 mask4) = Vector512.Widen(upperMask); + /// + /// Index of the first element among elements starting at that + /// does not beat under ; must be the + /// block's best element, so this is the first element tied with it (equal, or an equal-magnitude tie the operator does not order). + /// + [MethodImpl(MethodImplOptions.NoInlining)] // called once per search; keeps the caller within the inlining budget + private static int IndexOfFirstNotBeaten512(ref T xRef, int length, T value) + where T : INumber where TOperator : struct, IIndexOfMinMaxOperator + { + int count = Vector512.Count; + Vector512 best = Vector512.Create(value); + int i = 0; - // Update result and indices. - result = ElementWiseSelect(mask, current, result); - resultIndex1 = ElementWiseSelect(mask1.AsUInt32(), currentIndex, resultIndex1); - currentIndex += indexIncrement; - resultIndex2 = ElementWiseSelect(mask2.AsUInt32(), currentIndex, resultIndex2); - currentIndex += indexIncrement; - resultIndex3 = ElementWiseSelect(mask3.AsUInt32(), currentIndex, resultIndex3); - currentIndex += indexIncrement; - resultIndex4 = ElementWiseSelect(mask4.AsUInt32(), currentIndex, resultIndex4); - currentIndex += indexIncrement; + for (; i + count <= length; i += count) + { + var bits = (~TOperator.Compare(best, Vector512.LoadUnsafe(ref xRef, (nuint)i))).ExtractMostSignificantBits(); + if (bits != 0) + { + return i + BitOperations.TrailingZeroCount(bits); + } } + for (; i < length; i++) { - // Where result does not bitwise-equal the aggregate min/max value; replace indices with uint.MaxValue. Then find the min index. - T aggResult = TOperator.Aggregate(result); - Vector512 aggMask = ~Vector512.Equals(result.AsSByte(), Vector512.Create(aggResult).AsSByte()); - - (Vector512 lowerMask, Vector512 upperMask) = Vector512.Widen(aggMask); - (Vector512 mask1, Vector512 mask2) = Vector512.Widen(lowerMask); - (Vector512 mask3, Vector512 mask4) = Vector512.Widen(upperMask); - Vector512 aggIndex = resultIndex1 | mask1.AsUInt32(); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex2 | mask2.AsUInt32()); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex3 | mask3.AsUInt32()); - aggIndex = MinOperator.Invoke(aggIndex, resultIndex4 | mask4.AsUInt32()); - - return (int)HorizontalAggregate>(aggIndex); + if (!TOperator.Compare(value, Unsafe.Add(ref xRef, i))) + { + return i; + } } + + Debug.Fail("The block's best element was expected in the block."); + return -1; } + } } diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMax.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMax.cs index 96ff23100609e3..d078345a02443e 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMax.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMax.cs @@ -3,7 +3,6 @@ using System.Runtime.CompilerServices; using System.Runtime.Intrinsics; -using System.Runtime.Intrinsics.X86; namespace System.Numerics.Tensors { @@ -32,6 +31,14 @@ public static int IndexOfMax(ReadOnlySpan x) public static T Aggregate(Vector128 x) => HorizontalAggregate>(x); public static T Aggregate(Vector256 x) => HorizontalAggregate>(x); public static T Aggregate(Vector512 x) => HorizontalAggregate>(x); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static T Reduce(T x, T y) => MaxOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector128 Reduce(Vector128 x, Vector128 y) => MaxOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector256 Reduce(Vector256 x, Vector256 y) => MaxOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector512 Reduce(Vector512 x, Vector512 y) => MaxOperator.Invoke(x, y); [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool Compare(T x, T y) @@ -98,53 +105,5 @@ private static int IndexOfFirstMatch(Vector256 mask) => private static int IndexOfFirstMatch(Vector512 mask) => BitOperations.TrailingZeroCount(mask.ExtractMostSignificantBits()); - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private static unsafe Vector128 ElementWiseSelect(Vector128 mask, Vector128 left, Vector128 right) - { - if (Sse41.IsSupported) - { - if (typeof(T) == typeof(float)) return Sse41.BlendVariable(left.AsSingle(), right.AsSingle(), (~mask).AsSingle()).As(); - if (typeof(T) == typeof(double)) return Sse41.BlendVariable(left.AsDouble(), right.AsDouble(), (~mask).AsDouble()).As(); - - if (sizeof(T) == 1) return Sse41.BlendVariable(left.AsByte(), right.AsByte(), (~mask).AsByte()).As(); - if (sizeof(T) == 2) return Sse41.BlendVariable(left.AsUInt16(), right.AsUInt16(), (~mask).AsUInt16()).As(); - if (sizeof(T) == 4) return Sse41.BlendVariable(left.AsUInt32(), right.AsUInt32(), (~mask).AsUInt32()).As(); - if (sizeof(T) == 8) return Sse41.BlendVariable(left.AsUInt64(), right.AsUInt64(), (~mask).AsUInt64()).As(); - } - - return Vector128.ConditionalSelect(mask, left, right); - } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private static unsafe Vector256 ElementWiseSelect(Vector256 mask, Vector256 left, Vector256 right) - { - if (Avx2.IsSupported) - { - if (typeof(T) == typeof(float)) return Avx2.BlendVariable(left.AsSingle(), right.AsSingle(), (~mask).AsSingle()).As(); - if (typeof(T) == typeof(double)) return Avx2.BlendVariable(left.AsDouble(), right.AsDouble(), (~mask).AsDouble()).As(); - - if (sizeof(T) == 1) return Avx2.BlendVariable(left.AsByte(), right.AsByte(), (~mask).AsByte()).As(); - if (sizeof(T) == 2) return Avx2.BlendVariable(left.AsUInt16(), right.AsUInt16(), (~mask).AsUInt16()).As(); - if (sizeof(T) == 4) return Avx2.BlendVariable(left.AsUInt32(), right.AsUInt32(), (~mask).AsUInt32()).As(); - if (sizeof(T) == 8) return Avx2.BlendVariable(left.AsUInt64(), right.AsUInt64(), (~mask).AsUInt64()).As(); - } - - return Vector256.ConditionalSelect(mask, left, right); - } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - private static unsafe Vector512 ElementWiseSelect(Vector512 mask, Vector512 left, Vector512 right) - { - if (Avx512F.IsSupported) - { - if (typeof(T) == typeof(float)) return Avx512F.BlendVariable(left.AsSingle(), right.AsSingle(), (~mask).AsSingle()).As(); - if (typeof(T) == typeof(double)) return Avx512F.BlendVariable(left.AsDouble(), right.AsDouble(), (~mask).AsDouble()).As(); - - if (sizeof(T) == 4) return Avx512F.BlendVariable(left.AsUInt32(), right.AsUInt32(), (~mask).AsUInt32()).As(); - if (sizeof(T) == 8) return Avx512F.BlendVariable(left.AsUInt64(), right.AsUInt64(), (~mask).AsUInt64()).As(); - } - - return Vector512.ConditionalSelect(mask, left, right); - } } } diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMaxMagnitude.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMaxMagnitude.cs index f04ca6466f46e8..1c565b0228fda2 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMaxMagnitude.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMaxMagnitude.cs @@ -31,6 +31,14 @@ public static int IndexOfMaxMagnitude(ReadOnlySpan x) public static T Aggregate(Vector128 x) => HorizontalAggregate>(x); public static T Aggregate(Vector256 x) => HorizontalAggregate>(x); public static T Aggregate(Vector512 x) => HorizontalAggregate>(x); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static T Reduce(T x, T y) => MaxMagnitudeOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector128 Reduce(Vector128 x, Vector128 y) => MaxMagnitudeOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector256 Reduce(Vector256 x, Vector256 y) => MaxMagnitudeOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector512 Reduce(Vector512 x, Vector512 y) => MaxMagnitudeOperator.Invoke(x, y); [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool Compare(T x, T y) diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMin.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMin.cs index 4da71bcecabbff..c9cf847d7865a3 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMin.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMin.cs @@ -31,6 +31,14 @@ public static int IndexOfMin(ReadOnlySpan x) public static T Aggregate(Vector128 x) => HorizontalAggregate>(x); public static T Aggregate(Vector256 x) => HorizontalAggregate>(x); public static T Aggregate(Vector512 x) => HorizontalAggregate>(x); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static T Reduce(T x, T y) => MinOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector128 Reduce(Vector128 x, Vector128 y) => MinOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector256 Reduce(Vector256 x, Vector256 y) => MinOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector512 Reduce(Vector512 x, Vector512 y) => MinOperator.Invoke(x, y); [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool Compare(T x, T y) diff --git a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMinMagnitude.cs b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMinMagnitude.cs index fea0bd960ec6d0..59b76957e80d0d 100644 --- a/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMinMagnitude.cs +++ b/src/libraries/System.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMinMagnitude.cs @@ -31,6 +31,14 @@ public static int IndexOfMinMagnitude(ReadOnlySpan x) public static T Aggregate(Vector128 x) => HorizontalAggregate>(x); public static T Aggregate(Vector256 x) => HorizontalAggregate>(x); public static T Aggregate(Vector512 x) => HorizontalAggregate>(x); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static T Reduce(T x, T y) => MinMagnitudeOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector128 Reduce(Vector128 x, Vector128 y) => MinMagnitudeOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector256 Reduce(Vector256 x, Vector256 y) => MinMagnitudeOperator.Invoke(x, y); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public static Vector512 Reduce(Vector512 x, Vector512 y) => MinMagnitudeOperator.Invoke(x, y); [MethodImpl(MethodImplOptions.AggressiveInlining)] public static bool Compare(T x, T y) diff --git a/src/libraries/System.Numerics.Tensors/tests/TensorPrimitivesTests.cs b/src/libraries/System.Numerics.Tensors/tests/TensorPrimitivesTests.cs index f428af07f79d64..9cfe0ddc17d946 100644 --- a/src/libraries/System.Numerics.Tensors/tests/TensorPrimitivesTests.cs +++ b/src/libraries/System.Numerics.Tensors/tests/TensorPrimitivesTests.cs @@ -1356,6 +1356,173 @@ public void IndexOfMin_IndexAboveMaxValue() x.Span[size.Value - 1] = ConvertFromSingle(0); Assert.Equal(size.Value - 1, IndexOfMin(x)); } + + // Lengths that span several 256-element blocks of the block-minimum implementation (Helpers.TensorLengths stops at 256). + private static readonly int[] s_indexOfMinLongLengths = [255, 256, 257, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049, 4097, 65539]; + + private static IEnumerable IndexOfMinLongPositions(int tensorLength) => + new[] { 0, 1, 255, 256, 257, 511, 512, 513, tensorLength / 2, tensorLength - 2, tensorLength - 1 }.Where(i => i < tensorLength).Distinct(); + + [Fact] + public void IndexOfMin_LongLengths() + { + Assert.All(s_indexOfMinLongLengths, tensorLength => + { + foreach (int expected in IndexOfMinLongPositions(tensorLength)) + { + using BoundedMemory x = CreateAndFillTensor(tensorLength); + x[expected] = Enumerable.Min(MemoryMarshal.ToEnumerable(x.Memory)); + int actual = IndexOfMin(x.Span); + Assert.True(actual == expected || (actual < expected && x[actual].Equals(x[expected])), $"{tensorLength} {actual} {expected}"); + } + }); + } + + [Fact] + public void IndexOfMin_LongLengths_FirstOccurrenceReturned() + { + Assert.All(s_indexOfMinLongLengths, tensorLength => + { + foreach (int expected in IndexOfMinLongPositions(tensorLength)) + { + using BoundedMemory x = CreateTensor(tensorLength); + x.Span.Fill(ConvertFromSingle(1)); + x[expected] = ConvertFromSingle(0); + x[tensorLength - 1] = ConvertFromSingle(0); + Assert.Equal(expected, IndexOfMin(x.Span)); + } + }); + } + + [Fact] + public void IndexOfMin_LongLengths_FirstNaNReturned() + { + if (!IsFloatingPoint) return; + + Assert.All(s_indexOfMinLongLengths, tensorLength => + { + foreach (int expected in IndexOfMinLongPositions(tensorLength)) + { + using BoundedMemory x = CreateTensor(tensorLength); + x.Span.Fill(ConvertFromSingle(1)); + x[0] = ConvertFromSingle(-1); // a smaller value in an earlier block must not beat the NaN + x[expected] = ConvertFromSingle(float.NaN); + x[tensorLength - 1] = ConvertFromSingle(float.NaN); + Assert.Equal(expected, IndexOfMin(x.Span)); + } + }); + } + + [Fact] + public void IndexOfMin_LongLengths_Negative0LesserThanPositive0() + { + if (!IsFloatingPoint) return; + + Assert.All(s_indexOfMinLongLengths, tensorLength => + { + foreach (int expected in IndexOfMinLongPositions(tensorLength)) + { + using BoundedMemory x = CreateTensor(tensorLength); + x.Span.Fill(Zero); + x[expected] = NegativeZero; + x[tensorLength - 1] = NegativeZero; + Assert.Equal(expected, IndexOfMin(x.Span)); + } + }); + } + + // The same multi-block coverage for the other three searches that share the block-reduction core. + [Fact] + public void IndexOfMax_LongLengths_FirstOccurrenceReturned() => + AssertLongLengths(1, 2, IndexOfMax); + + [Fact] + public void IndexOfMinMagnitude_LongLengths_FirstOccurrenceReturned() => + AssertLongLengths(2, 1, IndexOfMinMagnitude); + + [Fact] + public void IndexOfMaxMagnitude_LongLengths_FirstOccurrenceReturned() => + AssertLongLengths(1, 2, IndexOfMaxMagnitude); + + [Fact] + public void IndexOfMax_LongLengths_FirstNaNReturned() + { + if (!IsFloatingPoint) return; + AssertLongLengthsNaN(1, 2, IndexOfMax); // a larger value in an earlier block must not beat the NaN + } + + [Fact] + public void IndexOfMinMagnitude_LongLengths_FirstNaNReturned() + { + if (!IsFloatingPoint) return; + AssertLongLengthsNaN(2, 1, IndexOfMinMagnitude); + } + + [Fact] + public void IndexOfMaxMagnitude_LongLengths_FirstNaNReturned() + { + if (!IsFloatingPoint) return; + AssertLongLengthsNaN(1, 2, IndexOfMaxMagnitude); + } + + [Fact] + public void IndexOfMax_LongLengths_Positive0GreaterThanNegative0() + { + if (!IsFloatingPoint) return; + AssertLongLengthsValues(NegativeZero, Zero, IndexOfMax); + } + + [Fact] + public void IndexOfMinMagnitude_LongLengths_Negative0LesserThanPositive0() + { + if (!IsFloatingPoint) return; + AssertLongLengthsValues(Zero, NegativeZero, IndexOfMinMagnitude); + } + + [Fact] + public void IndexOfMaxMagnitude_LongLengths_Positive0GreaterThanNegative0() + { + if (!IsFloatingPoint) return; + AssertLongLengthsValues(NegativeZero, Zero, IndexOfMaxMagnitude); + } + + private delegate int IndexOfSearch(ReadOnlySpan x); + + private void AssertLongLengths(float fill, float best, IndexOfSearch search) => + AssertLongLengthsValues(ConvertFromSingle(fill), ConvertFromSingle(best), search); + + /// Fills with , places at the expected index and at the end; the expected index must win. + private void AssertLongLengthsValues(T fill, T best, IndexOfSearch search) + { + Assert.All(s_indexOfMinLongLengths, tensorLength => + { + foreach (int expected in IndexOfMinLongPositions(tensorLength)) + { + using BoundedMemory x = CreateTensor(tensorLength); + x.Span.Fill(fill); + x[expected] = best; + x[tensorLength - 1] = best; + Assert.Equal(expected, search(x.Span)); + } + }); + } + + /// Fills with , puts first and NaN at the expected index and at the end; the first NaN must win. + private void AssertLongLengthsNaN(float fill, float better, IndexOfSearch search) + { + Assert.All(s_indexOfMinLongLengths, tensorLength => + { + foreach (int expected in IndexOfMinLongPositions(tensorLength)) + { + using BoundedMemory x = CreateTensor(tensorLength); + x.Span.Fill(ConvertFromSingle(fill)); + x[0] = ConvertFromSingle(better); + x[expected] = ConvertFromSingle(float.NaN); + x[tensorLength - 1] = ConvertFromSingle(float.NaN); + Assert.Equal(expected, search(x.Span)); + } + }); + } #endregion #region IndexOfMinMagnitude