kernels_v2.cu

#include "kernels_v2.h"
#include "kernels.h"
#include <cstdio>
#include <cuda_runtime.h>
#include <stdint.h>

#ifndef CUDA_CHECK
#define CUDA_CHECK(call) do {                                     \
  cudaError_t _e = (call);                                        \
  if (_e != cudaSuccess) {                                        \
    printf("[CUDA] %s:%d: %s\n", __FILE__, __LINE__,              \
           cudaGetErrorString(_e));                               \
  }                                                               \
} while(0)
#endif

// ------------------------------
// Tunables
// ------------------------------
static constexpr int TILE = 16;     // 16x16 output tile
static constexpr int KFIX = 3;      // fixed kernel
static constexpr int PAD  = 1;      // K=3 => pad=1
static constexpr int SHW  = TILE + (KFIX - 1); // 18

// =====================
// Conv2D
// =====================
// ------------------------------
// Optional constant memory weights
// NOTE: constant memory is limited (64KB total). We store weights only if small.
// Max floats we allow in const (leave room): 16384 floats ~ 64KB.
// ------------------------------
static constexpr int MAX_CONST_W = 16384;
__constant__ float c_weights[MAX_CONST_W];

// ---- Kernel: read weights from global memory ----
__global__ void conv2d_fwd_tiled_gmem_k3(const float* __restrict__ input,
                                        float* __restrict__ output,
                                        const float* __restrict__ weights,
                                        const float* __restrict__ bias,
                                        int H, int W, int C_in, int C_out)
{
  __shared__ float s_in[SHW][SHW];

  const int tx = threadIdx.x;
  const int ty = threadIdx.y;

  const int row_o = blockIdx.y * TILE + ty;
  const int col_o = blockIdx.x * TILE + tx;

  const int co = blockIdx.z; // one output channel per z-slice

  float sum = 0.0f;

  // Loop over input channels
  for (int ci = 0; ci < C_in; ++ci) {
    // Load shared tile (with halo)
    const int in_row0 = blockIdx.y * TILE - PAD;
    const int in_col0 = blockIdx.x * TILE - PAD;

    // cooperative load: cover 18x18 with 16x16 threads
    for (int i = ty; i < SHW; i += TILE) {
      for (int j = tx; j < SHW; j += TILE) {
        int r = in_row0 + i;
        int c = in_col0 + j;
        float v = 0.0f;
        if ((unsigned)r < (unsigned)H && (unsigned)c < (unsigned)W) {
          // NHWC/HWC: (r*W + c)*C_in + ci
          v = input[(r * W + c) * C_in + ci];
        }
        s_in[i][j] = v;
      }
    }
    __syncthreads();

    // Compute one output pixel per thread (within bounds)
    if ((unsigned)row_o < (unsigned)H && (unsigned)col_o < (unsigned)W) {
      const int w_base = ((co * C_in + ci) * KFIX) * KFIX; // [co][ci][3][3]
      // Unroll 3x3
      #pragma unroll
      for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
          sum += s_in[ty + kh][tx + kw] * weights[w_base + kh * 3 + kw];
        }
      }
    }

    __syncthreads();
  }

  if ((unsigned)row_o < (unsigned)H && (unsigned)col_o < (unsigned)W) {
    sum += bias[co];
    // NHWC/HWC output: (row_o*W + col_o)*C_out + co
    output[(row_o * W + col_o) * C_out + co] = sum;
  }
}

// ---- Kernel: read weights from constant memory ----
__global__ void conv2d_fwd_tiled_cmem_k3(const float* __restrict__ input,
                                        float* __restrict__ output,
                                        const float* __restrict__ bias,
                                        int H, int W, int C_in, int C_out)
{
  __shared__ float s_in[SHW][SHW];

  const int tx = threadIdx.x;
  const int ty = threadIdx.y;

  const int row_o = blockIdx.y * TILE + ty;
  const int col_o = blockIdx.x * TILE + tx;

  const int co = blockIdx.z;

  float sum = 0.0f;

  for (int ci = 0; ci < C_in; ++ci) {
    const int in_row0 = blockIdx.y * TILE - PAD;
    const int in_col0 = blockIdx.x * TILE - PAD;

    for (int i = ty; i < SHW; i += TILE) {
      for (int j = tx; j < SHW; j += TILE) {
        int r = in_row0 + i;
        int c = in_col0 + j;
        float v = 0.0f;
        if ((unsigned)r < (unsigned)H && (unsigned)c < (unsigned)W) {
          v = input[(r * W + c) * C_in + ci];
        }
        s_in[i][j] = v;
      }
    }
    __syncthreads();

    if ((unsigned)row_o < (unsigned)H && (unsigned)col_o < (unsigned)W) {
      const int w_base = ((co * C_in + ci) * KFIX) * KFIX;
      #pragma unroll
      for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
          sum += s_in[ty + kh][tx + kw] * c_weights[w_base + kh * 3 + kw];
        }
      }
    }

    __syncthreads();
  }

  if ((unsigned)row_o < (unsigned)H && (unsigned)col_o < (unsigned)W) {
    sum += bias[co];
    output[(row_o * W + col_o) * C_out + co] = sum;
  }
}

void launch_conv2d_forward_opt(float* d_input,
                               float* d_output,
                               const float* d_weights,
                               const float* d_bias,
                               int H, int W,
                               int C_in, int C_out,
                               int K)
{
  if (K != 3) {
    printf("[conv_opt] Only K=3 supported in this optimized forward (got K=%d)\n", K);
    return;
  }

  dim3 block(TILE, TILE);
  dim3 grid((W + TILE - 1) / TILE,
            (H + TILE - 1) / TILE,
            C_out);

  const int w_elems = C_out * C_in * 3 * 3;

  // Optional constant-memory path (auto)
  if (w_elems <= MAX_CONST_W) {
    CUDA_CHECK(cudaMemcpyToSymbol(c_weights, d_weights, w_elems * sizeof(float), 0,
                                 cudaMemcpyDeviceToDevice));
    conv2d_fwd_tiled_cmem_k3<<<grid, block>>>(d_input, d_output, d_bias, H, W, C_in, C_out);
  } else {
    conv2d_fwd_tiled_gmem_k3<<<grid, block>>>(d_input, d_output, d_weights, d_bias,
                                              H, W, C_in, C_out);
  }

  CUDA_CHECK(cudaGetLastError());
}

// ============================================================================
// Backward (K=3, NHWC)
//   d_grad_input  shape: [H][W][C_in]
//   d_grad_output shape: [H][W][C_out]
//   weights       shape: [C_out][C_in][3][3]
//   d_grad_weights same as weights
//   d_grad_bias   shape: [C_out]
//
// Strategies:
//  1) shared-memory tiling
//  2) unrolled K=3
//  3) optional constant-memory weights (same c_weights as forward)
//
// Notes:
//  - Optimized for batch=1 (pipeline hiện tại của bạn).
//  - d_grad_weights / d_grad_bias: mỗi block chỉ atomicAdd 9 lần (weights) + 1 lần (bias)
// ============================================================================

// ---- dInput kernel (weights from global memory) ----
__global__ void conv2d_dinput_tiled_gmem_k3(const float* __restrict__ d_go,
                                           float* __restrict__ d_gi,
                                           const float* __restrict__ w,
                                           int H, int W, int C_in, int C_out)
{
  __shared__ float s_go[SHW][SHW];

  const int tx = threadIdx.x;
  const int ty = threadIdx.y;
  const int row_i = blockIdx.y * TILE + ty;
  const int col_i = blockIdx.x * TILE + tx;
  const int ci = blockIdx.z; // one input channel per z-slice

  float sum = 0.0f;

  for (int co = 0; co < C_out; ++co) {
    const int go_row0 = blockIdx.y * TILE - PAD;
    const int go_col0 = blockIdx.x * TILE - PAD;

    // cooperative load 18x18 grad_output tile for this co
    for (int i = ty; i < SHW; i += TILE) {
      for (int j = tx; j < SHW; j += TILE) {
        int r = go_row0 + i;
        int c = go_col0 + j;
        float v = 0.0f;
        if ((unsigned)r < (unsigned)H && (unsigned)c < (unsigned)W) {
          v = d_go[(r * W + c) * C_out + co];
        }
        s_go[i][j] = v;
      }
    }
    __syncthreads();

    if ((unsigned)row_i < (unsigned)H && (unsigned)col_i < (unsigned)W) {
      const int w_base = ((co * C_in + ci) * KFIX) * KFIX;

      // Flip (kh,kw) for dInput:
      // use s_go[ty+(2-kh)][tx+(2-kw)] (because PAD=1)
      #pragma unroll
      for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
          sum += s_go[ty + (2 - kh)][tx + (2 - kw)] * w[w_base + kh * 3 + kw];
        }
      }
    }

    __syncthreads();
  }

  if ((unsigned)row_i < (unsigned)H && (unsigned)col_i < (unsigned)W) {
    d_gi[(row_i * W + col_i) * C_in + ci] = sum;
  }
}

// ---- dInput kernel (weights from constant memory) ----
__global__ void conv2d_dinput_tiled_cmem_k3(const float* __restrict__ d_go,
                                           float* __restrict__ d_gi,
                                           int H, int W, int C_in, int C_out)
{
  __shared__ float s_go[SHW][SHW];

  const int tx = threadIdx.x;
  const int ty = threadIdx.y;
  const int row_i = blockIdx.y * TILE + ty;
  const int col_i = blockIdx.x * TILE + tx;
  const int ci = blockIdx.z;

  float sum = 0.0f;

  for (int co = 0; co < C_out; ++co) {
    const int go_row0 = blockIdx.y * TILE - PAD;
    const int go_col0 = blockIdx.x * TILE - PAD;

    for (int i = ty; i < SHW; i += TILE) {
      for (int j = tx; j < SHW; j += TILE) {
        int r = go_row0 + i;
        int c = go_col0 + j;
        float v = 0.0f;
        if ((unsigned)r < (unsigned)H && (unsigned)c < (unsigned)W) {
          v = d_go[(r * W + c) * C_out + co];
        }
        s_go[i][j] = v;
      }
    }
    __syncthreads();

    if ((unsigned)row_i < (unsigned)H && (unsigned)col_i < (unsigned)W) {
      const int w_base = ((co * C_in + ci) * KFIX) * KFIX;
      #pragma unroll
      for (int kh = 0; kh < 3; ++kh) {
        #pragma unroll
        for (int kw = 0; kw < 3; ++kw) {
          sum += s_go[ty + (2 - kh)][tx + (2 - kw)] * c_weights[w_base + kh * 3 + kw];
        }
      }
    }

    __syncthreads();
  }

  if ((unsigned)row_i < (unsigned)H && (unsigned)col_i < (unsigned)W) {
    d_gi[(row_i * W + col_i) * C_in + ci] = sum;
  }
}

// ---- dW + dB kernel (tiled partial reduction; one (co,ci) per block.z) ----
// grid.z = C_out * C_in
__global__ void conv2d_dweight_dbias_tiled_k3(const float* __restrict__ d_go,
                                             const float* __restrict__ in,
                                             float* __restrict__ d_gw,
                                             float* __restrict__ d_gb,
                                             int H, int W, int C_in, int C_out)
{
  // Decode z -> (co, ci)
  const int z  = blockIdx.z;
  const int co = z / C_in;
  const int ci = z - co * C_in;

  __shared__ float s_in[SHW][SHW];
  // __shared__ float s_go[TILE][TILE];

  const int tx = threadIdx.x;
  const int ty = threadIdx.y;

  const int row_o = blockIdx.y * TILE + ty;
  const int col_o = blockIdx.x * TILE + tx;

  const int in_row0 = blockIdx.y * TILE - PAD;
  const int in_col0 = blockIdx.x * TILE - PAD;

  // Load grad_output tile for this co
  float go_val = 0.0f;
  if ((unsigned)row_o < (unsigned)H && (unsigned)col_o < (unsigned)W) {
    go_val = d_go[(row_o * W + col_o) * C_out + co];
  }
  // s_go[ty][tx] = go_val;

  // Load input tile with halo for this ci
  for (int i = ty; i < SHW; i += TILE) {
    for (int j = tx; j < SHW; j += TILE) {
      int r = in_row0 + i;
      int c = in_col0 + j;
      float v = 0.0f;
      if ((unsigned)r < (unsigned)H && (unsigned)c < (unsigned)W) {
        v = in[(r * W + c) * C_in + ci];
      }
      s_in[i][j] = v;
    }
  }
  __syncthreads();

  // Partial sums for 9 weights
  float p[9];
  #pragma unroll
  for (int t = 0; t < 9; ++t) p[t] = 0.0f;

  if ((unsigned)row_o < (unsigned)H && (unsigned)col_o < (unsigned)W) {
    #pragma unroll
    for (int kh = 0; kh < 3; ++kh) {
      #pragma unroll
      for (int kw = 0; kw < 3; ++kw) {
        p[kh * 3 + kw] = s_in[ty + kh][tx + kw] * go_val;
      }
    }
  }

  // Reduce p[9] across block
  __shared__ float s_red[9][TILE * TILE];
  const int lid = ty * TILE + tx;
  #pragma unroll
  for (int t = 0; t < 9; ++t) s_red[t][lid] = p[t];
  __syncthreads();

  for (int stride = (TILE * TILE) / 2; stride > 0; stride >>= 1) {
    if (lid < stride) {
      #pragma unroll
      for (int t = 0; t < 9; ++t) {
        s_red[t][lid] += s_red[t][lid + stride];
      }
    }
    __syncthreads();
  }

  if (lid == 0) {
    const int w_base = ((co * C_in + ci) * 3) * 3;
    #pragma unroll
    for (int t = 0; t < 9; ++t) {
      atomicAdd(&d_gw[w_base + t], s_red[t][0]);
    }
  }

  // Bias: reduce go tile & atomicAdd once
  __shared__ float s_b[TILE * TILE];
  s_b[lid] = go_val;
  __syncthreads();
  for (int stride = (TILE * TILE) / 2; stride > 0; stride >>= 1) {
    if (lid < stride) s_b[lid] += s_b[lid + stride];
    __syncthreads();
  }
  if (lid == 0) atomicAdd(&d_gb[co], s_b[0]);
}

void launch_conv2d_backward_opt(const float* d_grad_output,
                                const float* d_input,
                                const float* d_weights,
                                float* d_grad_input,
                                float* d_grad_weights,
                                float* d_grad_bias,
                                int H, int W,
                                int C_in, int C_out,
                                int K)
{
  if (K != 3) {
    printf("[conv_bwd_opt] Only K=3 supported (got K=%d)\n", K);
    return;
  }

  dim3 block(TILE, TILE);

  // dInput: grid.z = C_in
  dim3 grid_in((W + TILE - 1) / TILE,
               (H + TILE - 1) / TILE,
               C_in);

  const int w_elems = C_out * C_in * 3 * 3;

  const bool use_cmem = (w_elems <= MAX_CONST_W);
  if (use_cmem) {
    CUDA_CHECK(cudaMemcpyToSymbol(c_weights, d_weights, w_elems * sizeof(float), 0,
                                 cudaMemcpyDeviceToDevice));
    conv2d_dinput_tiled_cmem_k3<<<grid_in, block>>>(d_grad_output, d_grad_input, H, W, C_in, C_out);
  } else {
    conv2d_dinput_tiled_gmem_k3<<<grid_in, block>>>(d_grad_output, d_grad_input, d_weights, H, W, C_in, C_out);
  }
  CUDA_CHECK(cudaGetLastError());

  // dW + dB: grid.z = C_out * C_in
  dim3 grid_w((W + TILE - 1) / TILE,
              (H + TILE - 1) / TILE,
              C_out * C_in);

  conv2d_dweight_dbias_tiled_k3<<<grid_w, block>>>(d_grad_output, d_input,
                                                  d_grad_weights, d_grad_bias,
                                                  H, W, C_in, C_out);
  CUDA_CHECK(cudaGetLastError());
}

// =====================
// Max Pooling
// =====================

// --- helpers ---
static __device__ __forceinline__ float4 f4max(float4 a, float4 b) {
  a.x = (a.x > b.x) ? a.x : b.x;
  a.y = (a.y > b.y) ? a.y : b.y;
  a.z = (a.z > b.z) ? a.z : b.z;
  a.w = (a.w > b.w) ? a.w : b.w;
  return a;
}

static __device__ __forceinline__ uint8_t argmax4(float4 v0, float4 v1, float4 v2, float4 v3,
                                                  float4* out_max) {
  // lane-wise max across 4 candidates; return which candidate won per lane encoded in 2 bits each
  // We'll compute max, and separately winner id per lane.
  float4 m01 = f4max(v0, v1);
  float4 m23 = f4max(v2, v3);
  float4 m   = f4max(m01, m23);

  uint8_t w = 0;
  // x lane
  {
    float mx = m.x;
    uint8_t id = (v0.x == mx) ? 0 : (v1.x == mx) ? 1 : (v2.x == mx) ? 2 : 3;
    w |= (id & 0x3) << 0;
  }
  // y lane
  {
    float my = m.y;
    uint8_t id = (v0.y == my) ? 0 : (v1.y == my) ? 1 : (v2.y == my) ? 2 : 3;
    w |= (id & 0x3) << 2;
  }
  // z lane
  {
    float mz = m.z;
    uint8_t id = (v0.z == mz) ? 0 : (v1.z == mz) ? 1 : (v2.z == mz) ? 2 : 3;
    w |= (id & 0x3) << 4;
  }
  // w lane
  {
    float mw = m.w;
    uint8_t id = (v0.w == mw) ? 0 : (v1.w == mw) ? 1 : (v2.w == mw) ? 2 : 3;
    w |= (id & 0x3) << 6;
  }

  *out_max = m;
  return w;
}

// Forward: each thread handles one (h_out, w_out) and one float4 pack of channels
__global__ void maxpool2x2_forward_nhwc_f4(const float* __restrict__ in,
                                          float* __restrict__ out,
                                          int* __restrict__ indices,
                                          int H, int W, int C,
                                          int H_out, int W_out)
{
  // pack index over C/4
  const int c4 = C >> 2;
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  const int n = H_out * W_out * c4;
  if (tid >= n) return;

  const int pc = tid % c4;          // pack channel id
  const int t  = tid / c4;
  const int w_out_idx = t % W_out;
  const int h_out_idx = t / W_out;

  const int h0 = h_out_idx * 2;
  const int w0 = w_out_idx * 2;

  // base pointers
  const int c_base = pc * 4;

  // four input positions in 2x2 window
  // NHWC index = (h*W + w)*C + c
  const int base00 = (h0 * W + w0) * C + c_base;
  const int base01 = (h0 * W + (w0 + 1)) * C + c_base;
  const int base10 = ((h0 + 1) * W + w0) * C + c_base;
  const int base11 = ((h0 + 1) * W + (w0 + 1)) * C + c_base;

  // vector loads
  const float4 v00 = *reinterpret_cast<const float4*>(in + base00);
  const float4 v01 = *reinterpret_cast<const float4*>(in + base01);
  const float4 v10 = *reinterpret_cast<const float4*>(in + base10);
  const float4 v11 = *reinterpret_cast<const float4*>(in + base11);

  float4 vmax;
  // winner encoding per lane: 0=v00, 1=v01, 2=v10, 3=v11 (2 bits each lane)
  const uint8_t win = argmax4(v00, v01, v10, v11, &vmax);

  // write output float4
  const int out_base = ((h_out_idx * W_out + w_out_idx) * C + c_base);
  *reinterpret_cast<float4*>(out + out_base) = vmax;

  // store indices as absolute flat indices per element (int per channel)
  // decode per lane -> which base?? and write 4 ints
  if (indices) {
    int* idx_ptr = indices + out_base;

    const int base_ids[4] = { base00, base01, base10, base11 };

    uint8_t id0 = (win >> 0) & 0x3;
    uint8_t id1 = (win >> 2) & 0x3;
    uint8_t id2 = (win >> 4) & 0x3;
    uint8_t id3 = (win >> 6) & 0x3;

    idx_ptr[0] = base_ids[id0] + 0;
    idx_ptr[1] = base_ids[id1] + 1;
    idx_ptr[2] = base_ids[id2] + 2;
    idx_ptr[3] = base_ids[id3] + 3;
  }
}

// Backward: no atomic needed (stride=2 => no overlap), scatter grad_out to grad_in
__global__ void maxpool2x2_backward_scatter_noatomic(const float* __restrict__ go,
                                                     float* __restrict__ gi,
                                                     const int* __restrict__ indices,
                                                     int size_out)
{
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= size_out) return;

  int in_idx = indices[idx];
  // indices[idx] always valid in your forward path
  gi[in_idx] = go[idx];
}

void launch_maxpool_forward_opt(const float* d_input, float* d_output,
                                int* d_indices, int H, int W, int C)
{
  const int H_out = H >> 1;
  const int W_out = W >> 1;

  // Fast path: H,W even & C multiple of 4
  if (((H & 1) == 0) && ((W & 1) == 0) && ((C & 3) == 0)) {
    const int c4 = C >> 2;
    const int n = H_out * W_out * c4;

    constexpr int BS = 256;
    const int grid = (n + BS - 1) / BS;

    maxpool2x2_forward_nhwc_f4<<<grid, BS>>>(d_input, d_output, d_indices,
                                            H, W, C, H_out, W_out);
    return;
  }

  // Fallback: call your existing v1 (giữ đúng kết quả cho trường hợp lẻ)
  launch_maxpool_forward((float*)d_input, d_output, d_indices, H, W, C);
}

void launch_maxpool_backward_opt(const float* d_grad_output, float* d_grad_input,
                                 const int* d_indices, int size_out)
{
  constexpr int BS = 256;
  const int grid = (size_out + BS - 1) / BS;
  maxpool2x2_backward_scatter_noatomic<<<grid, BS>>>(
      d_grad_output, d_grad_input, d_indices, size_out);
}

// =====================
// ReLU
// =====================

__device__ __forceinline__ float relu(float x) { return x > 0.0f ? x : 0.0f; }

__global__ void relu_forward_f4(const float* __restrict__ in,
                               float* __restrict__ out,
                               int n)
{
  // float4
  int tid4 = blockIdx.x * blockDim.x + threadIdx.x;
  int n4 = n >> 2; // n/4

  for (int i = tid4; i < n4; i += blockDim.x * gridDim.x) {
    float4 v = reinterpret_cast<const float4*>(in)[i];
    v.x = relu(v.x);
    v.y = relu(v.y);
    v.z = relu(v.z);
    v.w = relu(v.w);
    reinterpret_cast<float4*>(out)[i] = v;
  }

  // tail (n % 4)
  int base = (n4 << 2);
  int tid = base + (blockIdx.x * blockDim.x + threadIdx.x);
  for (int j = tid; j < n; j += blockDim.x * gridDim.x) {
    float x = in[j];
    out[j] = relu(x);
  }
}

__global__ void relu_backward_f4(const float* __restrict__ go,
                                float* __restrict__ gi,
                                const float* __restrict__ in,
                                int n)
{
  int tid4 = blockIdx.x * blockDim.x + threadIdx.x;
  int n4 = n >> 2;

  for (int i = tid4; i < n4; i += blockDim.x * gridDim.x) {
    float4 g = reinterpret_cast<const float4*>(go)[i];
    float4 x = reinterpret_cast<const float4*>(in)[i];

    // grad passes only where input>0
    g.x = (x.x > 0.0f) ? g.x : 0.0f;
    g.y = (x.y > 0.0f) ? g.y : 0.0f;
    g.z = (x.z > 0.0f) ? g.z : 0.0f;
    g.w = (x.w > 0.0f) ? g.w : 0.0f;

    reinterpret_cast<float4*>(gi)[i] = g;
  }

  int base = (n4 << 2);
  int tid = base + (blockIdx.x * blockDim.x + threadIdx.x);
  for (int j = tid; j < n; j += blockDim.x * gridDim.x) {
    float x = in[j];
    gi[j] = (x > 0.0f) ? go[j] : 0.0f;
  }
}

void launch_relu_forward_opt(const float* d_input, float* d_output, int size)
{
  constexpr int BS = 256;
  // grid avoid small oversubscribe 
  int grid = (size / 4 + BS - 1) / BS;
  if (grid > 4096) grid = 4096;
  relu_forward_f4<<<grid, BS>>>(d_input, d_output, size);
}

void launch_relu_backward_opt(const float* d_grad_output, float* d_grad_input,
                              const float* d_input, int size)
{
  constexpr int BS = 256;
  int grid = (size / 4 + BS - 1) / BS;
  if (grid > 4096) grid = 4096;
  relu_backward_f4<<<grid, BS>>>(d_grad_output, d_grad_input, d_input, size);
}

// =============
// Upsamle
// =============
#include <cuda_runtime.h>
#include <stdint.h>

__global__ void upsample2x_forward_nhwc_f4(const float* __restrict__ in,
                                          float* __restrict__ out,
                                          int H, int W, int C,
                                          int H2, int W2)
{
  // each thread handles one input pixel and one float4 pack
  int c4 = C >> 2;
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int n = H * W * c4;
  if (tid >= n) return;

  int pc = tid % c4;
  int t  = tid / c4;
  int w0 = t % W;
  int h0 = t / W;

  int c_base = pc * 4;

  // load one float4 from input
  int in_idx = (h0 * W + w0) * C + c_base;
  float4 v = *reinterpret_cast<const float4*>(in + in_idx);

  // write to 4 positions in output (2x2)
  int oh = h0 * 2;
  int ow = w0 * 2;

  int o00 = (oh * W2 + ow) * C + c_base;
  int o01 = (oh * W2 + (ow + 1)) * C + c_base;
  int o10 = ((oh + 1) * W2 + ow) * C + c_base;
  int o11 = ((oh + 1) * W2 + (ow + 1)) * C + c_base;

  *reinterpret_cast<float4*>(out + o00) = v;
  *reinterpret_cast<float4*>(out + o01) = v;
  *reinterpret_cast<float4*>(out + o10) = v;
  *reinterpret_cast<float4*>(out + o11) = v;
}

__global__ void upsample2x_backward_sum4_nhwc_f4(const float* __restrict__ go,
                                                float* __restrict__ gi,
                                                int H, int W, int C,
                                                int H2, int W2)
{
  // each thread computes one input pixel grad and one float4 pack
  int c4 = C >> 2;
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int n = H * W * c4;
  if (tid >= n) return;

  int pc = tid % c4;
  int t  = tid / c4;
  int w0 = t % W;
  int h0 = t / W;

  int c_base = pc * 4;

  int oh = h0 * 2;
  int ow = w0 * 2;

  int o00 = (oh * W2 + ow) * C + c_base;
  int o01 = (oh * W2 + (ow + 1)) * C + c_base;
  int o10 = ((oh + 1) * W2 + ow) * C + c_base;
  int o11 = ((oh + 1) * W2 + (ow + 1)) * C + c_base;

  float4 g00 = *reinterpret_cast<const float4*>(go + o00);
  float4 g01 = *reinterpret_cast<const float4*>(go + o01);
  float4 g10 = *reinterpret_cast<const float4*>(go + o10);
  float4 g11 = *reinterpret_cast<const float4*>(go + o11);

  // unrolled sum of 4 grads
  float4 s;
  s.x = g00.x + g01.x + g10.x + g11.x;
  s.y = g00.y + g01.y + g10.y + g11.y;
  s.z = g00.z + g01.z + g10.z + g11.z;
  s.w = g00.w + g01.w + g10.w + g11.w;

  int in_idx = (h0 * W + w0) * C + c_base;
  *reinterpret_cast<float4*>(gi + in_idx) = s;
}

void launch_upsample_forward_opt(const float* d_input, float* d_output,
                                 int H, int W, int C)
{
  int H2 = H * 2;
  int W2 = W * 2;

  // fast path: C multiple of 4 (128/256 ok)
  if ((C & 3) == 0) {
    int c4 = C >> 2;
    int n = H * W * c4;
    constexpr int BS = 256;
    int grid = (n + BS - 1) / BS;
    upsample2x_forward_nhwc_f4<<<grid, BS>>>(d_input, d_output, H, W, C, H2, W2);
    return;
  }

  // fallback 
  launch_upsample_forward((float*)d_input, d_output, H, W, C);
}

void launch_upsample_backward_opt(const float* d_grad_output, float* d_grad_input,
                                  int H_in, int W_in, int C)
{
  int H2 = H_in * 2;
  int W2 = W_in * 2;

  if ((C & 3) == 0) {
    int c4 = C >> 2;
    int n = H_in * W_in * c4;
    constexpr int BS = 256;
    int grid = (n + BS - 1) / BS;
    upsample2x_backward_sum4_nhwc_f4<<<grid, BS>>>(d_grad_output, d_grad_input,
                                                   H_in, W_in, C, H2, W2);
    return;
  }

  // fallback 
  launch_upsample_backward((float*)d_grad_output, d_grad_input, H_in, W_in, C);
}


// -------------------------
// Warp + block reduction
// -------------------------
__device__ __forceinline__ float warpReduceSum(float v) {
  for (int offset = 16; offset > 0; offset >>= 1)
    v += __shfl_down_sync(0xffffffff, v, offset);
  return v;
}

__device__ __forceinline__ float blockReduceSum(float v) {
  // one warp result per warp
  __shared__ float shared[32]; // up to 1024 threads
  int lane = threadIdx.x & 31;
  int wid  = threadIdx.x >> 5;

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

  // final reduce by warp0
  v = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
  if (wid == 0) v = warpReduceSum(v);
  return v;
}

// -------------------------
// MSE loss optimized
// -------------------------
__global__ void mse_loss_reduce_f4(const float* __restrict__ out,
                                  const float* __restrict__ tgt,
                                  float* __restrict__ loss,
                                  int n)
{
  float sum = 0.0f;

  // vectorized part
  int n4 = n >> 2;
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int stride = blockDim.x * gridDim.x;

  for (int i = tid; i < n4; i += stride) {
    float4 o = reinterpret_cast<const float4*>(out)[i];
    float4 t = reinterpret_cast<const float4*>(tgt)[i];

    float dx = o.x - t.x; sum += dx * dx;
    float dy = o.y - t.y; sum += dy * dy;
    float dz = o.z - t.z; sum += dz * dz;
    float dw = o.w - t.w; sum += dw * dw;
  }

  // tail
  int base = (n4 << 2);
  for (int j = base + tid; j < n; j += stride) {
    float d = out[j] - tgt[j];
    sum += d * d;
  }

  float bsum = blockReduceSum(sum);
  if (threadIdx.x == 0) atomicAdd(loss, bsum); // 1 atomic per block
}

void launch_mse_loss_opt(const float* d_output, const float* d_target,
                         float* d_loss, int size)
{
  // must zero loss first (same behavior as naive)
  cudaMemset(d_loss, 0, sizeof(float));

  constexpr int BS = 256;

  int grid = (size / 4 + BS - 1) / BS;
  if (grid > 2048) grid = 2048;

  mse_loss_reduce_f4<<<grid, BS>>>(d_output, d_target, d_loss, size);
}

// -------------------------
// MSE backward optimized (vectorized, no atomics)
// dL/dOut = 2*(out-target)  
// -------------------------
__global__ void mse_loss_backward_f4(const float* __restrict__ out,
                                    const float* __restrict__ tgt,
                                    float* __restrict__ go,
                                    int n)
{
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int stride = blockDim.x * gridDim.x;

  int n4 = n >> 2;
  for (int i = tid; i < n4; i += stride) {
    float4 o = reinterpret_cast<const float4*>(out)[i];
    float4 t = reinterpret_cast<const float4*>(tgt)[i];

    float4 g;
    g.x = 2.0f * (o.x - t.x);
    g.y = 2.0f * (o.y - t.y);
    g.z = 2.0f * (o.z - t.z);
    g.w = 2.0f * (o.w - t.w);

    reinterpret_cast<float4*>(go)[i] = g;
  }

  int base = (n4 << 2);
  for (int j = base + tid; j < n; j += stride) {
    go[j] = 2.0f * (out[j] - tgt[j]);
  }
}

void launch_mse_loss_backward_opt(const float* d_output, const float* d_target,
                                  float* d_grad_input, int size)
{
  constexpr int BS = 256;
  int grid = (size / 4 + BS - 1) / BS;
  if (grid > 2048) grid = 2048;
  mse_loss_backward_f4<<<grid, BS>>>(d_output, d_target, d_grad_input, size);
}