kernels_v2.h

#pragma once
#include <cuda_runtime.h>

// Optimized Conv2D Forward (NHWC/HWC input/output)
// Strategies: shared-memory tiling + unrolled K=3 +
// optional constant-memory weights
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 /*expect 3*/);

// Optimized Conv2D Backward (NHWC/HWC) for K=3
// Computes d_grad_input, d_grad_weights, d_grad_bias.
// Strategies: shared-memory tiling + unrolled K=3 + optional constant-memory weights.
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 /*expect 3*/);

// Optimized MaxPool2D Forward/Backward (2x2, stride=2) for NHWC
// Fast-path optimized for even H,W and C multiple of 4 (128/256 fits your net).
void launch_maxpool_forward_opt(const float* d_input, float* d_output,
                                int* d_indices, int H, int W, int C);

void launch_maxpool_backward_opt(const float* d_grad_output, float* d_grad_input,
                                 const int* d_indices, int size_out);


void launch_relu_forward_opt(const float* d_input, float* d_output, int size);
void launch_relu_backward_opt(const float* d_grad_output, float* d_grad_input,
                              const float* d_input, int size);


// Upsample2D nearest-neighbor x2, NHWC
void launch_upsample_forward_opt(const float* d_input, float* d_output,
                                 int H, int W, int C);

// Backward: d_grad_input(h,w,c) = sum of 4 grads from output
void launch_upsample_backward_opt(const float* d_grad_output, float* d_grad_input,
                                  int H_in, int W_in, int C);

// Optimized MSE loss: sum((out-target)^2) over size elements
void launch_mse_loss_opt(const float* d_output, const float* d_target,
                         float* d_loss, int size);

// Backward giữ nguyên (nó đã O(N) đơn giản), nhưng có bản opt vectorized
void launch_mse_loss_backward_opt(const float* d_output, const float* d_target,
                                  float* d_grad_input, int size);