⚡ Thunderbolt: max — AVX2 vectorized 8x unroll reduction

.jules/thunderbolt.md

@@ -27,3 +27,10 @@
 **Evidence:** Microbenchmarking showed a 2x speedup (99ms -> 49ms) for max_v3 over max_v2 on L1-hot arrays. End-to-end framework benchmarks showed an 8% throughput increase (4.03 -> 4.36 GFLOP/s) on large fixed-memory allocations (N=6553600).
 
 **Action:** For reductions using instructions with >2 cycle latency (like max_ps or add_ps), default to 8x unrolling over 4x unrolling to fully saturate modern out-of-order execution engines.
+## 2025-02-27 - AVX2 Max Reduction Register Pressure
+
+**Learning:** While `_mm256_max_ps` has a 4-cycle latency, aggressively unrolling 16x to use all 16 YMM registers on AVX2 causes register spilling. This is because the load intrinsic requires temporary registers, leaving none available. An 8-way unroll perfectly covers the 4-cycle latency (given 0.5-cycle throughput) and shifts bottlenecks directly to L1/L2 cache bandwidth constraints without causing spills.
+
+**Evidence:** A 16-way unroll was initially implemented and passed tests, but code review pointed out that it forces the compiler to spill registers to the stack inside the innermost hot loop, defeating the purpose of perfect latency hiding. An 8-way unroll avoids this.
+
+**Action:** When unrolling AVX2 loops to hide latency, target an 8-way unroll (which perfectly matches 4-cycle latency ops with 0.5 cycle throughput) rather than exhausting all 16 YMM registers, ensuring temporary registers remain available for loads.

ml_kernels/include/ml_kernels/max.h

@@ -59,6 +59,67 @@ inline float max_v2(const float *input, std::size_t n) {
     return max_val;
 }
 
+
+// ⚡ Thunderbolt: AVX2 Vectorized Max Reduction (8x unroll)
+// Target: AVX2 (Haswell+)
+// Reason: `_mm256_max_ps` has a 4-cycle latency and 0.5-cycle throughput on most modern Intel uarchs. Simple vector reduction loops benefit from aggressive 8x unrolling to fully utilize all 16 YMM registers. A 16x unroll would cause register spilling because the load intrinsic requires temporary registers. An 8-way unroll perfectly covers the 4-cycle latency and shifts bottlenecks directly to L1/L2 cache bandwidth constraints without causing spills.
+// Expected gain: ~1.5x-2.0x throughput over 4x unroll (max_v2) on large arrays.
+inline float max_v4(const float *input, std::size_t n) {
+    if (n == 0) return 0.0f;
+
+    std::size_t i = 0;
+    __m256 max_v = _mm256_set1_ps(std::numeric_limits<float>::lowest());
+    __m256 m0 = max_v, m1 = max_v, m2 = max_v, m3 = max_v;
+    __m256 m4 = max_v, m5 = max_v, m6 = max_v, m7 = max_v;
+
+    // Unroll 8x for 64 elements per iteration
+    for (; i + 63 < n; i += 64) {
+        m0 = _mm256_max_ps(m0, _mm256_loadu_ps(input + i));
+        m1 = _mm256_max_ps(m1, _mm256_loadu_ps(input + i + 8));
+        m2 = _mm256_max_ps(m2, _mm256_loadu_ps(input + i + 16));
+        m3 = _mm256_max_ps(m3, _mm256_loadu_ps(input + i + 24));
+        m4 = _mm256_max_ps(m4, _mm256_loadu_ps(input + i + 32));
+        m5 = _mm256_max_ps(m5, _mm256_loadu_ps(input + i + 40));
+        m6 = _mm256_max_ps(m6, _mm256_loadu_ps(input + i + 48));
+        m7 = _mm256_max_ps(m7, _mm256_loadu_ps(input + i + 56));
+    }
+
+    // Reduce the 8 vectors into 1
+    m0 = _mm256_max_ps(m0, m4);
+    m1 = _mm256_max_ps(m1, m5);
+    m2 = _mm256_max_ps(m2, m6);
+    m3 = _mm256_max_ps(m3, m7);
+
+    m0 = _mm256_max_ps(m0, m1);
+    m2 = _mm256_max_ps(m2, m3);
+    m0 = _mm256_max_ps(m0, m2);
+
+    // Remainder loop for multiples of 8 elements
+    for (; i + 7 < n; i += 8) {
+        m0 = _mm256_max_ps(m0, _mm256_loadu_ps(input + i));
+    }
+
+    // In-register horizontal reduction
+    __m128 lo = _mm256_castps256_ps128(m0);
+    __m128 hi = _mm256_extractf128_ps(m0, 1);
+    lo = _mm_max_ps(lo, hi);
+
+    __m128 shuf = _mm_shuffle_ps(lo, lo, _MM_SHUFFLE(2, 3, 0, 1));
+    lo = _mm_max_ps(lo, shuf);
+    shuf = _mm_shuffle_ps(lo, lo, _MM_SHUFFLE(1, 0, 3, 2));
+    lo = _mm_max_ps(lo, shuf);
+
+    float max_val = _mm_cvtss_f32(lo);
+
+    // Scalar epilogue
+    for (; i < n; ++i) {
+        if (input[i] > max_val) {
+            max_val = input[i];
+        }
+    }
+    return max_val;
+}
+
 } // namespace ml_kernels
 
 // ⚡ Thunderbolt: AVX2 Vectorized Max Reduction (8x unroll)

ml_kernels/src/kernel_bench.cpp

@@ -518,3 +518,60 @@ class MaxV3Benchmark : public MaxBenchmarkBase {
     std::size_t current_idx_ = 0;
 };
 REGISTER_BENCHMARK(MaxV3Benchmark);
+
+
+class MaxV4Benchmark : public MaxBenchmarkBase {
+public:
+    const char *name() const override { return "max_v4"; }
+
+    void setup(int n) override {
+        size_t bytes_per_iteration = n * sizeof(float);
+        size_t target_pool_bytes = 100ULL * 1024 * 1024;
+        pool_size_ = g_use_pool ? std::max<std::size_t>(1, target_pool_bytes / bytes_per_iteration) : 1;
+
+        inputs_.resize(pool_size_);
+        std::mt19937 rng(12345);
+        std::uniform_real_distribution<float> dist(-4.0f, 4.0f);
+        for (std::size_t i = 0; i < pool_size_; ++i) {
+            inputs_[i].resize(n);
+            for (float &value : inputs_[i]) {
+                value = dist(rng);
+            }
+        }
+
+        result_ref_ = inputs_[0].size() == 0
+            ? 0.0f
+            : *std::max_element(inputs_[0].begin(), inputs_[0].end());
+        result_ = 0.0f;
+        current_idx_ = 0;
+    }
+
+    void run() override {
+        result_ = ml_kernels::max_v4(inputs_[current_idx_].data(), inputs_[current_idx_].size());
+        current_idx_ = (current_idx_ + 1) % pool_size_;
+    }
+
+    bool verify() override {
+        current_idx_ = 0;
+        run();
+        return std::fabs(result_ - result_ref_) <= 1e-6f;
+    }
+
+    void teardown() override {
+        inputs_.clear();
+        result_ = 0.0f;
+        result_ref_ = 0.0f;
+    }
+
+    double flops(int n) const override {
+        return static_cast<double>(n); // 1 comparison per element
+    }
+
+private:
+    std::vector<AlignedBuffer<float>> inputs_;
+    float result_;
+    float result_ref_;
+    std::size_t pool_size_;
+    std::size_t current_idx_ = 0;
+};
+REGISTER_BENCHMARK(MaxV4Benchmark);

ml_kernels/src/test_naive_ops.cpp

@@ -4,7 +4,7 @@
 #include <cmath>
 
 #include "ml_kernels/naive_ops.h"
-#include "ml_kernels/naive_ops.h"
+#include "ml_kernels/max.h"
 #include "ml_kernels/softmax.h"
 
 void test_max_naive() {
@@ -181,7 +181,43 @@ void test_softmax_v5() {
     std::cout << "test_softmax_v5 passed!" << std::endl;
 }
 
+
+void test_max_v4() {
+    std::cout << "Running test_max_v4..." << std::endl;
+    // Happy path (large enough for unrolled loop)
+    {
+        std::vector<float> input(150);
+        for (int i = 0; i < 150; ++i) input[i] = static_cast<float>(i);
+        input[145] = 1000.0f; // max value
+        float result = ml_kernels::max_v4(input.data(), input.size());
+        assert(result == 1000.0f);
+    }
+
+    // Negative values
+    {
+        std::vector<float> input = {-5.0f, -2.0f, -8.0f};
+        float result = ml_kernels::max_v4(input.data(), input.size());
+        assert(result == -2.0f);
+    }
+
+    // Single element
+    {
+        std::vector<float> input = {42.0f};
+        float result = ml_kernels::max_v4(input.data(), input.size());
+        assert(result == 42.0f);
+    }
+
+    // Empty array
+    {
+        float result = ml_kernels::max_v4(nullptr, 0);
+        assert(result == 0.0f);
+    }
+
+    std::cout << "test_max_v4 passed!" << std::endl;
+}
+
 int main() {
+    test_max_v4();
     test_relu_naive();
     test_max_naive();
     test_softmax_v3();