⚡ Thunderbolt: Softmax — 8x unroll via single-FMA exp256 optimization#68
⚡ Thunderbolt: Softmax — 8x unroll via single-FMA exp256 optimization#68bugparty 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+FMA exponential helper ChangesSoftmax v6 kernel and validation
Estimated code review effort: 3 (Moderate) | ~25 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: 1
🧹 Nitpick comments (4)
ml_kernels/src/kernel_bench.cpp (1)
337-342: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the new benchmark methods on brace-next-line style.
name()andrun()use inline function-body braces on the same line. Please switch them to brace-next-line form to match the repository's C/C++ 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 - 342, The new benchmark methods in the softmax_v6 benchmark class use same-line function-body braces, which violates the repository’s brace-next-line style. Update the inline overrides for name() and run() in kernel_bench.cpp so their opening and closing braces are on their own lines, matching the existing C/C++ formatting used by the benchmark classes.Source: Coding guidelines
ml_kernels/src/test_naive_ops.cpp (2)
183-183: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the new test functions on brace-next-line style.
test_softmax_v6()andmain()keep the opening brace on the signature line. Please move it to its own line for consistency with the repository's C/C++ style. As per coding guidelines,Keep braces on their own lines for function bodies.Also applies to: 216-216
🤖 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 183, The new test functions are using the opening brace on the signature line instead of the repository’s brace-next-line style. Update the function definitions for test_softmax_v6() and main() so their opening braces are placed on their own lines, matching the existing C/C++ formatting convention used in this test file.Source: Coding guidelines
185-210: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise the scalar tail too.
This 72-element case only hits the 64-wide unrolled body plus one 8-float vector remainder. The scalar tail in
softmax_v6never executes, so regressions in that path would still slip past this test. Add a second case with a non-multiple-of-8 size such as 69 or 71.🤖 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 185 - 210, The softmax comparison test only covers the 64-wide unrolled path plus an 8-float vector remainder, so the scalar tail in softmax_v6 is still untested. Update the existing test in test_naive_ops.cpp by adding a second input case with a size that is not a multiple of 8, using softmax_naive and softmax_v6 again to compare outputs and verify the sum; keep the new case alongside the current softmax_v6 test so both the vector remainder and scalar tail paths are exercised.ml_kernels/include/ml_kernels/softmax.h (1)
505-505: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMatch the repository brace style.
Both new function definitions keep the opening brace on the signature line. Move it to its own line so this header stays consistent with the C/C++ style used elsewhere in changed code. As per coding guidelines,
Keep braces on their own lines for function bodies.Also applies to: 541-541
🤖 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 new function definitions in softmax.h use a brace style that differs from the repository convention: move the opening brace for exp256_ps_v3 and the other affected AVX function to its own line. Keep the function signatures unchanged, and update the body formatting so these definitions match the existing C/C++ brace style used throughout the header.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 @.jules/thunderbolt.md:
- Around line 31-34: The entry under the 2024-05-24 heading in the thunderbolt
note is out of chronological order relative to the surrounding 2024-10-24 to
2024-10-26 entries. Update the note by either correcting the date if it is a
typo, or moving this "Single-FMA exp256 optimization" section to the proper
chronological position so the timeline reads forward; use the date heading and
the entry text as the unique anchors.
---
Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Line 505: The new function definitions in softmax.h use a brace style that
differs from the repository convention: move the opening brace for exp256_ps_v3
and the other affected AVX function to its own line. Keep the function
signatures unchanged, and update the body formatting so these definitions match
the existing C/C++ brace style used throughout the header.
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 337-342: The new benchmark methods in the softmax_v6 benchmark
class use same-line function-body braces, which violates the repository’s
brace-next-line style. Update the inline overrides for name() and run() in
kernel_bench.cpp so their opening and closing braces are on their own lines,
matching the existing C/C++ formatting used by the benchmark classes.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Line 183: The new test functions are using the opening brace on the signature
line instead of the repository’s brace-next-line style. Update the function
definitions for test_softmax_v6() and main() so their opening braces are placed
on their own lines, matching the existing C/C++ formatting convention used in
this test file.
- Around line 185-210: The softmax comparison test only covers the 64-wide
unrolled path plus an 8-float vector remainder, so the scalar tail in softmax_v6
is still untested. Update the existing test in test_naive_ops.cpp by adding a
second input case with a size that is not a multiple of 8, using softmax_naive
and softmax_v6 again to compare outputs and verify the sum; keep the new case
alongside the current softmax_v6 test so both the vector remainder and scalar
tail paths are exercised.
🪄 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: 9787adea-c29f-4539-b1e9-ab34c6707c5c
⛔ Files ignored due to path filters (1)
v6.logis excluded by!**/*.log
📒 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
| ## 2024-05-24 - Single-FMA exp256 optimization | ||
| **Learning:** In transcendental AVX2 SIMD approximations like exp256, combining constants for `r = x - n * ln(2)` into a single FMA instruction (rather than splitting `ln(2)` for exact precision) reduces register pressure and instruction count. This unlocks the ability to aggressively unroll the softmax loop 8x (processing 64 elements at once) without spilling registers, perfectly hiding latency and saturating the L1 cache bandwidth, all while staying within acceptable numerical ML tolerances. | ||
| **Evidence:** Implementation of `softmax_v6` showing significant reduction in instruction count and higher throughput compared to `softmax_v5`. | ||
| **Action:** When working on approximations where exact precision isn't strictly necessary, look for opportunities to combine mathematical constants to save registers and instructions, allowing for higher unrolling factors. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the note's date/order.
This entry is dated 2024-05-24 but is appended after the 2024-10-24 to 2024-10-26 entries, so the chronology now reads backward. If the date is intentional, move the section; otherwise fix the typo.
🤖 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 @.jules/thunderbolt.md around lines 31 - 34, The entry under the 2024-05-24
heading in the thunderbolt note is out of chronological order relative to the
surrounding 2024-10-24 to 2024-10-26 entries. Update the note by either
correcting the date if it is a typo, or moving this "Single-FMA exp256
optimization" section to the proper chronological position so the timeline reads
forward; use the date heading and the entry text as the unique anchors.
1 participant
💡 What: Added
softmax_v6optimized for AVX2, which unrolls the main computation loop 8x to process 64 elements at once. To avoid register spilling during this heavy unroll, the transcendental approximationexp256_ps_v3was modified to use a single FMA forr = x - n * ln(2)by combining constants.🎯 Why: The previous
softmax_v5implementation was bottlenecked by FMA instruction latency and register pressure. Reducing the instructions required for theexpapproximation frees up enough registers to allow aggressive 8x unrolling, fully saturating execution ports and hiding latency.🏗️ How:
exp256_psto computer = _mm256_fnmadd_ps(n, ln2_constant, x).max,exp/sum, andnormalizeloops insoftmax_v6by 8 to maintain 8 independent accumulators.test_naive_ops.cpp.📊 Impact:
On N=1,048,576 (Fixed Memory mode):
softmax_v5: 4.17 GFLOP/ssoftmax_v6: 4.57 GFLOP/s (~9.6% speedup)🖥️ Tested on: Intel Xeon @ 2.30GHz (AVX2-capable), Ubuntu Linux, GCC 13.3.0.
🔬 How to reproduce:
PR created automatically by Jules for task 11879846324715467560 started by @bugparty
Summary by CodeRabbit
New Features
Tests