Reduce 归约求和

#include <cuda.h>
#include <cuda_runtime.h>
#include <time.h>
#include <algorithm>
#include <iostream>
#include <cuda_fp16.h>
using namespace std;

#define PackSize 4
#define kWarpSize 32
#define N 32 * 1024 * 1024

// 【PackSize：一个 pack 装几个 float】4 个 float = 16 字节，是一次访存能拿的最大宽度。
#define PackSize 4
#define kWarpSize 32
#define N 32 * 1024 * 1024

constexpr int BLOCK_SIZE = 256;
constexpr int kBlockSize = 256;
// 【kNumWaves：要发几波】1 表示"只要发一波就够"（一波刚好把 SM 填满）。
constexpr int kNumWaves = 1;

// 【GetNumBlocks：按硬件规格算 block 数（OneFlow 的做法）】
int64_t GetNumBlocks(int64_t n) {
    // 【第 1 步：拿到当前设备号】cudaGetDevice 写进 dev。
    int dev;
    {
        cudaError_t err = cudaGetDevice(&dev);
        if (err != cudaSuccess) { return err; }
    }
    // 【第 2 步：SM 个数】cudaDevAttrMultiProcessorCount。A100 是 108。
    int sm_count;
    {
        cudaError_t err = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
        if (err != cudaSuccess) { return err; }
    }
    // 【第 3 步：每个 SM 最多常驻多少线程】cudaDevAttrMaxThreadsPerMultiProcessor。
   int tpm;
   {
        cudaError_t err = cudaDeviceGetAttribute(&tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev);
        if (err != cudaSuccess) { return err; }
   }

    int64_t num_blocks = std::max<int64_t>(1, std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
                                                   sm_count * tpm / kBlockSize * kNumWaves));
    return num_blocks;
}


//alignas(sizeof(T) * pack_size) —— 整段的开关
//含义：这个类型的每一个对象的起始地址，都必须是 N 的倍数。这里 N = sizeof(float) × 4 = 16。它把结构体的对齐要求抬到 16 字节，没有这句 alignas，结构体可能只有 4 字节对齐，"一次读 16 字节"就没法保证了。
template<typename T, int pack_size>
struct alignas(sizeof(T) * pack_size) Packed {

    __device__ Packed(T val){
        #pragma unroll
        for(int i = 0; i < pack_size; i++){
            elem[i] = val;
        }
      }

    __device__ Packed() {
    // do nothing
      }
    // 【匿名 union】union 的成员共用同一块内存，每个成员的偏移都是 0。
    //   这里只有 elem 一个成员，所以等价于直接写 `T elem[pack_size];`。
    //   匿名 → 没有名字，elem 直接就是 Packed 的成员，写 p.elem 即可。
   union {
     T elem[pack_size];
   };

     __device__ void operator+=(Packed<T, pack_size> packA){
        #pragma unroll
        for(int i = 0; i < pack_size; i++){
            elem[i] += packA.elem[i];
        }
      }
};

// 【PackReduce：把一个 pack 压成 1 个 T】就是把这 4 个 lane 加起来。
//   注意这一步是**串行的 3 次加法**（一条链），但那只是收尾一次，
//   前面几百次累加都是 4 路并行的。
template<typename T, int pack_size>
__device__ T PackReduce(Packed<T, pack_size> pack){
    T res = 0.0;
    #pragma unroll
    for(int i = 0; i < pack_size; i++){
        res += pack.elem[i];
    }
    return res;
}

// 【warp 内归约：这一版换成了循环写法】与 v7 的 5 个 if 相比：
//   - 没有模板参数、没有编译期门限，delta 从 16 一路除到 1，5 步恒定；
//   - 每次都用 0xffffffff 全掩码（32 个 lane 全参与）。
//   代价是"第二级只归约 8 个值"时也要走满 5 步 —— 好在 lane 8~31 手里是 0，
//   多走的那两步只是加了 0，不影响结果（见调用点那句三元表达式的注释）。
template<typename T>
__device__ T warpReduceSum(T val){
    for(int lane_mask = 16; lane_mask > 0; lane_mask /=2){
        val += __shfl_down_sync(0xffffffff, val, lane_mask);
    }
    return val;
}

__global__ void reduce_v8(float *g_idata,float *g_odata, unsigned int n){

     // 【本线程在网格里的全局编号
    unsigned int i = blockDim.x * blockIdx.x + threadIdx.x;
    
    Packed<float, PackSize> sum_pack(0.0);

    

    const auto* pack_ptr = reinterpret_cast<const Packed<float, PackSize>*>(g_idata);

        // 【网格步长循环】linear_index 是"pack 的下标"，上界是 n/PackSize（pack 的个数）。
    //   每次跨 blockDim.x*gridDim.x 个 pack，也就是"整个网格的线程数"。
    //   于是：每个线程读一个 pack，这和"相邻线程读相邻 pack"是同一件事 ——
    //   一个 warp 一次读 32 个 pack = 512 字节连续内存 = 4 条 128 字节的事务，完全合并。
    for(int32_t linear_index = i; linear_index < n / PackSize; linear_index+=blockDim.x * gridDim.x){

        Packed<float, PackSize> g_idata_load = pack_ptr[linear_index];

        sum_pack += g_idata_load;
    }

    float PackReduceVal = PackReduce<float, PackSize>(sum_pack);

    static __shared__ float warpLevelSums[kWarpSize];
    const int laneId = threadIdx.x % kWarpSize;
    const int warpId = threadIdx.x / kWarpSize;


    float sum = warpReduceSum<float>(PackReduceVal);

    __syncthreads();

    if(laneId == 0 )warpLevelSums[warpId] = sum;

    __syncthreads();

    sum = (threadIdx.x < blockDim.x / kWarpSize) ? warpLevelSums[laneId] : 0;

    if (warpId == 0) sum = warpReduceSum<float>(sum);

    if (threadIdx.x == 0) g_odata[blockIdx.x] = sum;
}

// ============================================================================
// check —— 精确相等比较（关于它为什么脆，见文件头那段）
// ============================================================================
bool check(float *out,float *res,int n){
    for(int i=0;i<n;i++){
        if(out[i]!=res[i])
            return false;
    }
    return true;
}

// ============================================================================
// main —— 两次启动：阶段1 出 block_num 个部分和，阶段2 收成 1 个
// ============================================================================
int main(){
    float *a=(float *)malloc(N*sizeof(float));
    float *d_a;
    cudaMalloc((void **)&d_a,N*sizeof(float));


    const int64_t block_num = GetNumBlocks(N);
    printf("Block num is: %ld \n", block_num);

    // 【out 只有 1 个元素】因为最后要的是总和
    float *out=(float *)malloc(sizeof(float));
    float *g_odata;
    cudaMalloc((void **)&g_odata,block_num*sizeof(float));
    
    // 【给阶段2 用的单元素输出缓冲】
    float *g_final_data;
    cudaMalloc((void **)&g_final_data,1*sizeof(float));

    // 【输入 a[i] = i】
    for(int i=0;i<N;i++){
        a[i]=i;
    }

    // 【标准答案：等差数列求和】N*(首项+末项)/2。这里的算式假设 a[i] 就是 i，
    float *res=(float *)malloc(sizeof(float));
    res[0] = N * (a[0] + a[N - 1]) / 2;


    cudaMemcpy(d_a,a,N*sizeof(float),cudaMemcpyHostToDevice);;

    // 【阶段1：864 个 block，每个 block 出一个部分和，共 864 个】
    reduce_v8<<<block_num, BLOCK_SIZE>>>(d_a, g_odata, N);
    // 【阶段2：只起 1 个 block，把 864 个部分和收成 1 个】
    //   同一个 kernel 复用：此时"数据"是 g_odata，n = block_num。
    //   1 个 block 的 8 个 warp 先各自归约、再在 0 号 warp 里收成 1 个 ——
    //   和阶段1 内部的机制完全一样，所以不用另写一个 kernel。
    reduce_v8<<<1, BLOCK_SIZE>>>(g_odata, g_final_data, block_num);

    cudaMemcpy(out,g_final_data,1*sizeof(float),cudaMemcpyDeviceToHost);

    // 【对拍】只有一个数，直接比。
    if(check(out,res,1))printf("the answer is right\n");
    else{
        printf("the answer is wrong\n");
        printf("out[0]: %lf\n",out[0]);
        printf("res[0]: %lf\n", res[0]);
    }

    cudaFree(d_a);
    cudaFree(g_odata);
    free(a);
    free(out);
    return 0;
}