⚡ Thunderbolt: Softmax — Combine FMA for exp range reduction#50
⚡ Thunderbolt: Softmax — Combine FMA for exp range reduction#50bugparty wants to merge 1 commit into
Conversation
Combined `ln(2)` splitting into a single FMA instruction in AVX2 Softmax `exp` range reduction to improve throughput by ~6% while retaining necessary precision. 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. |
📝 WalkthroughWalkthroughThis PR introduces a new softmax kernel variant (softmax_v6) with an optimized AVX2 exp approximation (exp256_ps_v3) that uses single-FMA range reduction to improve instruction count and latency. The implementation includes vectorized max/sum reduction, 32-wide and 8-wide processing loops, test validation against the naive baseline, and performance benchmarking infrastructure. ChangesSoftmax V6 Kernel and Optimization
Sequence DiagramsequenceDiagram
participant Caller
participant softmax_v6
participant ReduceMax as Reduce Max
participant exp256_ps_v3
participant ReduceSum as Reduce Sum
participant Output
Caller->>softmax_v6: call softmax_v6(input, output, n)
softmax_v6->>ReduceMax: Vectorized loop (32-wide + 8-wide + scalar tail)
ReduceMax->>softmax_v6: max value
softmax_v6->>exp256_ps_v3: For each element: exp(x - max) vectorized
exp256_ps_v3->>exp256_ps_v3: Range reduction, polynomial, scaling
exp256_ps_v3->>Output: Write unnormalized exp(x - max)
softmax_v6->>ReduceSum: Vectorized sum reduction (32-wide + 8-wide + scalar)
ReduceSum->>softmax_v6: sum value
softmax_v6->>Output: Normalize: y[i] *= 1.0/sum (vectorized + scalar)
Output->>Caller: Softmax output
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)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Infer (1.2.0)ml_kernels/src/test_naive_ops.cppml_kernels/src/test_naive_ops.cpp:6:10: fatal error: 'ml_kernels/naive_ops.h' file not found ... [truncated 1112 characters] ... l/lib/clang/18/include" ml_kernels/src/kernel_bench.cppml_kernels/src/kernel_bench.cpp:14:10: fatal error: 'aligned_buffer.h' file not found ... [truncated 1089 characters] ... all/lib/clang/18/include" 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)
184-184: ⚡ Quick winUse newline opening brace for
test_softmax_v6function body.Please move the opening
{to its own line for this newly added function.Proposed formatting patch
-void test_softmax_v6() { +void test_softmax_v6() +{As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}must “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 184, The function declaration for test_softmax_v6 places the opening brace on the same line; update the function definition for test_softmax_v6 so the `{` is moved to its own line (i.e., use a newline before the function body brace) to comply with the project's brace style for function bodies.Source: Coding guidelines
ml_kernels/src/kernel_bench.cpp (1)
335-343: ⚡ Quick winApply function brace placement rule in
SoftmaxV6Benchmarkmethods.The new method definitions use same-line
{. Please switch to next-line opening braces to match project C/C++ style.Proposed formatting patch
- const char *name() const override { return "softmax_v6"; } + const char *name() const override + { + return "softmax_v6"; + } @@ - void run() override { + 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_; }As per coding guidelines,
**/*.{c,cpp,cc,h,hpp}must “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 - 343, The brace placement for SoftmaxV6Benchmark's method definitions violates project style: move the opening braces for name() and run() to their own lines; update the class SoftmaxV6Benchmark so name() const override and run() override have the function body opening brace on the next line (i.e., change "const char *name() const override { ... }" and "void run() override { ... }" to use newline before the "{"), preserving the existing bodies and current_idx_/pool_size_ logic.Source: Coding guidelines
ml_kernels/include/ml_kernels/softmax.h (1)
504-530: ⚡ Quick winMove function opening braces to their own lines in new softmax_v6 additions.
exp256_ps_v3andsoftmax_v6currently keep{on the declaration line. Please align these new function bodies with the repo’s brace rule.Proposed formatting patch
-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}must “Keep braces on their own lines for function bodies.”Also applies to: 536-632
🤖 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 - 530, The function definitions exp256_ps_v3 and softmax_v6 violate the repo brace style by keeping the opening brace on the same line as the declaration; update each function declaration so the `{` is moved to its own line (i.e., change "inline __m256 exp256_ps_v3(__m256 x) {" to have the brace on the following line, and do the same for softmax_v6) to comply with the "Keep braces on their own lines for function bodies" rule; ensure indentation for the brace and body matches surrounding code style.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`:
- Around line 504-530: The function definitions exp256_ps_v3 and softmax_v6
violate the repo brace style by keeping the opening brace on the same line as
the declaration; update each function declaration so the `{` is moved to its own
line (i.e., change "inline __m256 exp256_ps_v3(__m256 x) {" to have the brace on
the following line, and do the same for softmax_v6) to comply with the "Keep
braces on their own lines for function bodies" rule; ensure indentation for the
brace and body matches surrounding code style.
In `@ml_kernels/src/kernel_bench.cpp`:
- Around line 335-343: The brace placement for SoftmaxV6Benchmark's method
definitions violates project style: move the opening braces for name() and run()
to their own lines; update the class SoftmaxV6Benchmark so name() const override
and run() override have the function body opening brace on the next line (i.e.,
change "const char *name() const override { ... }" and "void run() override {
... }" to use newline before the "{"), preserving the existing bodies and
current_idx_/pool_size_ logic.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Line 184: The function declaration for test_softmax_v6 places the opening
brace on the same line; update the function definition for test_softmax_v6 so
the `{` is moved to its own line (i.e., use a newline before the function body
brace) to comply with the project's brace style for function bodies.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 95645985-e70d-4f6e-b021-8ca69687bcd8
📒 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:
Combined the two
_mm256_fnmadd_psinstructions handling range reduction (r = x - n * ln(2)) in the AVX2expfunction approximation into a single_mm256_fnmadd_psusing the combinedln(2)float constant (0.6931471805599453f). This is implemented insoftmax_v6.🎯 Why:$x \leq 0$ ), this extreme precision splitting is unneeded and unnecessarily lengthens the critical FMA execution chain latency. Removing it improves overall throughput without sacrificing numerical tolerances (1e-4).
The previous implementation mathematically split
ln(2)into two FMA operations to retain 53-bit precision and avoid catastrophic cancellation. For typical ML softmax kernels where the inputs are already pre-shifted bymax_val(🏗️ How:
exp256_ps_v2toexp256_ps_v3.__m256 r = _mm256_fnmadd_ps(n, _mm256_set1_ps(0.6931471805599453f), x);softmax_v6insoftmax.hreferencing the new exp logic.kernel_bench.cppto register the new variant andtest_naive_ops.cppwith a new, full unroll boundary correctness test suite forsoftmax_v6.📊 Impact:
Measured throughput improvements on N=1048576 arrays via microbenchmarks (Fixed Memory):
softmax_v5): 4.21 GFLOP/ssoftmax_v6): 4.46 GFLOP/s🖥️ Tested on:
x86-64AVX2 target, compiled with GCC 13.🔬 How to reproduce:
PR created automatically by Jules for task 17176520017703929985 started by @bugparty
Summary by CodeRabbit
New Features
Tests
Chores