kernels.h

#pragma once
#include <cuda_runtime.h>

// 1. ReLU: f(x) = max(0, x)
// size = H * W * C
void launch_relu_forward(float *d_input, float *d_output, int size);

// 2. MaxPool2D: 2x2 pooling
// Input: (H, W, C) -> Output: (H/2, W/2, C)
void launch_maxpool_forward(float *d_input, float *d_output, int *d_indices,
                            int H, int W, int C);

// 3. Upsample2D: 2x2 nearest neighbor
// Input: (H, W, C) -> Output: (H*2, W*2, C)
void launch_upsample_forward(float *d_input, float *d_output, int H, int W,
                             int C);

// 4. MSE Loss (Naive atomicAdd)
// Tính tổng bình phương sai số
void launch_mse_loss(float *d_output, float *d_target, float *d_loss, int size);

// 5. Naive Conv2D Forward
// Input: HWC layout
// Weights: [C_out][C_in][K][K] layout (phẳng hóa)
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);

// --- BACKWARD KERNELS ---

// 6. MSE Loss Backward
// Tính dL/dOutput = 2/N * (Output - Target)
void launch_mse_loss_backward(float *d_output, float *d_target,
                              float *d_grad_input, int size);

// 7. ReLU Backward
// Nếu input > 0 thì grad_in = grad_out, ngược lại = 0
void launch_relu_backward(float *d_grad_output, float *d_grad_input,
                          float *d_input, int size);

// 8. MaxPool2D Backward
// Truyền gradient về đúng vị trí Max (dùng indices đã lưu lúc forward)
void launch_maxpool_backward(float *d_grad_output, float *d_grad_input,
                             int *d_indices, int size_out);

// 9. Upsample2D Backward
// Cộng dồn gradient từ 4 pixel output về 1 pixel input
void launch_upsample_backward(float *d_grad_output, float *d_grad_input,
                              int H_in, int W_in, int C);

// 10. Conv2D Backward (Naive with Atomics)
// Tính dL/dInput và dL/dWeights
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);

// 11. Adam Update Weights
// Cập nhật weights và bias theo công thức Adam
void launch_adam_update(float *weights, float *grad, float *m, float *v,
                        int size, int step, float alpha, float beta1,
                        float beta2, float epsilon);