main.cpp

#include "gemv_bf16.h"

int main(int argc, char **argv) {
  bool do_verify = false;
  float alpha = 1.0f;
  float beta = 0.0f;
  int M = 11264;
  int K = 4096;
  // int N = 1;  // Unused
  int lda = K;
  int block_size = 256;

  if (char *value = getCmdOption(argv, argv + argc, "--verify")) {
    do_verify = std::stoi(value) == 1;
  }

  if (char *value = getCmdOption(argv, argv + argc, "--alpha")) {
    alpha = std::stof(value);
  }

  if (char *value = getCmdOption(argv, argv + argc, "-M")) {
    M = std::stoi(value);
  }

  if (char *value = getCmdOption(argv, argv + argc, "-K")) {
    K = std::stoi(value);
    lda = K;
  }

  if (char *value = getCmdOption(argv, argv + argc, "--lda")) {
    lda = std::stoi(value);
  }

  if (char *value = getCmdOption(argv, argv + argc, "-B")) {
    block_size = std::stoi(value);
  }

  // transA=T，因此是行优先
  size_t count_A = (size_t)M * lda;
  size_t size_A = count_A * sizeof(hip_bfloat16);
  size_t size_x = (size_t)K * sizeof(hip_bfloat16);
  size_t size_y = (size_t)M * sizeof(hip_bfloat16);

  // Host 内存分配
  std::vector<hip_bfloat16> h_A(count_A);
  std::vector<hip_bfloat16> h_x(K);
  std::vector<hip_bfloat16> h_y(M);

  // 随机初始数据
  const float rand_max = static_cast<float>(RAND_MAX);
  for (int i = 0; i < count_A; i++)
    h_A[i] = hip_bfloat16(static_cast<float>(rand()) / rand_max);
  for (int i = 0; i < K; i++)
    h_x[i] = hip_bfloat16(static_cast<float>(rand()) / rand_max);
  for (int i = 0; i < M; i++)
    h_y[i] = hip_bfloat16(0.0f);

  // Device 内存分配
  hip_bfloat16 *d_A, *d_x, *d_y;
  checkHipErrors(hipMalloc(&d_A, size_A));
  checkHipErrors(hipMalloc(&d_x, size_x));
  checkHipErrors(hipMalloc(&d_y, size_y));

  checkHipErrors(hipMemcpy(d_A, h_A.data(), size_A, hipMemcpyHostToDevice));
  checkHipErrors(hipMemcpy(d_x, h_x.data(), size_x, hipMemcpyHostToDevice));
  checkHipErrors(hipMemcpy(d_y, h_y.data(), size_y, hipMemcpyHostToDevice));

  // Kernel 注册表
  std::vector<KernelCase> kernels;

  constexpr bool NTL = true;
  constexpr int UNROLL = 4;
  constexpr int TILE_K = calculate_tile_k<8>(4);
  constexpr int ROWS_PER_WARP = 2;

  kernels.push_back(
      {"naive", [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
                    const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int grid = (M + block_size - 1) / block_size;
         gemv_bf16_TN_naive<<<grid, block_size>>>(M, K, alpha, A, lda, x, beta,
                                                  y);
       }});

  kernels.push_back(
      {"vec8", [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
                   const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int grid = (M + block_size - 1) / block_size;
         gemv_bf16_TN_vec<<<grid, block_size>>>(M, K, alpha, A, lda, x, beta,
                                                y);
       }});

  kernels.push_back(
      {"vec8_ntl",
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int grid = (M + block_size - 1) / block_size;
         gemv_bf16_TN_vec<NTL>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"warp", [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
                   const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_warp<<<grid, block_size>>>(M, K, alpha, A, lda, x, beta,
                                                 y);
       }});

  kernels.push_back(
      {"vec8+warp",
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp<<<grid, block_size>>>(M, K, alpha, A, lda, x,
                                                     beta, y);
       }});

  kernels.push_back(
      {"vec8_ntl+warp",
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp<NTL>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8+warp_mr" + std::to_string(ROWS_PER_WARP),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid =
             ((M + ROWS_PER_WARP - 1) / ROWS_PER_WARP + warps_per_block - 1) /
             warps_per_block;
         gemv_bf16_TN_vec_warp_mr<!NTL, ROWS_PER_WARP>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8_ntl+warp_mr" + std::to_string(ROWS_PER_WARP),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid =
             ((M + ROWS_PER_WARP - 1) / ROWS_PER_WARP + warps_per_block - 1) /
             warps_per_block;
         gemv_bf16_TN_vec_warp_mr<NTL, ROWS_PER_WARP>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8+warp+unroll" + std::to_string(UNROLL),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp_unroll<!NTL, UNROLL>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8_ntl+warp+unroll" + std::to_string(UNROLL),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp_unroll<NTL, UNROLL>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8+warp+shm" + std::to_string(TILE_K),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp_shm<!NTL, TILE_K>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8_ntl+warp+shm" + std::to_string(TILE_K),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp_shm<NTL, TILE_K>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8+warp+unroll" + std::to_string(UNROLL) + "+shm" +
           std::to_string(TILE_K),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp_unroll_shm<!NTL, UNROLL, TILE_K>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8_ntl+warp+unroll" + std::to_string(UNROLL) + "+shm" +
           std::to_string(TILE_K),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid = (M + warps_per_block - 1) / warps_per_block;
         gemv_bf16_TN_vec_warp_unroll_shm<NTL, UNROLL, TILE_K>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8+warp_mr" + std::to_string(ROWS_PER_WARP) + "+shm" +
           std::to_string(TILE_K),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid =
             ((M + ROWS_PER_WARP - 1) / ROWS_PER_WARP + warps_per_block - 1) /
             warps_per_block;
         gemv_bf16_TN_vec_warp_mr_shm<!NTL, TILE_K, ROWS_PER_WARP>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  kernels.push_back(
      {"vec8_ntl+warp_mr" + std::to_string(ROWS_PER_WARP) + "+shm" +
           std::to_string(TILE_K),
       [&](int M, int K, float alpha, const hip_bfloat16 *A, int lda,
           const hip_bfloat16 *x, float beta, hip_bfloat16 *y) {
         int warps_per_block = block_size / WARP_SIZE;
         int grid =
             ((M + ROWS_PER_WARP - 1) / ROWS_PER_WARP + warps_per_block - 1) /
             warps_per_block;
         gemv_bf16_TN_vec_warp_mr_shm<NTL, TILE_K, ROWS_PER_WARP>
             <<<grid, block_size>>>(M, K, alpha, A, lda, x, beta, y);
       }});

  // 运行所有测试
  run_benchmark(kernels, M, K, alpha, d_A, lda, d_x, beta, d_y, do_verify);

  // 清理
  checkHipErrors(hipFree(d_A));
  checkHipErrors(hipFree(d_x));
  checkHipErrors(hipFree(d_y));

  return 0;
}