最近邻 2 倍上采样

// ============================================================================
// upsample_nearest_2d —— OneFlow 的 2 倍最近邻上采样前后向 kernel：
//   用 Pack2X<T> 把"一次搬 1 个元素"改成"一次搬水平相邻的 2 个元素"
// ============================================================================#
// include <cuda.h>          // driver API，本文件没用到（kernel 启动走的是 runtime）
#include <cuda_runtime.h>  // 真正需要的：__global__、cudaMalloc/cudaMemcpy、dim3
#include <time.h>          // 【未使用】main 里没有计时代码，耗时是 README 里用 ncu 外部测的
#include <algorithm>       // 【未使用】
#include <iostream>        // 【未使用】打印用的是 printf（随 cuda_runtime 间接引入）
#include <cuda_fp16.h>     // 本文件 main 只实例化 float，这个 include 是给 T=__half 准备的
using namespace std;       // 【本文件其实没用到 std 里的东西】reinterpret_cast 是关键字，不是 std 成员
#define N 32*1024*1024   // 输入元素个数 33,554,432（= 32 个 1024x1024 平面）；输出是 4N
#define kBlockSize 256   // 每 block 线程数


#define CUDA_1D_KERNEL_LOOP(i, n)                                                                 \
  for (int32_t i = blockIdx.x * blockDim.x + threadIdx.x, step = blockDim.x * gridDim.x; i < (n); \
       i += step)  // 【循环体末尾的 i += step】就是"跳到下一个归本线程的元素"


template<typename T>
struct alignas(2 * sizeof(T)) Pack2X {
    T x;
    T y;
};


// 【正向】线程按**输入**下标划分：每线程读 1 个 in 元素，把它复制成输出上的一个 2x2 块。
// 4 个输出位置 = 2 个 Pack2X（上下两行各一个），所以是 2 条宽存储。 
// 【形参】in_elem_cnt：输入元素总数（= N）；in_height/in_width：输入的 H、W（输出是 2 倍）；
//   in_dptr：输入显存地址（只读）；out_dptr：输出显存地址（4 倍元素数，写）。
template<typename T>
__global__ void UpsampleNearest2D2XForward(const int32_t in_elem_cnt, const T* in_dptr,
                                           const int32_t in_height, const int32_t in_width,
                                           T* out_dptr){
  const int32_t in_hw_size = in_width * in_height;  // 一个 (n,c) 平面的输入元素数；输出平面是它的 2 倍（pack 数）
  CUDA_1D_KERNEL_LOOP(index, in_elem_cnt) {        // index 是**输入**的线性下标：index = ((n*C + c)*H + h)*W + w
    const T in_value = in_dptr[index];             // 1 条标量读（float 时 PTX 是 ld.global.b32）
    const int32_t nc_idx = index / in_hw_size;     // 第几个 (n,c) 平面。除数是运行期形参 → 走倒数近似，见文件头"整数除法的代价"
    const int32_t hw_off = index - nc_idx * in_hw_size;  // 平面内偏移；用商反推余数，省掉一次 `%`
    const int32_t h = hw_off / in_width;           // 输入行号（第二次除法）
    const int32_t w = hw_off - h * in_width;       // 输入列号（同样用商反推余数）
    Pack2X<T> out_value{in_value, in_value};       // 造出一个"两个相同值"的 pack：{x,y} = {v,v}
    Pack2X<T>* out_pack_dptr = reinterpret_cast<Pack2X<T>*>(out_dptr);  // 换成 pack 视图（输出的一行有 2*in_width 个元素 = in_width 个 pack，(n,c) 平面的 pack 数是 2*in_hw_size）；要求 out_dptr 8B(float)/4B(half) 对齐，见文件头"main 的测例…"第 3 条
    out_pack_dptr[nc_idx * in_hw_size * 2 + h * 2 * in_width + w] = out_value;       // 写输出的第 2h 行：覆盖 (2h,2w) 和 (2h,2w+1)。一个 warp 内 lane 只改变 w，32 个地址连续（float 时 256B 完整事务）
    out_pack_dptr[nc_idx * in_hw_size * 2 + (h * 2 + 1) * in_width + w] = out_value; // 写输出的第 2h+1 行：覆盖 (2h+1,2w) 和 (2h+1,2w+1)。与上一条相隔 2*in_width 个元素（in_width=1024 时 8KB），也是独立的完整事务
  }
}



// 【反向】梯度回传，线程按**输入（dx，小图）**下标划分：每线程收自己那 2x2 块的 4 个 dy。
// 4 个 dy = 2 个 Pack2X（上下两行各一个），所以是 2 条宽读；dx 只有 1 个值，标量写。
// 【形参】命名沿用正向的套路但含义换了：in_elem_cnt 这里指 dx 的元素数（= 正向的输入数），
//   dy_dptr 是上游梯度（大图，4 倍），dx_dptr 是要算的输入梯度（小图）。
template<typename T>
__global__ void UpsampleNearest2D2XBackward(const int32_t in_elem_cnt, const T* dy_dptr, 
                                            const int32_t dx_height, const int32_t dx_width,
                                            T* dx_dptr) {
  const int32_t dx_hw_size = dx_height * dx_width;
  CUDA_1D_KERNEL_LOOP(index, in_elem_cnt) {
    T dx_value = 0.0;
    const int32_t nc_idx = index / dx_hw_size;     // 与正向同样的两次除法
    const int32_t dx_hw_off = index - nc_idx * dx_hw_size;
    const int32_t dx_h = dx_hw_off / dx_width;
    const int32_t dx_w = dx_hw_off - dx_h * dx_width;
    const Pack2X<T>* dy_pack_dptr = reinterpret_cast<const Pack2X<T>*>(dy_dptr);  // dy 用同一个 pack 视图，因为正向就是这么写下去的
    const Pack2X<T> dy_pack_value1 =
        dy_pack_dptr[nc_idx * dx_hw_size * 2 + dx_h * 2 * dx_width + dx_w];  // dy 第 2*dx_h 行的那一个 pack：含 (2h,2w) 与 (2h,2w+1) 两个梯度
    const Pack2X<T> dy_pack_value2 =
        dy_pack_dptr[nc_idx * dx_hw_size * 2 + (dx_h * 2 + 1) * dx_width + dx_w];  // dy 第 2*dx_h+1 行：含 (2h+1,2w) 与 (2h+1,2w+1)
    dx_value += dy_pack_value1.x;  // 4 项相加：正向里同一个 in_value 复制到 4 个位置，
    dx_value += dy_pack_value1.y;  // 反向按链式法则就是把这 4 个位置的梯度加起来。这里是普通
    dx_value += dy_pack_value2.x;  // 浮点加法（不是 atomicAdd）：4 项都在本线程手里，
    dx_value += dy_pack_value2.y;  // 别的线程不会碰同一个 dx 槽位
    dx_dptr[index] = dx_value;     // 标量写（float 时 PTX 是 st.global.b32）
  }
}


// 【测例】1024x1024 的输入、32 个 (n,c) 平面、倍率 2，只跑正向并打印前 50 个输出值。
int main(){
    float *input_host = (float*)malloc(N*sizeof(float));    // 主机输入缓冲 128MB
    float *input_device;
    cudaMalloc((void **)&input_device, N*sizeof(float));    // 设备输入缓冲 128MB；cudaMalloc 的地址 256B 对齐，满足 Pack2X<float> 的 8B 要求
    for (int i = 0; i < N; i++) input_host[i] = 1.0;       // 【全 1 数据】只能验证不崩/不越界，验证不了下标映射（见文件头"main 的测例…"第 2 条）
    cudaMemcpy(input_device, input_host, N*sizeof(float), cudaMemcpyHostToDevice);

    float *output_host = (float*)malloc(N * 4 * sizeof(float));  // 输出 4 倍元素数 = 512MB
    float *output_device;
    cudaMalloc((void **)&output_device, N * 4 * sizeof(float));

    // 【启动配置】grid = 33,554,432/256 = 131,072 个 block，线程总数恰好 = N，
    //   于是 grid-stride 循环每线程只走一圈。本次 kernel 的访存量：读 128MB + 写 512MB = 640MB，
    //   写读比 4:1 —— 属于纯带宽受限，所以"每条指令搬几字节"直接决定性能。
    dim3 grid(N / kBlockSize, 1);
    dim3 block(kBlockSize, 1);
    UpsampleNearest2D2XForward<<<grid, block>>>(N, input_device, 1024, 1024, output_device);  // 显式传 1024x1024；N/1024/1024 = 32 个平面
    cudaMemcpy(output_host, output_device, N * 4 * sizeof(float), cudaMemcpyDeviceToHost);
    for(int i = 0; i < 50; i++) {
        printf("%.5f\n", output_host[i]);   // 本机（RTX 5070 / sm_120）实测打印 50 行 1.00000，与原版一致
    }

    // 【原样保留】host 的 malloc 与 device 的 cudaMalloc 都没有释放，进程退出后由 OS 回收。
    // 反向 kernel UpsampleNearest2D2XBackward 这里没调用，也就不会被实例化进二进制。
    return 0;
}