main.cpp

#include <iostream>
#include <cstdlib>
#include <vector>
#include <cmath>
#include <numeric>
#include <algorithm>
#include <tbb/parallel_for.h>
#include <tbb/parallel_reduce.h>
#include <tbb/parallel_scan.h>
#include <tbb/blocked_range.h>
#include "ticktock.h"

#include <atomic>
#include "pod.h"

// 1. fill - 使用 parallel_for 直接并行
template <class T, class Func>
std::vector<T> fill(std::vector<T> &arr, Func const &func) {
    TICK(fill);
    tbb::parallel_for(tbb::blocked_range<size_t>(0, arr.size()), [&](auto const &r) {
        for (size_t i = r.begin(); i < r.end(); i++) {
            arr[i] = func(i);
        }
    });
    TOCK(fill);
    return arr;
}

// 2. saxpy - 使用 parallel_for 直接并行
template <class T>
void saxpy(T a, std::vector<T> &x, std::vector<T> const &y) {
    TICK(saxpy);
    tbb::parallel_for(tbb::blocked_range<size_t>(0, x.size()), [&](auto const &r) {
        for (size_t i = r.begin(); i < r.end(); i++) {
            x[i] = a * x[i] + y[i];
        }
    });
    TOCK(saxpy);
}

// 3. sqrtdot - 使用 parallel_reduce 进行并行规约
template <class T>
T sqrtdot(std::vector<T> const &x, std::vector<T> const &y) {
    TICK(sqrtdot);
    size_t n = std::min(x.size(), y.size());
    T ret = tbb::parallel_reduce(
        tbb::blocked_range<size_t>(0, n), (T)0,
        [&](auto const &r, T local_res) {
            for (size_t i = r.begin(); i < r.end(); i++) {
                local_res += x[i] * y[i];
            }
            return local_res;
        },
        [](T a, T b) { return a + b; }
    );
    ret = std::sqrt(ret);
    TOCK(sqrtdot);
    return ret;
}

// 4. minvalue - 使用 parallel_reduce 找最小值
template <class T>
T minvalue(std::vector<T> const &x) {
    TICK(minvalue);
    T ret = tbb::parallel_reduce(
        tbb::blocked_range<size_t>(0, x.size()), x[0],
        [&](auto const &r, T local_min) {
            for (size_t i = r.begin(); i < r.end(); i++) {
                if (x[i] < local_min) local_min = x[i];
            }
            return local_min;
        },
        [](T a, T b) { return std::min(a, b); }
    );
    TOCK(minvalue);
    return ret;
}

// 5. magicfilter - 使用 原子索引 + 预分配空间
template <class T>
std::vector<T> magicfilter(std::vector<T> const &x, std::vector<T> const &y) {
    TICK(magicfilter);
    size_t n = std::min(x.size(), y.size());
    
    // 1. 预分配一个足够大的临时空间 (最大可能长度是 2n)
    // 使用 pod<T> 可以避免不必要的构造函数开销
    std::vector<pod<T>> tmp(n * 2); 
    std::atomic<size_t> counter{0};

    // 2. 并行写入，每个线程通过 atomic 获取写入位置
    tbb::parallel_for(tbb::blocked_range<size_t>(0, n), [&](auto const &r) {
        // 局部缓冲区，减少对全局原子变量的竞争
        std::vector<T> local;
        local.reserve(r.size()); 
        
        for (size_t i = r.begin(); i < r.end(); i++) {
            if (x[i] > y[i]) {
                local.push_back(x[i]);
            } else if (y[i] > x[i] && y[i] > 0.5f) {
                local.push_back(y[i]);
                local.push_back(x[i] * y[i]);
            }
        }
        
        // 一次性申请位置并拷贝，效率极高
        if (!local.empty()) {
            size_t base = counter.fetch_add(local.size());
            for (size_t j = 0; j < local.size(); j++) {
                tmp[base + j] = local[j];
            }
        }
    });

    // 3. 裁剪回最终大小并转回 std::vector
    size_t final_size = counter.load();
    std::vector<T> res(final_size);
    tbb::parallel_for(tbb::blocked_range<size_t>(0, final_size), [&](auto const &r) {
        for (size_t i = r.begin(); i < r.end(); i++) {
            res[i] = tmp[i];
        }
    });

    TOCK(magicfilter);
    return res;
}

// 6. scanner - 使用 parallel_scan 实现并行前缀和
template <class T>
T scanner(std::vector<T> &x) {
    TICK(scanner);
    T total_sum = tbb::parallel_scan(
        tbb::blocked_range<size_t>(0, x.size()), (T)0,
        [&](auto const &r, T sum, bool is_final) {
            for (size_t i = r.begin(); i < r.end(); i++) {
                sum += x[i];
                if (is_final) x[i] = sum;
            }
            return sum;
        },
        [](T a, T b) { return a + b; }
    );
    TOCK(scanner);
    return total_sum;
}

// 下面是测试逻辑，必须保留！
int main() {
    size_t n = 1 << 26;
    std::vector<float> x(n);
    std::vector<float> y(n);

    fill(x, [](size_t i) { return std::sin(i); });
    fill(y, [](size_t i) { return std::cos(i); });

    saxpy(0.5f, x, y);

    std::cout << sqrtdot(x, y) << std::endl;
    std::cout << minvalue(x) << std::endl;

    auto res = magicfilter(x, y);
    std::cout << res.size() << std::endl;

    scanner(x);
    std::cout << x.back() << std::endl;

    return 0;
}