kernels.cu

#include "kernels.h"
#include <algorithm>
#include <cstdio>

#define BLOCK_SIZE 256

// --------------------------------------------------------------------------
// 1. ReLU KERNEL
// Logic: Mỗi thread xử lý 1 phần tử. Rất dễ song song hóa.
// --------------------------------------------------------------------------
__global__ void relu_forward_kernel(float *input, float *output, int size) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size) {
    output[idx] = fmaxf(0.0f, input[idx]);
  }
}

void launch_relu_forward(float *d_input, float *d_output, int size) {
  int grid_size = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;
  relu_forward_kernel<<<grid_size, BLOCK_SIZE>>>(d_input, d_output, size);
}

// --------------------------------------------------------------------------
// 2. MaxPool2D KERNEL
// Logic: Mỗi thread tính 1 pixel của OUTPUT (H_out, W_out, C).
//        Nó sẽ quét 4 pixel tương ứng ở Input để tìm Max.
// --------------------------------------------------------------------------
__global__ void maxpool_forward_kernel(float *input, float *output,
                                       int *indices, int H, int W, int C,
                                       int H_out, int W_out) {
  // Ánh xạ Thread ID sang tọa độ Output (h_out, w_out, c)
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int size_out = H_out * W_out * C;

  if (idx < size_out) {
    // Giải mã chỉ số phẳng thành 3D (c, h_out, w_out)
    // Layout bộ nhớ giả định: [H][W][C] (như code CPU cũ)
    // Để khớp với CPU code: index = (h * W + w) * C + c

    int c = idx % C;
    int tmp = idx / C;
    int w_out = tmp % W_out;
    int h_out = tmp / W_out;

    int h_in_start = h_out * 2;
    int w_in_start = w_out * 2;

    float max_val = -1e30f; // -Infinity
    int max_idx = -1;

    // Quét cửa sổ 2x2
    for (int kh = 0; kh < 2; ++kh) {
      for (int kw = 0; kw < 2; ++kw) {
        int h_in = h_in_start + kh;
        int w_in = w_in_start + kw;

        if (h_in < H && w_in < W) {
          int in_idx = (h_in * W + w_in) * C + c;
          float val = input[in_idx];
          if (val > max_val) {
            max_val = val;
            max_idx = in_idx;
          }
        }
      }
    }

    output[idx] = max_val;
    if (indices)
      indices[idx] = max_idx; // Lưu vị trí để dùng cho Backward
  }
}

void launch_maxpool_forward(float *d_input, float *d_output, int *d_indices,
                            int H, int W, int C) {
  int H_out = H / 2;
  int W_out = W / 2;
  int size_out = H_out * W_out * C;
  int grid_size = (size_out + BLOCK_SIZE - 1) / BLOCK_SIZE;

  maxpool_forward_kernel<<<grid_size, BLOCK_SIZE>>>(
      d_input, d_output, d_indices, H, W, C, H_out, W_out);
}

// --------------------------------------------------------------------------
// 3. Upsample2D KERNEL
// Logic: Mỗi thread tính 1 pixel của OUTPUT.
//        Nó lấy giá trị từ pixel input tương ứng (tọa độ chia 2).
// --------------------------------------------------------------------------
__global__ void upsample_forward_kernel(float *input, float *output, int H_out,
                                        int W_out, int C, int H_in, int W_in) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int size_out = H_out * W_out * C;

  if (idx < size_out) {
    // Output layout: (h_out, w_out, c)
    int c = idx % C;
    int tmp = idx / C;
    int w_out = tmp % W_out;
    int h_out = tmp / W_out;

    // Nearest Neighbor: chia 2 lấy phần nguyên
    int h_in = h_out / 2;
    int w_in = w_out / 2;

    int in_idx = (h_in * W_in + w_in) * C + c;
    output[idx] = input[in_idx];
  }
}

void launch_upsample_forward(float *d_input, float *d_output, int H_in,
                             int W_in, int C) {
  int H_out = H_in * 2;
  int W_out = W_in * 2;
  int size_out = H_out * W_out * C;
  int grid_size = (size_out + BLOCK_SIZE - 1) / BLOCK_SIZE;

  upsample_forward_kernel<<<grid_size, BLOCK_SIZE>>>(d_input, d_output, H_out,
                                                     W_out, C, H_in, W_in);
}

// --------------------------------------------------------------------------
// 4. MSE Loss KERNEL (Naive atomicAdd)
// Logic: Mỗi thread tính sai số bình phương của 1 pixel, sau đó cộng dồn vào
// biến global loss. Lưu ý: atomicAdd lên biến global rất chậm, nhưng đây là
// "Naive Implementation". Ở Phase 3 (Optimized), ta sẽ dùng Parallel Reduction.
// --------------------------------------------------------------------------
__global__ void mse_loss_kernel(const float *output, const float *target,
                                float *d_loss, int size) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size) {
    float diff = output[idx] - target[idx];
    atomicAdd(d_loss, diff * diff);
  }
}

void launch_mse_loss(float *d_output, float *d_target, float *d_loss,
                     int size) {
  cudaMemset(d_loss, 0, sizeof(float));

  int grid = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;
  mse_loss_kernel<<<grid, BLOCK_SIZE>>>(d_output, d_target, d_loss, size);
}

// --------------------------------------------------------------------------
// 5. Conv2D Naive KERNEL
// Strategy: One thread per Output Pixel (h, w, co)
// --------------------------------------------------------------------------
__global__ void conv2d_naive_kernel(float *input, float *output, float *weights,
                                    float *bias, int H, int W, int C_in,
                                    int C_out, int K) {

  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int size_out = H * W * C_out;

  if (idx < size_out) {
    // Mapping: idx -> (h, w, co)
    // Layout Output: [H][W][C_out]
    int co = idx % C_out;
    int tmp = idx / C_out;
    int w = tmp % W;
    int h = tmp / W;

    float sum = bias[co];

    // Convolution
    // Receptive Field
    int pad = K / 2; // padding = 1, K=3

    // Base pointer
    // Weight Layout: [C_out][C_in][K][K]
    int weight_co_offset = co * C_in * K * K;

    // Loop Input Channels
    for (int ci = 0; ci < C_in; ++ci) {

      // Offset weights of input current channel
      int weight_ci_offset = weight_co_offset + ci * K * K;

      // Loop Kernel
      for (int kh = -pad; kh <= pad; ++kh) {
        for (int kw = -pad; kw <= pad; ++kw) {

          int h_in = h + kh;
          int w_in = w + kw;

          // Padding check
          if (h_in >= 0 && h_in < H && w_in >= 0 && w_in < W) {

            // Layout Input: [H][W][C_in]
            int input_idx = (h_in * W + w_in) * C_in + ci;
            float val_in = input[input_idx];

            // Kernel local coordinate : (kh+pad, kw+pad) => (0..K-1)
            int k_row = kh + pad;
            int k_col = kw + pad;
            int weight_idx = weight_ci_offset + k_row * K + k_col;
            float val_w = weights[weight_idx];

            sum += val_in * val_w;
          }
        }
      }
    }

    output[idx] = sum;
  }
}

void launch_conv2d_naive(float *d_input, float *d_output, float *d_weights,
                         float *d_bias, int H, int W, int C_in, int C_out,
                         int K) {
  int size_out = H * W * C_out;
  int grid_size = (size_out + BLOCK_SIZE - 1) / BLOCK_SIZE;

  conv2d_naive_kernel<<<grid_size, BLOCK_SIZE>>>(d_input, d_output, d_weights,
                                                 d_bias, H, W, C_in, C_out, K);

  cudaError_t err = cudaGetLastError();
  if (err != cudaSuccess) {
    printf("Conv2D Kernel Error: %s\n", cudaGetErrorString(err));
  }
}

// ==========================================================================
// BACKWARD KERNELS IMPLEMENTATION
// ==========================================================================

// 6. MSE LOSS BACKWARD
__global__ void mse_loss_backward_kernel(float *output, float *target,
                                         float *grad_input, int size) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size) {
    // dL/dx = 2/N * (x - y)
    // Lưu ý: Việc chia cho N (size) có thể làm ở đây hoặc gộp vào learning
    // rate. Ở đây ta tính raw gradient: 2 * (output - target)
    grad_input[idx] = 2.0f * (output[idx] - target[idx]) / (float)size;
  }
}

void launch_mse_loss_backward(float *d_output, float *d_target,
                              float *d_grad_input, int size) {
  int grid = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;
  mse_loss_backward_kernel<<<grid, BLOCK_SIZE>>>(d_output, d_target,
                                                 d_grad_input, size);
}

// 7. RELU BACKWARD
__global__ void relu_backward_kernel(float *grad_out, float *grad_in,
                                     float *input, int size) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size) {
    // Nếu input > 0, gradient đi qua, ngược lại bị chặn (0)
    grad_in[idx] = (input[idx] > 0.0f) ? grad_out[idx] : 0.0f;
  }
}

void launch_relu_backward(float *d_grad_output, float *d_grad_input,
                          float *d_input, int size) {
  int grid = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;
  relu_backward_kernel<<<grid, BLOCK_SIZE>>>(d_grad_output, d_grad_input,
                                             d_input, size);
}

// 8. MAXPOOL BACKWARD
// Cần indices từ bước Forward
__global__ void maxpool_backward_kernel(float *grad_out, float *grad_in,
                                        int *indices, int size_out) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size_out) {
    int max_idx = indices[idx];
    // Chỉ truyền gradient về đúng vị trí max
    // atomicAdd phòng trường hợp MaxPool có overlap (dù ở đây stride=2
    // không overlap) Dùng atomicAdd để an toàn tuyệt đối.
    if (max_idx >= 0) {
      atomicAdd(&grad_in[max_idx], grad_out[idx]);
    }
  }
}

void launch_maxpool_backward(float *d_grad_output, float *d_grad_input,
                             int *d_indices, int size_out) {
  // Lưu ý: Trước khi gọi hàm này, d_grad_input phải được reset về 0
  // (cudaMemset)
  int grid = (size_out + BLOCK_SIZE - 1) / BLOCK_SIZE;
  maxpool_backward_kernel<<<grid, BLOCK_SIZE>>>(d_grad_output, d_grad_input,
                                                d_indices, size_out);
}

// 9. UPSAMPLE BACKWARD
// Forward: 1 -> 4. Backward: 4 -> 1 (Sum)
__global__ void upsample_backward_kernel(float *grad_out, float *grad_in,
                                         int H_in, int W_in, int C, int H_out,
                                         int W_out) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int size_in = H_in * W_in * C;

  if (idx < size_in) {
    // idx là tọa độ của grad_in (Input của Forward)
    int c = idx % C;
    int tmp = idx / C;
    int w_in = tmp % W_in;
    int h_in = tmp / W_in;

    // Vùng tương ứng trên Output (của Forward) là 2x2
    int h_out_start = h_in * 2;
    int w_out_start = w_in * 2;

    float sum_grad = 0.0f;

    // Cộng dồn 4 ô
    sum_grad += grad_out[((h_out_start)*W_out + w_out_start) * C + c];
    sum_grad += grad_out[((h_out_start)*W_out + (w_out_start + 1)) * C + c];
    sum_grad += grad_out[((h_out_start + 1) * W_out + w_out_start) * C + c];
    sum_grad +=
        grad_out[((h_out_start + 1) * W_out + (w_out_start + 1)) * C + c];

    grad_in[idx] = sum_grad;
  }
}

void launch_upsample_backward(float *d_grad_output, float *d_grad_input,
                              int H_in, int W_in, int C) {
  int H_out = H_in * 2;
  int W_out = W_in * 2;
  int size_in = H_in * W_in * C;
  int grid = (size_in + BLOCK_SIZE - 1) / BLOCK_SIZE;
  upsample_backward_kernel<<<grid, BLOCK_SIZE>>>(d_grad_output, d_grad_input,
                                                 H_in, W_in, C, H_out, W_out);
}

// 10. CONV2D BACKWARD (NAIVE with ATOMICS)
// Đây là kernel nặng nhất. Mỗi thread xử lý 1 pixel của Grad_Output và cộng dồn
// vào Grad_Input và Grad_Weights.
__global__ void conv2d_backward_naive_kernel(float *grad_out, float *input,
                                             float *weights, float *grad_in,
                                             float *grad_w, float *grad_b,
                                             int H, int W, int C_in, int C_out,
                                             int K) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  int size_out = H * W * C_out; // Kích thước của Output (cũng là Grad_Out)

  if (idx < size_out) {
    // 1. Phân giải tọa độ Grad_Output
    int co = idx % C_out;
    int tmp = idx / C_out;
    int w = tmp % W;
    int h = tmp / W;

    float g_out = grad_out[idx];

    // 2. Tích lũy Gradient cho Bias (Atomic)
    atomicAdd(&grad_b[co], g_out);

    // 3. Vòng lặp tính dW và dInput
    int pad = K / 2;
    int weight_co_offset = co * C_in * K * K;

    for (int ci = 0; ci < C_in; ++ci) {
      int weight_ci_offset = weight_co_offset + ci * K * K;

      for (int kh = -pad; kh <= pad; ++kh) {
        for (int kw = -pad; kw <= pad; ++kw) {
          int h_in = h + kh;
          int w_in = w + kw;

          if (h_in >= 0 && h_in < H && w_in >= 0 && w_in < W) {
            int in_idx = (h_in * W + w_in) * C_in + ci;

            // --- Tính Gradient cho Weights (dW) ---
            // dW += Input * Grad_Out
            int k_row = kh + pad;
            int k_col = kw + pad;
            int w_idx = weight_ci_offset + k_row * K + k_col;

            float val_in = input[in_idx];
            atomicAdd(&grad_w[w_idx], val_in * g_out);

            // --- Tính Gradient cho Input (dX) ---
            // dX += Weights * Grad_Out
            float val_w = weights[w_idx];
            atomicAdd(&grad_in[in_idx], val_w * g_out);
          }
        }
      }
    }
  }
}

void launch_conv2d_backward(float *d_grad_output, float *d_input,
                            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) {
  // Quan trọng: Phải memset d_grad_input, d_grad_weights, d_grad_bias về 0
  // trước khi gọi hàm này! (Việc memset sẽ được làm ở gpu_autoencoder.cpp)

  int size_out = H * W * C_out;
  int grid = (size_out + BLOCK_SIZE - 1) / BLOCK_SIZE;

  conv2d_backward_naive_kernel<<<grid, BLOCK_SIZE>>>(
      d_grad_output, d_input, d_weights, d_grad_input, d_grad_weights,
      d_grad_bias, H, W, C_in, C_out, K);
}

// 11. ADAM UPDATE KERNEL
__global__ void adam_update_kernel(float *weights, float *grad, float *m,
                                   float *v, int size, int step, float alpha,
                                   float beta1, float beta2, float epsilon) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < size) {
    float g = grad[idx];

    // Update moments
    float m_t = beta1 * m[idx] + (1.0f - beta1) * g;
    float v_t = beta2 * v[idx] + (1.0f - beta2) * g * g;

    m[idx] = m_t;
    v[idx] = v_t;

    // Bias correction
    float m_hat = m_t / (1.0f - powf(beta1, step));
    float v_hat = v_t / (1.0f - powf(beta2, step));

    // Update weight
    weights[idx] -= alpha * m_hat / (sqrtf(v_hat) + epsilon);
  }
}

void launch_adam_update(float *weights, float *grad, float *m, float *v,
                        int size, int step, float alpha, float beta1,
                        float beta2, float epsilon) {
  int grid = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;
  adam_update_kernel<<<grid, BLOCK_SIZE>>>(weights, grad, m, v, size, step,
                                           alpha, beta1, beta2, epsilon);
}