⚡ Thunderbolt: Softmax — AVX2 8x unroll with single-FMA exp256#67
⚡ Thunderbolt: Softmax — AVX2 8x unroll with single-FMA exp256#67bugparty wants to merge 1 commit into
Conversation
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. |
📝 WalkthroughWalkthroughAdds a new AVX2 softmax kernel variant ( ChangesSoftmax v6 AVX2 kernel
Sequence Diagram(s)sequenceDiagram
participant Test as test_softmax_v6
participant V6 as softmax_v6
participant Exp as exp256_ps_v3
participant Naive as softmax_naive
Test->>V6: run on input vector
V6->>Exp: compute exponentials (8x unrolled)
Exp-->>V6: exp values
V6-->>Test: output probabilities
Test->>Naive: run on same input
Naive-->>Test: reference output
Test->>Test: compare outputs, check sum ≈ 1
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 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 (3)
ml_kernels/src/test_naive_ops.cpp (1)
184-220: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove these function-body braces onto their own lines.
The changed function definitions still keep
{on the signature line. As per coding guidelines, "Keep braces on their own lines for function bodies."Also applies to: 222-228
🤖 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 184 - 220, The function-body brace style in test_softmax_v6 (and the related test functions noted in the comment) still places the opening brace on the signature line, which violates the project’s brace convention. Update the affected function definitions so the opening and closing braces are on their own lines, matching the existing style used elsewhere in ml_kernels/src/test_naive_ops.cpp, and keep the function bodies unchanged.Source: Coding guidelines
ml_kernels/src/kernel_bench.cpp (1)
337-342: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReformat these function bodies to match the C/C++ style guide.
The opening braces for
name()andrun()are still on the signature line. As per coding guidelines, "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 in the benchmark kernel class are not following the C/C++ brace style guide because the opening braces for name() and run() are on the same lines as their signatures. Reformat these overrides in kernel_bench.cpp so the braces for the softmax_v6 benchmark’s name() and run() methods are placed on their own lines, matching the project’s function-body style used elsewhere in the class.Source: Coding guidelines
ml_kernels/include/ml_kernels/softmax.h (1)
504-531: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the new function-body braces onto their own lines.
Lines 504 and 537 keep the opening brace on the signature line, which does not match the repo's C/C++ style rule. As per coding guidelines, "Keep braces on their own lines for function bodies."
Also applies to: 537-627
🤖 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 504 - 531, The new function bodies in softmax.h do not follow the repo’s brace style, since the opening brace is kept on the signature line for exp256_ps_v3 and the other affected functions. Move the function-body opening braces onto their own lines for exp256_ps_v3 and the related functions in the same block, keeping the existing function names and structure unchanged.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.
Inline comments:
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 335-344: SoftmaxV6Benchmark still inherits SoftmaxBenchmark’s
stricter verify() behavior, so it can fail even though softmax_v6 now only
guarantees the looser 1e-4f contract. Update SoftmaxV6Benchmark (or its verify
override in kernel_bench.cpp) to match the new kernel accuracy tolerance used by
softmax_v6, keeping the change scoped to the benchmark class and its
verification logic.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 184-220: The test_softmax_v6() input currently leaves an 8-element
tail, so softmax_v6 only exercises the vectorized remainder path and never the
scalar cleanup loop. Adjust the test data in test_softmax_v6 so the total input
size leaves a remainder smaller than the vector-tail width after the unrolled
body, ensuring the final scalar for (; i < n; ++i) path in softmax_v6 is
actually hit while still comparing against softmax_naive.
---
Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Around line 504-531: The new function bodies in softmax.h do not follow the
repo’s brace style, since the opening brace is kept on the signature line for
exp256_ps_v3 and the other affected functions. Move the function-body opening
braces onto their own lines for exp256_ps_v3 and the related functions in the
same block, keeping the existing function names and structure unchanged.
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 337-342: The function bodies in the benchmark kernel class are not
following the C/C++ brace style guide because the opening braces for name() and
run() are on the same lines as their signatures. Reformat these overrides in
kernel_bench.cpp so the braces for the softmax_v6 benchmark’s name() and run()
methods are placed on their own lines, matching the project’s function-body
style used elsewhere in the class.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 184-220: The function-body brace style in test_softmax_v6 (and the
related test functions noted in the comment) still places the opening brace on
the signature line, which violates the project’s brace convention. Update the
affected function definitions so the opening and closing braces are on their own
lines, matching the existing style used elsewhere in
ml_kernels/src/test_naive_ops.cpp, and keep the function bodies 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: c3f922d4-e064-404b-9e30-0371e7e886af
📒 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
| class SoftmaxV6Benchmark : public SoftmaxBenchmark { | ||
| public: | ||
| const char *name() const override { return "softmax_v6"; } | ||
|
|
||
| void run() override { | ||
| ml_kernels::softmax_v6(inputs_[current_idx_].data(), outputs_[current_idx_].data(), inputs_[0].size()); | ||
| current_idx_ = (current_idx_ + 1) % pool_size_; | ||
| } | ||
| }; | ||
| REGISTER_BENCHMARK(SoftmaxV6Benchmark); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Relax softmax_v6's benchmark verifier to the new kernel contract.
Because Line 335 subclasses SoftmaxBenchmark unchanged, this benchmark still inherits the 1e-5f verify() threshold, while the new kernel/test/docs only claim 1e-4f accuracy. That can make the benchmark report a false failure for a kernel that is behaving as designed.
Proposed fix
class SoftmaxV6Benchmark : public SoftmaxBenchmark {
public:
const char *name() const override { return "softmax_v6"; }
void run() override {
ml_kernels::softmax_v6(inputs_[current_idx_].data(), outputs_[current_idx_].data(), inputs_[0].size());
current_idx_ = (current_idx_ + 1) % pool_size_;
}
+
+ bool verify() override {
+ current_idx_ = 0;
+ run();
+ constexpr float tol = 1e-4f;
+ for (std::size_t i = 0; i < outputs_[0].size(); ++i) {
+ if (std::fabs(outputs_[0][i] - output_ref_[i]) > tol) {
+ return false;
+ }
+ }
+ return true;
+ }
};📝 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.
| class SoftmaxV6Benchmark : public SoftmaxBenchmark { | |
| public: | |
| const char *name() const override { return "softmax_v6"; } | |
| void run() override { | |
| ml_kernels::softmax_v6(inputs_[current_idx_].data(), outputs_[current_idx_].data(), inputs_[0].size()); | |
| current_idx_ = (current_idx_ + 1) % pool_size_; | |
| } | |
| }; | |
| REGISTER_BENCHMARK(SoftmaxV6Benchmark); | |
| class SoftmaxV6Benchmark : public SoftmaxBenchmark { | |
| public: | |
| const char *name() const override { return "softmax_v6"; } | |
| void run() override { | |
| ml_kernels::softmax_v6(inputs_[current_idx_].data(), outputs_[current_idx_].data(), inputs_[0].size()); | |
| current_idx_ = (current_idx_ + 1) % pool_size_; | |
| } | |
| bool verify() override { | |
| current_idx_ = 0; | |
| run(); | |
| constexpr float tol = 1e-4f; | |
| for (std::size_t i = 0; i < outputs_[0].size(); ++i) { | |
| if (std::fabs(outputs_[0][i] - output_ref_[i]) > tol) { | |
| return false; | |
| } | |
| } | |
| return true; | |
| } | |
| }; | |
| REGISTER_BENCHMARK(SoftmaxV6Benchmark); |
🤖 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 335 - 344, SoftmaxV6Benchmark
still inherits SoftmaxBenchmark’s stricter verify() behavior, so it can fail
even though softmax_v6 now only guarantees the looser 1e-4f contract. Update
SoftmaxV6Benchmark (or its verify override in kernel_bench.cpp) to match the new
kernel accuracy tolerance used by softmax_v6, keeping the change scoped to the
benchmark class and its verification logic.
| void test_softmax_v6() { | ||
| std::cout << "Running test_softmax_v6..." << std::endl; | ||
| 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, | ||
| -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, | ||
| 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f // To test remainder | ||
| }; | ||
|
|
||
| std::vector<float> output_naive(input.size(), 0.0f); | ||
| std::vector<float> output_v6(input.size(), 0.0f); | ||
|
|
||
| ml_kernels::softmax_naive(input.data(), output_naive.data(), input.size()); | ||
| ml_kernels::softmax_v6(input.data(), output_v6.data(), input.size()); | ||
|
|
||
| float sum = 0.0f; | ||
| for (std::size_t i = 0; i < input.size(); ++i) { | ||
| assert(std::fabs(output_naive[i] - output_v6[i]) < 1e-4f); | ||
| sum += output_v6[i]; | ||
| } | ||
| assert(std::fabs(sum - 1.0f) < 1e-4f); | ||
|
|
||
| std::cout << "test_softmax_v6 passed!" << std::endl; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
This test never reaches softmax_v6's scalar remainder loop.
Line 203 adds exactly 8 trailing values, so after the 64-element unrolled body the i + 7 < n vector-tail loop consumes the remainder too. The final scalar for (; i < n; ++i) path is still untested even though the comment says otherwise.
Proposed fix
- 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f // To test remainder
+ 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f // Forces a scalar tail after the 8-lane remainder loop📝 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.
| void test_softmax_v6() { | |
| std::cout << "Running test_softmax_v6..." << std::endl; | |
| 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, | |
| -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, | |
| 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f // To test remainder | |
| }; | |
| std::vector<float> output_naive(input.size(), 0.0f); | |
| std::vector<float> output_v6(input.size(), 0.0f); | |
| ml_kernels::softmax_naive(input.data(), output_naive.data(), input.size()); | |
| ml_kernels::softmax_v6(input.data(), output_v6.data(), input.size()); | |
| float sum = 0.0f; | |
| for (std::size_t i = 0; i < input.size(); ++i) { | |
| assert(std::fabs(output_naive[i] - output_v6[i]) < 1e-4f); | |
| sum += output_v6[i]; | |
| } | |
| assert(std::fabs(sum - 1.0f) < 1e-4f); | |
| std::cout << "test_softmax_v6 passed!" << std::endl; | |
| } | |
| 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f // Forces a scalar tail after the 8-lane remainder loop |
🤖 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 184 - 220, The
test_softmax_v6() input currently leaves an 8-element tail, so softmax_v6 only
exercises the vectorized remainder path and never the scalar cleanup loop.
Adjust the test data in test_softmax_v6 so the total input size leaves a
remainder smaller than the vector-tail width after the unrolled body, ensuring
the final scalar for (; i < n; ++i) path in softmax_v6 is actually hit while
still comparing against softmax_naive.
1 participant
💡 What: Added
softmax_v6toml_kernelsusing AVX2. This implementation simplifies theexp256approximation by using a single FMA forrcalculation, and unrolls all phases (max, exp+sum, scaling) 8x.🎯 Why: To improve throughput over
softmax_v5by reducing instruction count and register pressure during the exponential calculation, allowing deeper unrolling to better saturate execution ports.🏗️ How: Replaced the two-part subtraction of
ln(2)in the range reduction step ofexp256with a single FMA operation using a combined constant. Applied 8x unrolling to the entire Softmax kernel.📊 Impact:
softmax_v5at 5.55 GFLOP/s (~12% gain)softmax_v5at 4.06 GFLOP/s (~13% gain)🖥️ Tested on: GCC 13.3.0, Ubuntu Linux container.
🔬 How to reproduce:
PR created automatically by Jules for task 13459054292008055690 started by @bugparty
Summary by CodeRabbit
New Features
Bug Fixes
Tests