⚡ Thunderbolt: Softmax — single FMA range reduction#55
Conversation
Replaced split ln(2) precision constants with a combined single-FMA step in `exp256_ps_v3`. 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 Changessoftmax_v6 kernel, benchmark, and tests
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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.
🧹 Nitpick comments (3)
ml_kernels/src/test_naive_ops.cpp (1)
156-156: ⚡ Quick winUse the repository brace style for the new test function definition.
test_softmax_v6currently places{on the same line as the signature.As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}requires: “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 156, The function `test_softmax_v6` currently has its opening brace on the same line as the function signature, which violates the repository's C++ brace style guidelines. Move the opening brace to its own line by placing it on a new line immediately after the function signature `void test_softmax_v6()` to conform to the requirement that braces for function bodies must be on their own lines.Source: Coding guidelines
ml_kernels/src/kernel_bench.cpp (1)
335-344: ⚡ Quick winApply the project brace style in the new benchmark class methods.
SoftmaxV6Benchmarkmethods use inline opening braces on signature lines; this is inconsistent with the C/C++ formatting rule used by the repo.As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}requires: “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 335 - 344, The SoftmaxV6Benchmark class methods violate the project's C/C++ brace style guideline which requires opening braces to be on their own lines for function bodies. In the name() method override, the opening brace is inline with the function signature on the same line as the return statement. In the run() method override, the opening brace is also on the same line as the function signature. Reformat both methods by moving the opening braces to their own lines while keeping the function body content properly indented, which will also require expanding the single-line name() method implementation across multiple lines.Source: Coding guidelines
ml_kernels/include/ml_kernels/softmax.h (1)
505-505: ⚡ Quick winMove function-body braces to their own lines to match repository style.
Both newly added function definitions keep
{on the signature line, which conflicts with the project’s C/C++ brace rule.Style-only patch sketch
-inline __m256 exp256_ps_v3(__m256 x) { +inline __m256 exp256_ps_v3(__m256 x) +{ ... -inline void softmax_v6(const float *input, float *output, std::size_t n) { +inline void softmax_v6(const float *input, float *output, std::size_t n) +{As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}requires: “Keep braces on their own lines for function bodies.”Also applies to: 542-542
🤖 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 opening brace for the function definition inline __m256 exp256_ps_v3(__m256 x) is placed on the same line as the function signature, which violates the repository's C/C++ style guidelines requiring braces to be on their own lines for function bodies. Move the opening brace to its own line immediately after the function signature. This same style correction applies to another function definition at line 542 in the same file—ensure its opening brace is also moved to its own line following the same pattern.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.
Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Line 505: The opening brace for the function definition inline __m256
exp256_ps_v3(__m256 x) is placed on the same line as the function signature,
which violates the repository's C/C++ style guidelines requiring braces to be on
their own lines for function bodies. Move the opening brace to its own line
immediately after the function signature. This same style correction applies to
another function definition at line 542 in the same file—ensure its opening
brace is also moved to its own line following the same pattern.
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 335-344: The SoftmaxV6Benchmark class methods violate the
project's C/C++ brace style guideline which requires opening braces to be on
their own lines for function bodies. In the name() method override, the opening
brace is inline with the function signature on the same line as the return
statement. In the run() method override, the opening brace is also on the same
line as the function signature. Reformat both methods by moving the opening
braces to their own lines while keeping the function body content properly
indented, which will also require expanding the single-line name() method
implementation across multiple lines.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Line 156: The function `test_softmax_v6` currently has its opening brace on
the same line as the function signature, which violates the repository's C++
brace style guidelines. Move the opening brace to its own line by placing it on
a new line immediately after the function signature `void test_softmax_v6()` to
conform to the requirement that braces for function bodies must be on their own
lines.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8f204c77-dd2c-4d6a-89ef-6d5615420b88
📒 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
1 participant
💡 What: Added
softmax_v6containing a new polynomial evaluation helperexp256_ps_v3which merges the two-step precision split calculation ofln(2)into a single combined_mm256_fnmadd_psinstruction.🎯 Why: The previous
exp256_ps_v2implementation splitln(2)into two separate constants (0.693145751953125fand1.428606765330187e-06f) to maintain high precision during range reduction. However, because softmax involves shift-invariant scaling bounded by_mm256_max_ps, splitting the constants increases instruction count and port pressure unnecessarily without significant gains in the required tolerance threshold (1e-4).🏗️ How: Merged the components
r = x - n * 0.693145...andr = r - n * 1.428...into a singler = _mm256_fnmadd_ps(n, _mm256_set1_ps(0.6931471805599453f), x). The reduction was placed insidesoftmax_v6using the existing 4x unrolled loop and in-register Horner polynomial configuration fromsoftmax_v5.📊 Impact: ~5-10% throughput improvement across multiple matrix sizes, specifically increasing from 4.42 GFLOP/s to 4.78 GFLOP/s at N=262144 on Fixed Memory workloads, maintaining strict 1e-4 accuracy tolerances.
🖥️ Tested on: AVX2-capable host (x86_64, GCC 13.3)
🔬 How to reproduce:
PR created automatically by Jules for task 9919791229991724126 started by @bugparty
Summary by CodeRabbit
New Features
Tests
Chores