__global__ void Kernel_A(int *d_s, int *d_o)
{
    int tid = threadIdx.x;
    int *b_s = d_s + blockIdx.x * num_per_block;

    // 1) 每个线程先将负责的部分的数据规约到自己的寄存器中
    int sum = 0;
    for(int i=tid; i<num_per_block; i+=blockDim.x)
    {
        sum += b_s[i];
    }
    
    // 2) 每个块中都分配共享内存, 共享内存负责存储每个线程的变量的值
    __shared__ int tmp_sum[thread_per_block];
    tmp_sum[tid] = sum;
    __syncthreads();

    // 3) 这样规约的好处, 1) 避免了线程束的分歧 2) 不存在 bank 冲突 3) 最后一个 warp 内不需要同步, 避免了同步造成的影响
    for(int s=blockDim.x/2; s>16; s>>=1)
    {
        if(tid<s)
        {
            tmp_sum[tid] += tmp_sum[tid+s];
        }
        __syncthreads();  // 这个不能放到 if 里面
    }

    // 一个 warp 中的所有线程无论什么时候, 都是处在同一种状态, SIMD 的特点
    if(tid<16)  // 因为是同一个 warp 所以不需要同步
    {
        tmp_sum[tid] += tmp_sum[tid+16];
        tmp_sum[tid] += tmp_sum[tid+8];
        tmp_sum[tid] += tmp_sum[tid+4];
        tmp_sum[tid] += tmp_sum[tid+2];
        tmp_sum[tid] += tmp_sum[tid+1];
    }

    // 4) 将每个块规约之后的值赋值到输出
    if(tid==0) d_o[blockIdx.x] = tmp_sum[0];
}

#define FINAL_MASK 0xffffffff

template <typename T>
__inline__ __device__
T warpReduceSum(T val)
{
  for(int mask = 16; mask > 0; mask >>= 1)
    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
  return val;
}

template <typename T>
__inline__ __device__
T blockReduceSum(T val)
{
  static __shared__ T shared[32]; 
  int lane = threadIdx.x & 0x1f;  // 相当于取余
  int wid = threadIdx.x >> 5;  

  val = warpReduceSum<T>(val);

  if(lane == 0)
    shared[wid] = val;
  __syncthreads();

  val = (threadIdx.x < (blockDim.x >> 5 )) ? shared[lane] : (T)0.0f;
  val = warpReduceSum<T>(val);

  return val;  // 返回是都要返回的, 只不过只有部分线程的值为想要的值
}

// 每个 warp 进行规约, 规约后的值在 warp 中的 0 号线程
// __shfl_down_sync 相比于 __shfl_down 允许开发者指定一个线程掩码, 确保只有在指定的线程都完成数据交换后, 才继续执行后续的操作
__device__ int warpReduceSum(int sum)
{
    sum += __shfl_down_sync(0xffffffff, sum, 16);
    sum += __shfl_down_sync(0xffffffff, sum, 8);
    sum += __shfl_down_sync(0xffffffff, sum, 4);
    sum += __shfl_down_sync(0xffffffff, sum, 2);
    sum += __shfl_down_sync(0xffffffff, sum, 1);
    
    // sum += __shfl_xor_sync(0xffffffff, sum, 16);
    // sum += __shfl_xor_sync(0xffffffff, sum, 8);
    // sum += __shfl_xor_sync(0xffffffff, sum, 4);
    // sum += __shfl_xor_sync(0xffffffff, sum, 2);
    // sum += __shfl_xor_sync(0xffffffff, sum, 1);

    return sum;
}

__device__ int blockReduceSum(int sum)
{
    int laneId = threadIdx.x % warpSize;  // 每个 warp 中第几个线程
    int warpId = threadIdx.x / warpSize;  // 属于第几个 warp

    // 填充进去
    sum = warpReduceSum(sum);
    __shared__ int reduceSum[32];         // 最多也就 32 个 warp
    if(laneId==0)
    {
        reduceSum[warpId] = sum;
    }
    __syncthreads();

    bool pred = (warpId==0) && (laneId < blockDim.x / warpSize);
    int value =  (int)pred * reduceSum[laneId];  // 在范围内的为正常的数字, 否则为 0
    if(warpId==0) value = warpReduceSum(value);
    if(warpId==0 && laneId==0) sum=value;

    return sum;
}


__global__ void Kernel_B(int *d_s, int *d_o)
{
    int tid = threadIdx.x;  // 每个块内线程的 id
    int *b_s = d_s + blockIdx.x * num_per_block;

    // 1) 每个线程先将负责的部分的数据规约到自己的寄存器中
    int sum = 0;
    for(int i=tid; i<num_per_block; i+=blockDim.x)
    {
        sum += b_s[i];
    }
    sum = blockReduceSum(sum);

    if(tid==0) d_o[blockIdx.x] = sum;
}

// 每个线程负责的元素个数不能被线程数整除
__global__ void kernel_C(int *d_s, int *d_o){
    int tid = threadIdx.x;  // 每个块内的线程 id
    
    // 网格跨步法处理能够向量化读取的部分
    for(int i=tid; i < num_per_block/2; i+=blockDim.x)
    {
        reinterpret_cast<int2*>(d_o)[i] = reinterpret_cast<int2*>(d_s)[i];
    }

    // 处理不够 2 的部分
    int remainder = num_per_block % 2;  // 余数
    int quotient = num_per_block / 2;  // 商
    if(tid < remainder)
    {
        d_o[2*quotient+tid] = d_s[2*quotient+tid];
    }

}

// 网格块跨步法, 依据全局内存中 arr 的元素个数, 每个线程迭代几次, 将数据从全局加载在到其寄存器中
for (i=threadIdx.x; i<imax; i+= BLOCKDIMX) {
    double locvar = arr[i];
    /* 接下来很多指令使用 locvar 变量*/
}

// 每个线程都引入了计数器 ctr, 来记录当前线程迭代的步数, 这里的预取距离(PDIST)是 4, 这里必须假设每个线程的迭代步数都能够被 4 整除
// 这个是批量预取到寄存器中, 通常预取的值越多, 方法越有效
double v0, v1, v2, v3;
for (i=threadIdx.x, ctr=0; i<imax; i+= BLOCKDIMX, ctr++) {
    int ctr_mod = ctr%4;
    if (ctr_mod==0) {  // 一个线程每迭代四步, 就填充 buffer
        v0=arr[i+0* BLOCKDIMX]; 
        v1=arr[i+1* BLOCKDIMX]; 
        v2=arr[i+2* BLOCKDIMX]; 
        v3=arr[i+3* BLOCKDIMX];
    }
    switch (ctr_mod) { // 依据当前的迭代步数, 从预取的寄存器中得到值
        case 0: double locvar = v0; break;
        case 1: double locvar = v1; break;
        case 2: double locvar = v2; break;
        case 3: double locvar = v3; break;
    }
    /* 接下来很多指令使用 locvar 变量*/
}

// 当寄存器数组中元素个数比较少的时候, 才是放到寄存器中
double v[4];
for (int i=threadIdx.x, ctr=0; i<imax; i+= BLOCKDIMX, ctr++) {
    int ctr_mod = ctr%4;
    if (ctr_mod==0) {  // 一个线程每迭代四步, 就填充 buffer
        for(int k=0; k<4; ++k){
            v[i] = arr[i+k* BLOCKDIMX];
        }
    }
    double locvar = v[ctr_mod];
    /* 接下来很多指令使用 locvar 变量*/
}

/*
    这种对共享内存的操作, 不需要加共享内存同步指令
    1) 每个线程只会访问自己所需要的共享内存数据, 不会出现数据竞争问题
    2) 每个线程都是写完才读的, 不存在其他线程写的, 自己来读这种情况
*/
// 批量预取到共享内存中, 也要保证每个线程的迭代步数能够被 PDIST 整除
constexpr int PDIST = 4;  // PDIST 是预取举例
// 用共享变量来存储预取的数据, index 是线程迭代次数 % PDIST 后的结果.
#define vsmem(index)  v[index+PDIST*threadIdx.x]

__shared__ double v[PDIST* BLOCKDIMX];  // 分配空间, 每个线程预留 PDIST 个空间
for (int i=threadIdx.x, int ctr=0; i<imax; i+= BLOCKDIMX, ctr++) {
    int ctr_mod = ctr % PDIST;
    if (ctr_mod==0) {
        for (int k=0; k<PDIST; ++k) vsmem(k) = arr[i+k* BLOCKDIMX];
    }
    double locvar = vsmem(ctr_mod);
    /* 接下来很多指令使用 locvar 变量*/
}


/* 以上的批量预取到寄存器和批量预取到共享内存, 都有一个问题, 就是需要每个线程的迭代步数, 都能整除 PDIST, 接下来的滚动预取, 能够克服这一点*/

/*

*/
// 仍然是每个线程预取 PDIST 的空间, 但是不再要求每个线程的迭代步数能够整除 PDIST
constexpr int PDIST = 4;
#define vsmem(index)  v[index+PDIST*threadIdx.x]
__shared__ double v[PDIST* BLOCKDIMX];  // 仍然是每个线程预取 PDIST 的空间

for (int k=0; k<PDIST; ++k) vsmem(k) = arr[threadIdx.x+k* BLOCKDIMX];  // 要求第一次预取能够达到
for (int i=threadIdx.x, ctr=0; i<imax; i+= BLOCKDIMX, ctr++) {
    int ctr_mod= ctr%PDIST;
    double locvar = vsmem(ctr_mod);
    // 这个判断条件可能 i+PDIST*BLOCKDIMX < imax 能更好理解一些, 判断这个线程是否还有需要的缓冲元素
    if (i<imax-PDIST* BLOCKDIMX) vsmem(ctr_mod) = arr[i+PDIST* BLOCKDIMX];

    /* 接下来很多指令使用 locvar 变量*/
}

/*以上可能共享内存会存在 bank 冲突, 在共享内存中填充(padding)数组大小，以避免错误的跨步*/
#define vsmem(index) v[index+(PDIST+PADDING)*threadIdx.x]

// 分配双倍的共享存储空间
__shared__ float As[2][BLOCK_SIZE_K][BLOCK_SIZE_M];
__shared__ float Bs[2][BLOCK_SIZE_K][BLOCK_SIZE_N];

// 分配双倍的寄存器存储空间
float frag_a[2][THREAD_SIZE_Y];
float frag_b[2][THREAD_SIZE_X];

// 分配双倍的寄存器存储空间, 要先将全局存储数据搬运到寄存器中再搬运到共享内存中
float ldg_a_reg[4*ldg_num_a];
float ldg_b_reg[4*ldg_num_b];

// 每个寄存器的结果, 迭代完 256 个大迭代后, 将结果写入全局内存中
float accum[THREAD_SIZE_Y][THREAD_SIZE_X] = {0};  // 这个最后直接赋值到输出矩阵 C 中

// 把第一个大迭代块所需要的数据从全局内存读入到共享内存
// 把第一个小迭代块所需要的数据从共享内存读入到寄存器中
for k in 256 大迭代:
    // 将下一个大迭代块所需要的数据从共享内存预取到寄存器中
    for k in 7 小迭代:
        // 将下一个迭代所需要的数据从共享内存预取到寄存器中
        // 依靠本次寄存器中的数据进行本次迭代计算, 结果写到 accum 中
    // 将下一个大迭代块所需要的数据从寄存器取到共享内存中
    // 计算最后一个小迭代, 结果写到 accum 中

// 完成 256 次大迭代后, 每个线程将结果写到全局内存中

/*
    以下的内核代码都是一维网格一维线程模型来启动的
*/

// 访问数组元素有偏移(偏移可能会造成不对齐的情况发生)
template<typename T>
__global__ void offset(T* a, int s)
{
  int i = blockDim.x * blockIdx.x + threadIdx.x + s;
  a[i] = a[i] + 1;
}

// 这个是跨步访问, s 是跨步的步长。这个导致了同一个 warp 中的线程访问的全局内存地址是不连续的
template<typename T>
__global__ void stride(T* a, int s)
{
  int i = (blockDim.x * blockIdx.x + threadIdx.x) * s;
  a[i] = a[i] + 1;
}

// 32*16 大小的线程块要操作 block_size_m*block_size_n 大小的数据块
template<int block_size_m, int block_size_n> 
__global__ void Kernel_A(
    float* __restrict__ d_i,   // 输入矩阵大小为 M*N
    float* __restrict__ d_o,   // 转置输入矩阵大小为 N*M
    const int M,
    const int N
){
    /*
        第一步将全局内存中的数据读入到共享内存中
    */
    __shared__ float sm_data[block_size_m][block_size_n+1];  // 每个共享内存块都要读入这么多的数据

    int tx = threadIdx.x;  // laneId
    int ty = threadIdx.y;  // warpId
    int bx = blockIdx.x;   // x 方向是第几个块
    int by = blockIdx.y;   // y 方向是第几个块

    for(int i=0; i<block_size_m; i+=blockDim.y)
    {
        int idx_glo = (by*block_size_m+i+ty)*N+bx*block_size_n;  // 定位到第几行, 第几列的起始位置
        reinterpret_cast<float2*>(&sm_data[i+ty][tx*2])[0] = reinterpret_cast<float2*>(&d_i[idx_glo])[tx];
    }

    __syncthreads();  // 不同的线程要操作相同的共享内存, 所以要先进行同步

    /*
        将共享内存中的数据逐列读取, 逐行写入到全局内存中
    */
    for(int i=0; i<block_size_n; i+=blockDim.y) 
    {
        int idx_glo = (bx*block_size_n+ty+i)*M + by*block_size_m+tx;
        d_o[idx_glo] = sm_data[tx][ty+i];
    }
}

int input_arr[] = {3, 1,  7,  0,  4,  1,  6,  3};
// 包含扫描和为    [3, 4, 11, 11, 15, 16, 22, 25]
// 独占扫描和为    [0, 3,  4, 11, 11, 15, 16, 22]

// input_arr 末尾填充一些 0
int input_arr[] = {3, 1,  7,  0,  4,  1,  6,  3,  0,  0,  0};
// 包含扫描和为    [3, 4, 11, 11, 15, 16, 22, 25, 25, 25, 25]
// 独占扫描和为    [0, 3,  4, 11, 11, 15, 16, 22, 25, 25, 25]

__global__ void parallel_block_scan_kernel(int *data, int *prefix_sum, int N)  // 这里的 N 是填充前, 真实元素的个数
{
    extern __shared__ int tmp[];
    int tid = threadIdx.x;
    int leaf_num = blockDim.x * 2;  // equals to length of tmp

    // 这里其实是循环展开了, 每个线程负责运输两个数据, 用 0 来补充
    tmp[tid * 2] = tid * 2 < N ? data[tid * 2] : 0;
    tmp[tid * 2 + 1] = tid * 2 + 1 < N ? data[tid * 2 + 1] : 0;
    __syncthreads();  // 同步不能忘记了

    // 上扫过程, 每个线程处理两个元素
    int offset = 1;  // 相隔多少个元素取一个元素
    // 先确定要使用的线程, 再给这些线程分配操纵的数据
    for (int d = leaf_num >> 1; d > 0; d >>= 1)
    {
        if (tid < d)  // 先筛选出需要的线程个数, 选好后的每个线程总是负责两个数据
        {
            int ai = offset * (2 * tid + 1) - 1;  // 2*tid 是每个线程处理 2*tid*offset 个元素, +1 是指取出第一个块的后一个元素
            int bi = offset * (2 * tid + 2) - 1;
            tmp[bi] += tmp[ai];
        }
        offset *= 2;
        __syncthreads();
    }

    // 每个块都使用 tid 为 0 的线程, 将共享存储中最后一个元素置为0, 只有完成这一步才能下扫
    if (tid == 0)
    {
        tmp[leaf_num - 1] = 0;
    }
    __syncthreads();

    // 下扫过程, 每个线程处理两个元素
    for (int d = 1; d < leaf_num; d *= 2)
    {
        offset >>= 1;
        if (tid < d)
        {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;

            float v = tmp[ai];
            tmp[ai] = tmp[bi];
            tmp[bi] += v;
        }
        __syncthreads();
    }

    // 每个线程处理两个元素, N 的个数可能为奇数, 也可能为偶数, 这里其实也是展开了
    // 先搬所有偶数, 再搬所有奇数
    if (tid * 2 < N)
    {
        prefix_sum[tid * 2] = tmp[tid * 2];
    }
    if (tid * 2 + 1 < N)
    {
        prefix_sum[tid * 2 + 1] = tmp[tid * 2 + 1];
    }
}

#define NUM_BANKS 32
#define LOG_NUM_BANKS 5
#define CONFLICT_FREE_OFFSET(n) (((n) >> LOG_NUM_BANKS) + ((n) >> (2 * LOG_NUM_BANKS)))

// 有 bank 填充的时候, tmp 分配的空间更多了, 这些空间既不会读到, 也不会写到。
// 仍然是正常的逻辑来处理元素, 只不过放入共享内存, 或者从共享内存读取的时候, 需要依据其所在元素的位置计算偏移
__global__ void parallel_block_scan_bcao_kernel(int *data, int *prefix_sum, int N)
{
    extern __shared__ int tmp[];
    int tid = threadIdx.x;
    int leaf_num = blockDim.x * 2; // not equals to length of tmp, 填充到 2 的幂次的大小, 但是没有填充 bank

    int ai = tid;  // 前一半中的元素
    int bi = tid + (leaf_num >> 1);  // 后一半中对应的元素
    int offset_ai = CONFLICT_FREE_OFFSET(ai);  // ai 前有多少个元素的填充
    int offset_bi = CONFLICT_FREE_OFFSET(bi);  // bi 前有多少个元素的填充

    tmp[ai + offset_ai] = ai < N ? data[ai] : 0;
    tmp[bi + offset_bi] = bi < N ? data[bi] : 0;
    __syncthreads();

    // 加 bank 的 offset 是是先算出来原本的索引, 再依据原本的索引添加 offset
    int offset = 1;
    for (int d = leaf_num >> 1; d > 0; d >>= 1)
    {
        if (tid < d)
        {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            ai += CONFLICT_FREE_OFFSET(ai);
            bi += CONFLICT_FREE_OFFSET(bi);
            tmp[bi] += tmp[ai];
        }
        offset *= 2;
        __syncthreads();
    }

    if (tid == 0)
    {
        tmp[leaf_num - 1 + CONFLICT_FREE_OFFSET(leaf_num - 1)] = 0;
    }
    __syncthreads();

    for (int d = 1; d < leaf_num; d *= 2)
    {
        offset >>= 1;
        if (tid < d)
        {
            int ai = offset * (2 * tid + 1) - 1;
            int bi = offset * (2 * tid + 2) - 1;
            ai += CONFLICT_FREE_OFFSET(ai);
            bi += CONFLICT_FREE_OFFSET(bi);

            float v = tmp[ai];
            tmp[ai] = tmp[bi];
            tmp[bi] += v;
        }
        __syncthreads();
    }

    if (ai < N)
    {
        prefix_sum[ai] = tmp[ai + offset_ai];
    }
    if (bi < N)
    {
        prefix_sum[bi] = tmp[bi + offset_bi];
    }
}

void recursive_scan(int *d_data, int *d_prefix_sum, int N, bool bcao)
{
    // 这两行是一个向上取整的操作
    int block_num = N / MAX_ELEMENTS_PER_BLOCK;  // 分配几个块
    if (N % MAX_ELEMENTS_PER_BLOCK != 0)
    {
        block_num += 1;
    }

    // 将每个段的和再次规约
    int *d_sums, *d_sums_prefix_sum;  // 数组中存放的是每个块的数据
    CUDA_CHECK(cudaMalloc(&d_sums, block_num * sizeof(int)));
    CUDA_CHECK(cudaMalloc(&d_sums_prefix_sum, block_num * sizeof(int)));

    // 依据当前数据所需要的块来进行前缀扫描和
    if (bcao)
    {
        parallel_large_scan_bcao_kernel<<<block_num, MAX_THREADS_PER_BLOCK>>>(d_data, d_prefix_sum, N, d_sums);
    }
    else
    {
        parallel_large_scan_kernel<<<block_num, MAX_THREADS_PER_BLOCK>>>(d_data, d_prefix_sum, N, d_sums);
    }
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());

    // 当 block_num 的个数不为 1 的时候, 就需要再次递归, 转化为子问题。
    // 当 block_num 的个数为 1 的时候, 就不需要递归了, 函数就结束了(递归结束的条件, 数学归纳法的初始条件)。
    if (block_num != 1)  // 判断块的个数是否为1
    {
        recursive_scan(d_sums, d_sums_prefix_sum, block_num, bcao);
        add_kernel<<<block_num, MAX_THREADS_PER_BLOCK>>>(d_prefix_sum, d_sums_prefix_sum, N);
        CUDA_CHECK(cudaGetLastError());
        CUDA_CHECK(cudaDeviceSynchronize());
    }

    CUDA_CHECK(cudaFree(d_sums));
    CUDA_CHECK(cudaFree(d_sums_prefix_sum));
}

__global__ void kernel_scalarProduct(half2 *vec1, half2 *vec2, float* result, int num_ele){

    extern __shared__ half2 sm_value[];
    const int tid = threadIdx.x;
    const int id = tid + blockIdx.x * blockDim.x;

    // 每个线程先将自己负责的数据归约到共享内存中
    for(int i=0; i<(num_ele>>1); i+=gridDim.x*blockDim.x)
    {
        sm_value[tid] = __hadd2(sm_value[tid], __hmul2(vec1[i+id], vec2[i+id]));
    }
    __syncthreads();

    for(int s=(blockDim.x>>1); s>=32; s>>=1)
    {
        if(tid<s)
        {
            sm_value[tid] = __hadd2(sm_value[tid], sm_value[tid+s]);
        }
        __syncthreads();
    }

    if(tid<32)
    {
        sm_value[tid] = __hadd2(sm_value[tid], sm_value[tid+16]);
        sm_value[tid] = __hadd2(sm_value[tid], sm_value[tid+8]);
        sm_value[tid] = __hadd2(sm_value[tid], sm_value[tid+4]);
        sm_value[tid] = __hadd2(sm_value[tid], sm_value[tid+2]);
        sm_value[tid] = __hadd2(sm_value[tid], sm_value[tid+1]);
    }
    __syncthreads();

    if(tid==0)
    {
        result[blockIdx.x] = __high2float(sm_value[0]) + __low2float(sm_value[0]);
    }
}

// host 端调用, 依据传入的 cudaError_t 类型的变量来打印错误所在位置和错误信息
// #var 是一种预处理器操作，被称为字符串化操作符。将宏参数 val 转换为一个字符串。也就是说，它会把 val 所代表的实际值转换为字符串形式。
#define checkCudaErrors(val) check((val), #val, __FILE__, __LINE__)

// host 端调用, 主要是打印 kernel 内发生的错误信息。因为 kernel 并不能使用 checkCudaErrors。但是有几点注意事项
// 1) 其主要是用于 kernel 的检查。因为核函数是异步的, 也没有任何返回值。所以必须在核函数启动之后调用 cudaGetLastError 来检索核函数是否启动成功
// 2) 我们要确保核函数启动之前 cudaError_t 类型的变量是 cudaSuccess, 排除核函数以外的错误信息。
// 3) 由于核函数的启动是异步的, 所以必须在调用 cudaGetLastError() 前同步核函数(其实也好理解, 只有核函数执行完, 才能得到 cudaError_t 类型的变量)。
#define getLastCudaError(msg) __getLastCudaError(msg, __FILE__, __LINE__)

template <typename T>
void check(T result, char const *const func, const char *const file,
           int const line) {
  if (result) {
    fprintf(stderr, "CUDA error at %s:%d code=%d(%s) \"%s\" \n", file, line,
            static_cast<unsigned int>(result), cudaGetErrorName(result), func);
    exit(EXIT_FAILURE);
  }
}

inline void __getLastCudaError(const char *errorMessage, const char *file,
                               const int line) {
  cudaError_t err = cudaGetLastError();

  if (cudaSuccess != err) {
    fprintf(stderr,
            "%s(%i) : getLastCudaError() CUDA error :"
            " %s : (%d) %s.\n",
            file, line, errorMessage, static_cast<int>(err),
            cudaGetErrorString(err));
    exit(EXIT_FAILURE);
  }
}

// 使用多个流, 实现 copy 与 kernel 的并行操作
checkCudaErrors(cudaEventRecord(start_event, 0));  // default 上这个执行完了, 才会执行非默认流中的操作
for(int i=0; i<numStream; ++i)
{
    int cur = i * streamSize;  // 当前处理的元素个数
    checkCudaErrors(cudaMemcpyAsync(&d_i[cur], &h_i[cur], streamBytes, cudaMemcpyHostToDevice, streamVec[i]));
    kernel_A<<<streamSize/blockSize, blockSize, 0, streamVec[i]>>>(&d_i[cur], &d_o[cur], 0);
    checkCudaErrors(cudaMemcpyAsync(&h_o[cur], &d_o[cur], streamBytes, cudaMemcpyDeviceToHost, streamVec[i]));
}
checkCudaErrors(cudaEventRecord(stop_event, 0));
checkCudaErrors(cudaEventSynchronize(stop_event));
checkCudaErrors(cudaEventElapsedTime(&letancy, start_event, stop_event));
std::cout << "multiple stream letancy: " << letancy << " ms" << std::endl;

// CPU 版本的算法流程如下:
// maxBits 是输入元素的最大位数; numBits 是将多少个位做为一组来对比, k = 1 << numBits; 
for(int i=0; i<maxBits; i+=numBits)
{
    // 这个是额外的存储空间, 用来统计每个 k 的每个状态的元素个数和元素个数的前缀和
    std::vector<int> binHistogram(numBins, 0);
    std::vector<int> binScan(numBins, 0);

    // step 1 遍历 inVec, inVec 中第 i 位值为 [0, numBins-1] 的各有多少, 其值分别放到 binHistogram
    // step 2 遍历 binHistogram, 计算 binHistogram 的独占式累加和, 存入 binScan 中
    // step 3 遍历 inVec, 计算其值所属的桶, 得到前缀和, 放到输出的位置 outVec[binScan[bin]++] = inVec[j];
    // step 4 swap(intVec, outVec)
}

#include <iostream>
#include <vector>

// 基数排序是从最低位看起, 把这一位数字值相同的按照扫描顺序放入同一个桶里面, 值小的桶在前面。
// 当所有数字都扫描完，再使用高一位，循环上述步骤，直至达到所有数字所有的最高位数，最后输出的就是排序后的答案。
void radixSort(std::vector<unsigned int>& inputVals, int numBits = 2) {
    int n = inputVals.size();    // 输入元素的个数
    int numBins = 1 << numBits;  // numBins = 2^numBits, 桶的个数
    int maxBits = 32;  // Assuming unsigned int (32-bit)

    std::vector<unsigned int> outputVals(n);  // 存放排序后的结果
    for (int i = 0; i < maxBits; i += numBits){
        std::vector<int> binHistogram(numBins, 0);
        std::vector<int> binScan(numBins, 0);

        // Step 1: 计算统计图
        for (int j = 0; j < n; ++j) {  // 遍历所有元素, 统计每个桶中元素个数
            int bin = (inputVals[j] >> i) & (numBins - 1);
            binHistogram[bin]++;
        }

        // Step 2: 计算前缀和
        for (int j = 1; j < numBins; ++j) {  // 遍历所有桶, 统计元素个数的前缀和
            binScan[j] = binScan[j - 1] + binHistogram[j - 1];
        }

        // Step 3: 基于当前数字来排序
        for (int j = 0; j < n; ++j) {  // 遍历所有元素, 将输入元素赋值到输出对应的位置
            int bin = (inputVals[j] >> i) & (numBins - 1);  // 判断当前数字属于那个 bin, 调整其位置
            outputVals[binScan[bin]++] = inputVals[j];
            // binScan[bin]++;
        }

        // Step 4: 下一轮的排序基于当前排序的结果
        inputVals.swap(outputVals);
    }
}

int main() {
    std::vector<unsigned int> inputVals = {170, 45, 75, 90, 802, 24, 2, 66};
    radixSort(inputVals, 2);

    for (unsigned int val : inputVals) {
        std::cout << val << " ";
    }
    std::cout << std::endl;

    return 0;
}

unsigned int *data_temp = (unsigned int *)(&src_data[i]);
*data_temp = (*data_temp >> 31 & 0x1)? ~(*data_temp): (*data_temp) | 0x80000000; 

void radix_sort(unsigned int* const d_out,
    unsigned int* const d_in,
    unsigned int d_in_len)  // d_in_len 要排序是数据元素总个数
{
    unsigned int block_sz = MAX_BLOCK_SZ;  // 一个 block 最大线程个数
    unsigned int max_elems_per_block = block_sz;
    unsigned int grid_sz = d_in_len / max_elems_per_block;  // 每个线程对应一个元素,总共分成多少个 block
    // Take advantage of the fact that integer division drops the decimals
    if (d_in_len % max_elems_per_block != 0)  // 向上取整
        grid_sz += 1;

    unsigned int* d_prefix_sums;  // 存放相对位置索引
    unsigned int d_prefix_sums_len = d_in_len;
    checkCudaErrors(cudaMalloc(&d_prefix_sums, sizeof(unsigned int) * d_prefix_sums_len));
    checkCudaErrors(cudaMemset(d_prefix_sums, 0, sizeof(unsigned int) * d_prefix_sums_len));

    unsigned int* d_block_sums;  // 统计每个线程块的各个桶的数量
    unsigned int d_block_sums_len = 4 * grid_sz; // 4-way split
    checkCudaErrors(cudaMalloc(&d_block_sums, sizeof(unsigned int) * d_block_sums_len));
    checkCudaErrors(cudaMemset(d_block_sums, 0, sizeof(unsigned int) * d_block_sums_len));

    unsigned int* d_scan_block_sums;  // 每个线程块各个桶的数量的前缀和
    checkCudaErrors(cudaMalloc(&d_scan_block_sums, sizeof(unsigned int) * d_block_sums_len));
    checkCudaErrors(cudaMemset(d_scan_block_sums, 0, sizeof(unsigned int) * d_block_sums_len));

    // shared memory consists of 3 arrays the size of the block-wise input
    // and 2 arrays the size of n in the current n-way split (4)
    unsigned int s_data_len = max_elems_per_block;  // 存放数据, 每个线程对应一个数据
    unsigned int s_mask_out_len = max_elems_per_block + 1;  // 
    unsigned int s_merged_scan_mask_out_len = max_elems_per_block;
    unsigned int s_mask_out_sums_len = 4; // 4-way split
    unsigned int s_scan_mask_out_sums_len = 4;
    unsigned int shmem_sz = (s_data_len 
                            + s_mask_out_len
                            + s_merged_scan_mask_out_len
                            + s_mask_out_sums_len
                            + s_scan_mask_out_sums_len)
                            * sizeof(unsigned int);


    // for every 2 bits from LSB to MSB:
    //  block-wise radix sort (write blocks back to global memory)
    for (unsigned int shift_width = 0; shift_width <= 30; shift_width += 2)
    {
        // 每个 block 上进行排序得到最终的结果
        gpu_radix_sort_local<<<grid_sz, block_sz, shmem_sz>>>(d_out, 
                                                            d_prefix_sums, 
                                                            d_block_sums, 
                                                            shift_width, 
                                                            d_in, 
                                                            d_in_len, 
                                                            max_elems_per_block);


        // scan global block sum array, 四路分别进行规约
        sum_scan_blelloch(d_scan_block_sums, d_block_sums, d_block_sums_len);

        // scatter/shuffle block-wise sorted array to final positions
        gpu_glbl_shuffle<<<grid_sz, block_sz>>>(d_in,
                                                d_out, 
                                                d_scan_block_sums, 
                                                d_prefix_sums, 
                                                shift_width, 
                                                d_in_len, 
                                                max_elems_per_block);
    }

    checkCudaErrors(cudaMemcpy(d_out, d_in, sizeof(unsigned int) * d_in_len, cudaMemcpyDeviceToDevice));
    checkCudaErrors(cudaFree(d_scan_block_sums));
    checkCudaErrors(cudaFree(d_block_sums));
    checkCudaErrors(cudaFree(d_prefix_sums));
}

std::vector<unsigned int> quickSort(std::vector<unsigned int>&nums, int start, int end)
{
    // 注意这里的 start 和 end 都是左闭右闭
    // start >= end 时候, 这个部分就排序结束啦
    if (start >= end) return nums;

    int base = nums[start];  // 选择最开始的当作 base

    int i = start;
    int j = end;
    while (i < j)  // i 和 j 重合以后, 这段就排序完了
    {
        while (i < j && nums[j] >= base) j--; //从右往左，寻找比base小的数
        std::swap(nums[i], nums[j]);          //找到比base小的数，即与base交换位置, 即使 i == j 也不影响。
        while (i < j && nums[i] <= base) i++; //从左往右，寻找比base大的数
        std::swap(nums[i], nums[j]);          //找到比base大的数，即与base交换位置
    }
    quickSort(nums, start, i - 1);  // 分成两个部分来再次排序
    quickSort(nums, i + 1, end);    // 分成两个部分来再次排序

    return nums;
}

#include <iostream>
#include <vector>

// 建立大根堆, 升序排序
// 将下标为 k 的节点进行下沉
// 整体思路是 k 必须为非叶子节点。
// 然后判断其值是否大于等于其左右子节点的最大值, 不大于等于的话就进行进行下沉操作, 直至 k 为叶子节点, 或者 k 的值大于其左右子节点的最大值
void HeapAdjust(std::vector<int>& nums, int heapSize, int k){
	while(k<heapSize/2)  // 最后一个非叶子节点索引为 heapSize/2-1
	{
		int largest = 2*k+1;  // 因为 k 不是叶子节点, 所以其左节点一定存在
		int right = 2*k+2;
		if(right<heapSize && nums[largest]<nums[right]) largest = right;
		if(nums[k] > nums[largest]) break;  // 这里应该是 break, 而不是 continue;
		std::swap(nums[k], nums[largest]);
		k = largest;
	}
}

// 建立小根堆, 降序排序
/*
void HeapAdjust(std::vector<int>& nums, int heapSize, int k){
	while(k<heapSize/2)  // 最后一个非叶子节点索引为 heapSize/2-1
	{
		int small = 2*k+1;  // 因为 k 不是叶子节点, 所以其左节点一定存在
		int right = 2*k+2;
		if(right<heapSize && nums[small]>nums[right]) small = right;
		if(nums[k] < nums[small]) break;  // 这里应该是 break, 而不是 continue;
		std::swap(nums[k], nums[small]);
		k = small;
	}
}
*/

void HeapSort(std::vector<int>& nums)
{
	int heapSize = nums.size();       // 最初的堆的大小
	// 建立堆。从最后一个非叶子节点开始, 依次将所有的非叶子节点进行下沉操作。
	// 建立堆需要将无序数组构建成一个堆, 时间复杂度为 O(n)。而每个节点下沉操作最多需要 O(log(n))。建堆时间复杂度为 O(n)
	for(int i=heapSize/2-1; i>=0; --i)
	{
		HeapAdjust(nums, heapSize, i);
	}

	// 堆排序。从最后一个元素开始, 每次将最后一个元素与第一元素进行交换, 然后调整第一个元素, 注意此时的 heapSize 已经减少了 1。
	// 需要执行 n-1 次, 每次调整是时间复杂度为 O(log(n)), 这部分时间复杂度为 O(nlog(n))。
	for(int i=heapSize-1; i>0; --i)
	{
		std::swap(nums[0], nums[i]);  // 将堆中最后一个元素与第一个元素进行交换
		HeapAdjust(nums, i, 0);       // 对堆中的第一个元素进行调整
	}
	// 由于有栈空间的支出。所以最好和平均的空间复杂度均为 O(log2(n)), 最坏为 O(n)[当元素均为逆序的时候, 每次划分的都不均衡]
}

int main()
{
	std::vector<int> arr {9, 5, 6, 3, 5, 3, 1, 0, 96, 66};
	std::cout << "排序前的结果为: " << std::endl;
	for(int i=0; i<arr.size(); ++i)
	{
		std::cout << arr[i] << " ";
	}
	std::cout << std::endl;

	HeapSort(arr);

	std::cout << "排序后的结果为: " << std::endl;
	for(int i=0; i<arr.size(); ++i)
	{
		std::cout << arr[i] << " ";
	}
	std::cout << std::endl;

	return 0;
}

#include <iostream>
#include <vector>

// 合并
void merge(std::vector<int>& arr, std::vector<int>& tempArr, int left, int mid, int right)
{
    // 标记左半区第一个未排序的元素
    int l_pos = left;
    // 标记右半区第一个未排序的元素
    int r_pos = mid + 1;
    // 临时数组元素的下标
    int pos = left;  // 当前存到哪里了

    // 合并
    while (l_pos <= mid && r_pos <= right)
    {
        if (arr[l_pos] < arr[r_pos])  // 左半区第一个剩余元素更小, 这里小于号保证了算法的稳定性, 如果是小于等于算法就是不稳定了
            tempArr[pos++] = arr[l_pos++];
        else  // 右半区第一个剩余元素更小
            tempArr[pos++] = arr[r_pos++];
    }

    // 合并左半区剩余的元素
    while (l_pos <= mid)
        tempArr[pos++] = arr[l_pos++];

    // 合并右半区剩余的元素
    while (r_pos <= right)
        tempArr[pos++] = arr[r_pos++];

    // 把临时数组中合并后的元素复制回原来的数组, 是在这里将排序后的数组拷贝到原始的输入数组中
    while (left <= right)
    {
        arr[left] = tempArr[left];
        left++;
    }
}

// 归并排序
void msort(std::vector<int>& arr, std::vector<int>& tempArr, int left, int right)
{
    // 如果只有一个元素，那么不需要继续划分
    // 只有一个元素的区域，本身就是有序的，只需要被归并即可
    if (left < right)
    {
        // 找中间点
        int mid = left + (right-left) / 2;
        // 递归划分左半区
        msort(arr, tempArr, left, mid);
        // 递归划分右半区
        msort(arr, tempArr, mid + 1, right);

        // 执行到这里的时候, [left, mid] [mid+1, right] 这两个区间已经有序了(某个区间只有一个元素也是有序), 合并已经排序的部分
        merge(arr, tempArr, left, mid, right);
    }
}

// 归并排序入口
void merge_sort(std::vector<int>& arr)
{
    // 分配一个辅助数组, 空间复杂度就在这里 O(n)
    std::vector<int> tempArr(arr.size());
    // 调用实际的归并排序
    msort(arr, tempArr, 0, arr.size() - 1);
}

int main()
{
    std::vector<int> arr = {38, 27, 43, 3, 9, 82, 10};
    
    std::cout << "排序前: ";
    for (int i : arr)
        std::cout << i << " ";
    std::cout << std::endl;

    // 调用排序函数
    merge_sort(arr);

    std::cout << "排序后: ";
    for (int i : arr)
        std::cout << i << " ";
    std::cout << std::endl;

    return 0;
}

#include <iostream>
#include <cuda_runtime.h>

#define THREADS_PER_BLOCK 256  // 每个block中的线程数

// CUDA 核函数, 并行合并两个有序的子数组
__global__ void parallelMerge(int *arr, int *temp, int subArraySize, int n) {
    int threadId = blockIdx.x * blockDim.x + threadIdx.x;  // 全局线程ID
    int start = threadId * (2 * subArraySize);  // 当前线程负责的起始位置, 每个线程负责 2*subArraySize 的大小的数据

    // 保证在数组边界内
    if (start >= n) return;  // 因为 block 是向上取整分配的

    // 记录负责的这部分数据起始位置和终止位置
    int mid = min(start + subArraySize, n);      // 中间位置
    int end = min(start + 2 * subArraySize, n);  // 结束位置

    // 双指针合并两个子数组 [start, mid) 和 [mid, end)
    int i = start, j = mid, k = start;  // k 就是记录的写到 temp 的位置

    while (i < mid && j < end) {
        if (arr[i] <= arr[j]) {
            temp[k++] = arr[i++];
        } else {
            temp[k++] = arr[j++];
        }
    }

    while (i < mid) {
        temp[k++] = arr[i++];
    }

    while (j < end) {
        temp[k++] = arr[j++];
    }

    // 将合并结果复制回原数组
    for (i = start; i < end; i++) {
        arr[i] = temp[i];
    }
}

// 主程序
void mergeSortParallel(int *h_arr, int n) {
    // 设备端数组
    int *d_arr, *d_temp;  // d_temp 是临时内存, 其实就是空间复杂度是 O(n)
    cudaMalloc((void**)&d_arr, n * sizeof(int));
    cudaMalloc((void**)&d_temp, n * sizeof(int));

    // 复制数据到设备
    cudaMemcpy(d_arr, h_arr, n * sizeof(int), cudaMemcpyHostToDevice);

    // 确定每次合并的子数组大小, 初始为1, 逐步翻倍
    for (int subArraySize = 1; subArraySize < n; subArraySize *= 2) {
        // 相当于一个线程负责 2*subArraySize 大小的数据的合并, 其实是负责 2 个大小为 subArraySize 大小的数组的合并, 合并之后的结果是 2*subArraySize 的大小。
        int numBlocks = (n + (2 * subArraySize * THREADS_PER_BLOCK) - 1) / (2 * subArraySize * THREADS_PER_BLOCK);  // 分配的块数是向上取整
        parallelMerge<<<numBlocks, THREADS_PER_BLOCK>>>(d_arr, d_temp, subArraySize, n);
        cudaDeviceSynchronize();  // 因为是异步执行, 所以每一次 for 循环之后都需要进行一个同步
    }

    // 复制排序结果回主机
    cudaMemcpy(h_arr, d_arr, n * sizeof(int), cudaMemcpyDeviceToHost);

    // 释放设备内存
    cudaFree(d_arr);
    cudaFree(d_temp);
}

int main() {
    const int n = 100000;  // 输入数组长度
    int* h_arr = new int[n];

    // 随机初始化数据
    for (int i = 0; i < n; i++) {
        h_arr[i] = rand() % 1000000;
    }

    // 调用并行归并排序
    mergeSortParallel(h_arr, n);

    // 打印排序后的前 N 个元素
    std::cout << "Sorted array (first N elements): ";
    for (int i = 0; i < 30; i++) {
        std::cout << h_arr[i] << " ";
    }
    std::cout << std::endl;

    // 释放主机内存
    delete[] h_arr;

    return 0;
}

#include <iostream>
#include <cuda_runtime.h>

#define THREADS_PER_BLOCK 256  // 每个block中的线程数

// CUDA 核函数, 并行合并两个有序数组
__global__ void parallelMerge(int *A, int m, int *B, int n, int *C) {
    int L = m + n;
    int totalThreads = gridDim.x * blockDim.x;  // 总线程数
    int threadId = blockIdx.x * blockDim.x + threadIdx.x;  // 全局线程ID
    int chunk_size = (L + totalThreads - 1) / totalThreads;  // 每个线程处理的元素数量
    int s = threadId * chunk_size;  // 当前线程负责的输出数组起始位置
    int e = min(s + chunk_size, L);  // 当前线程负责的输出数组结束位置

    // 在 A 和 B 中寻找划分点 i 和 j
    int i_lower = max(0, s - n);
    int i_upper = min(s, m-1);
    int i, j;

    // 二分法查找 i 和 j
    while (i_lower <= i_upper) {
        i = (i_lower + i_upper) / 2;
        j = s - i;

        int A_i_1 = (i > 0) ? A[i - 1] : INT_MIN;
        int A_i = (i < m) ? A[i] : INT_MAX;
        int B_j_1 = (j > 0) ? B[j - 1] : INT_MIN;
        int B_j = (j < n) ? B[j] : INT_MAX;

        if (A_i_1 <= B_j && B_j_1 <= A_i) {
            break; // 找到合适的划分点
        } else if (A_i_1 > B_j) {
            i_upper = i - 1;
        } else {
            i_lower = i + 1;
        }
    }

    // 现在 i 和 j 是输入数组的起始位置
    int i_start = i;
    int j_start = s - i_start;

    // 合并操作，确保只处理当前线程的输出区间 [s, e)
    int idx = s;
    while (i_start < m && j_start < n && idx < e) {
        if (A[i_start] <= B[j_start]) {
            C[idx++] = A[i_start++];
        } else {
            C[idx++] = B[j_start++];
        }
    }

    while (i_start < m && idx < e) {
        C[idx++] = A[i_start++];
    }

    while (j_start < n && idx < e) {
        C[idx++] = B[j_start++];
    }
}

// CPU 端的归并函数：归并两个有序数组
void merge(int* A, int m, int* B, int n, int* C) {
    int i = 0, j = 0, k = 0;

    while (i < m && j < n) {
        if (A[i] <= B[j]) {
            C[k++] = A[i++];
        } else {
            C[k++] = B[j++];
        }
    }

    // 复制剩余元素
    while (i < m) {
        C[k++] = A[i++];
    }
    while (j < n) {
        C[k++] = B[j++];
    }
}

// 递归归并排序函数
void mergeSort(int* arr, int n) {
    if (n <= 1) {
        return; // 已经是有序的
    }

    int mid = n / 2;
    int* left = new int[mid];
    int* right = new int[n - mid];

    // 分割数组
    for (int i = 0; i < mid; i++) {
        left[i] = arr[i];
    }
    for (int i = mid; i < n; i++) {
        right[i - mid] = arr[i];
    }

    // 递归排序
    mergeSort(left, mid);
    mergeSort(right, n - mid);

    // 归并排序
    merge(left, mid, right, n - mid, arr);

    // 释放内存
    delete[] left;
    delete[] right;
}

// 主程序
int main() {
    const int n = 100000;  // 数据量10万
    int* h_A = new int[n];  // 数组A
    int* h_B = new int[n];  // 数组B
    int* h_C = new int[2 * n];  // 用于存放结果的数组C

    // 初始化数据
    for (int i = 0; i < n; i++) {
        h_A[i] = rand() % (n * 2);  // 随机数填充数组A
        h_B[i] = rand() % (n * 2);  // 随机数填充数组B
    }

    // CPU 归并排序
    mergeSort(h_A, n);
    mergeSort(h_B, n);

    // 打印归并后的结果（可以去掉，避免大数据量下打印影响性能）
    std::cout << "Array A sorted (first 10 elements): ";
    for (int i = 0; i < 20; i++) {
        std::cout << h_A[i] << " ";
    }
    std::cout << std::endl;

    std::cout << "Array B sorted (first 10 elements): ";
    for (int i = 0; i < 20; i++) {
        std::cout << h_B[i] << " ";
    }
    std::cout << std::endl;

    // 分配设备内存
    int *d_A, *d_B, *d_C;
    cudaMalloc((void**)&d_A, n * sizeof(int));
    cudaMalloc((void**)&d_B, n * sizeof(int));
    cudaMalloc((void**)&d_C, 2 * n * sizeof(int));

    // 复制数据到设备
    cudaMemcpy(d_A, h_A, n * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, n * sizeof(int), cudaMemcpyHostToDevice);

    // 启动 CUDA 核函数进行并行归并
    int numBlocks = (2 * n + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;  // 计算需要的块数
    parallelMerge<<<numBlocks, THREADS_PER_BLOCK>>>(d_A, n, d_B, n, d_C);

    // 复制结果回主机
    cudaMemcpy(h_C, d_C, 2 * n * sizeof(int), cudaMemcpyDeviceToHost);

    // 打印结果（可以去掉，避免大数据量下打印影响性能）
    std::cout << "Merged array (first 10 elements): ";
    for (int i = 0; i < 30; i++) {
        std::cout << h_C[i] << " ";
    }
    std::cout << std::endl;

    // 释放设备内存
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    // 释放主机内存
    delete[] h_A;
    delete[] h_B;
    delete[] h_C;

    return 0;
}

// 全局内存分配
// count : 所分配全局内存的字节 
// devptr：指针返回该内存的地址
cudaError_t cudaMalloc(void **devPtr, size_t count);

// 全局内存的释放
cudaError_t cudaFree(void *devPtr); 

// 同步传输
cudaMemcpy(void *dst, const void *src, size_t count, enum cudaMemcpyKind kind);

// 异步传输
cudaMemcpyAsync(void *dst, const void *src, size_t count, enum cudaMemcpyKind kind, cudaStream_t stream);

/*
    * 函数从内存位置 src 复制 count 字节到内存位置 dst
    * kind 表示内存传输方向(要与src到dst保持一致)
        kind取值：
        cudaMemcpyHostToHost
        cudaMemcpyHostToDevice
        cudaMemcpyDeviceToHost
        cudaMemcpyDeviceToDevice
    * 注意同步传输的时候, 主机内存可以是分页内存和页锁定内存。而异步传输的时候，主机内存必须是页锁定内存。
*/

// 页锁定内存分配
cudaError_t cudaMallocHost(void **devPtr, size_t count);
cudaError_t cudaHostAlloc(T **ptr, size_t size, unsigned int flags)

// 页锁定内存释放
cudaError_t cudaFreeHost(void *ptr);

// 零拷贝内存的分配
cudaError_t cudaHostAlloc(void **pHost, size_t count, unsigned int flags);
/*
    flags 有以下情况:
    cudaHostAllocDefault: 默认分配策略, 默认分配普通的页锁定内存, 不会为 GPU 提供直接访问能力。
    cudaHostAllocPortable: 分配的内存可以跨不同的 GPU 设备, 意味着在一个设备上分配的内存可以在多个设备之间使用, 适用于多 GPU 环境。但并不意味着分配的主机内存能够被多个 GPU 直接访问。它的实际含义是，这段内存可以跨多个 CUDA 上下文（contexts）共享使用。
    cudaHostAllocMapped: 分配的主机内存是映射的, 允许 GPU 直接访问, 减少了主机和设备之间的数据拷贝开销。
*/

// 零拷贝内存的释放, 其实就是页锁定内存的释放函数
cudaError_t cudaFreeHost(void *ptr); 

// 就是需要注意, 零拷贝虽然不需要显示地传递到设备上, 但是设备不能通过 pHost 直接访问对应的内存地址, 设备需要访问主机上的零拷贝内存, 需要先获得另外一个地址,这个地址帮助设备访问到主机对应的内存。
cudaError_t cudaHostGetDevicePointer(void ** pDevice,void * pHost,unsigned flags);  //这里将flags设置为0

// 多线程环境的上下文管理, 每个线程独立地与 GPU 交互。问了避免上下文冲突, 需要确保每个线程正确的管理自己的 CUDA 上下文。
// 每个线程调用 cudaSetDevice(device), 设置其使用的 GPU 设备, 其实就是会把对应 Device 的 primary context 设置为栈顶 context。
#include <cuda_runtime.h>
#include <thread>
#include <vector>
#include <iostream>

void gpuTask(int device) {
    cudaSetDevice(device);

    // 在该线程中执行 GPU 操作
    // 分配内存、启动内核等
    // ...

    // 线程结束时，CUDA 资源会自动释放
}

int main() {
    int deviceCount = 0;
    cudaGetDeviceCount(&deviceCount);

    std::vector<std::thread> threads;
    for (int device = 0; device < deviceCount; ++device) {
        threads.emplace_back(gpuTask, device);
    }

    // 等待所有线程完成
    for (auto& t : threads) {
        t.join();
    }

    return 0;
}

// 同一个线程管理多个设备, 让每个设备执行不同的任务。但是上下文切换会带来性能开销，应尽量减少切换次数。
// 上下文创建和销毁成本高：应尽量重用上下文，避免频繁的创建和销毁操作。
// cuDeviceGet(CUdevice *device, int ordinal) 返回 ordinal 指定的设备的函数句柄。
// cuCtxCreate() 是宏定义, cuCtxCreate(CUcontext *pctx, unsigned int flags, CUdevice dev), 创建 a CUDA Context。
#include <cuda.h>
#include <iostream>

int main() {
    CUdevice cuDevice1, cuDevice2;
    CUcontext cuContext1, cuContext2;

    cuInit(0);
    cuDeviceGet(&cuDevice1, 0);
    cuDeviceGet(&cuDevice2, 1);

    cuCtxCreate(&cuContext1, 0, cuDevice1);
    cuCtxCreate(&cuContext2, 0, cuDevice2);

    // 切换到上下文1
    cuCtxSetCurrent(cuContext1);
    // 在设备1上执行操作
    // ...

    // 切换到上下文2
    cuCtxSetCurrent(cuContext2);
    // 在设备2上执行操作
    // ...

    // 清理资源
    cuCtxDestroy(cuContext1);
    cuCtxDestroy(cuContext2);

    return 0;
}

// 想要使用 CUDA Driver API, 必须包含头文件 <cuda.h>,  <cuda_runtime.h> 是运行时 API
// 显示的上下文管理, 通常用在一个设备上需要多个上下文, 而 CUDA Runtime API 隐式的管理上下文不太满足要求。
// CUDA Driver API 显示地创建和管理上下文, 提供更细粒度的控制。
// cuInit(0); 初始化 CUDA Driver API, 必须在其他 CUDA Driver API 调用之前调用, 其中的 flags 参数必须为 0。
// cuCtxCreate 显示创建 CUDA Context, 注意每个线程是通过 Context stack 来管理的; cuCtxDestroy 销毁上下文, 释放资源。
#include <cuda.h>
#include <iostream>

int main() {
    CUdevice cuDevice;
    CUcontext cuContext;
    int deviceCount = 0;

    cuInit(0);
    cuDeviceGetCount(&deviceCount);

    for (int device = 0; device < deviceCount; ++device) {
        cuDeviceGet(&cuDevice, device);
        cuCtxCreate(&cuContext, 0, cuDevice);

        // 使用 Driver API 进行操作，如内存分配、内核加载等
        // ...

        cuCtxDestroy(cuContext);
    }

    return 0;
}

// 1.
// Cooperative Groups 中的基本数据类型是 thread_group, 是一组线程的函数句柄。该函数句柄只能由它组内的线程成员访问。
// unsigned size();  // 获取线程组中线程的个数
// thread_rank();    // 组中查找调用线程的索引 (介于 0 到 size()-1)
// 线程组提供了在组中所有线程之间执行集体操作的能力(分成线程组的意义), 集体操作是指需要在一组指定的线程之间进行同步或者通信的操作, 最简单的集体操作就是一个屏障, 不传输任何数据, 只是同步线程组中的线程。
// 第一个集体操作调用 cooperative_groups::sync() 来同步线程组, 此时组内的所有线程都会进行一个同步。g.sync(); cg::syncthronize(g);

using namespace cooperative_groups;
using cooperative_groups::thread_group;

// 开发者可以自由指定多少个线程为一组进行规约求和, 如果线程组的大小为 warpSize, 则等价于 warpReduceSum。如果线程组的大小为 block size, 就等价于 blockReduceSum
__device__ int reduce_sum(thread_group g, int *temp, int val)
{
    int lane = g.thread_rank();  // 得到线程组中线程的索引序号

    // i 的初始值为线程组中线程的个数, val 对每个线程私有
    for (int i = g.size() / 2; i > 0; i /= 2)
    {
        // 这种思路是利用 shared_memory 做一个线程间信息交换的介质
        temp[lane] = val;
        g.sync();  // wait for all threads to store
        if(lane<i) val += temp[lane + i];
        g.sync();  // wait for all threads to load
    }
    return val;  // note: only thread 0 will return full sum
}

// 2.
// thread blocks。Cooperative Groups 引入了新的数据类型 thread_block, thread_block block = this_thread_block(); 每个线程都会得到线程块的实列。
// 其实线程块组的意思是一个 block 中的所有线程分成一个组。
// __syncthreads();  block.sync();  cg::synchronize(block);  都是同一个意思

// 3.
// Partitioning Groups。Cooperative Groups 为开发者提供了通过划分现有组来创建新组的方法。cg::tiled_partition() 就负责将一个线程块组划分成多个组。
// thread_group tile32 = cg::partition(this_block_thread(), 32);  // 每个线程都将获得一个 32 线程组的函数句柄。

// 4.
// Cooperative Groups 的真正用途在于模块化, 将组做为显示参数传递给函数, 利用模块化统一接口。有效的减少了竞争和死锁的情况。
// 每个 block 中只有一部分的线程会调用 sum() 函数, 并非块中的所有线程都会达到 __syncthreads(), 因此存在死锁情况。
__device__ int sum(int *x, int n) 
{
    ...
    __syncthreads();
    ...
    return total;
}

__global__ void parallel_kernel(float *x, int n)
{
    if (threadIdx.x < blockDim.x / 2)
        sum(x, count);  // error: half of threads in block skip
                        // __syncthreads() => deadlock
}

// 使用协作组之后
__device__ int sum(thread_block block, int *x, int n) 
{
    ...
    block.sync();
    ...
    return total;
}

__global__ void parallel_kernel(float *x, int n)
{
    sum(this_thread_block(), x, count); // no divergence around call
}

// 5.
// Cooperative Groups 提供了 cg::tiled_partition() 的一个替代版本，它将分组大小作为模板参数，返回一个静态大小的组，称为 thread_block_tile. 在编译时了解分组大小可以提供更好的优化机会。
thread_block_tile<32> tile32 = tiled_partition<32>(this_thread_block());
thread_block_tile<4>  tile4  = tiled_partition<4> (this_thread_block());

// thread Block Tiles 也提供了一些 warp-level 原语操作。
// .shfl()  .shfl_down()  .shfl_up()  .shfl_xor()

// CPU 版本的实现
void softmax_cpu(float* input, float* output, const int M, const int N) {
    // 一行一行的处理
    for (int m = 0; m < M; ++m) {
        // 先获取这一行的最大值
        float maxval = -INFINITY;
        const float* x = input + m * N;
        for (int n = 0; n < N; ++n) {
            maxval = x[n] > maxval ? x[n] : maxval;
        }

        // 依据最大值获取这行的指数和
        float s = 0.0f;
        for (int n = 0; n < N; ++n) {
            s += exp(x[n] - maxval);
        }

        // 依据这行的最大值和指数和得到输出
        float* y = output + m * N;
        for (int n = 0; n < N; ++n) {
            y[n] = exp(x[n] - maxval) / s;
        }
    }
}

// online 的 softmax cpu 版本
void online_softmax_cpu(float* input, float* output, const int M, const int N) {
    // 一行一行的处理
    for (int m = 0; m < M; ++m) {
        const float* x = input + m * N;
        // 两个循环合并成一个(合并求最大值和求指数和)
        float maxval = -INFINITY;  // 这里初始化为 负无穷也是有讲究的, 对于第一次计算 s, 前一项为 0
        float s = 0.0f;
        for (int n = 0; n < N; ++n) {
            float bigger = std::max(maxval, x[n]);          // 历史最大值和当前值选个大的
            s = s * exp(maxval-bigger) + exp(x[n]-bigger);  // 历史最大值是否更迭, 都能用这个公式。假如没有更迭, 前一项就是 s, 更迭的话前一项就会更新, 后一项是没有任何问题的。
            maxval = bigger;                                // 始终每一步都取得最大值
        }

        // 第二个循环, 依据最大值和指数和输出得到结果
        float* y = output + m * N;
        for (int n = 0; n < N; ++n) {
            y[n] = exp(x[n] - maxval) / s;
        }
    }
}

// blockSize = 128, 每个 warp 处理一个行。
// dim3 grid((M*32 + blockSize-1)/blockSize, 1, 1);  dim3 block(blockSize, 1, 1);
__global__ void online_softmax_kernel3(float* __restrict__ input, float* __restrict__ output, const int M, const int N) {
    const int tid = threadIdx.x;                         // 块中的线程索引 id
    const int warpId = tid / warpSize;                    // 此线程位于块内的第几个 warp
    const int laneId = tid % warpSize;                    // 此线程位于一个 warp 中的线程索引
    const int warpsPerBlock = blockDim.x / warpSize;      // 一个 block 的 warp 个数
    const int numWarps = warpsPerBlock * gridDim.x;       // 一个 grid 中总的 warp 个数
    const int idx = warpsPerBlock * blockIdx.x + warpId;  // 当前这个 warp 的 id

    for (int m = idx; m < M; m += numWarps) {  // 这里其实就已经把可能多分配的 warp 给排除了
        // 1) 找对每个 warp 对应的输入和输出的位置
        const float* x = input + m * N;
        float* const y = output + m * N;

        // 2) 先将值汇总到每个线程的寄存器中。经过这里的步骤之后, 每个线程中保留了其负责的那部分数据的最大的值和指数和。
        float maxval = -INFINITY, sum = 0.0f;  // 
        for (int i = laneId; i < N; i += warpSize) {
            float bigger = fmaxf(maxval, x[i]);  // 得到更大的值
            sum = sum * expf(maxval - bigger) + expf(x[i] - bigger);  // 这里就没有判断了, 直接做就可以啦
            maxval = bigger;
        }

        // 3) warp 内的规约, 因为 warp 中的每个线程都保留了其负责那部分数据的最大值和指数和, 这里要对这些部分进行规约。
        for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
            // 在一个 warp 内部强制同步，确保在调用该函数之前的所有指令已经被所有参与同步的线程执行完毕，随后这些线程才继续执行后续的代码。它类似于 CPU 编程中的屏障（barrier），但仅限于 warp 内部。通常用于一个 warp 中有条件分支, 可能会有某些线程会提前执行到某一行代码, 而其他线程还没有执行完之前的部分。
            __syncwarp();  // warp 内的同步操作
            // 先得到与自己对应的那个线程的最大值和指数和。
            float offsetMax = __shfl_xor_sync(0xFFFFFFFF, maxval, offset);  // 得到其他局部块的最大值
            float offsetSum = __shfl_xor_sync(0xFFFFFFFF, sum, offset);     // 得到其他局部块的指数和
            // 一定是根据这两个线程对应的数据, 更新最大值和指数和
            if (maxval < offsetMax) {  // 如果本线程的比另一个线程的最大值小, 更新自己的最大值和指数和
                sum *= expf(maxval - offsetMax);  // 将自己的 sum 乘以 exp
                maxval = offsetMax;               // 将自己的 maxval 更新
            } else {  // 本线程的比另一个线程的最大值大, 更新另一个线程指数和
                offsetSum *= expf(offsetMax - maxval);  // 更新另一个块的
            }
            sum += offsetSum;  // 最后肯定是 更新自己块+其他块   或  自己块+更新其他块
        }

        // 4) 这样操作一遍, warp 中的每个线程的最大值和指数和都是全局的了, 直接遍历一次得到最终结果
        for (int i = laneId; i < N; i += warpSize) {
            y[i] = expf(x[i] - maxval) / sum;
        }
    }
}

// 每个 block 处理一行, dim3 grid(M, 1, 1);  dim3 block(blockSize, 1, 1);
// shared_memory 的大小为 numwarp*2*sizeof(float), 前 numwarp 个 float 存放的是每个 warp 对应数据的最大值;
// 后 numwarp 个 float 存放的是每个 warp 对应的数据的指数和
__global__ void online_softmax_kernel4(float*__restrict__ input, float*__restrict__ output, const int M, const int N) {
    // further, each block handles one row
    if (blockIdx.x < M) {
        extern __shared__ float shared[];
        const int laneId = threadIdx.x % warpSize;                // warp 中第几个线程
        const int warpId = threadIdx.x / warpSize;                // block 中第几个 warp
        // 计算每个 block 中的 warp 个数和每个 warp 负责的数据总个数
        const int warpsPerBlock = ceilDiv(blockDim.x, warpSize);  // 每个 block 中 warp 的个数
        const int dataPerWarp = ceilDiv(N, warpsPerBlock);        // 每个 warp 负责的数据总个数
        // 得到每个 warp 负责数据的起始位置和终止位置
        const int start = dataPerWarp * warpId;                   // 每个 warp 负责的数据开始的位置
        const int end = min((warpId + 1) * dataPerWarp, N);       // 每个 warp 负责的数据结束的位置
        const float* x = input + blockIdx.x * N;
        float* const y = output + blockIdx.x * N;

        // 获取最大值起始位置的共享内存
        float* const maxVals = shared;                  // 先存放每个 warp 的最大值
        // 获取指数和起始位置的共享内存
        float* const sumVals = shared + warpsPerBlock;  // 再存放每个 warp 的指数和

        // 每个 warp 先将其负责的部分数据归约到自己的寄存器中
        float maxval = -INFINITY, sumval = 0.f;
        for (int i = start + laneId; i < end; i += warpSize) {
            float newLaneMax = fmaxf(maxval, x[i]);  // 每检索一个 x, 就求一次最大值
            sumval = sumval * expf(maxval - newLaneMax) + expf(x[i] - newLaneMax);  // 依据最大值进行更新
            maxval = newLaneMax;
        }

        // 每个 warp 内进行规约, 得到 warp 内最大值和指数和, 其实这里就没有是求 max 和 求 sum 进行同时处理
        float maxvalOld = maxval;            // 保留每个线程中旧的 maxVal
        maxval = warpReduceMax(maxval);      // warp 内进行 MaxReduce, 得到一个 warp 中的最大的 maxval
        sumval *= expf(maxvalOld - maxval);  // 更新每个线程中指数和
        sumval = warpReduceSum(sumval);      // 得到一个 warp 内指数和的 Reduce

        // 将warp 中的数据放到共享内存中, 为了后续的 warp 间的规约
        if (laneId == 0) {
            maxVals[warpId] = maxval;
            sumVals[warpId] = sumval;
        }
        __syncthreads();
        // re-init for the next block reduce process

        // 在一个 block 中, 将每个 warp 内的结果放倒共享内存中
        if (warpId == 0) {
            // 通过 shared_memory 将数据取到第 0 个 warp 的线程的寄存器中
            if (laneId < warpsPerBlock){
                maxval = maxVals[laneId];
                sumval = sumVals[laneId];
            } else {
                maxval = -INFINITY;  // 没有的默认值
                sumval = 0.f;
            }
            // 对第 0 warp 进行再次规约操作, 这时候 Max 和 sum 就是同时来做了
            for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
                float otherMax = __shfl_xor_sync(0xFFFFFFFF, maxval, offset);
                float otherSum = __shfl_xor_sync(0xFFFFFFFF, sumval, offset);
                if (maxval < otherMax) {
                    sumval = sumval * expf(maxval - otherMax) + otherSum;
                    maxval = otherMax;
                } else {
                    sumval += otherSum * expf(otherMax - maxval);
                }
            }
        }
        __syncthreads();

        // 利用共享内存做为中间介质, 将整个 block 的最大值和指数和发送给每个线程
        // update the blockwise sum and max to smem
        if (warpId == 0 && laneId == 0) {
            maxVals[0] = maxval;
            sumVals[0] = sumval;
        }
        __syncthreads();

        // 将值传递给每个线程
        // warps take the blockwise max and sum and do the final calculation
        float blockSum = sumVals[0], blockMax = maxVals[0];
        for (int i = start + laneId; i < end; i += 32) {
            y[i] = expf(x[i] - blockMax) / blockSum;
        }
    }
}

// batchSize 为 2, seq_len 为 64, num_head 为 32, head_dim 为 128
#define NUM_PAIRS 64  // head_dim / 2
#define SEQ_LEN 64

// Precomputed θ_i values stored in constant memory, 存储角度值
__constant__ float theta[NUM_PAIRS];

// Precomputed cosine and sine values for θ_i * p, 存储 m*theta 的 sin 和 cos 值, 形状是 [seq_len, num_pairs], 先存 seq_len 的所有 num_pair
__constant__ float cos_theta_p[SEQ_LEN * NUM_PAIRS];
__constant__ float sin_theta_p[SEQ_LEN * NUM_PAIRS];

__global__ void RoPE(float *x, int batchSize, int seq_len, int num_head, int head_dim) {
    // Each block processes one (batch_id, seq_id, head_id)
    int batch_id = blockIdx.z;
    int seq_id = blockIdx.y;
    int head_id = blockIdx.x;
    int tid = threadIdx.x;

    int num_pairs = head_dim / 2;
    int total_threads = blockDim.x;    // 一个 block 中的总的线程个数

    int p = seq_id;
    int theta_offset = p * num_pairs;  // 找 theta 的偏移

    for (int i = tid; i < num_pairs; i += total_threads) {
        int idx = 2 * i;

        // Retrieve precomputed cosine and sine values
        float cos_theta = cos_theta_p[theta_offset + i];
        float sin_theta = sin_theta_p[theta_offset + i];

        // Compute offset
        int offset = batch_id * seq_len * num_head * head_dim 
                        + seq_id * num_head * head_dim + 
                            head_id * head_dim + idx;

        float x1 = x[offset];
        float x2 = x[offset + 1];

        // Apply rotation
        float x1_new = x1 * cos_theta - x2 * sin_theta;
        float x2_new = x1 * sin_theta + x2 * cos_theta;

        // Write back
        x[offset] = x1_new;
        x[offset + 1] = x2_new;
    }
}

#include <cuda_runtime.h>
#include <iostream>

__constant__ float constData[256];  // 定义常量内存

__global__ void kernel() {
    // 设备端代码可以读取常量内存
    printf("Thread %d: constData[0] = %f\n", threadIdx.x, constData[0]);
}

int main() {
    float hostData[256];
    
    // 初始化数据
    for (int i = 0; i < 256; ++i) {
        hostData[i] = static_cast<float>(i);
    }

    // 将数据从主机端拷贝到设备端的常量内存
    cudaMemcpyToSymbol(constData, hostData, sizeof(hostData), 0, cudaMemcpyHostToDevice);

    // 启动内核，读取常量内存中的数据
    kernel<<<1, 10>>>();
    cudaDeviceSynchronize();

    // 从常量内存拷贝回主机
    float resultData[256];
    cudaMemcpyFromSymbol(resultData, constData, sizeof(resultData), cudaMemcpyDeviceToHost);

    // 打印拷贝回来的数据，确认是否正确
    std::cout << "Data copied from constant memory back to host: " << resultData[0] << std::endl;

    return 0;
}

#include <cuda_runtime.h>

// 输入矩阵的大小为 (B, T, C) = (8, 1024, 768)  B 为 batchSize, T 为 sequence 的长度, C 是 embedding 的维度
constexpr int B = 8;
constexpr int T = 1024;
constexpr int C = 768;
constexpr int eps = 1e-5;
constexpr int blockSize = 128;

// 这个没什么心意, 就是注意这里可以进行 unroll
template <typename T1>
__device__ T1 warpReduceSum(T1 val) {
    #pragma unroll
    for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    }
    return val;
}

// 1. 利用了个公式 Var(X) = E(x^2) - E(X)^2, 这样就能节省一个循环
// 2. 每个 warp 处理一行
// 3. 因为 输入数据会被多次用到, 所以利用 shared_memory 来进行缓存输入数据
__global__ void layernorm_forward_kernel6(float *input, float *out, float *weight,
                                          float *bias, float eps, int B, int T, int C) {
    // one warp one row
    // use smem to store the shift (x - mean) values
    // use D(X) = E(X^2) - E(X)^2
    assert((C % warpSize) == 0);
    extern __shared__ float sharedX[];
    int warpsPerBlock = blockDim.x / warpSize;           // 每个 block 中的 warp 个数
    int warpId = threadIdx.x / warpSize;                 // 当前线程的 warp id
    int laneId = threadIdx.x % warpSize;                 // 当前线程的 lane id
    int numWarps = gridDim.x * warpsPerBlock;            // grid 的所有 warp 个数
    float *const xSharedWarp = sharedX + warpId * C;     // 当前 warp 用到的 SM 的起始地址

    for (int row = blockIdx.x * warpsPerBlock + warpId; row < B * T; row += numWarps)
        if (row < B * T) {
            // 定位到当前 warp 处理的全局内存的起始地址
            float *const x = input + row * C;
            float *const y = out + row * C;

            float partialSum = 0.0f, partialSum2 = 0.0f;
            for (int i = laneId; i < C; i += warpSize) {
                float xi = x[i];     // 先从共享内存读取到寄存器中, 这样再从寄存器到其他的寄存器和 SM 就快了
                xSharedWarp[i] = xi;
                partialSum += xi;
                partialSum += xi * xi;
            }

            float mean = warpReduceSum(partialSum) / C;    // warp 规约求 E(X)
            float mean2 = warpReduceSum(partialSum2) / C;  // warp 规约求 E(X^2)
            float var = (mean2 - mean * mean);             // 依据公式求得方差
            float inv_std = 1.0f / sqrt(var / C + eps);

            // 写到输出中
            for (int i = laneId; i < C; i += warpSize) {
                y[i] = weight[i] * (sharedX[i] - mean) * inv_std + bias[i];
            }
        }
}

int main(){

    float *inputGPU = nullptr;
    float *outputGPU = nullptr;
    float *weightGPU = nullptr;
    float *biasGPU = nullptr;
    cudaMalloc(&inputGPU, B * T * C * sizeof(float));
    cudaMalloc(&outputGPU, B * T * C * sizeof(float));
    // weight 和 bias 都是在 C 这个维度展开的
    cudaMalloc(&weightGPU, C * sizeof(float));
    cudaMalloc(&biasGPU, C * sizeof(float));

    const int smemSize = blockSize / 32 * C * sizeof(float);  // 每个 warp 处理一个 C, 那就需要 一个block中 warp 个数*C*sizeof(float)
    layernorm_forward_kernel6<<<B * T * 32 / blockSize, blockSize, smemSize>>>(inputGPU, outputGPU, weightGPU, biasGPU, eps, B, T, C); 

    return 0;
}

// Dropout kernel function
__global__ void DropoutKernel(const float* __restrict__ x,
                              uint8_t* __restrict__ mask,
                              float* __restrict__ y,
                              size_t n,
                              float dropout_prob,
                              bool is_upscale_in_train,
                              unsigned int seed) {

    // Calculate global thread ID
    int idx = blockIdx.x * blockDim.x + threadIdx.x;  // 计算线程的 id
    int stride = gridDim.x * blockDim.x;              // 所有线程的总数

    // Initialize the random number generator state
    curandStatePhilox4_32_10_t state;
    curand_init(seed, idx, 0, &state);

    float inv_prob = 1.0f / (1.0f - dropout_prob);  // 提前计算出来

    // Loop over the elements, processing four elements per thread for efficiency
    for (size_t i = idx * 4; i < n; i += stride * 4) {
        // Generate 4 random numbers
        float4 rand = curand_uniform4(&state);  // 一次得到 float4 的随机数

        // Load 4 input values, 因为输入不一定是 4 的整数倍, 所以可以直接展开, 超的直接置为 0 就可以了
        float4 x_val;
        x_val.x = (i + 0 < n) ? x[i + 0] : 0.0f;
        x_val.y = (i + 1 < n) ? x[i + 1] : 0.0f;
        x_val.z = (i + 2 < n) ? x[i + 2] : 0.0f;
        x_val.w = (i + 3 < n) ? x[i + 3] : 0.0f;

        // Initialize output values and masks, 保存输出值和mask
        float4 y_val;
        uint8_t m_val[4];

        // Apply dropout, 小于 dropout 概率, 都置为 0, 还需要依据 is_upscale_in_train 来判断是否需要缩放
        y_val.x = (rand.x >= dropout_prob) ? (is_upscale_in_train ? x_val.x * inv_prob : x_val.x) : 0.0f;
        y_val.y = (rand.y >= dropout_prob) ? (is_upscale_in_train ? x_val.y * inv_prob : x_val.y) : 0.0f;
        y_val.z = (rand.z >= dropout_prob) ? (is_upscale_in_train ? x_val.z * inv_prob : x_val.z) : 0.0f;
        y_val.w = (rand.w >= dropout_prob) ? (is_upscale_in_train ? x_val.w * inv_prob : x_val.w) : 0.0f;

        // Set mask values, 存储 mask
        m_val[0] = (rand.x >= dropout_prob);
        m_val[1] = (rand.y >= dropout_prob);
        m_val[2] = (rand.z >= dropout_prob);
        m_val[3] = (rand.w >= dropout_prob);

        // Store the results back to global memory, 展开着写回, 防止超越边界
        if (i + 0 < n) {
            y[i + 0] = y_val.x;
            mask[i + 0] = m_val[0];
        }
        if (i + 1 < n) {
            y[i + 1] = y_val.y;
            mask[i + 1] = m_val[1];
        }
        if (i + 2 < n) {
            y[i + 2] = y_val.z;
            mask[i + 2] = m_val[2];
        }
        if (i + 3 < n) {
            y[i + 3] = y_val.w;
            mask[i + 3] = m_val[3];
        }
    }
}

// curand_init(seed, thread_id, offset, &state) 是 cuRAND 中用于初始化随机数生成器状态的函数, 它为每个线程生成独立的随机数序列。
// seed: 用于初始化随机数生成器的种子，它决定了生成的随机数序列的起点。相同的种子会生成相同的随机数序列
// thread_id: 是用于区分不同线程的标识符，它确保每个线程生成不同的随机数序列。通常将 thread_id 设为每个线程的唯一标识符（如 threadIdx.x + blockIdx.x * blockDim.x），确保每个线程拥有独立的随机数生成器状态，从而生成不同的随机数。
// offset: 用于偏移随机数序列的起点，它决定了每个线程从随机数序列中的哪个位置开始生成随机数。用于在同一个线程中跳过随机数序列的前几个数。通常可以设置为 0。
// state: state 是一个指向 curandState 类型的指针，表示线程的随机数生成器状态。

// Calculate global thread ID
int idx = blockIdx.x * blockDim.x + threadIdx.x;  // 计算线程的 id
int stride = gridDim.x * blockDim.x;              // 所有线程的总数

// Initialize the random number generator state
curandStatePhilox4_32_10_t state;
curand_init(seed, idx, 0, &state);
float rand = curand_uniform(&state);    // 生成 0-1 之间的均匀分布
float4 rand = curand_uniform4(&state);  // 生成四个