[SIMD] Reductions: Add SIMD for Prod, ArgMax/ArgMin, and Axis Reductions

[Nucs]

Overview

Extend the IL kernel generator's SIMD reduction support beyond the current Sum/Max/Min element-wise operations.

Parent issue: #545

Current State

Reduction	Element-wise (axis=null)	Axis Reduction
Sum	✅ SIMD	❌ NDIterator
Max	✅ SIMD	❌ NDIterator
Min	✅ SIMD	❌ NDIterator
Prod	❌ Scalar	❌ NDIterator
ArgMax	❌ Scalar	❌ NDIterator
ArgMin	❌ Scalar	❌ NDIterator

Existing SIMD code: ILKernelGenerator.cs:3996-4095

Task List

Tier 1: Element-wise Improvements

• SIMD Prod (Product reduction)

  • Challenge: No Vector256.HorizontalMultiply() in .NET
  • Solution: Lane-wise multiply accumulator, horizontal product at end

  • Vector256<T> prodVec = Vector256.Create((T)1);
    for (i = 0; i <= vectorEnd; i += vectorCount)
        prodVec = prodVec * Vector256.Load(input + i);
    result = prodVec[0] * prodVec[1] * ... * prodVec[N-1];

  • Expected: 1.5-2× speedup

• SIMD ArgMax/ArgMin

  • Challenge: Need to track indices alongside values
  • Solution: Dual-accumulator (value + index vectors)

  • Vector256<T> maxVec = Vector256.Create(T.MinValue);
    Vector256<int> idxVec = Vector256.Create(0, 1, 2, ...);
    for (i = 0; i <= vectorEnd; i += vectorCount) {
        var vec = Vector256.Load(input + i);
        var mask = Vector256.GreaterThan(vec, maxVec);
        maxVec = Vector256.ConditionalSelect(mask, vec, maxVec);
        idxVec = Vector256.ConditionalSelect(mask.AsInt32(), currentIdx, idxVec);
    }

  • Expected: 2× speedup

Tier 2: Axis Reductions

• SIMD for inner-axis reductions
  • Current: Default.Reduction.AMax.cs uses slow NDIterator loops
  • Opportunity: For axis=-1 (inner), the reduction dimension is contiguous

  • // Current (slow):
    for each outer_coord:
        iter = arr[outer_coord, :].AsIterator()
        while iter.HasNext(): val = iter.MoveNext()  // Virtual calls!
    
    // With SIMD:
    for each outer_coord:
        T* rowPtr = basePtr + outer_coord * outerStride
        result[outer_coord] = SimdReduce(rowPtr, innerSize)  // Vector256 loop

  • Expected: 3-5× speedup for large axis reductions

Files to Modify

File	Changes
ILKernelGenerator.cs	Add Prod/ArgMax/ArgMin to element-wise SIMD paths
ReductionKernel.cs	Add AxisReductionKernel delegate if needed
Default.Reduction.*.cs	Refactor axis reductions to use IL kernels
DefaultEngine.ReductionOp.cs	Wire up new kernels

Benchmarks to Add

[Benchmark] public double Prod_10M() => np.prod(_array);
[Benchmark] public int ArgMax_10M() => np.argmax(_array);
[Benchmark] public NDArray Sum_Axis0() => np.sum(_matrix, axis: 0);
[Benchmark] public NDArray Max_Axis1() => np.amax(_matrix, axis: 1);

Success Criteria

1. Prod element-wise: ≥1.5× faster than current
2. ArgMax/ArgMin element-wise: ≥1.5× faster than current
3. Inner-axis reductions: ≥2× faster than current
4. All existing reduction tests pass

[Nucs]

Implemented: SIMD Axis Reductions (PR pending)

Completed implementation of SIMD axis reductions with two major optimizations:

1. AVX2 Gather for Strided Float/Double

When reducing along non-contiguous axes (e.g., np.sum(arr, axis=0) on a 2D array), the current implementation used scalar loops. Now uses AVX2 gather instructions for float and double types:

// Float: 8 elements per gather using Vector256
var indices = Vector256.Create(0, stride, stride*2, ..., stride*7);
var gathered = Avx2.GatherVector256(data + i * stride, indices, 4);

// Double: 4 elements per gather using Vector256
var indices = Vector128.Create(0, stride, stride*2, stride*3);
var gathered = Avx2.GatherVector256(data + i * stride, indices, 8);

Expected speedup: ~2-3x for strided float/double reductions.

2. Parallel Outer Loop for Large Outputs

When the output array has >1000 elements, the outer loop now uses Parallel.For:

if (outputSize > AxisReductionParallelThreshold) // threshold = 1000
{
    Parallel.For(0, outputSize, outIdx => {
        ReduceAxisElement(input, output, ...);
    });
}

This is especially beneficial for operations like np.sum(large_array, axis=1) where you're computing many independent row sums.

3. Scalar Loop Unrolling

For integer types (no gather support), added 4x loop unrolling for better instruction-level parallelism:

for (; i <= unrollEnd; i += 4)
{
    T v0 = data[i * stride];
    T v1 = data[(i + 1) * stride];
    T v2 = data[(i + 2) * stride];
    T v3 = data[(i + 3) * stride];
    result = Combine(result, v0, v1, v2, v3);
}

Files Changed

• ILKernelGenerator.Reduction.Axis.Simd.cs
  • Added AxisReductionParallelThreshold = 1000
  • Modified AxisReductionSimdHelper<T>() to use parallel loop for large outputs
  • Added ReduceAxisElement<T>() helper for parallel execution
  • Added ReduceStridedAxisGatherFloat() with AVX2 gather
  • Added ReduceStridedAxisGatherDouble() with AVX2 gather
  • Added ReduceStridedAxisScalar<T>() with 4x unrolling

Tests Added

17 new tests in AxisReductionBenchmarkTests.cs:

• Strided access correctness (float, double, int32, int64)
• Parallel path correctness (Sum, Mean, Max)
• Edge cases (small axis, NaN propagation, 3D arrays)
• AVX2 feature detection reporting

All 17 tests pass. No regressions in existing tests (85 pre-existing failures unchanged).

[Nucs]

Following up on my "PR pending" note here, shipped in 0.70.0 (#628). argmax/argmin are a single-pass SIMD tournament, and the axis reductions are vectorized:

// SIMD argmax
Console.WriteLine(np.argmax(np.array(new[] { 1.0, 9, 3, 2 })));   // 1

// vectorized axis reduction
var m = np.array(new[,] { { 1, 2, 3 }, { 4, 5, 6 } });
Console.WriteLine(np.sum(m, axis: 0));   // [5 7 9]

(prod, std, var, and the nan* reductions are vectorized along axes too.) Closing.