gpu_autoencoder.cpp

#include "gpu_autoencoder.h"
#include "kernels.h"
#include "kernels_v2.h"
#include <fstream>
#include <iostream>
#include <random>
#include <vector>

// Macro rút gọn
#define CUDA_MALLOC(ptr, size)                                                 \
  CUDA_CHECK(cudaMalloc((void **)&ptr, size * sizeof(float)))
#define CUDA_MALLOC_INT(ptr, size)                                             \
  CUDA_CHECK(cudaMalloc((void **)&ptr, size * sizeof(int)))
#define CUDA_ZERO(ptr, size)                                                   \
  CUDA_CHECK(cudaMemset(ptr, 0, size * sizeof(float)))

// Helper init weight
void init_weight_buffer(float *d_ptr, size_t size) {
  std::vector<float> host_data(size);
  std::mt19937 gen(42);
  std::normal_distribution<float> dist(0.0f, 0.02f);
  for (size_t i = 0; i < size; ++i)
    host_data[i] = dist(gen);
  CUDA_CHECK(cudaMemcpy(d_ptr, host_data.data(), size * sizeof(float),
                        cudaMemcpyHostToDevice));
}

void GPUAutoencoder::initialize() {
  // Init weights only (Bias = 0 set in allocate via CUDA_ZERO not needed, but
  // good practice)
  init_weight_buffer(d_conv1_w, 256 * 3 * 9);
  init_weight_buffer(d_conv2_w, 128 * 256 * 9);
  init_weight_buffer(d_conv3_w, 128 * 128 * 9);
  init_weight_buffer(d_conv4_w, 256 * 128 * 9);
  init_weight_buffer(d_conv5_w, 3 * 256 * 9);
  std::cout << "GPU Initialized." << std::endl;
}

GPUAutoencoder::GPUAutoencoder() { allocate_memory(); }
GPUAutoencoder::~GPUAutoencoder() { free_memory(); }
void GPUAutoencoder::allocate_memory() {
  // 1. Weights & Gradients & Adam Moments
  auto alloc_layer = [&](float *&w, float *&b, float *&gw, float *&gb,
                         float *&mw, float *&vw, float *&mb, float *&vb,
                         int c_out, int c_in, int k) {
    int w_size = c_out * c_in * k * k;
    int b_size = c_out;
    CUDA_MALLOC(w, w_size);
    CUDA_MALLOC(b, b_size);
    CUDA_MALLOC(gw, w_size);
    CUDA_MALLOC(gb, b_size);
    CUDA_MALLOC(mw, w_size);
    CUDA_MALLOC(vw, w_size);
    CUDA_MALLOC(mb, b_size);
    CUDA_MALLOC(vb, b_size);

    // Zero out gradients & moments
    CUDA_ZERO(gw, w_size);
    CUDA_ZERO(gb, b_size);
    CUDA_ZERO(mw, w_size);
    CUDA_ZERO(vw, w_size);
    CUDA_ZERO(mb, b_size);
    CUDA_ZERO(vb, b_size);
  };

  // Conv1: 3 -> 256
  alloc_layer(d_conv1_w, d_conv1_b, d_grad_conv1_w, d_grad_conv1_b, d_m_conv1_w,
              d_v_conv1_w, d_m_conv1_b, d_v_conv1_b, 256, 3, 3);
  // Conv2: 256 -> 128
  alloc_layer(d_conv2_w, d_conv2_b, d_grad_conv2_w, d_grad_conv2_b, d_m_conv2_w,
              d_v_conv2_w, d_m_conv2_b, d_v_conv2_b, 128, 256, 3);
  // Conv3: 128 -> 128
  alloc_layer(d_conv3_w, d_conv3_b, d_grad_conv3_w, d_grad_conv3_b, d_m_conv3_w,
              d_v_conv3_w, d_m_conv3_b, d_v_conv3_b, 128, 128, 3);
  // Conv4: 128 -> 256
  alloc_layer(d_conv4_w, d_conv4_b, d_grad_conv4_w, d_grad_conv4_b, d_m_conv4_w,
              d_v_conv4_w, d_m_conv4_b, d_v_conv4_b, 256, 128, 3);
  // Conv5: 256 -> 3
  alloc_layer(d_conv5_w, d_conv5_b, d_grad_conv5_w, d_grad_conv5_b, d_m_conv5_w,
              d_v_conv5_w, d_m_conv5_b, d_v_conv5_b, 3, 256, 3);

  // 2. Activations & Gradients & Indices
  auto alloc_act = [&](float *&act, float *&grad, int size) {
    CUDA_MALLOC(act, size);
    CUDA_MALLOC(grad, size);
  };

  alloc_act(d_input, d_grad_input, 32 * 32 * 3);

  // Encoder
  alloc_act(d_act1, d_grad_act1, 32 * 32 * 256);
  alloc_act(d_act2, d_grad_act2, 32 * 32 * 256);
  alloc_act(d_act3, d_grad_act3, 16 * 16 * 256);
  CUDA_MALLOC_INT(d_indices1, 16 * 16 * 256); // Indices Pool1

  alloc_act(d_act4, d_grad_act4, 16 * 16 * 128);
  alloc_act(d_act5, d_grad_act5, 16 * 16 * 128);
  alloc_act(d_act6, d_grad_act6, 8 * 8 * 128);
  CUDA_MALLOC_INT(d_indices2, 8 * 8 * 128); // Indices Pool2

  // Decoder
  alloc_act(d_act7, d_grad_act7, 8 * 8 * 128);
  alloc_act(d_act8, d_grad_act8, 8 * 8 * 128);
  alloc_act(d_act9, d_grad_act9, 16 * 16 * 128);

  alloc_act(d_act10, d_grad_act10, 16 * 16 * 256);
  alloc_act(d_act11, d_grad_act11, 16 * 16 * 256);
  alloc_act(d_act12, d_grad_act12, 32 * 32 * 256);

  // Output & Loss
  alloc_act(d_output, d_grad_output, 32 * 32 * 3);
  CUDA_MALLOC(d_target, 32 * 32 * 3);
  CUDA_MALLOC(d_loss_val, 1);
}

void GPUAutoencoder::free_memory() {
  // 1. Free Weights & Gradients & Adam Moments
  auto free_layer = [&](float *w, float *b, float *gw, float *gb, float *mw,
                        float *vw, float *mb, float *vb) {
    if (w)
      cudaFree(w);
    if (b)
      cudaFree(b);
    if (gw)
      cudaFree(gw);
    if (gb)
      cudaFree(gb);
    if (mw)
      cudaFree(mw);
    if (vw)
      cudaFree(vw);
    if (mb)
      cudaFree(mb);
    if (vb)
      cudaFree(vb);
  };

  free_layer(d_conv1_w, d_conv1_b, d_grad_conv1_w, d_grad_conv1_b, d_m_conv1_w,
             d_v_conv1_w, d_m_conv1_b, d_v_conv1_b);
  free_layer(d_conv2_w, d_conv2_b, d_grad_conv2_w, d_grad_conv2_b, d_m_conv2_w,
             d_v_conv2_w, d_m_conv2_b, d_v_conv2_b);
  free_layer(d_conv3_w, d_conv3_b, d_grad_conv3_w, d_grad_conv3_b, d_m_conv3_w,
             d_v_conv3_w, d_m_conv3_b, d_v_conv3_b);
  free_layer(d_conv4_w, d_conv4_b, d_grad_conv4_w, d_grad_conv4_b, d_m_conv4_w,
             d_v_conv4_w, d_m_conv4_b, d_v_conv4_b);
  free_layer(d_conv5_w, d_conv5_b, d_grad_conv5_w, d_grad_conv5_b, d_m_conv5_w,
             d_v_conv5_w, d_m_conv5_b, d_v_conv5_b);

  // 2. Free Activations
  auto free_act = [&](float *act, float *grad) {
    if (act)
      cudaFree(act);
    if (grad)
      cudaFree(grad);
  };

  free_act(d_input, d_grad_input);
  free_act(d_act1, d_grad_act1);
  free_act(d_act2, d_grad_act2);
  free_act(d_act3, d_grad_act3);
  if (d_indices1)
    cudaFree(d_indices1);

  free_act(d_act4, d_grad_act4);
  free_act(d_act5, d_grad_act5);
  free_act(d_act6, d_grad_act6);
  if (d_indices2)
    cudaFree(d_indices2);

  free_act(d_act7, d_grad_act7);
  free_act(d_act8, d_grad_act8);
  free_act(d_act9, d_grad_act9);
  free_act(d_act10, d_grad_act10);
  free_act(d_act11, d_grad_act11);
  free_act(d_act12, d_grad_act12);

  free_act(d_output, d_grad_output);
  if (d_target)
    cudaFree(d_target);
  if (d_loss_val)
    cudaFree(d_loss_val);
}

void GPUAutoencoder::update_weights(int step) {
  float alpha = 0.001f;
  float beta1 = 0.9f;
  float beta2 = 0.999f;
  float eps = 1e-8f;

  // Helper lambda
  auto update_layer = [&](float *w, float *b, float *gw, float *gb, float *mw,
                          float *vw, float *mb, float *vb, int w_size,
                          int b_size) {
    launch_adam_update(w, gw, mw, vw, w_size, step, alpha, beta1, beta2, eps);
    launch_adam_update(b, gb, mb, vb, b_size, step, alpha, beta1, beta2, eps);
  };

  update_layer(d_conv1_w, d_conv1_b, d_grad_conv1_w, d_grad_conv1_b,
               d_m_conv1_w, d_v_conv1_w, d_m_conv1_b, d_v_conv1_b, 256 * 3 * 9,
               256);

  update_layer(d_conv2_w, d_conv2_b, d_grad_conv2_w, d_grad_conv2_b,
               d_m_conv2_w, d_v_conv2_w, d_m_conv2_b, d_v_conv2_b,
               128 * 256 * 9, 128);

  update_layer(d_conv3_w, d_conv3_b, d_grad_conv3_w, d_grad_conv3_b,
               d_m_conv3_w, d_v_conv3_w, d_m_conv3_b, d_v_conv3_b,
               128 * 128 * 9, 128);

  update_layer(d_conv4_w, d_conv4_b, d_grad_conv4_w, d_grad_conv4_b,
               d_m_conv4_w, d_v_conv4_w, d_m_conv4_b, d_v_conv4_b,
               256 * 128 * 9, 256);

  update_layer(d_conv5_w, d_conv5_b, d_grad_conv5_w, d_grad_conv5_b,
               d_m_conv5_w, d_v_conv5_w, d_m_conv5_b, d_v_conv5_b, 3 * 256 * 9,
               3);

  cudaDeviceSynchronize();
}

void GPUAutoencoder::zero_grad() {
  // Reset toàn bộ Gradient của Trọng số và Bias về 0
  // (Copy các dòng này từ hàm backward cũ sang)
  CUDA_ZERO(d_grad_conv1_w, 256 * 3 * 9);
  CUDA_ZERO(d_grad_conv1_b, 256);
  CUDA_ZERO(d_grad_conv2_w, 128 * 256 * 9);
  CUDA_ZERO(d_grad_conv2_b, 128);
  CUDA_ZERO(d_grad_conv3_w, 128 * 128 * 9);
  CUDA_ZERO(d_grad_conv3_b, 128);
  CUDA_ZERO(d_grad_conv4_w, 256 * 128 * 9);
  CUDA_ZERO(d_grad_conv4_b, 256);
  CUDA_ZERO(d_grad_conv5_w, 3 * 256 * 9);
  CUDA_ZERO(d_grad_conv5_b, 3);
}

void GPUAutoencoder::set_input(const std::vector<float> &host_input) {
  CUDA_CHECK(cudaMemcpy(d_input, host_input.data(),
                        host_input.size() * sizeof(float),
                        cudaMemcpyHostToDevice));
  CUDA_CHECK(cudaMemcpy(d_target, host_input.data(),
                        host_input.size() * sizeof(float),
                        cudaMemcpyHostToDevice));
}

float GPUAutoencoder::get_loss() {
  float h_loss;
  CUDA_CHECK(
      cudaMemcpy(&h_loss, d_loss_val, sizeof(float), cudaMemcpyDeviceToHost));
  return h_loss / (32.0f * 32.0f * 3.0f);
}

// Helper save vector to file
void write_buffer(std::ofstream &out, float *d_ptr, size_t size) {
  std::vector<float> host_data(size);
  // Copy Device -> Host
  CUDA_CHECK(cudaMemcpy(host_data.data(), d_ptr, size * sizeof(float),
                        cudaMemcpyDeviceToHost));

  size_t n = host_data.size();
  out.write(reinterpret_cast<const char *>(&n), sizeof(n));
  out.write(reinterpret_cast<const char *>(host_data.data()),
            n * sizeof(float));
}

void GPUAutoencoder::save_weights(const std::string &filename) {
  std::ofstream out(filename, std::ios::binary);
  if (!out)
    throw std::runtime_error("Cannot open file for saving weights");

  write_buffer(out, d_conv1_w, 256 * 3 * 9);
  write_buffer(out, d_conv1_b, 256);
  write_buffer(out, d_conv2_w, 128 * 256 * 9);
  write_buffer(out, d_conv2_b, 128);
  write_buffer(out, d_conv3_w, 128 * 128 * 9);
  write_buffer(out, d_conv3_b, 128);
  write_buffer(out, d_conv4_w, 256 * 128 * 9);
  write_buffer(out, d_conv4_b, 256);
  write_buffer(out, d_conv5_w, 3 * 256 * 9);
  write_buffer(out, d_conv5_b, 3);

  std::cout << "Weights saved to " << filename << std::endl;
}

void GPUAutoencoder::load_weights(const std::string &filename) {
  std::ifstream in(filename, std::ios::binary);
  if (!in)
    throw std::runtime_error("Cannot open file for loading weights");

  auto read_buffer = [&](float *d_ptr, size_t expected_size) {
    size_t n;
    in.read(reinterpret_cast<char *>(&n), sizeof(n));
    if (n != expected_size)
      throw std::runtime_error("Weight size mismatch in load_weights");
    std::vector<float> host_data(n);
    in.read(reinterpret_cast<char *>(host_data.data()), n * sizeof(float));
    CUDA_CHECK(cudaMemcpy(d_ptr, host_data.data(), n * sizeof(float), cudaMemcpyHostToDevice));
  };

  read_buffer(d_conv1_w, 256 * 3 * 9);
  read_buffer(d_conv1_b, 256);
  read_buffer(d_conv2_w, 128 * 256 * 9);
  read_buffer(d_conv2_b, 128);
  read_buffer(d_conv3_w, 128 * 128 * 9);
  read_buffer(d_conv3_b, 128);
  read_buffer(d_conv4_w, 256 * 128 * 9);
  read_buffer(d_conv4_b, 256);
  read_buffer(d_conv5_w, 3 * 256 * 9);
  read_buffer(d_conv5_b, 3);

  std::cout << "Weights loaded from " << filename << std::endl;
}

// ============================================
// GPU Auntoencoder version 1
// ============================================
void GPUAutoencoder::forward() {
  // Encoder
  launch_conv2d_naive(d_input, d_act1, d_conv1_w, d_conv1_b, 32, 32, 3, 256, 3);
  launch_relu_forward(d_act1, d_act2, 32 * 32 * 256);

  launch_maxpool_forward(d_act2, d_act3, d_indices1, 32, 32, 256);

  launch_conv2d_naive(d_act3, d_act4, d_conv2_w, d_conv2_b, 16, 16, 256, 128,
                      3);
  launch_relu_forward(d_act4, d_act5, 16 * 16 * 128);
  
  launch_maxpool_forward(d_act5, d_act6, d_indices2, 16, 16, 128);

  // Decoder
  launch_conv2d_naive(d_act6, d_act7, d_conv3_w, d_conv3_b, 8, 8, 128, 128, 3);
  launch_relu_forward(d_act7, d_act8, 8 * 8 * 128);
  launch_upsample_forward(d_act8, d_act9, 8, 8, 128);

  launch_conv2d_naive(d_act9, d_act10, d_conv4_w, d_conv4_b, 16, 16, 128, 256,
                      3);
  launch_relu_forward(d_act10, d_act11, 16 * 16 * 256);
  launch_upsample_forward(d_act11, d_act12, 16, 16, 256);

  launch_conv2d_naive(d_act12, d_output, d_conv5_w, d_conv5_b, 32, 32, 256, 3,
                      3);

  // Compute Loss 
  launch_mse_loss(d_output, d_target, d_loss_val, 32 * 32 * 3);
  cudaDeviceSynchronize();
}

void GPUAutoencoder::backward() {
  // 0. Loss Backward
  launch_mse_loss_backward(d_output, d_target, d_grad_output, 32 * 32 * 3);

  // --- DECODER BACKWARD ---

  // 1. Conv5 Backward
  CUDA_ZERO(d_grad_act12, 32 * 32 * 256);
  launch_conv2d_backward(d_grad_output, d_act12, d_conv5_w, d_grad_act12,
                         d_grad_conv5_w, d_grad_conv5_b, 32, 32, 256, 3, 3);

  // 2. Up2 Backward
  launch_upsample_backward(d_grad_act12, d_grad_act11, 16, 16, 256);

  // 3. ReLU4 Backward
  launch_relu_backward(d_grad_act11, d_grad_act10, d_act10, 16 * 16 * 256);

  // 4. Conv4 Backward
  CUDA_ZERO(d_grad_act9, 16 * 16 * 128);
  launch_conv2d_backward(d_grad_act10, d_act9, d_conv4_w, d_grad_act9,
                         d_grad_conv4_w, d_grad_conv4_b, 16, 16, 128, 256, 3);

  // 5. Up1 Backward
  launch_upsample_backward(d_grad_act9, d_grad_act8, 8, 8, 128);

  // 6. ReLU3 Backward
  launch_relu_backward(d_grad_act8, d_grad_act7, d_act7, 8 * 8 * 128);

  // 7. Conv3 Backward
  CUDA_ZERO(d_grad_act6, 8 * 8 * 128);
  launch_conv2d_backward(d_grad_act7, d_act6, d_conv3_w, d_grad_act6,
                         d_grad_conv3_w, d_grad_conv3_b, 8, 8, 128, 128, 3);

  // --- ENCODER BACKWARD ---

  // 8. Pool2 Backward
  CUDA_ZERO(d_grad_act5, 16 * 16 * 128);
  launch_maxpool_backward(d_grad_act6, d_grad_act5, d_indices2, 8 * 8 * 128);

  // 9. ReLU2 Backward
  launch_relu_backward(d_grad_act5, d_grad_act4, d_act4, 16 * 16 * 128);

  // 10. Conv2 Backward
  CUDA_ZERO(d_grad_act3, 16 * 16 * 256);
  launch_conv2d_backward(d_grad_act4, d_act3, d_conv2_w, d_grad_act3,
                         d_grad_conv2_w, d_grad_conv2_b, 16, 16, 256, 128, 3);

  // 11. Pool1 Backward
  CUDA_ZERO(d_grad_act2, 32 * 32 * 256);
  launch_maxpool_backward(d_grad_act3, d_grad_act2, d_indices1, 16 * 16 * 256);

  // 12. ReLU1 Backward
  launch_relu_backward(d_grad_act2, d_grad_act1, d_act1, 32 * 32 * 256);

  // 13. Conv1 Backward
  CUDA_ZERO(d_grad_input, 32 * 32 * 3);
  launch_conv2d_backward(d_grad_act1, d_input, d_conv1_w, d_grad_input,
                         d_grad_conv1_w, d_grad_conv1_b, 32, 32, 3, 256, 3);

  cudaDeviceSynchronize();
}

// Extract Features 
void GPUAutoencoder::extract_features(const std::vector<float> &host_input,
                                      std::vector<float> &host_latent) {
  set_input(host_input);

  launch_conv2d_naive(d_input, d_act1, d_conv1_w, d_conv1_b, 32, 32, 3, 256, 3);
  launch_relu_forward(d_act1, d_act2, 32 * 32 * 256);
  launch_maxpool_forward(d_act2, d_act3, nullptr, 32, 32, 256);

  launch_conv2d_naive(d_act3, d_act4, d_conv2_w, d_conv2_b, 16, 16, 256, 128,
                      3);
  launch_relu_forward(d_act4, d_act5, 16 * 16 * 128);
  launch_maxpool_forward(d_act5, d_act6, nullptr, 16, 16,
                         128); // d_act6 là Latent (8x8x128)

  cudaDeviceSynchronize();

  int latent_size = 8 * 8 * 128;
  host_latent.resize(latent_size);
  CUDA_CHECK(cudaMemcpy(host_latent.data(), d_act6, latent_size * sizeof(float),
                        cudaMemcpyDeviceToHost));
}

// ============================================
// GPU Auntoencoder version 2 (NEED UPDATEEE)
// ============================================
void GPUAutoencoder::forward_v2() {
  // Encoder
  launch_conv2d_forward_opt(d_input, d_act1, d_conv1_w, d_conv1_b, 32, 32, 3, 256, 3);
  launch_relu_forward_opt(d_act1, d_act2, 32 * 32 * 256);

  launch_maxpool_forward_opt(d_act2, d_act3, d_indices1, 32, 32, 256);

  launch_conv2d_forward_opt(d_act3, d_act4, d_conv2_w, d_conv2_b, 16, 16, 256, 128,
                      3);
  launch_relu_forward_opt(d_act4, d_act5, 16 * 16 * 128);

  launch_maxpool_forward_opt(d_act5, d_act6, d_indices2, 16, 16, 128);

  // Decoder
  launch_conv2d_forward_opt(d_act6, d_act7, d_conv3_w, d_conv3_b, 8, 8, 128, 128, 3);
  launch_relu_forward_opt(d_act7, d_act8, 8 * 8 * 128);
  launch_upsample_forward_opt(d_act8, d_act9, 8, 8, 128);

  launch_conv2d_forward_opt(d_act9, d_act10, d_conv4_w, d_conv4_b, 16, 16, 128, 256,
                      3);
  launch_relu_forward_opt(d_act10, d_act11, 16 * 16 * 256);
  launch_upsample_forward_opt(d_act11, d_act12, 16, 16, 256);

  launch_conv2d_forward_opt(d_act12, d_output, d_conv5_w, d_conv5_b, 32, 32, 256, 3,
                      3);

  // Compute Loss
  launch_mse_loss_opt(d_output, d_target, d_loss_val, 32 * 32 * 3);
  cudaDeviceSynchronize();
}

void GPUAutoencoder::backward_v2() {
  // 0. Loss Backward
  launch_mse_loss_backward_opt(d_output, d_target, d_grad_output, 32 * 32 * 3);

  // --- DECODER BACKWARD ---

  // 1. Conv5 Backward
  CUDA_ZERO(d_grad_act12, 32 * 32 * 256);
  launch_conv2d_backward_opt(d_grad_output, d_act12, d_conv5_w, d_grad_act12,
                         d_grad_conv5_w, d_grad_conv5_b, 32, 32, 256, 3, 3);

  // 2. Up2 Backward
  launch_upsample_backward_opt(d_grad_act12, d_grad_act11, 16, 16, 256);

  // 3. ReLU4 Backward
  launch_relu_backward_opt(d_grad_act11, d_grad_act10, d_act10, 16 * 16 * 256);

  // 4. Conv4 Backward
  CUDA_ZERO(d_grad_act9, 16 * 16 * 128);
  launch_conv2d_backward_opt(d_grad_act10, d_act9, d_conv4_w, d_grad_act9,
                         d_grad_conv4_w, d_grad_conv4_b, 16, 16, 128, 256, 3);

  // 5. Up1 Backward
  launch_upsample_backward_opt(d_grad_act9, d_grad_act8, 8, 8, 128);

  // 6. ReLU3 Backward
  launch_relu_backward_opt(d_grad_act8, d_grad_act7, d_act7, 8 * 8 * 128);

  // 7. Conv3 Backward
  CUDA_ZERO(d_grad_act6, 8 * 8 * 128);
  launch_conv2d_backward_opt(d_grad_act7, d_act6, d_conv3_w, d_grad_act6,
                         d_grad_conv3_w, d_grad_conv3_b, 8, 8, 128, 128, 3);

  // --- ENCODER BACKWARD ---

  // 8. Pool2 Backward
  CUDA_ZERO(d_grad_act5, 16 * 16 * 128);
  launch_maxpool_backward_opt(d_grad_act6, d_grad_act5, d_indices2, 8 * 8 * 128);

  // 9. ReLU2 Backward
  launch_relu_backward_opt(d_grad_act5, d_grad_act4, d_act4, 16 * 16 * 128);

  // 10. Conv2 Backward
  CUDA_ZERO(d_grad_act3, 16 * 16 * 256);
  launch_conv2d_backward_opt(d_grad_act4, d_act3, d_conv2_w, d_grad_act3,
                         d_grad_conv2_w, d_grad_conv2_b, 16, 16, 256, 128, 3);

  // 11. Pool1 Backward
  CUDA_ZERO(d_grad_act2, 32 * 32 * 256);
  launch_maxpool_backward_opt(d_grad_act3, d_grad_act2, d_indices1, 16 * 16 * 256);

  // 12. ReLU1 Backward
  launch_relu_backward_opt(d_grad_act2, d_grad_act1, d_act1, 32 * 32 * 256);

  // 13. Conv1 Backward
  CUDA_ZERO(d_grad_input, 32 * 32 * 3);
  launch_conv2d_backward_opt(d_grad_act1, d_input, d_conv1_w, d_grad_input,
                         d_grad_conv1_w, d_grad_conv1_b, 32, 32, 3, 256, 3);

  cudaDeviceSynchronize();
}

void GPUAutoencoder::extract_features_v2(const std::vector<float> &host_input,
                                         std::vector<float> &host_latent) {
  set_input(host_input);

  // Conv1
  launch_conv2d_forward_opt(d_input, d_act1, d_conv1_w, d_conv1_b, 32, 32, 3,
                              256, 3);
  launch_relu_forward_opt(d_act1, d_act2, 32 * 32 * 256);
  launch_maxpool_forward_opt(d_act2, d_act3, nullptr, 32, 32, 256);

  // Conv2
  launch_conv2d_forward_opt(d_act3, d_act4, d_conv2_w, d_conv2_b, 16, 16, 256,
                              128, 3);
  launch_relu_forward_opt(d_act4, d_act5, 16 * 16 * 128);
  launch_maxpool_forward_opt(d_act5, d_act6, nullptr, 16, 16,
                         128); // d_act6 là Latent

  cudaDeviceSynchronize();

  int latent_size = 8 * 8 * 128;
  host_latent.resize(latent_size);
  CUDA_CHECK(cudaMemcpy(host_latent.data(), d_act6, latent_size * sizeof(float),
                        cudaMemcpyDeviceToHost));
}