Fix integer overflow in TensorPrimitives IndexOf methods

[pine919]

Fixes #124233

Description

IndexOfMax, IndexOfMin, IndexOfMaxMagnitude, and IndexOfMinMagnitude were using Vector<T> for index tracking, where T is the element type. This caused integer overflow when the array length exceeded the maximum value that T could represent.

For example:

• A byte array with 258 elements would incorrectly return index 1 instead of 257
• A short array with 32770 elements would incorrectly return index 1 instead of 32769

Changes

Changed index tracking to use Vector<int> throughout the implementation, preventing overflow regardless of array size.

• Updated IIndexOfOperator<T> interface signatures to use Vector128/256/512<int> for index parameters
• Modified all four IndexOf operators to use Vector<int> for index tracking
• Simplified IndexLessThan helper methods to return Vector<int> directly
• Removed redundant type conversion helper functions
• Added regression test covering byte, sbyte, short, and ushort types

Testing

Added IndexOfMax_NoIntegerOverflow test that verifies correct behavior for arrays exceeding the element type's maximum value:

• byte[258] → expects index 257
• sbyte[258] → expects index 257
• short[32770] → expects index 32769
• ushort[65538] → expects index 65537

cc: @dotnet/area-system-numerics-tensors

[pine919]

@dotnet-policy-service agree

[pine919]

@dhartglassMSFT Please review my PR

[tannergooding]

There is a lot of extra churn here that makes it harder to see what is fix vs what is "cleanup" and which will make it harder to backport. Ideally you separate out such changes, particularly for bug fixes, so that we have the minimum churn possible.

Additionally, this particular approach is going to be more expensive and halve or quarter the throughput for the small integer types. It would likely be better to have an additional ref int baseIndex parameter that can be used to adjust resultIndex in the final step based on the T.MaxValue wraparound. This way the core of the logic doesn't change at all and instead we only have a correcting step at the end to account for the mod wraparound

[pine919]

Thanks for your time

[stephentoub]

🤖 Copilot Code Review — PR #124274

Holistic Assessment

Motivation: The bug is real and well-documented in issue #124233. The vectorized IndexOf* methods store indices in Vector<T> where T is the element type, which overflows when array length exceeds byte.MaxValue (256) or short.MaxValue (32768) / ushort.MaxValue (65536).

Approach: The PR takes a scalar-fallback approach rather than fixing the vectorized path to use wider index types. While this correctly fixes the bug, it abandons vectorization for larger arrays of small element types — a potentially significant performance regression for common use cases.

Summary: ⚠️ Needs Human Review. The fix is correct for the reported bug, but the approach may have unintended performance consequences that warrant maintainer input. There are also several issues with the implementation details.

---

Detailed Findings

⚠️ Performance Regression — Scalar fallback for large arrays

The fix forces scalar processing for arrays larger than 256 elements (byte/sbyte) or 65536 elements (short/ushort). This is a significant performance cliff:

• byte[] arrays > 256 elements will use scalar path — this is a very small threshold
• short[] arrays > 65536 elements will use scalar path

For context, a 65KB byte[] is common in many workloads. The performance difference between vectorized and scalar paths can be 4-16x depending on hardware. Consider instead widening the index vector to Vector<int> or Vector<long> to maintain vectorization while supporting larger indices.

Location: TensorPrimitives.IndexOfMax.cs lines 152-157, 243-248, 334-339

⚠️ Incomplete Fix — Byte case not fully addressed

The PR adds unsigned comparison for short types (Vector*.LessThan(indices1.AsUInt16(), indices2.AsUInt16())), but for byte it continues using signed AsByte(). While byte is unsigned by nature in C#, the IndexOfFinalAggregate method's return at line 98 still does:

return resultIndex.As<T, byte>().ToScalar();

This only returns a byte (0-255), but the indices in the vector could have wrapped around. For consistency with the short fix, this should cast to int:

return (int)resultIndex.As<T, byte>().ToScalar();

Although byte is unsigned so this won't produce negative results, it's inconsistent with the short path fix.

Location: TensorPrimitives.IIndexOfOperator.cs line 98

⚠️ Off-by-one threshold — 256 should be 255

The check uses > 256 for byte types, but a byte can represent values 0-255 (256 distinct values). An array of length 256 would have indices 0-255, which fit in a byte. The threshold should be > 255 (or equivalently >= 256).

Similarly, for unsigned short the threshold > 65536 is correct (indices 0-65535 fit), but the comment says "byte and short types" — signed short can only represent indices 0-32767 before going negative. The code treats short and ushort identically with threshold 65536, which is only correct for ushort.

For signed short, index 32768 would be stored as -32768, and unsigned comparison would treat it incorrectly. The threshold for short should be > 32767.

Location: TensorPrimitives.IndexOfMax.cs lines 152-157

⚠️ Test structure — Should use [Theory] pattern

The test uses a single [ConditionalFact] with runtime type checks instead of the preferred [Theory] with [InlineData] or [MemberData]. The existing test patterns in this file use [Fact] per type-parameterized test class, so the runtime if (typeof(T) == ...) checks effectively make this test run only for specific type instantiations while silently passing for others.

Additionally, the test only covers IndexOfMax but the bug affects IndexOfMin, IndexOfMaxMagnitude, and IndexOfMinMagnitude as well.

Location: TensorPrimitivesTests.cs lines 1137-1181

💡 Suggestion — Comment accuracy

The added comment says "For short/ushort, use unsigned comparison to handle indices up to 65535" but this is in the sizeof(T) == 2 block which also handles Half (floating-point). The comment should clarify this is for the index comparison, not the value comparison.

Location: TensorPrimitives.IIndexOfOperator.cs line 56

---

Summary

The PR correctly identifies and addresses a real bug, but the chosen approach (scalar fallback) has meaningful performance implications that should be discussed. A fix that maintains vectorization by widening the index storage type would be preferable if feasible. The implementation also has threshold accuracy issues (256 vs 255, signed vs unsigned short handling).

I recommend a maintainer weigh in on whether:

1. The scalar fallback performance trade-off is acceptable, or
2. A more complex fix that preserves vectorization is warranted

If the scalar fallback approach is accepted, the threshold values need correction.

[dotnet-policy-service]

Tagging subscribers to this area: @dotnet/area-system-numerics-tensors
See info in area-owners.md if you want to be subscribed.