From b9145029671802970a5f4d4b10720caf70bebe47 Mon Sep 17 00:00:00 2001 From: Niklas Schilli Date: Tue, 15 Sep 2026 19:06:59 +0200 Subject: [PATCH 1/2] Use a block-reduction search in TensorPrimitives.IndexOfMin/Max/MinMagnitude/MaxMagnitude IndexOfMinMaxCore tracked a running best value and a vector of indices per lane, which costs a compare and two selects per element and needs per-element-size variants so that indices fit in lanes. On AVX-512 hardware the blends of the running value also go through mask-register conversions, and IndexOfMin over 65536 elements runs at about 5 GElem/s on a Zen 4 box, no faster than the .NET 10 package. Replace it with a two-pass block reduction: - Pass 1 reduces each block of 32 vectors to its best element with a pure vector loop (one load and one lane-wise Reduce per vector, two accumulators), then a horizontal Aggregate, and keeps the first block whose result beats the running result under the operator's strict Compare. - Pass 2 scans only the winning block for the first element the result does not beat, which is the operator's own tie rule (first equal element; -0 before +0 for min; the sign preference for equal magnitudes). IIndexOfMinMaxOperator gains Reduce (scalar and per vector width); each operator forwards it to its aggregation operator (Vector.Min/Max/MinMagnitude/MaxMagnitude), the same operators the previous code used for its final horizontal step and bitwise match, so the tie behavior is unchanged. Because those reductions implement IEEE minimum/maximum and propagate NaN, a block containing a NaN reduces to NaN and the first NaN of that block is returned; the per-vector NaN checks are no longer needed. Indices never live in vector lanes, so the Size4Plus/Size2/Size1 variants and ElementWiseSelect are removed. Measured with BenchmarkDotNet on an AMD Ryzen 7 7800X3D (.NET 10.0.5, N = 65536 ints): IndexOfMin goes from 5.1 GElem/s to 35.6-36.7 GElem/s on increasing, random, decreasing and sparse-hit inputs, which is the L2 bandwidth ceiling for that working set. The hot loop is two vpminsd with memory operands per iteration. Tests: adds multi-block length, NaN and signed-zero cases for all four searches (the existing lengths stop at 256 elements, one block). The IndexOf* tests pass for all element types on the 512-, 256- and 128-bit and scalar paths. Co-Authored-By: Claude Fable 5.1 --- ...TensorPrimitives.IIndexOfMinMaxOperator.cs | 887 +++++++----------- .../netcore/TensorPrimitives.IndexOfMax.cs | 57 +- .../TensorPrimitives.IndexOfMaxMagnitude.cs | 8 + .../netcore/TensorPrimitives.IndexOfMin.cs | 8 + .../TensorPrimitives.IndexOfMinMagnitude.cs | 8 + .../tests/TensorPrimitivesTests.cs | 167 ++++ 6 files changed, 537 insertions(+), 598 deletions(-) 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..809bef096a1ba5 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 + public static 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,461 @@ private static int IndexOfMinMaxFallback(ReadOnlySpan x) return resultIndex; } - private static int IndexOfMinMaxVectorized128Size4Plus(ReadOnlySpan x) - where T : INumber where TOperator : struct, IIndexOfMinMaxOperator where TInt : IBinaryInteger + /// See . + private static int IndexOfMinMaxBlocks128(ReadOnlySpan x) + 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)); + Debug.Assert(Vector128.IsHardwareAccelerated && Vector128.IsSupported); + Debug.Assert(x.Length >= Vector128.Count); - // 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); - } - } + int blockSize = BlockVectors * Vector128.Count; + int length = x.Length; + ref T xRef = ref MemoryMarshal.GetReference(x); - // Initialize indices. - Vector128 indexIncrement = Vector128.Create(TInt.CreateTruncating(Vector128.Count)); - Vector128 resultIndex = Vector128.Indices; - Vector128 currentIndex = resultIndex + indexIncrement; - ReadOnlySpan span = x.Slice(Vector128.Count); - - while (!span.IsEmpty) + // 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); - - // 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)); - } + // 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); - - // 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); - } - } + Debug.Assert(length >= 1); - // 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); + nuint count = (nuint)Vector128.Count; + nuint end = (nuint)length; + T result; + nuint i; - while (!span.IsEmpty) + if (end >= 2 * count) { - Vector128 current; - if (span.Length >= Vector128.Count) + // 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) { - 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; + 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); + + 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); - - // 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)); - } + // 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) + // 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) { - 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; + 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) + if (!TOperator.Compare(value, 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(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); + + 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); - - // 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)); - } + // 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) - { - current = Vector512.Create(span); - span = span.Slice(Vector512.Count); - } - else + if (T.IsNaN(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((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 From 46e30e09bd8d3b5fb081819e3abf8637e6f0ccbc Mon Sep 17 00:00:00 2001 From: Niklas Schilli Date: Tue, 15 Sep 2026 19:27:26 +0200 Subject: [PATCH 2/2] Assert the block-reduction contracts in debug builds Add debug-only checks to IndexOfMinMaxCore: after each block reduction, no element of the block beats the reduction result under the operator's Compare and no NaN survived the reduction (the properties the final scan and the NaN handling rely on), the element size is one the lane extraction supports, and pass 1 selected a block. Co-Authored-By: Claude Fable 5.1 --- ...TensorPrimitives.IIndexOfMinMaxOperator.cs | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) 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 809bef096a1ba5..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 @@ -8,7 +8,7 @@ namespace System.Numerics.Tensors { - public static partial class TensorPrimitives + public static unsafe partial class TensorPrimitives { /// The operations an IndexOfMin/Max-style search needs from its ordering. /// @@ -107,12 +107,33 @@ private static int IndexOfMinMaxFallback(ReadOnlySpan x) return resultIndex; } + /// + /// 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 + { + for (int i = 0; i < length; i++) + { + T element = Unsafe.Add(ref xRef, i); + if (T.IsNaN(element) || TOperator.Compare(element, value)) + { + return false; + } + } + + return true; + } + /// 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; @@ -134,6 +155,9 @@ private static int IndexOfMinMaxBlocks128(ReadOnlySpan x) } } + 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."); + if (resultBlock < 0 || TOperator.Compare(blockResult, result)) { result = blockResult; @@ -141,6 +165,8 @@ private static int IndexOfMinMaxBlocks128(ReadOnlySpan x) } } + 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); } @@ -265,6 +291,7 @@ private static int IndexOfMinMaxBlocks256(ReadOnlySpan x) { 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; @@ -286,6 +313,9 @@ private static int IndexOfMinMaxBlocks256(ReadOnlySpan x) } } + 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."); + if (resultBlock < 0 || TOperator.Compare(blockResult, result)) { result = blockResult; @@ -293,6 +323,8 @@ private static int IndexOfMinMaxBlocks256(ReadOnlySpan x) } } + 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); } @@ -417,6 +449,7 @@ private static int IndexOfMinMaxBlocks512(ReadOnlySpan x) { 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; @@ -438,6 +471,9 @@ private static int IndexOfMinMaxBlocks512(ReadOnlySpan x) } } + 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."); + if (resultBlock < 0 || TOperator.Compare(blockResult, result)) { result = blockResult; @@ -445,6 +481,8 @@ private static int IndexOfMinMaxBlocks512(ReadOnlySpan x) } } + 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); }