⚡ Thunderbolt: Softmax — AVX2 8x Unrolling#80
Conversation
💡 What: Implemented `softmax_v6` utilizing an 8x unrolled structure to parallelize compute. 🎯 Why: Max reduction and exp approximations are bounded by instruction latencies (like the 4-cycle max and add) rather than throughput when unrolled only 4x. 8x perfectly saturates these. 🏗️ How: By interleaving 8 independent accumulators and Horner chains for exp across all three kernel phases without spilling YMM. 📊 Impact: ~6% throughput over `softmax_v5` on large workloads (N=1048576, Fixed Memory). 🖥️ Tested on: Standard test runner VM w/ GCC 13. 🔬 How to reproduce: `cd build && DISABLE_CPU_BINDING=1 ./ml_kernels/ml_kernel_bench --filter "softmax_v[56]"` 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
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Benchmark
participant softmax_v6
participant Test
participant softmax_naive
Benchmark->>softmax_v6: Run benchmark on input/output buffers
Test->>softmax_naive: Compute reference output
Test->>softmax_v6: Compute optimized output
softmax_v6-->>Test: Return vectorized output
Test->>Test: Compare elements and output sum
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.
🧹 Nitpick comments (2)
ml_kernels/src/test_naive_ops.cpp (1)
196-205: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd elements to test the scalar tail.
The input array currently has exactly 96 elements (a multiple of 8). This completely bypasses the scalar tail loops (
i < n) in thesoftmax_v6implementation. Adding a few extra elements will ensure that these boundary condition fallback paths are properly exercised.🧪 Proposed fix to increase coverage
1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, - 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f + 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, + // Add elements to trigger the scalar tail (not a multiple of 8) + 2.0f, -1.0f, 0.5f };🤖 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 196 - 205, Extend the input array in the test around softmax_v6 beyond its current 96 elements by adding a few additional values, ensuring the test length is not divisible by the vector width of 8 and exercises the scalar tail loop while preserving the existing test setup..jules/thunderbolt.md (1)
36-36: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFix typographical error.
The phrase "map-reduce map math kernels" contains a redundant "map". Consider simplifying it to "map-reduce math kernels" for clarity.
🔤 Proposed fix
-**Action:** Default to 8x unrolling across all phases for multi-pass map-reduce map math kernels (like softmax) when using `_mm256_max_ps` and `_mm256_add_ps` to perfectly match their 4-cycle latency and hide FMA chains, without risking register spilling since YMM allows 16 registers. +**Action:** Default to 8x unrolling across all phases for multi-pass map-reduce math kernels (like softmax) when using `_mm256_max_ps` and `_mm256_add_ps` to perfectly match their 4-cycle latency and hide FMA chains, without risking register spilling since YMM allows 16 registers.🤖 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 at line 36, Correct the wording in the multi-pass kernel guidance by changing “map-reduce map math kernels” to “map-reduce math kernels,” leaving the rest of the unrolling and instruction details unchanged.
🤖 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 @.jules/thunderbolt.md:
- Line 36: Correct the wording in the multi-pass kernel guidance by changing
“map-reduce map math kernels” to “map-reduce math kernels,” leaving the rest of
the unrolling and instruction details unchanged.
In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 196-205: Extend the input array in the test around softmax_v6
beyond its current 96 elements by adding a few additional values, ensuring the
test length is not divisible by the vector width of 8 and exercises the scalar
tail loop while preserving the existing test setup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8be6b1e7-82a6-46b1-88ea-318b9b45f729
📒 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: Implemented
softmax_v6utilizing an 8x unrolled structure to parallelize compute.🎯 Why: Max reduction and exp approximations are bounded by instruction latencies (like the 4-cycle max and add) rather than throughput when unrolled only 4x. 8x perfectly saturates these.
🏗️ How: By interleaving 8 independent accumulators and Horner chains for exp across all three kernel phases without spilling YMM.
📊 Impact: ~6% throughput over
softmax_v5on large workloads (N=1048576, Fixed Memory).🖥️ Tested on: Standard test runner VM w/ GCC 13.
🔬 How to reproduce:
cd build && DISABLE_CPU_BINDING=1 ./ml_kernels/ml_kernel_bench --filter "softmax_v[56]"PR created automatically by Jules for task 17337070163629464084 started by @bugparty
Summary by CodeRabbit
New Features
Bug Fixes
Documentation