mm32.cu

// Copyright (c) 2024-present, Junyeol Ryu, junyeol@aces.snu.ac.kr

#include <stdio.h>
#include <time.h>
#include <cassert>
#include <cublas_v2.h>

#define EPS 1e-3 
#define CHECK_CUDA(e) \
  if ((e) != cudaSuccess) { \
    printf("[%s:%d CudaError]: %s\n", \
        __FILE__, __LINE__, cudaGetErrorString(e)); \
    exit(EXIT_FAILURE);                          \
  } 
#define CHECK_CUBLAS(e)                                  \
  if ((e) != CUBLAS_STATUS_SUCCESS) {                    \
    printf("[%s:%d CublasError]\n", __FILE__, __LINE__); \
    exit(EXIT_FAILURE);                                  \
  }

#define MAX(a, b) (((a) < (b)) ? (b) : (a))
#define MIN(a, b) (((a) < (b)) ? (a) : (b))

#define WARMUP

float *A, *B, *C, *C_ans;
float *A_cuda, *B_cuda, *C_cuda, *C_cublas;

constexpr int M = 4096;
constexpr int K = 4096;
constexpr int N = 4096;

constexpr int BLOCK_M = 16;
constexpr int BLOCK_N = 16;
constexpr int TILE_M = 64;
constexpr int TILE_N = 64;
constexpr int TILE_K = 32;

/* SMEM load */
constexpr int A_N = MAX(MIN(TILE_K, BLOCK_N), 1); // A tile row is loaded by BLOCK_N threads, and can take multiple iterations
constexpr int A_M = (BLOCK_M * BLOCK_N) / A_N;    // Number of A tile rows loaded by a thread block in a single iteration
constexpr int B_N = MAX(MIN(TILE_N, BLOCK_N), 1); // B tile row is loaded by BLOCK_N threads, and can take multiple iterations
constexpr int B_M = (BLOCK_M * BLOCK_N) / B_N;    // Number of A tile rows loaded by a thread block in a single iteration

/* a thread computes REG_M * REG_N C elements */
constexpr int REG_M = (TILE_M + BLOCK_M - 1) / BLOCK_M; // Number of C row elements computed per thread
constexpr int REG_N = (TILE_N + BLOCK_N - 1) / BLOCK_N; // number of C col elements computer per thread

__global__ void mm32(float *A, float *B, float *C, const int M, const int K, const int N)
{
  if (blockIdx.x * TILE_N >= N || blockIdx.y * TILE_M >= M) return;

  const float ZERO = { 0.f };

  __shared__ float A_shared[TILE_M][TILE_K];
  __shared__ float B_shared[TILE_K][TILE_N];

  const int ay = (blockDim.x * threadIdx.y + threadIdx.x) / A_N;
  const int ax = (blockDim.x * threadIdx.y + threadIdx.x) % A_N;
  const int by = (blockDim.x * threadIdx.y + threadIdx.x) / B_N;
  const int bx = (blockDim.x * threadIdx.y + threadIdx.x) % B_N;

  float c_reg[REG_M][REG_N];

  // init c_reg
  for (int i = 0; i < REG_M; ++i) {
    for (int j = 0; j < REG_N; ++j) {
      c_reg[i][j] = ZERO;
    }
  }

  for (int tk = 0; tk < K; tk += TILE_K) {
    // load A
    for (int ii = 0; ii < TILE_M; ii += A_M) {
      int li = ii + ay; // which row in shared A
      int Ai = TILE_M * blockIdx.y + li;  // which row in A
      for (int kk = 0; kk < TILE_K; kk += A_N) {  // load A row iteratively
        int lk = kk + ax; // which col in shared A
        int Ak = tk + lk;
        A_shared[li][lk] = (Ai < M && Ak < K) ? A[Ai * K + Ak] : ZERO;
      }
    }

    // load B
    for (int kk = 0; kk < TILE_K; kk += B_M) {
      int lk = kk + by; // which row in shared B
      int Bk = tk + lk;  // which row in B
      for (int jj = 0; jj < TILE_N; jj += B_N) {  // load B row iteratively
        int lj = jj + bx; // which col in shared B
        int Bj = blockIdx.x * TILE_N + lj;  // which col in B
        B_shared[lk][lj] = (Bk < K && Bj < N) ? B[Bk * N + Bj] : ZERO;
      }
    }
    // sync after load
    __syncthreads();

    for (int y = 0; y < REG_M; ++y) {
      int i = y * BLOCK_M + threadIdx.y;
      for (int x = 0; x < REG_N; ++x) {
        int j = x * BLOCK_N + threadIdx.x;
        for (int k = 0; k < TILE_K; ++k) {
          c_reg[y][x] += A_shared[i][k] * B_shared[k][j];
        }
      }
    }
    // sync after use
    __syncthreads();
  }

  // copy back
  for (int y = 0; y < REG_M; ++y) {
    int i = blockIdx.y * TILE_M + y * BLOCK_M + threadIdx.y;
    for (int x = 0; x < REG_N; ++x) {
      int j = blockIdx.x * TILE_N + x * BLOCK_N + threadIdx.x;
      C[i * N + j] = c_reg[y][x];
    }
  }
}

int main(int argc, char *argv[]) {
  A = (float *)malloc(M * K * sizeof(float));
  B = (float *)malloc(K * N * sizeof(float));
  C = (float *)malloc(M * N * sizeof(float));
  C_ans = (float *)malloc(M * N * sizeof(float));

  CHECK_CUDA(cudaMalloc(&A_cuda, M * K * sizeof(float)));
  CHECK_CUDA(cudaMalloc(&B_cuda, K * N * sizeof(float)));
  CHECK_CUDA(cudaMalloc(&C_cuda, M * N * sizeof(float)));
  CHECK_CUDA(cudaMalloc(&C_cublas, M * N * sizeof(float)));

  for (int i = 0; i < M * K; ++i) {
    A[i] = 2 * (rand() / (double)RAND_MAX);
  }
  for (int i = 0; i < K * N; ++i) {
    B[i] = 2 * (rand() / (double)RAND_MAX);
  }

  CHECK_CUDA(cudaMemcpy(A_cuda, A, M * K * sizeof(float), cudaMemcpyHostToDevice));
  CHECK_CUDA(cudaMemcpy(B_cuda, B, K * N * sizeof(float), cudaMemcpyHostToDevice));

  printf("Running kernel\n");
#ifdef WARMUP
  {
    for (int ii = 0; ii < 10; ++ii) {
      dim3 blockDim{ BLOCK_N, BLOCK_M, 1 };
      dim3 gridDim{ (unsigned int)(N + TILE_N - 1) / TILE_N, (unsigned int)(M + TILE_M - 1) / TILE_M, 1 };
      mm32 << < gridDim, blockDim >> > (A_cuda, B_cuda, C_cuda, M, K, N);
      CHECK_CUDA(cudaGetLastError());
    }
    CHECK_CUDA(cudaDeviceSynchronize());
  }
#endif

  struct timespec s, e;
  clock_gettime(CLOCK_MONOTONIC, &s);
  {
    dim3 blockDim{ BLOCK_N, BLOCK_M, 1 };
    dim3 gridDim{ (unsigned int)(N + TILE_N - 1) / TILE_N, (unsigned int)(M + TILE_M - 1) / TILE_M, 1 };
    mm32 << < gridDim, blockDim >> > (A_cuda, B_cuda, C_cuda, M, K, N);
    CHECK_CUDA(cudaGetLastError());
    CHECK_CUDA(cudaDeviceSynchronize());
  }

  clock_gettime(CLOCK_MONOTONIC, &e);
  double elapsed = (e.tv_sec - s.tv_sec) + ((double)e.tv_nsec - s.tv_nsec) / 1000000000.;
  double bw = 2.0 * M * K * N / 1000000000. / elapsed;
  printf("elapsed time: %lfs, bandwidth: %lf GB/s\n", elapsed, bw);
  CHECK_CUDA(cudaMemcpy(C, C_cuda, M * N * sizeof(float), cudaMemcpyDeviceToHost));

  // cublas verify
  if (argc == 2) {
    cublasHandle_t handle;
    CHECK_CUBLAS(cublasCreate(&handle));

    struct timespec s, e;
    clock_gettime(CLOCK_MONOTONIC, &s);
    {
      printf("Running cublas\n");
      float alpha = 1.f;
      float beta = 0.f;
      CHECK_CUBLAS(cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B_cuda, N, A_cuda, K, &beta, C_cublas, N));
      CHECK_CUDA(cudaDeviceSynchronize());
    }
    clock_gettime(CLOCK_MONOTONIC, &e);
    double elapsed = (e.tv_sec - s.tv_sec) + ((double)e.tv_nsec - s.tv_nsec) / 1000000000.;
    double cublas_bw = 2.0 * M * K * N / 1000000000. / elapsed;
    printf("elapsed time: %lfs, bandwidth: %lf GB/s\n", elapsed, cublas_bw);

    printf("Kernel / cuBlas = %lf / %lf = %lf %%\n", bw, cublas_bw, bw / cublas_bw * 100);
    CHECK_CUDA(cudaMemcpy(C_ans, C_cublas, sizeof(float) * M * N, cudaMemcpyDeviceToHost));

    for (int i = 0; i < M; ++i) {
      for (int j = 0; j < N; ++j) {
        if (fabs((C[i * N + j] - C_ans[i * N + j]) / C[i * N + j]) >= EPS) {
          printf("Validation Failed! C[%d, %d]: %f %f\n", i, j, C[i * N + j], C_ans[i * N + j]);
          exit(1);
        }
      }
    }
    printf("Verification Success!\n");
  }
}