⚡ Thunderbolt: softmax_v6 — 8x unrolled Max and Normalization passes - #89
⚡ Thunderbolt: softmax_v6 — 8x unrolled Max and Normalization passes#89bugparty wants to merge 1 commit into
Conversation
💡 What: Implemented `softmax_v6` with 8x unrolled Max reduction and Normalization passes, while retaining 4x unroll for the exponential pass. 🎯 Why: Max reduction (`_mm256_max_ps`) has a 4-cycle latency which is perfectly hidden by an 8x unroll. Normalization is entirely memory bound and benefits from an 8x unroll to saturate the store buffers. However, the exp pass requires many registers for polynomial constants, so keeping it at 4x avoids register spilling. 🏗️ How: Unrolled passes 1 and 3 8x (handling 64 elements per iteration). Fallbacks for 32, 8, and 1 elements were updated appropriately. 📊 Impact: Benchmark GFLOPS improved from ~2.35 to ~2.41 for N=1M elements. 🖥️ Tested on: AVX2-capable CPU via `ml_kernel_bench` and `ml_kernel_test`. 🔬 How to reproduce: Run `./build/ml_kernels/ml_kernel_bench --filter "softmax_v6"` and compare it against `softmax_v5`. 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 the AVX2 ChangesAVX2 softmax_v6
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
🚥 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: 2
🧹 Nitpick comments (1)
ml_kernels/src/test_naive_ops.cpp (1)
187-215: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover each fallback path.
The 72-element case executes the 64-element and 8-element paths. It does not execute the scalar fallback or the 32-element normalization fallback. Add cases such as 65 and 96 elements to cover those paths.
🤖 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 187 - 215, Expand the softmax comparison test around softmax_naive and softmax_v6 to include input sizes that exercise every fallback path, specifically add cases for 65 elements to cover the scalar remainder and 96 elements to cover the 32-element normalization fallback. Preserve the existing value-comparison and sum-to-one assertions for each case.
🤖 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/include/ml_kernels/softmax.h`:
- Line 510: Move each function body opening brace onto its own line for
softmax_v6 in ml_kernels/include/ml_kernels/softmax.h:510-510,
SoftmaxV6Benchmark::run in ml_kernels/src/kernel_bench.cpp:339-339,
test_softmax_v6 in ml_kernels/src/test_naive_ops.cpp:185-185, and main in
ml_kernels/src/test_naive_ops.cpp:220-220; make no other changes.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Line 226: Adjust the indentation of the test_softmax_v6() call in main so it
uses 4 leading spaces, matching the surrounding test calls.
---
Nitpick comments:
In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 187-215: Expand the softmax comparison test around softmax_naive
and softmax_v6 to include input sizes that exercise every fallback path,
specifically add cases for 65 elements to cover the scalar remainder and 96
elements to cover the 32-element normalization fallback. Preserve the existing
value-comparison and sum-to-one assertions for each case.
🪄 Autofix
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 Plus
Run ID: 24615759-5fb7-4b40-adc4-fb7c90b57b9b
📒 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
| // Reason: Avoids `round_ps` by leveraging `cvtps_epi32` rounding mode, and replaces Estrin's scheme with Horner's. | ||
| // Max reduction pass (pass 1) and Normalization pass (pass 3) are now 8x unrolled to maximize execution unit saturation for max_ps and memory bandwidth utilization. | ||
| // Expected gain: ~2-5% over softmax_v5 due to slightly better throughput. | ||
| inline void softmax_v6(const float *input, float *output, std::size_t n) { |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Move function opening braces to their own lines.
ml_kernels/include/ml_kernels/softmax.h#L510-L510: Place thesoftmax_v6opening brace on the next line.ml_kernels/src/kernel_bench.cpp#L339-L339: Place theSoftmaxV6Benchmark::runopening brace on the next line.ml_kernels/src/test_naive_ops.cpp#L185-L185: Place thetest_softmax_v6opening brace on the next line.ml_kernels/src/test_naive_ops.cpp#L220-L220: Place themainopening brace on the next line.
As per coding guidelines, "Keep braces on their own lines for function bodies."
📍 Affects 3 files
ml_kernels/include/ml_kernels/softmax.h#L510-L510(this comment)ml_kernels/src/kernel_bench.cpp#L339-L339ml_kernels/src/test_naive_ops.cpp#L185-L185ml_kernels/src/test_naive_ops.cpp#L220-L220
🤖 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, Move each function body
opening brace onto its own line for softmax_v6 in
ml_kernels/include/ml_kernels/softmax.h:510-510, SoftmaxV6Benchmark::run in
ml_kernels/src/kernel_bench.cpp:339-339, test_softmax_v6 in
ml_kernels/src/test_naive_ops.cpp:185-185, and main in
ml_kernels/src/test_naive_ops.cpp:220-220; make no other changes.
Source: Coding guidelines
| test_softmax_v3(); | ||
| test_softmax_v4(); | ||
| test_softmax_v5(); | ||
| test_softmax_v6(); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use 4-space indentation.
Line 226 has 8 leading spaces. Align this call with the other calls in main.
Proposed fix
- test_softmax_v6();
+ test_softmax_v6();As per coding guidelines, "Use 4-space indentation."
📝 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.
| test_softmax_v6(); | |
| test_softmax_v6(); |
🤖 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` at line 226, Adjust the indentation of the
test_softmax_v6() call in main so it uses 4 leading spaces, matching the
surrounding test calls.
Source: Coding guidelines
1 participant
💡 What: Implemented
softmax_v6with 8x unrolled Max reduction and Normalization passes, while retaining 4x unroll for the exponential pass.🎯 Why: Max reduction (
_mm256_max_ps) has a 4-cycle latency which is perfectly hidden by an 8x unroll. Normalization is entirely memory bound and benefits from an 8x unroll to saturate the store buffers. However, the exp pass requires many registers for polynomial constants, so keeping it at 4x avoids register spilling.🏗️ How: Unrolled passes 1 and 3 8x (handling 64 elements per iteration). Fallbacks for 32, 8, and 1 elements were updated appropriately.
📊 Impact: Benchmark GFLOPS improved from ~2.35 to ~2.41 for N=1M elements.
🖥️ Tested on: AVX2-capable CPU via
ml_kernel_benchandml_kernel_test.🔬 How to reproduce: Run
./build/ml_kernels/ml_kernel_bench --filter "softmax_v6"and compare it againstsoftmax_v5.PR created automatically by Jules for task 7106996948143984687 started by @bugparty
Summary by CodeRabbit
New Features
Tests
Documentation