⚡ Thunderbolt: Softmax — 8x Unroll Max and Norm Phases#44
Open
bugparty wants to merge 1 commit into
Open
Conversation
This commit introduces `softmax_v6`, heavily optimizing `softmax_v5` via differentiated loop unrolling strategies in AVX2. 💡 **What:** Separates the unroll factors of the softmax phases. The max reduction and normalization phases are now unrolled 8x (64 elements/iteration). The exponential calculation remains unrolled at 4x (32 elements/iteration). 🎯 **Why:** Simple pointwise phases like max and division/normalization scale perfectly to 8x unrolling because they have extremely short dependency chains and low register pressure, effectively saturating Execution Ports. Conversely, unrolling the complex FMA chains within the `exp256` polynomial evaluation to 8x forces YMM register spilling and drastically reduces throughput. 🏗️ **How:** - Implemented an 8x unrolled loop in the max-finding phase utilizing 8 parallel accumulators (`max0` through `max7`). - Implemented an 8x unrolled loop in the normalization phase. - Left the FMA-heavy `exp` evaluation at a 4x unroll. - Registered the new kernel as `SoftmaxV6Benchmark`. - Expanded test bounds in `test_naive_ops.cpp` to >64 elements to explicitly verify correctness of both the 8x primary blocks and the scalar remainder loops. 📊 **Impact:** Microbenchmarks demonstrated fixed-memory GFLOP/s increasing from 3.89 to 4.16 (+6.9%) over `softmax_v5` on large 1M element inputs. 🖥️ **Tested on:** Haswell+ architecture AVX2 nodes using GCC 13. 🔬 **How to reproduce:** ```bash make -j$(nproc) ml_kernel_bench DISABLE_CPU_BINDING=1 ./build/ml_kernels/ml_kernel_bench --filter softmax ``` Co-authored-by: bugparty <1510776+bugparty@users.noreply.github.com>
Contributor
|
👋 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 adds ChangesAVX2 Softmax v6 Implementation and Validation
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~25 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 (4)
ml_kernels/src/test_naive_ops.cpp (1)
185-185: ⚡ Quick winMove the function-body opening brace onto its own line.
The new test uses same-line brace placement, which is out of repo style for C/C++ function bodies.
As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}:Keep braces on their own lines for function bodies.🤖 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 185, The function definition for test_softmax_v6 currently places the opening brace on the same line; change it so the opening brace is on its own line to match repo C/C++ style (i.e., convert "void test_softmax_v6() {" to "void test_softmax_v6()\n{"). Update the function declaration for test_softmax_v6 accordingly, ensuring the rest of the function body remains unchanged.ml_kernels/src/kernel_bench.cpp (1)
337-342: ⚡ Quick winMove the function-body opening braces onto their own lines.
The new benchmark methods use same-line brace placement, which is out of repo style for C/C++ function bodies.
As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}:Keep braces on their own lines for function bodies.🤖 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/kernel_bench.cpp` around lines 337 - 342, The function bodies for the methods name() and run() use same-line opening braces and must follow the repo C/C++ style of putting the opening brace on its own line; update the methods (the name() override returning "softmax_v6" and the run() override that calls ml_kernels::softmax_v6 and updates current_idx_) so the brace for each function is moved to its own line before the body, preserving existing function signatures and bodies exactly.ml_kernels/include/ml_kernels/softmax.h (2)
597-623: ⚡ Quick winReuse the loaded YMM registers in the 8x normalization loop.
This hot loop creates
o0..o7and then another full set ofm0..m7. Reusingo0..o7avoids extra register pressure in the exact path this PR is trying to speed up.♻️ Suggested simplification
- __m256 m0 = _mm256_mul_ps(o0, inv_sum_v); - __m256 m1 = _mm256_mul_ps(o1, inv_sum_v); - __m256 m2 = _mm256_mul_ps(o2, inv_sum_v); - __m256 m3 = _mm256_mul_ps(o3, inv_sum_v); - __m256 m4 = _mm256_mul_ps(o4, inv_sum_v); - __m256 m5 = _mm256_mul_ps(o5, inv_sum_v); - __m256 m6 = _mm256_mul_ps(o6, inv_sum_v); - __m256 m7 = _mm256_mul_ps(o7, inv_sum_v); - - _mm256_storeu_ps(output + i, m0); - _mm256_storeu_ps(output + i + 8, m1); - _mm256_storeu_ps(output + i + 16, m2); - _mm256_storeu_ps(output + i + 24, m3); - _mm256_storeu_ps(output + i + 32, m4); - _mm256_storeu_ps(output + i + 40, m5); - _mm256_storeu_ps(output + i + 48, m6); - _mm256_storeu_ps(output + i + 56, m7); + o0 = _mm256_mul_ps(o0, inv_sum_v); + o1 = _mm256_mul_ps(o1, inv_sum_v); + o2 = _mm256_mul_ps(o2, inv_sum_v); + o3 = _mm256_mul_ps(o3, inv_sum_v); + o4 = _mm256_mul_ps(o4, inv_sum_v); + o5 = _mm256_mul_ps(o5, inv_sum_v); + o6 = _mm256_mul_ps(o6, inv_sum_v); + o7 = _mm256_mul_ps(o7, inv_sum_v); + + _mm256_storeu_ps(output + i, o0); + _mm256_storeu_ps(output + i + 8, o1); + _mm256_storeu_ps(output + i + 16, o2); + _mm256_storeu_ps(output + i + 24, o3); + _mm256_storeu_ps(output + i + 32, o4); + _mm256_storeu_ps(output + i + 40, o5); + _mm256_storeu_ps(output + i + 48, o6); + _mm256_storeu_ps(output + i + 56, o7);🤖 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` around lines 597 - 623, The 8x normalization loop in the softmax code creates both o0..o7 (loaded YMM registers) and m0..m7 (results), which is unnecessary and increases register pressure; instead perform the multiply in-place by reusing/overwriting the loaded registers (e.g., replace m0 = _mm256_mul_ps(o0, inv_sum_v) with o0 = _mm256_mul_ps(o0, inv_sum_v) and similarly for o1..o7) and then store o0..o7 back to output; update the stores to write the reused registers so the loop uses only the original o* variables and removes m* temporaries.
509-509: ⚡ Quick winMove the function-body opening brace onto its own line.
The new function uses same-line brace placement, which is out of repo style for C/C++ function bodies.
As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}:Keep braces on their own lines for function bodies.🤖 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 509, The function softmax_v6 currently places the opening brace on the same line as the signature; change it to place the opening brace on its own line to match repo C/C++ style (i.e., transform "inline void softmax_v6(const float *input, float *output, std::size_t n) {" into a two-line form with the brace on the next line) so the function body conforms with the repository's brace placement rule for functions.
🤖 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 207-210: The test code uses std::exit(1) but doesn't include the
proper header; add the missing include for <cstdlib> at the top of the file so
std::exit is declared (avoid relying on transitive includes). Locate the use of
std::exit in the mismatch check loop (where output_ref and output_v6 are
compared) and add the include directive for <cstdlib> near other standard
headers so the compilation is robust and portable.
---
Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Around line 597-623: The 8x normalization loop in the softmax code creates
both o0..o7 (loaded YMM registers) and m0..m7 (results), which is unnecessary
and increases register pressure; instead perform the multiply in-place by
reusing/overwriting the loaded registers (e.g., replace m0 = _mm256_mul_ps(o0,
inv_sum_v) with o0 = _mm256_mul_ps(o0, inv_sum_v) and similarly for o1..o7) and
then store o0..o7 back to output; update the stores to write the reused
registers so the loop uses only the original o* variables and removes m*
temporaries.
- Line 509: The function softmax_v6 currently places the opening brace on the
same line as the signature; change it to place the opening brace on its own line
to match repo C/C++ style (i.e., transform "inline void softmax_v6(const float
*input, float *output, std::size_t n) {" into a two-line form with the brace on
the next line) so the function body conforms with the repository's brace
placement rule for functions.
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 337-342: The function bodies for the methods name() and run() use
same-line opening braces and must follow the repo C/C++ style of putting the
opening brace on its own line; update the methods (the name() override returning
"softmax_v6" and the run() override that calls ml_kernels::softmax_v6 and
updates current_idx_) so the brace for each function is moved to its own line
before the body, preserving existing function signatures and bodies exactly.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Line 185: The function definition for test_softmax_v6 currently places the
opening brace on the same line; change it so the opening brace is on its own
line to match repo C/C++ style (i.e., convert "void test_softmax_v6() {" to
"void test_softmax_v6()\n{"). Update the function declaration for
test_softmax_v6 accordingly, ensuring the rest of the function body remains
unchanged.
🪄 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: 1860f408-b85a-4cb0-8909-cef4f05116cd
📒 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
Comment on lines
+207
to
+210
| for (size_t i = 0; i < input.size(); ++i) { | ||
| if (std::abs(output_ref[i] - output_v6[i]) > 1e-4) { | ||
| std::cerr << "Mismatch at index " << i << ": expected " << output_ref[i] << ", got " << output_v6[i] << std::endl; | ||
| std::exit(1); |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="$(fd -p 'test_naive_ops.cpp' | head -n1)"
echo "Inspecting ${file}"
rg -n 'std::exit|`#include` <cstdlib>' "${file}"Repository: bugparty/cpu_math_kernels_pri
Length of output: 148
Include <cstdlib> for std::exit.
ml_kernels/src/test_naive_ops.cpp calls std::exit(1) (line 210) but doesn’t include <cstdlib>, relying on transitive includes that aren’t guaranteed to provide std::exit.
🛠️ Minimal fix
`#include` <cassert>
+#include <cstdlib>
`#include` <iostream>🤖 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 207 - 210, The test code uses
std::exit(1) but doesn't include the proper header; add the missing include for
<cstdlib> at the top of the file so std::exit is declared (avoid relying on
transitive includes). Locate the use of std::exit in the mismatch check loop
(where output_ref and output_v6 are compared) and add the include directive for
<cstdlib> near other standard headers so the compilation is robust and portable.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR introduces
softmax_v6, heavily optimizingsoftmax_v5via differentiated loop unrolling strategies in AVX2.💡 What:
Separates the unroll factors of the softmax phases. The max reduction and normalization phases are now unrolled 8x (64 elements/iteration). The exponential calculation remains unrolled at 4x (32 elements/iteration).
🎯 Why:
Simple pointwise phases like max and division/normalization scale perfectly to 8x unrolling because they have extremely short dependency chains and low register pressure, effectively saturating Execution Ports. Conversely, unrolling the complex FMA chains within the
exp256polynomial evaluation to 8x forces YMM register spilling and drastically reduces throughput.🏗️ How:
max0throughmax7).expevaluation at a 4x unroll.SoftmaxV6Benchmark.test_naive_ops.cppto >64 elements to explicitly verify correctness of both the 8x primary blocks and the scalar remainder loops.📊 Impact:
Microbenchmarks demonstrated fixed-memory GFLOP/s increasing from 3.89 to 4.16 (+6.9%) over
softmax_v5on large 1M element inputs.🖥️ Tested on:
Haswell+ architecture AVX2 nodes using GCC 13.
🔬 How to reproduce:
make -j$(nproc) ml_kernel_bench DISABLE_CPU_BINDING=1 ./build/ml_kernels/ml_kernel_bench --filter softmaxPR created automatically by Jules for task 478666488249728123 started by @bugparty
Summary by CodeRabbit
Refactor
Tests
Documentation