⚡ Thunderbolt: Softmax — 8x Unrolling + Single-FMA Range Reduction#64
⚡ Thunderbolt: Softmax — 8x Unrolling + Single-FMA Range Reduction#64bugparty 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. |
📝 WalkthroughWalkthroughAdded a new AVX2 softmax variant with a new exponential helper, benchmark registration, correctness coverage, and a documentation note about the range-reduction and unrolling approach. ChangesAVX2 softmax v6
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 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: 1
🧹 Nitpick comments (4)
ml_kernels/src/kernel_bench.cpp (1)
337-341: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the new function-body braces onto their own lines.
The new
name()andrun()definitions use same-line opening braces, which doesn't match the repository's C/C++ function style. 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 - 341, The new softmax_v6 benchmark methods in the unnamed class use same-line function-body braces, which violates the project’s C/C++ style. Update the `name()` and `run()` overrides in `kernel_bench.cpp` so the opening braces for these function bodies are on their own lines, matching the surrounding function formatting used elsewhere in this file.Source: Coding guidelines
ml_kernels/src/test_naive_ops.cpp (2)
184-184: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the new function-body braces onto their own lines.
test_softmax_v6andmainkeep the opening brace on the signature line, which doesn't match the repository's C/C++ function style. As per coding guidelines, "Keep braces on their own lines for function bodies".Also applies to: 219-219
🤖 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 184, The function-body braces for test_softmax_v6 and main are attached to the signature line, which does not match the C/C++ style used here. Move each opening brace onto its own line for those function definitions, following the same brace placement pattern used elsewhere in this file.Source: Coding guidelines
184-214: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a 1-7 element tail case for
softmax_v6.This fixture has 72 values, so it covers the 64-wide unrolled body and the 8-lane vector remainder, but it never reaches the scalar cleanup loop in
softmax_v6. A 65-71 element case would exercise the new tail path as well.🤖 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 - 214, The softmax_v6 test only covers the 64-wide unrolled path plus the 8-lane remainder, so it misses the scalar cleanup tail. Update test_softmax_v6 to include a separate 65-71 element input case that forces the final scalar loop in ml_kernels::softmax_v6. Compare that result against softmax_naive the same way as the existing assertions, so all execution paths are validated.ml_kernels/include/ml_kernels/softmax.h (1)
505-505: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the new function-body braces onto their own lines.
exp256_ps_v3andsoftmax_v6keep the opening brace on the signature line, which doesn't match the repository's C/C++ function style. As per coding guidelines, "Keep braces on their own lines for function bodies".Also applies to: 539-539
🤖 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 505, The function-body opening braces in exp256_ps_v3 and softmax_v6 should follow the repository’s C/C++ style and be placed on their own lines instead of staying on the signature lines. Update the function declarations in softmax.h so each opening brace is moved to its own line for these two functions, matching the existing brace style used elsewhere in the file.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 is still using the inherited
SoftmaxBenchmark::verify() tolerance, which is stricter than the published
softmax_v6 kernel contract. Override verify() in SoftmaxV6Benchmark and use the
kernel’s expected 1e-4f tolerance so the benchmark matches the registered
softmax_v6 behavior.
---
Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Line 505: The function-body opening braces in exp256_ps_v3 and softmax_v6
should follow the repository’s C/C++ style and be placed on their own lines
instead of staying on the signature lines. Update the function declarations in
softmax.h so each opening brace is moved to its own line for these two
functions, matching the existing brace style used elsewhere in the file.
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 337-341: The new softmax_v6 benchmark methods in the unnamed class
use same-line function-body braces, which violates the project’s C/C++ style.
Update the `name()` and `run()` overrides in `kernel_bench.cpp` so the opening
braces for these function bodies are on their own lines, matching the
surrounding function formatting used elsewhere in this file.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Line 184: The function-body braces for test_softmax_v6 and main are attached
to the signature line, which does not match the C/C++ style used here. Move each
opening brace onto its own line for those function definitions, following the
same brace placement pattern used elsewhere in this file.
- Around line 184-214: The softmax_v6 test only covers the 64-wide unrolled path
plus the 8-lane remainder, so it misses the scalar cleanup tail. Update
test_softmax_v6 to include a separate 65-71 element input case that forces the
final scalar loop in ml_kernels::softmax_v6. Compare that result against
softmax_naive the same way as the existing assertions, so all execution paths
are validated.
🪄 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: 307a9fe9-20c4-4dbd-95dc-a0d0b307d6ad
📒 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
Align SoftmaxV6Benchmark::verify() with the kernel's published tolerance.
This subclass inherits SoftmaxBenchmark::verify(), which still hard-codes 1e-5f, while the new kernel is documented and tested at 1e-4f. That makes the benchmark contract stricter than the implementation you're registering here.
Suggested localized 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
is still using the inherited SoftmaxBenchmark::verify() tolerance, which is
stricter than the published softmax_v6 kernel contract. Override verify() in
SoftmaxV6Benchmark and use the kernel’s expected 1e-4f tolerance so the
benchmark matches the registered softmax_v6 behavior.
1 participant
💡 What:
Introduced
softmax_v6using 8x unrolling and a fasterexp256_ps_v3approximation that leverages a single-FMA range reduction instead of splittingln(2)for exact precision. Added the new kernel to the benchmark suite (SoftmaxV6Benchmark) and verified numerical correctness (test_softmax_v6).🎯 Why:
$e^{x_i - \max}$ ) rather than exact IEEE outputs of individual $e^x$ , we can safely sacrifice minor precision during range reduction without exceeding ML numerical tolerance (verified
softmax_v5heavily relied on the exact precision split forln(2)inexp256. However, because Softmax relies on relative differences (1e-4f). This simplification reduces register pressure and instruction count, allowing us to safely increase unrolling to 8x to perfectly match execution port throughput and hide instruction latency.🏗️ How:
exp256_ps_v3which replaces the two-stepln(2)subtraction with a single_mm256_fnmadd_ps.softmax_v6which implements an 8x unroll (processing 64 elements per loop iteration) for finding the max, computing📊 Impact:
Measured on AVX2:
softmax_v5's 3.81 GFLOP/s (~14.4% throughput gain).🖥️ Tested on:
GNU C++ 13.3.0, Ubuntu 22.04 (Docker), AVX2/FMA target.
🔬 How to reproduce:
PR created automatically by Jules for task 48235933610456447 started by @bugparty
Summary by CodeRabbit
New Features
Bug Fixes
Documentation