Skip to content

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

Closed
Closed
[SIMD] Reductions: Add SIMD for Prod, ArgMax/ArgMin, and Axis Reductions#576
Labels
coreInternal engine: Shape, Storage, TensorEngine, iteratorsenhancementNew feature or requestperformancePerformance improvements or optimizations

Description

@Nucs
Nucs
opened on Feb 21, 2026

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

Activity

  1. Nucs
    added
    enhancementNew feature or request
    performancePerformance improvements or optimizations
    coreInternal engine: Shape, Storage, TensorEngine, iterators
    on Feb 21, 2026

  2. Nucs commented on Mar 14, 2026

    @Nucs
    Nucs
    on Mar 14, 2026
    MemberAuthor

    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).

  3. Nucs
    added a commit that references this issue on Mar 14, 2026

    perf: SIMD axis reductions with AVX2 gather and parallel outer loop (#…

    e5b26c4
  4. Nucs
    added a commit that references this issue on Mar 23, 2026

    perf: SIMD axis reductions with AVX2 gather and parallel outer loop (#…

    9734ec0

  5. Nucs commented on Sep 7, 2026

    @Nucs
    Nucs
    on Sep 7, 2026
    MemberAuthor

    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.

  6. Nucs
    closed this as completedon Sep 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    coreInternal engine: Shape, Storage, TensorEngine, iteratorsenhancementNew feature or requestperformancePerformance improvements or optimizations

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions