Skip to content

Use a block-reduction search in TensorPrimitives.IndexOfMin/Max/MinMagnitude/MaxMagnitude #133969

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mrnikbobjeff wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Use a block-reduction search in TensorPrimitives.IndexOfMin/Max/MinMagnitude/MaxMagnitude #133969

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b914502
Use a block-reduction search in TensorPrimitives.IndexOfMin/Max/MinMa…
Sep 15, 2026
6bba49a
Merge branch 'main' into argmin-blocks
Mrnikbobjeff Sep 15, 2026
46e30e0
Assert the block-reduction contracts in debug builds
Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
909 changes: 368 additions & 541 deletions ...ors/src/System/Numerics/Tensors/netcore/Common/TensorPrimitives.IIndexOfMinMaxOperator.cs
View file
Open in desktop

Large diffs are not rendered by default.

57 changes: 8 additions & 49 deletions ...ystem.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMax.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

using System.Runtime.CompilerServices;
using System.Runtime.Intrinsics;
using System.Runtime.Intrinsics.X86;

namespace System.Numerics.Tensors
{
Expand Down Expand Up @@ -32,6 +31,14 @@ public static int IndexOfMax<T>(ReadOnlySpan<T> x)
public static T Aggregate(Vector128<T> x) => HorizontalAggregate<T, MaxOperator<T>>(x);
public static T Aggregate(Vector256<T> x) => HorizontalAggregate<T, MaxOperator<T>>(x);
public static T Aggregate(Vector512<T> x) => HorizontalAggregate<T, MaxOperator<T>>(x);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static T Reduce(T x, T y) => MaxOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector128<T> Reduce(Vector128<T> x, Vector128<T> y) => MaxOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector256<T> Reduce(Vector256<T> x, Vector256<T> y) => MaxOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector512<T> Reduce(Vector512<T> x, Vector512<T> y) => MaxOperator<T>.Invoke(x, y);

[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static bool Compare(T x, T y)
Expand Down Expand Up @@ -98,53 +105,5 @@ private static int IndexOfFirstMatch<T>(Vector256<T> mask) =>
private static int IndexOfFirstMatch<T>(Vector512<T> mask) =>
BitOperations.TrailingZeroCount(mask.ExtractMostSignificantBits());

[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static unsafe Vector128<T> ElementWiseSelect<T>(Vector128<T> mask, Vector128<T> left, Vector128<T> right)
{
if (Sse41.IsSupported)
{
if (typeof(T) == typeof(float)) return Sse41.BlendVariable(left.AsSingle(), right.AsSingle(), (~mask).AsSingle()).As<float, T>();
if (typeof(T) == typeof(double)) return Sse41.BlendVariable(left.AsDouble(), right.AsDouble(), (~mask).AsDouble()).As<double, T>();

if (sizeof(T) == 1) return Sse41.BlendVariable(left.AsByte(), right.AsByte(), (~mask).AsByte()).As<byte, T>();
if (sizeof(T) == 2) return Sse41.BlendVariable(left.AsUInt16(), right.AsUInt16(), (~mask).AsUInt16()).As<ushort, T>();
if (sizeof(T) == 4) return Sse41.BlendVariable(left.AsUInt32(), right.AsUInt32(), (~mask).AsUInt32()).As<uint, T>();
if (sizeof(T) == 8) return Sse41.BlendVariable(left.AsUInt64(), right.AsUInt64(), (~mask).AsUInt64()).As<ulong, T>();
}

return Vector128.ConditionalSelect(mask, left, right);
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static unsafe Vector256<T> ElementWiseSelect<T>(Vector256<T> mask, Vector256<T> left, Vector256<T> right)
{
if (Avx2.IsSupported)
{
if (typeof(T) == typeof(float)) return Avx2.BlendVariable(left.AsSingle(), right.AsSingle(), (~mask).AsSingle()).As<float, T>();
if (typeof(T) == typeof(double)) return Avx2.BlendVariable(left.AsDouble(), right.AsDouble(), (~mask).AsDouble()).As<double, T>();

if (sizeof(T) == 1) return Avx2.BlendVariable(left.AsByte(), right.AsByte(), (~mask).AsByte()).As<byte, T>();
if (sizeof(T) == 2) return Avx2.BlendVariable(left.AsUInt16(), right.AsUInt16(), (~mask).AsUInt16()).As<ushort, T>();
if (sizeof(T) == 4) return Avx2.BlendVariable(left.AsUInt32(), right.AsUInt32(), (~mask).AsUInt32()).As<uint, T>();
if (sizeof(T) == 8) return Avx2.BlendVariable(left.AsUInt64(), right.AsUInt64(), (~mask).AsUInt64()).As<ulong, T>();
}

return Vector256.ConditionalSelect(mask, left, right);
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static unsafe Vector512<T> ElementWiseSelect<T>(Vector512<T> mask, Vector512<T> left, Vector512<T> right)
{
if (Avx512F.IsSupported)
{
if (typeof(T) == typeof(float)) return Avx512F.BlendVariable(left.AsSingle(), right.AsSingle(), (~mask).AsSingle()).As<float, T>();
if (typeof(T) == typeof(double)) return Avx512F.BlendVariable(left.AsDouble(), right.AsDouble(), (~mask).AsDouble()).As<double, T>();

if (sizeof(T) == 4) return Avx512F.BlendVariable(left.AsUInt32(), right.AsUInt32(), (~mask).AsUInt32()).As<uint, T>();
if (sizeof(T) == 8) return Avx512F.BlendVariable(left.AsUInt64(), right.AsUInt64(), (~mask).AsUInt64()).As<ulong, T>();
}

return Vector512.ConditionalSelect(mask, left, right);
}
}
}
8 changes: 8 additions & 0 deletions ...erics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMaxMagnitude.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,14 @@ public static int IndexOfMaxMagnitude<T>(ReadOnlySpan<T> x)
public static T Aggregate(Vector128<T> x) => HorizontalAggregate<T, MaxMagnitudeOperator<T>>(x);
public static T Aggregate(Vector256<T> x) => HorizontalAggregate<T, MaxMagnitudeOperator<T>>(x);
public static T Aggregate(Vector512<T> x) => HorizontalAggregate<T, MaxMagnitudeOperator<T>>(x);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static T Reduce(T x, T y) => MaxMagnitudeOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector128<T> Reduce(Vector128<T> x, Vector128<T> y) => MaxMagnitudeOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector256<T> Reduce(Vector256<T> x, Vector256<T> y) => MaxMagnitudeOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector512<T> Reduce(Vector512<T> x, Vector512<T> y) => MaxMagnitudeOperator<T>.Invoke(x, y);

[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static bool Compare(T x, T y)
Expand Down
8 changes: 8 additions & 0 deletions ...ystem.Numerics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMin.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,14 @@ public static int IndexOfMin<T>(ReadOnlySpan<T> x)
public static T Aggregate(Vector128<T> x) => HorizontalAggregate<T, MinOperator<T>>(x);
public static T Aggregate(Vector256<T> x) => HorizontalAggregate<T, MinOperator<T>>(x);
public static T Aggregate(Vector512<T> x) => HorizontalAggregate<T, MinOperator<T>>(x);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static T Reduce(T x, T y) => MinOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector128<T> Reduce(Vector128<T> x, Vector128<T> y) => MinOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector256<T> Reduce(Vector256<T> x, Vector256<T> y) => MinOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector512<T> Reduce(Vector512<T> x, Vector512<T> y) => MinOperator<T>.Invoke(x, y);

[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static bool Compare(T x, T y)
Expand Down
8 changes: 8 additions & 0 deletions ...erics.Tensors/src/System/Numerics/Tensors/netcore/TensorPrimitives.IndexOfMinMagnitude.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,14 @@ public static int IndexOfMinMagnitude<T>(ReadOnlySpan<T> x)
public static T Aggregate(Vector128<T> x) => HorizontalAggregate<T, MinMagnitudeOperator<T>>(x);
public static T Aggregate(Vector256<T> x) => HorizontalAggregate<T, MinMagnitudeOperator<T>>(x);
public static T Aggregate(Vector512<T> x) => HorizontalAggregate<T, MinMagnitudeOperator<T>>(x);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static T Reduce(T x, T y) => MinMagnitudeOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector128<T> Reduce(Vector128<T> x, Vector128<T> y) => MinMagnitudeOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector256<T> Reduce(Vector256<T> x, Vector256<T> y) => MinMagnitudeOperator<T>.Invoke(x, y);
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static Vector512<T> Reduce(Vector512<T> x, Vector512<T> y) => MinMagnitudeOperator<T>.Invoke(x, y);

[MethodImpl(MethodImplOptions.AggressiveInlining)]
public static bool Compare(T x, T y)
Expand Down
167 changes: 167 additions & 0 deletions src/libraries/System.Numerics.Tensors/tests/TensorPrimitivesTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int> 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<T> x = CreateAndFillTensor(tensorLength);
x[expected] = Enumerable.Min(MemoryMarshal.ToEnumerable<T>(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<T> 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<T> 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<T> 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<T> x);

private void AssertLongLengths(float fill, float best, IndexOfSearch search) =>
AssertLongLengthsValues(ConvertFromSingle(fill), ConvertFromSingle(best), search);

/// <summary>Fills with <paramref name="fill"/>, places <paramref name="best"/> at the expected index and at the end; the expected index must win.</summary>
private void AssertLongLengthsValues(T fill, T best, IndexOfSearch search)
{
Assert.All(s_indexOfMinLongLengths, tensorLength =>
{
foreach (int expected in IndexOfMinLongPositions(tensorLength))
{
using BoundedMemory<T> x = CreateTensor(tensorLength);
x.Span.Fill(fill);
x[expected] = best;
x[tensorLength - 1] = best;
Assert.Equal(expected, search(x.Span));
}
});
}

/// <summary>Fills with <paramref name="fill"/>, puts <paramref name="better"/> first and NaN at the expected index and at the end; the first NaN must win.</summary>
private void AssertLongLengthsNaN(float fill, float better, IndexOfSearch search)
{
Assert.All(s_indexOfMinLongLengths, tensorLength =>
{
foreach (int expected in IndexOfMinLongPositions(tensorLength))
{
using BoundedMemory<T> 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
Expand Down
Loading