⚡ Thunderbolt: Softmax — 8x unroll and single-FMA exp approximation

[bugparty]

💡 What:
Implemented softmax_v6, a highly optimized AVX2 SIMD kernel that computes softmax. It replaces the exact two-FMA constant representation of x - n * ln(2) with a single fused-multiply-add using combined constants, and aggressively unrolls the main processing loops 8x instead of 4x.

🎯 Why:
In shift-invariant operations like Softmax, transcendental accuracy bounds on exp256_ps are significantly less sensitive to intermediate precision variations (since input elements are shifted by max_val to avoid overflow). This allows combining the range reduction constants into a single FMA to reduce instruction count and relieve execution port pressure. The remaining FMA latency bottleneck is hidden by scaling up to an 8x unroll scheme (using 8 independent accumulators/horner chains).

🏗️ How:

• Implemented exp256_ps_v3 combining n * ln(2) constants into a single _mm256_fnmadd_ps.
• Created softmax_v6 performing 64 elements (8 YMM registers) per iteration.
• Registered standard tests and microbenchmarks within the testing harness.

📊 Impact:

• softmax_v5 (4x unroll): ~5.70 GFLOP/s at N=65536
• softmax_v6 (8x unroll + 1-FMA): ~6.29 GFLOP/s at N=65536
• Improvement: ~10% throughput increase on fixed memory allocations.

🖥️ Tested on:
x86_64 CPU (AVX2), GCC 13.3.0, Ubuntu 24.04 Docker Environment

🔬 How to reproduce:

cd build && make -j$(nproc) ml_kernel_bench
./ml_kernels/ml_kernel_bench -f "softmax_v5"
./ml_kernels/ml_kernel_bench -f "softmax_v6"

---

PR created automatically by Jules for task 7445448542612547065 started by @bugparty

Summary by CodeRabbit

• New Features

  • Added a new, faster softmax implementation.
  • Introduced an updated exponential approximation used by the softmax routine.

• Bug Fixes

  • Improved numerical stability and precision in softmax calculations.
  • Added handling for empty inputs and zero-sum edge cases.

• Tests

  • Added coverage comparing the new softmax output against the baseline and checking that results sum to 1.

• Chores

  • Added a new benchmark entry to measure the updated softmax performance.

[google-labs-jules]

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

---

For security, I will only act on instructions from the user who triggered this task.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a new AVX2 exponential approximation function exp256_ps_v3 and a new softmax implementation softmax_v6 using 8x-unrolled max-reduction, exp/sum, and normalization loops. Registers a corresponding benchmark and adds a correctness test comparing against the naive softmax. Documents the technique in notes.

Changes

Softmax v6 Kernel Implementation

Layer / File(s)	Summary
AVX2 exp approximation and softmax_v6 kernel ml_kernels/include/ml_kernels/softmax.h	Adds exp256_ps_v3 (range-clamped AVX2 exponent approximation using combined-constant FMA) and softmax_v6, which uses 8× unrolled __m256 loops for max-reduction, exp/sum accumulation via exp256_ps_v3, staged reduction, and normalization with early returns for n==0 and zero sum.
Benchmark registration and correctness test ml_kernels/src/kernel_bench.cpp, ml_kernels/src/test_naive_ops.cpp	Adds SoftmaxV6Benchmark registered as "softmax_v6" calling softmax_v6, and a test_softmax_v6() routine comparing softmax_v6 output against softmax_naive with a 1e-4 tolerance, wired into main().
Optimization notes documentation .jules/thunderbolt.md	Documents the combined-constant single-FMA range-reduction approach with benchmark evidence comparing softmax_v6 to softmax_v5.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• bugparty/cpu_math_kernels_pri#28: Both PRs add AVX2 exp256_* approximation functions wired into new softmax variants (softmax_v6/exp256_ps_v3 vs softmax_v4/exp256_ps_estrin).
• bugparty/cpu_math_kernels_pri#31: Both PRs add new exp256_ps_* range-reduction implementations and corresponding softmax kernels (softmax_v6 vs softmax_v5) with the same max→exp+sum→normalize structure.
• bugparty/cpu_math_kernels_pri#7: Both PRs extend test coverage around ml_kernels::softmax_naive, used here as the reference for validating softmax_v6.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: a softmax optimization with 8x unrolling and a single-FMA exp approximation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch thunderbolt/softmax-8x-fma-7445448542612547065

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

ml_kernels/include/ml_kernels/softmax.h (1)

510-510: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Brace placement deviates from the repo style guideline.

exp256_ps_v3 and softmax_v6 open the function body brace on the declaration line. The guideline requires the opening brace on its own line for function bodies. This matches the surrounding file's existing style, so treat as optional.

As per coding guidelines: "Keep braces on their own lines for function bodies".

Also applies to: 540-540

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ml_kernels/include/ml_kernels/softmax.h` at line 510, The function body
braces in exp256_ps_v3 and softmax_v6 should follow the repo’s brace style by
placing the opening brace on its own line instead of on the declaration line.
Update the function declarations for these symbols so the body starts with a
standalone brace, matching the surrounding style in softmax.h and the coding
guideline for function bodies.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Line 510: The function body braces in exp256_ps_v3 and softmax_v6 should
follow the repo’s brace style by placing the opening brace on its own line
instead of on the declaration line. Update the function declarations for these
symbols so the body starts with a standalone brace, matching the surrounding
style in softmax.h and the coding guideline for function bodies.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: b5e1a571-da4a-4e15-ac11-2c5da4493ee5

📥 Commits

Reviewing files that changed from the base of the PR and between acca01e and 2d8aaf8.

📒 Files selected for processing (4)

• .jules/thunderbolt.md
• ml_kernels/include/ml_kernels/softmax.h
• ml_kernels/src/kernel_bench.cpp
• ml_kernels/src/test_naive_ops.cpp