⚡ Thunderbolt: softmax_v6 — AVX2 8x Unrolled Softmax#75
Conversation
💡 What Implemented `softmax_v6`, an AVX2-vectorized Softmax kernel that unrolls the main computational loops 8x and utilizes `exp256_ps_v2` for transcendental approximation. Updated `kernel_bench.cpp` to benchmark it and added functional correctness checks in `test_naive_ops.cpp`. 🎯 Why The `_mm256_max_ps` instruction and the FMA chains inside `exp256_ps_v2` have significant instruction latency (4+ cycles). A 4x unroll (as seen in `softmax_v5`) only maintains 4 independent streams, which is insufficient to completely hide the latency on modern out-of-order execution engines. Unrolling 8x maintains 8 streams, fully saturating the FMA execution units and shifting the bottleneck from instruction latency directly to memory bandwidth. 🏗️ How The three main passes of Softmax (find max, compute exp/sum, normalize) were unrolled by a factor of 8 (handling 64 elements per loop iteration). Eight independent accumulators/streams were utilized across all phases, combined via tree-reductions before the scalar remainder loop. 📊 Impact Benchmarking results (N=1048576, Fixed Memory mode) showed `softmax_v6` achieved 3.46 GFLOP/s compared to `softmax_v5`'s 3.17 GFLOP/s, a ~9% throughput improvement for large arrays due to better FMA execution port saturation. 🖥️ Tested on Local runner (Linux, GCC). 🔬 How to reproduce 1. `mkdir build && cd build && cmake .. && make -j$(nproc) ml_kernel_bench` 2. `DISABLE_CPU_BINDING=1 ./ml_kernels/ml_kernel_bench --filter "softmax"` 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. |
📝 WalkthroughWalkthroughAdded an 8x-unrolled AVX2 ChangesSoftmax v6
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant TestOrBenchmark
participant softmax_v6
participant exp256_ps_v2
participant Output
TestOrBenchmark->>softmax_v6: Provide input and output buffers
softmax_v6->>softmax_v6: Compute maximum and subtract
softmax_v6->>exp256_ps_v2: Evaluate exponentials
exp256_ps_v2-->>softmax_v6: Return vector exponentials
softmax_v6->>Output: Write normalized probabilities
Output-->>TestOrBenchmark: Return results for validation or timing
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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
🤖 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.
Inline comments:
In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 186-199: Add a separate test case based on the existing 72-element
input in the relevant test function, extending it to 73 elements with one
additional value and updating the expected result/length assertions. Ensure the
test invokes the same operation so the 64-element unrolled loop, 8-element tail,
and final scalar remainder path are all exercised.
🪄 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: 550a5e0d-08e5-4882-90e1-a63e7cffa0c5
📒 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
| // Length is 72 to test the 8x unrolled loop (64 elements) and remainder handling | ||
| 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, | ||
|
|
||
| 0.1f, 0.2f, 0.3f, 0.4f, -0.1f, -0.2f, -0.3f, -0.4f, | ||
| 2.5f, 3.5f, 4.5f, 5.5f, -2.5f, -3.5f, -4.5f, -5.5f, | ||
| 15.0f, 25.0f, 35.0f, 45.0f, -15.0f, -25.0f, -35.0f, -45.0f, | ||
| 8.0f, 9.0f, 10.0f, 11.0f, -8.0f, -9.0f, -10.0f, -11.0f, | ||
|
|
||
| 12.0f, 13.0f, 14.0f, 15.0f, -12.0f, -13.0f, -14.0f, -15.0f | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add a test case exercising the scalar remainder path.
The 72-element input tests the 64-element unrolled loop and the 8-element tail, but the scalar remainder loop (for (; i < n; ++i) at lines 614-618 and 657-659) is never reached since 72 is divisible by 8. Consider adding a test with n not divisible by 8 (e.g., 73 elements) to cover all three code paths.
💡 Suggested addition
Append one extra element to the input vector to make it 73 elements, exercising the scalar tail:
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,
0.1f, 0.2f, 0.3f, 0.4f, -0.1f, -0.2f, -0.3f, -0.4f,
2.5f, 3.5f, 4.5f, 5.5f, -2.5f, -3.5f, -4.5f, -5.5f,
15.0f, 25.0f, 35.0f, 45.0f, -15.0f, -25.0f, -35.0f, -45.0f,
8.0f, 9.0f, 10.0f, 11.0f, -8.0f, -9.0f, -10.0f, -11.0f,
- 12.0f, 13.0f, 14.0f, 15.0f, -12.0f, -13.0f, -14.0f, -15.0f
+ 12.0f, 13.0f, 14.0f, 15.0f, -12.0f, -13.0f, -14.0f, -15.0f,
+ 7.0f
};📝 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.
| // Length is 72 to test the 8x unrolled loop (64 elements) and remainder handling | |
| 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, | |
| 0.1f, 0.2f, 0.3f, 0.4f, -0.1f, -0.2f, -0.3f, -0.4f, | |
| 2.5f, 3.5f, 4.5f, 5.5f, -2.5f, -3.5f, -4.5f, -5.5f, | |
| 15.0f, 25.0f, 35.0f, 45.0f, -15.0f, -25.0f, -35.0f, -45.0f, | |
| 8.0f, 9.0f, 10.0f, 11.0f, -8.0f, -9.0f, -10.0f, -11.0f, | |
| 12.0f, 13.0f, 14.0f, 15.0f, -12.0f, -13.0f, -14.0f, -15.0f | |
| }; | |
| // Length is 72 to test the 8x unrolled loop (64 elements) and remainder handling | |
| 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, | |
| 0.1f, 0.2f, 0.3f, 0.4f, -0.1f, -0.2f, -0.3f, -0.4f, | |
| 2.5f, 3.5f, 4.5f, 5.5f, -2.5f, -3.5f, -4.5f, -5.5f, | |
| 15.0f, 25.0f, 35.0f, 45.0f, -15.0f, -25.0f, -35.0f, -45.0f, | |
| 8.0f, 9.0f, 10.0f, 11.0f, -8.0f, -9.0f, -10.0f, -11.0f, | |
| 12.0f, 13.0f, 14.0f, 15.0f, -12.0f, -13.0f, -14.0f, -15.0f, | |
| 7.0f | |
| }; |
🤖 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/src/test_naive_ops.cpp` around lines 186 - 199, Add a separate
test case based on the existing 72-element input in the relevant test function,
extending it to 73 elements with one additional value and updating the expected
result/length assertions. Ensure the test invokes the same operation so the
64-element unrolled loop, 8-element tail, and final scalar remainder path are
all exercised.
1 participant
💡 What
Implemented
softmax_v6, an AVX2-vectorized Softmax kernel that unrolls the main computational loops 8x and utilizesexp256_ps_v2for transcendental approximation. Updatedkernel_bench.cppto benchmark it and added functional correctness checks intest_naive_ops.cpp.🎯 Why
The
_mm256_max_psinstruction and the FMA chains insideexp256_ps_v2have significant instruction latency (4+ cycles). A 4x unroll (as seen insoftmax_v5) only maintains 4 independent streams, which is insufficient to completely hide the latency on modern out-of-order execution engines. Unrolling 8x maintains 8 streams, fully saturating the FMA execution units and shifting the bottleneck from instruction latency directly to memory bandwidth.🏗️ How
The three main passes of Softmax (find max, compute exp/sum, normalize) were unrolled by a factor of 8 (handling 64 elements per loop iteration). Eight independent accumulators/streams were utilized across all phases, combined via tree-reductions before the scalar remainder loop.
📊 Impact
Benchmarking results (N=1048576, Fixed Memory mode) showed
softmax_v6achieved 3.46 GFLOP/s compared tosoftmax_v5's 3.17 GFLOP/s, a ~9% throughput improvement for large arrays due to better FMA execution port saturation.🖥️ Tested on
Local runner (Linux, GCC).
🔬 How to reproduce
mkdir build && cd build && cmake .. && make -j$(nproc) ml_kernel_benchDISABLE_CPU_BINDING=1 ./ml_kernels/ml_kernel_bench --filter "softmax"PR created automatically by Jules for task 15516429383615574720 started by @bugparty
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Documentation