Modified AVX2 Quicksort with three element rule versus Standard AVX2 Quicksort

#include <iostream>
#include <vector>
#include <algorithm>
#include <random>
#include <chrono>
#include <immintrin.h> // Required for AVX2 Intrinsics

// Modified AVX2 Quicksort with three element rule versus Standard AVX2 Quicksort
// Contact Email [email protected]
// Created by I.A

const int NUM_ELEMENTS = 100000;
const int NUM_TRIALS = 10000;

// ---------------------------------------------------------
// 1. Core AVX2 Partition Logic
// ---------------------------------------------------------
int avx2_partition(std::vector<int>& arr, int low, int high) {
    int pivot_val = arr[low];
    __m256i pivot = _mm256_set1_epi32(pivot_val);
    int size = high - low;

    alignas(32) std::vector<int> left_buf(size);
    alignas(32) std::vector<int> right_buf(size);

    int left_count = 0;
    int right_count = 0;

    int i = low + 1;
    for (; i <= high - 7; i += 8) {
        __m256i data = _mm256_loadu_si256((__m256i*)&arr[i]);
        __m256i cmp = _mm256_cmpgt_epi32(data, pivot);

        alignas(32) int data_extracted[8];
        alignas(32) int mask_extracted[8];
        _mm256_storeu_si256((__m256i*)data_extracted, data);
        _mm256_storeu_si256((__m256i*)mask_extracted, cmp);

        for (int k = 0; k < 8; ++k) {
            if (mask_extracted[k] == 0) {
                left_buf[left_count++] = data_extracted[k];
            } else {
                right_buf[right_count++] = data_extracted[k];
            }
        }
    }

    // Scalar fallback for remaining elements
    for (; i <= high; ++i) {
        if (arr[i] <= pivot_val) {
            left_buf[left_count++] = arr[i];
        } else {
            right_buf[right_count++] = arr[i];
        }
    }

    std::copy(left_buf.begin(), left_buf.begin() + left_count, arr.begin() + low);

    int pivot_idx = low + left_count;
    arr[pivot_idx] = pivot_val;

    std::copy(right_buf.begin(), right_buf.begin() + right_count, arr.begin() + pivot_idx + 1);

    return pivot_idx;
}

// ---------------------------------------------------------
// 2. Modified AVX2 Quicksort (With Three-Element Rule)
// ---------------------------------------------------------
void apply_three_element_rule(std::vector<int>& arr, int low, int high) {
    if (high - low == 2) {
        int pivot = arr[low];
        int third_num = arr[high];

        if (third_num <= pivot) {
            int temp_mid = arr[low + 1];
            arr[high] = temp_mid;
            arr[low + 1] = pivot;
            arr[low] = third_num;
        } else {
            int temp_mid = arr[low + 1];
            arr[low] = temp_mid;
            arr[low + 1] = pivot;
            arr[high] = third_num;
        }
    }
}

void quickSortModified(std::vector<int>& arr, int low, int high) {
    if (low < high) {
        if (high - low == 2) {
            apply_three_element_rule(arr, low, high);
        }

        int p;
        if (high - low >= 16) {
            p = avx2_partition(arr, low, high);
            quickSortModified(arr, low, p - 1);
            quickSortModified(arr, p + 1, high);
        } else {
            int pivot = arr[low];
            int i = low - 1;
            int j = high + 1;
            while (true) {
                do { i++; } while (arr[i] < pivot);
                do { j--; } while (arr[j] > pivot);
                if (i >= j) {
                    p = j;
                    break;
                }
                std::swap(arr[i], arr[j]);
            }
            quickSortModified(arr, low, p);
            quickSortModified(arr, p + 1, high);
        }
    }
}

// ---------------------------------------------------------
// 3. Standard AVX2 Quicksort (Without Modifications)
// ---------------------------------------------------------
void quickSortStandard(std::vector<int>& arr, int low, int high) {
    if (low < high) {
        int p;
        if (high - low >= 16) {
            p = avx2_partition(arr, low, high);
            quickSortStandard(arr, low, p - 1);
            quickSortStandard(arr, p + 1, high);
        } else {
            int pivot = arr[low];
            int i = low - 1;
            int j = high + 1;
            while (true) {
                do { i++; } while (arr[i] < pivot);
                do { j--; } while (arr[j] > pivot);
                if (i >= j) {
                    p = j;
                    break;
                }
                std::swap(arr[i], arr[j]);
            }
            quickSortStandard(arr, low, p);
            quickSortStandard(arr, p + 1, high);
        }
    }
}


// ---------------------------------------------------------
// 4. Benchmark Runner
// ---------------------------------------------------------
int main() {
    std::cout << "Starting AVX2 Benchmark (Standard vs Modified)..." << std::endl;
    std::cout << "Elements per array: " << NUM_ELEMENTS << std::endl;
    std::cout << "Total Trials: " << NUM_TRIALS << "\n" << std::endl;

    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> distrib(1, 10000000);

    std::chrono::duration<double> total_duration_mod(0);
    std::chrono::duration<double> total_duration_std(0);

    for (int t = 1; t <= NUM_TRIALS; ++t) {
        // Generate base data once per trial
        std::vector<int> base_data(NUM_ELEMENTS);
        for (int i = 0; i < NUM_ELEMENTS; ++i) {
            base_data[i] = distrib(gen);
        }

        // Copy data so both algorithms sort the exact same numbers
        std::vector<int> data_mod = base_data;
        std::vector<int> data_std = base_data;

        if (t % 2 == 0) {
            // --- EVEN TRIALS: STANDARD RUNS FIRST ---

            // 1. Standard

            auto start_std = std::chrono::high_resolution_clock::now();
            quickSortStandard(data_std, 0, NUM_ELEMENTS - 1);
            auto end_std = std::chrono::high_resolution_clock::now();
            total_duration_std += (end_std - start_std);

            // 2. Modified

            auto start_mod = std::chrono::high_resolution_clock::now();
            quickSortModified(data_mod, 0, NUM_ELEMENTS - 1);
            auto end_mod = std::chrono::high_resolution_clock::now();
            total_duration_mod += (end_mod - start_mod);

        } else {
            // --- ODD TRIALS: MODIFIED RUNS FIRST ---

            // 1. Modified

            auto start_mod = std::chrono::high_resolution_clock::now();
            quickSortModified(data_mod, 0, NUM_ELEMENTS - 1);
            auto end_mod = std::chrono::high_resolution_clock::now();
            total_duration_mod += (end_mod - start_mod);

            // 2. Standard

            auto start_std = std::chrono::high_resolution_clock::now();
            quickSortStandard(data_std, 0, NUM_ELEMENTS - 1);
            auto end_std = std::chrono::high_resolution_clock::now();
            total_duration_std += (end_std - start_std);
        }

        // Verification Checks (Kept identical)
        if (!std::is_sorted(data_mod.begin(), data_mod.end())) {
            std::cerr << "Error: Modified Array not sorted on trial " << t << "!" << std::endl;
            return 1;
        }
        if (!std::is_sorted(data_std.begin(), data_std.end())) {
            std::cerr << "Error: Standard Array not sorted on trial " << t << "!" << std::endl;
            return 1;
        }

        std::cout << "Completed Trial " << t << "/" << NUM_TRIALS << "\r" << std::flush;
    }

    std::cout << "\n\n--- Benchmark Complete ---" << std::endl;

    double avg_mod = total_duration_mod.count() / NUM_TRIALS;
    double avg_std = total_duration_std.count() / NUM_TRIALS;

    std::cout << "**Modified Version (With 3-Element Rule):**" << std::endl;
    std::cout << "Total Time:   " << total_duration_mod.count() << " seconds" << std::endl;
    std::cout << "Average Time: " << avg_mod << " seconds\n" << std::endl;

    std::cout << "**Standard Version (Without Modifications):**" << std::endl;
    std::cout << "Total Time:   " << total_duration_std.count() << " seconds" << std::endl;
    std::cout << "Average Time: " << avg_std << " seconds\n" << std::endl;

    // Calculate percent difference relative to the standard baseline
    double percent_diff = ((avg_std - avg_mod) / avg_std) * 100.0;

    std::cout << "--------------------------------------" << std::endl;
    if (percent_diff > 0) {
        std::cout << "Result: Modified version is " << percent_diff << "% FASTER than Standard." << std::endl;
    } else if (percent_diff < 0) {
        std::cout << "Result: Modified version is " << -percent_diff << "% SLOWER than Standard." << std::endl;
    } else {
        std::cout << "Result: Both versions performed identically." << std::endl;
    }
    std::cout << "--------------------------------------" << std::endl;

    return 0;
}