⚡ Thunderbolt: max_v4 — 16x unrolled AVX2 max reduction#51
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. |
📝 WalkthroughWalkthroughThis PR introduces Changesmax_v4 AVX2 Kernel Implementation
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 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 (1)
ml_kernels/include/ml_kernels/max.h (1)
128-133: ⚡ Quick winClarify the conditions under which max_v4 achieves >2x speedup.
The comment states "Expected gain: >2x throughput vs max_v3 on large arrays," but the documentation in
thunderbolt.md(line 33) says the gains occur "on sizes fitting in L2/L3 cache" while also noting the kernel "remains bottlenecked by DRAM bandwidth on very large arrays." The example N=104,857,600 (~400MB) cited in the evidence far exceeds typical L2/L3 cache sizes (8-64MB).Consider revising to clarify the actual conditions, perhaps: "Expected gain: ~2x throughput vs max_v3 when not bottlenecked by DRAM bandwidth" or provide more specific guidance on array size ranges.
🤖 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/max.h` around lines 128 - 133, Update the header comment for the AVX2 kernel to clarify the conditions for the "Expected gain: >2x throughput vs max_v3" claim: mention that the >2x improvement for max_v4 compared to max_v3 applies when the working set fits in CPU caches (L2/L3) and the kernel is not DRAM-bandwidth bound, and either change the phrasing to something like "Expected gain: ~2x throughput vs max_v3 when not bottlenecked by DRAM (working set fits in L2/L3 cache)" or add an explicit size range note and cross-reference thunderbolt.md; edit the comment block containing the Thunderbolt/AVX2 description (the one describing max_v4) so it references max_v3 and thunderbolt.md for detailed benchmarks.
🤖 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/max.h`:
- Around line 128-133: Update the header comment for the AVX2 kernel to clarify
the conditions for the "Expected gain: >2x throughput vs max_v3" claim: mention
that the >2x improvement for max_v4 compared to max_v3 applies when the working
set fits in CPU caches (L2/L3) and the kernel is not DRAM-bandwidth bound, and
either change the phrasing to something like "Expected gain: ~2x throughput vs
max_v3 when not bottlenecked by DRAM (working set fits in L2/L3 cache)" or add
an explicit size range note and cross-reference thunderbolt.md; edit the comment
block containing the Thunderbolt/AVX2 description (the one describing max_v4) so
it references max_v3 and thunderbolt.md for detailed benchmarks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 681cbfaa-36cf-4af0-a215-cefac497a499
📒 Files selected for processing (4)
.jules/thunderbolt.mdml_kernels/include/ml_kernels/max.hml_kernels/src/kernel_bench.cppml_kernels/src/test_naive_ops.cpp
1 participant
💡 What: Added
max_v4, an AVX2-vectorized max reduction kernel explicitly unrolled 16x to use all 16 available YMM registers in x86-64.🎯 Why: The
_mm256_max_psinstruction has a 4-cycle latency and 0.5-cycle throughput. Whilemax_v3was unrolled 8x, it did not fully utilize the architectural registers or completely saturate the execution ports for a simple reduction operation. By unrolling 16x, we keep 16 independent dependency chains alive, which perfectly hides the 4-cycle instruction latency and transitions the performance bottleneck entirely to L1/L2 cache memory bandwidth.🏗️ How:
📊 Impact:
🖥️ Tested on: Built and tested with GCC (
make ml_kernel_testandml_kernel_bench).🔬 How to reproduce:
DISABLE_CPU_BINDING=1 ./build/ml_kernels/ml_kernel_bench --filter 'max' --sizes 100000000PR created automatically by Jules for task 18357408306933147473 started by @bugparty
Summary by CodeRabbit
New Features
Tests