⚡ Thunderbolt: Softmax — Optimized exp256 range reduction and polynomial eval#31
Conversation
…ial eval Co-authored-by: bugparty <1510776+bugparty@users.noreply.github.com>
|
👋 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 New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
📝 WalkthroughWalkthroughThis PR introduces new optimized AVX2 implementations for the exponential and softmax functions, featuring an improved exp256_ps_v2 that uses Horner polynomial evaluation and direct rounding-to-nearest-even via integer conversion instead of floating-point rounding, and a new softmax_v5 kernel that integrates this optimization with aggressive loop unrolling for improved throughput. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
ml_kernels/include/ml_kernels/softmax.h (1)
403-502: Consider templatizing the v3/v4/v5 softmax bodies on theexpfunctor.
softmax_v5is byte-for-byte identical tosoftmax_v4(andsoftmax_v3) except for theexp256_ps*call on lines 441–444 and 462. That's ~100 lines duplicated three times; every future bug fix (numerical edge cases, tail-loop bounds, sum==0 handling) has to be applied in three places.A small template like the following would eliminate the duplication without any runtime cost (calls inline just like today):
♻️ Sketch
template <__m256 (*ExpFn)(__m256)> inline void softmax_impl(const float *input, float *output, std::size_t n) { // ...existing body, calling ExpFn(x) instead of exp256_ps/exp256_ps_estrin/exp256_ps_v2... } inline void softmax_v3(const float *in, float *out, std::size_t n) { softmax_impl<exp256_ps>(in, out, n); } inline void softmax_v4(const float *in, float *out, std::size_t n) { softmax_impl<exp256_ps_estrin>(in, out, n); } inline void softmax_v5(const float *in, float *out, std::size_t n) { softmax_impl<exp256_ps_v2>(in, out, n); }Non-blocking since the intent seems to be to preserve each historical variant verbatim for benchmarking/archival, but worth considering as the v-series grows.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ml_kernels/include/ml_kernels/softmax.h` around lines 403 - 502, softmax_v3/softmax_v4/softmax_v5 duplicate ~100 lines with only the exp function differing; refactor by extracting the shared body into a templated helper (e.g., softmax_impl) that takes the exp function as a template parameter (pointer-to-function or functor) and call that from softmax_v3/softmax_v4/softmax_v5 with exp256_ps, exp256_ps_estrin, and exp256_ps_v2 respectively; update all internal calls that currently invoke exp256_ps*/exp (locations around the e0/e1/e2/e3 and scalar tail) to call the templated ExpFn so behavior remains identical and inlining/no-runtime-cost is preserved.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 157-166: The test uses a 32-element input vector named input in
test_softmax_v5 which only exercises the 32-wide main loop and never hits the
8-wide tail or scalar remainder in softmax_v5; change the input vector length to
e.g. 41 or 45 (add extra float values) so that the code executes the 32-wide
main loop once, the 8-wide tail (i + 7 < n) at least once, and the scalar
remainder (< n) at least once, leaving test_softmax_v5 and the input variable
name unchanged so the new tail code is covered by the test.
---
Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Around line 403-502: softmax_v3/softmax_v4/softmax_v5 duplicate ~100 lines
with only the exp function differing; refactor by extracting the shared body
into a templated helper (e.g., softmax_impl) that takes the exp function as a
template parameter (pointer-to-function or functor) and call that from
softmax_v3/softmax_v4/softmax_v5 with exp256_ps, exp256_ps_estrin, and
exp256_ps_v2 respectively; update all internal calls that currently invoke
exp256_ps*/exp (locations around the e0/e1/e2/e3 and scalar tail) to call the
templated ExpFn so behavior remains identical and inlining/no-runtime-cost is
preserved.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: cfa6d598-c9cb-4cfc-ba53-757753fc7193
📒 Files selected for processing (4)
.jules/thunderbolt.mdml_kernels/include/ml_kernels/softmax.hml_kernels/src/kernel_bench.cppml_kernels/src/test_naive_ops.cpp
| std::vector<float> input = { | ||
| -2.0f, -0.5f, 1.0f, 3.0f, | ||
| 0.0f, 0.0f, 0.0f, 0.0f, | ||
| 100.0f, 100.0f, -100.0f, -100.0f, | ||
| 5.0f, -5.0f, 2.0f, -2.0f, | ||
| 1.1f, 1.2f, 1.3f, 1.4f, | ||
| -1.1f, -1.2f, -1.3f, -1.4f, | ||
| 10.0f, 20.0f, 30.0f, 40.0f, | ||
| -10.0f, -20.0f, -30.0f, -40.0f | ||
| }; |
There was a problem hiding this comment.
Test size misses the tail loops in softmax_v5.
The input is exactly 32 floats, so execution enters the 32-wide main loop exactly once and both tails (i + 7 < n 8-wide loop and the scalar < n remainder) are never exercised. Given this PR's whole point is new code in the hot path that feeds into those same tails, and given test_softmax_v3/v4 use a 40-element vector (1 iteration of the 8-wide tail), test_softmax_v5 is actually a regression in tail coverage.
Consider bumping the input to a size like 41 or 45 so all three loop phases run at least once:
Proposed additional coverage
std::vector<float> input = {
-2.0f, -0.5f, 1.0f, 3.0f,
0.0f, 0.0f, 0.0f, 0.0f,
100.0f, 100.0f, -100.0f, -100.0f,
5.0f, -5.0f, 2.0f, -2.0f,
1.1f, 1.2f, 1.3f, 1.4f,
-1.1f, -1.2f, -1.3f, -1.4f,
10.0f, 20.0f, 30.0f, 40.0f,
- -10.0f, -20.0f, -30.0f, -40.0f
+ -10.0f, -20.0f, -30.0f, -40.0f,
+ // 8-wide tail
+ 0.25f, -0.25f, 0.75f, -0.75f, 1.5f, -1.5f, 2.5f, -2.5f,
+ // scalar tail
+ 0.1f, -0.1f, 0.3f
};📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| std::vector<float> input = { | |
| -2.0f, -0.5f, 1.0f, 3.0f, | |
| 0.0f, 0.0f, 0.0f, 0.0f, | |
| 100.0f, 100.0f, -100.0f, -100.0f, | |
| 5.0f, -5.0f, 2.0f, -2.0f, | |
| 1.1f, 1.2f, 1.3f, 1.4f, | |
| -1.1f, -1.2f, -1.3f, -1.4f, | |
| 10.0f, 20.0f, 30.0f, 40.0f, | |
| -10.0f, -20.0f, -30.0f, -40.0f | |
| }; | |
| std::vector<float> input = { | |
| -2.0f, -0.5f, 1.0f, 3.0f, | |
| 0.0f, 0.0f, 0.0f, 0.0f, | |
| 100.0f, 100.0f, -100.0f, -100.0f, | |
| 5.0f, -5.0f, 2.0f, -2.0f, | |
| 1.1f, 1.2f, 1.3f, 1.4f, | |
| -1.1f, -1.2f, -1.3f, -1.4f, | |
| 10.0f, 20.0f, 30.0f, 40.0f, | |
| -10.0f, -20.0f, -30.0f, -40.0f, | |
| // 8-wide tail | |
| 0.25f, -0.25f, 0.75f, -0.75f, 1.5f, -1.5f, 2.5f, -2.5f, | |
| // scalar tail | |
| 0.1f, -0.1f, 0.3f | |
| }; |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@ml_kernels/src/test_naive_ops.cpp` around lines 157 - 166, The test uses a
32-element input vector named input in test_softmax_v5 which only exercises the
32-wide main loop and never hits the 8-wide tail or scalar remainder in
softmax_v5; change the input vector length to e.g. 41 or 45 (add extra float
values) so that the code executes the 32-wide main loop once, the 8-wide tail (i
+ 7 < n) at least once, and the scalar remainder (< n) at least once, leaving
test_softmax_v5 and the input variable name unchanged so the new tail code is
covered by the test.
1 participant
What:
Added a new AVX2 kernel
softmax_v5and a companionexp256_ps_v2function that optimizes the exponential approximation._mm256_round_psinstruction with the sequence_mm256_cvtepi32_ps(_mm256_cvtps_epi32(x))to achieve round-to-nearest-even with lower latency.Why:
While Estrin's scheme breaks the FMA dependency chain for a single exponential evaluation,
softmax_v4is explicitly unrolled 4x. In a 4x unrolled loop, multiple independent Horner FMA chains interleave perfectly, saturating the execution ports and hiding latency naturally. Estrin's scheme in this context creates unnecessary instruction overhead and port pressure, acting as a bottleneck. Additionally,round_psis a slow instruction.How:
Implemented
exp256_ps_v2using thecvtps_epi32rounding trick and Horner's FMA chain, then integrated it intosoftmax_v5while maintaining the 4x unroll and shuffle-based horizontal reduction.Impact:
softmax_v5achieves ~5.10 GFLOP/s vssoftmax_v4's 4.48 GFLOP/s in Fixed Memory mode (N=16384). This is a solid ~13.8% throughput improvement.Tested on:
Linux / GCC 13.3.0 / AVX2 (CI runner environment).
How to reproduce:
DISABLE_CPU_BINDING=1 ./build/ml_kernels/ml_kernel_bench --filter "softmax" --sizes 16384,65536PR created automatically by Jules for task 18282112880023903289 started by @bugparty
Summary by CodeRabbit
Documentation
New Features
Tests