CUDA矩阵乘法的优化

#include <cstdio>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cutlass/gemm/device/gemm.h>
#include <cutlass/util/host_tensor.h>
#include <cutlass/util/reference/device/tensor_fill.h>
#include <functional>

///////////////////////////////////////////////////////////////////////////////////////////////////
//
// This function defines a CUTLASS GEMM kernel instantiation, constructs its
// parameters object, and launches it on the CUDA device.
//
///////////////////////////////////////////////////////////////////////////////////////////////////

/// Define a CUTLASS GEMM template and launch a GEMM kernel.
cudaError_t CutlassSgemmNN(int M, int N, int K, float alpha, float const *A,
                           int lda, float const *B, int ldb, float beta,
                           float *C, int ldc) {

  // Define type definition for single-precision CUTLASS GEMM with column-major
  // input matrices and 128x128x8 threadblock tile size (chosen by default).
  //
  // To keep the interface manageable, several helpers are defined for plausible
  // compositions including the following example for single-precision GEMM.
  // Typical values are used as default template arguments. See
  // `cutlass/gemm/device/default_gemm_configuration.h` for more details.
  //
  // To view the full gemm device API interface, see
  // `cutlass/gemm/device/gemm.h`

  using ColumnMajor = cutlass::layout::ColumnMajor;

  using CutlassGemm =
      cutlass::gemm::device::Gemm<float,        // Data-type of A matrix
                                  ColumnMajor,  // Layout of A matrix
                                  float,        // Data-type of B matrix
                                  ColumnMajor,  // Layout of B matrix
                                  float,        // Data-type of C matrix
                                  ColumnMajor>; // Layout of C matrix

  // Define a CUTLASS GEMM type
  CutlassGemm gemm_operator;

  // Construct the CUTLASS GEMM arguments object.
  //
  // One of CUTLASS's design patterns is to define gemm argument objects that
  // are constructible in host code and passed to kernels by value. These may
  // include pointers, strides, scalars, and other arguments needed by Gemm and
  // its components.
  //
  // The benefits of this pattern are (1.) a structured, composable strategy for
  // passing host-constructible arguments to kernels and (2.) minimized
  // initialization overhead on kernel entry.
  //
  CutlassGemm::Arguments args(
      {M, N, K}, // Gemm Problem dimensions
      {A, lda},  // Tensor-ref for source matrix A
      {B, ldb},  // Tensor-ref for source matrix B
      {C, ldc},  // Tensor-ref for source matrix C
      {C, ldc},  // Tensor-ref for destination matrix D (may be different memory
                 // than source C matrix)
      {alpha, beta}); // Scalars used in the Epilogue

  //
  // Launch the CUTLASS GEMM kernel.
  //

  cutlass::Status status = gemm_operator(args);

  //
  // Return a cudaError_t if the CUTLASS GEMM operator returned an error code.
  //

  if (status != cutlass::Status::kSuccess) {
    return cudaErrorUnknown;
  }

  // Return success, if no errors were encountered.
  return cudaSuccess;
}

namespace Alcanderian // https://github.com/Alcanderian/CUDA-tutorial/sgemm
{
__global__ void cuda_kernel_sgemm_1(float *a, float *b, float *c, size_t N,
                                    size_t M, size_t K, float alpha,
                                    float beta) {
  int tr = threadIdx.x;                   // row idx in block
  int tc = threadIdx.y;                   // col idx in block
  int ir = blockIdx.x * 32 + threadIdx.x; // row idx in global
  int ic = blockIdx.y * 32 + threadIdx.y; // col idx in global

  __shared__ float a_sub[32][32];
  __shared__ float b_sub[32][32];

  int load_size = K / 32;
  if (K % 32 != 0) {
    load_size += 1;
  }
  float acc = 0.0f;
  int a_ir = ir;
  int b_ic = ic;
#define idx(ri, ci, nc) ((ri) * (nc) + (ci))
  for (int l = 0; l < load_size; ++l) {
    int a_ic = l * 32 + tc;
    int b_ir = l * 32 + tr;
    a_sub[tr][tc] = 0.0f;
    b_sub[tr][tc] = 0.0f;
    if (a_ir < M && a_ic < K)
      a_sub[tr][tc] = a[idx(a_ir, a_ic, K)];
    if (b_ir < K && b_ic < N)
      b_sub[tr][tc] = b[idx(b_ir, b_ic, N)];

    __syncthreads();

#pragma unroll
    for (int k = 0; k < 32; ++k) {
      acc += a_sub[tr][k] * b_sub[k][tc];
    }

    __syncthreads();
  }

  if (ir < M && ic < N)
    c[idx(ir, ic, N)] = alpha * acc + beta * c[idx(ir, ic, N)];
#undef idx
}

// use __ldg & avoid bank conflict
__global__ void cuda_kernel_sgemm_2(float *__restrict__ a,
                                    float *__restrict__ b,
                                    float *__restrict__ c, size_t N, size_t M,
                                    size_t K, float alpha, float beta) {
  int tr = threadIdx.x;                   // row idx in block
  int tc = threadIdx.y;                   // col idx in block
  int ir = blockIdx.x * 32 + threadIdx.x; // row idx in global
  int ic = blockIdx.y * 32 + threadIdx.y; // col idx in global

  __shared__ float a_sub[32][32 + 1]; // avoid bank conflict
  __shared__ float b_sub[32][32 + 1];

  int load_size = K / 32;
  if (K % 32 != 0) {
    load_size += 1;
  }
  float acc = 0.0f;
  int a_ir = ir;
  int b_ic = ic;
#define idx(ri, ci, nc) ((ri) * (nc) + (ci))
  for (int l = 0; l < load_size; ++l) {
    int a_ic = l * 32 + tc;
    int b_ir = l * 32 + tr;
    a_sub[tr][tc] = 0.0f;
    b_sub[tr][tc] = 0.0f;
    if (a_ir < M && a_ic < K)
      a_sub[tr][tc] = a[idx(a_ir, a_ic, K)];
    if (b_ir < K && b_ic < N)
      b_sub[tr][tc] = b[idx(b_ir, b_ic, N)];

    __syncthreads();

#pragma unroll
    for (int k = 0; k < 32; ++k) {
      acc += a_sub[tr][k] * b_sub[k][tc];
    }

    __syncthreads();
  }

  if (ir < M && ic < N)
    c[idx(ir, ic, N)] = alpha * acc + beta * c[idx(ir, ic, N)];
#undef idx
}

// use __ldg without avoiding bank conflict
__global__ void cuda_kernel_sgemm_3(float *__restrict__ a,
                                    float *__restrict__ b,
                                    float *__restrict__ c, size_t N, size_t M,
                                    size_t K, float alpha, float beta) {
  int tr = threadIdx.x;                   // row idx in block
  int tc = threadIdx.y;                   // col idx in block
  int ir = blockIdx.x * 32 + threadIdx.x; // row idx in global
  int ic = blockIdx.y * 32 + threadIdx.y; // col idx in global

  __shared__ float a_sub[32][32]; // avoid bank conflict
  __shared__ float b_sub[32][32];

  int load_size = K / 32;
  if (K % 32 != 0) {
    load_size += 1;
  }
  float acc = 0.0f;
  int a_ir = ir;
  int b_ic = ic;
#define idx(ri, ci, nc) ((ri) * (nc) + (ci))
  for (int l = 0; l < load_size; ++l) {
    int a_ic = l * 32 + tc;
    int b_ir = l * 32 + tr;
    a_sub[tr][tc] = 0.0f;
    b_sub[tr][tc] = 0.0f;
    if (a_ir < M && a_ic < K)
      a_sub[tr][tc] = a[idx(a_ir, a_ic, K)];
    if (b_ir < K && b_ic < N)
      b_sub[tr][tc] = b[idx(b_ir, b_ic, N)];

    __syncthreads();

#pragma unroll
    for (int k = 0; k < 32; ++k) {
      acc += a_sub[tr][k] * b_sub[k][tc];
    }

    __syncthreads();
  }

  if (ir < M && ic < N)
    c[idx(ir, ic, N)] = alpha * acc + beta * c[idx(ir, ic, N)];
#undef idx
}
}; // namespace Alcanderian

namespace fynv // https://github.com/fynv/optimal_sgemm_cuda_c
{
struct f8 {
  float4 a, b;
  __device__ inline f8() { memset(this, 0, sizeof(f8)); }
};

struct f88 {
  f8 a, b, c, d, e, f, g, h;
};

__device__ inline void d_load8(const float *p, f8 &c) {
  c.a = ((float4 *)p)[0];
  c.b = ((float4 *)p)[16];
}

__device__ inline void d_store8(float *p, const f8 &c) {
  ((float4 *)p)[0] = c.a;
  ((float4 *)p)[16] = c.b;
}

__device__ inline void d_mult8v(f8 &c, const f8 &a, float b) {
  c.a.x += a.a.x * b;
  c.a.y += a.a.y * b;
  c.a.z += a.a.z * b;
  c.a.w += a.a.w * b;
  c.b.x += a.b.x * b;
  c.b.y += a.b.y * b;
  c.b.z += a.b.z * b;
  c.b.w += a.b.w * b;
}

template <typename T> __device__ inline void Swap(T &a, T &b) {
  T t = a;
  a = b;
  b = t;
}

__global__ __launch_bounds__(
    256) // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#launch-bounds
    void g_fgemm(float *d_C, const float *d_A, const float *d_B, int n, int lda,
                 int ldb, int ldc) {
  int x_a = threadIdx.x & 31;
  int y_a = threadIdx.x >> 5;

  int x_b = threadIdx.x & 1;
  int y_b = threadIdx.x >> 1;

  int x_c = threadIdx.x & 15;
  int y_c = threadIdx.x >> 4;

  __shared__ float smem[4096];
  float *s_A1 = smem;
  float *s_A2 = smem + 1024;
  float *s_B1 = smem + 2048;
  float *s_B2 = smem + 3072;

  f88 l_C;

  const float *p_A = d_A + (blockIdx.x << 7);
  const float *p_B = d_B + (blockIdx.y << 7) * ldb;

  float4 p, q;
  p = ((float4 *)p_A)[y_a * (lda >> 2) + x_a];
  q = ((float4 *)p_B)[y_b * (ldb >> 2) + x_b];

  for (int i = 0; i < n; i += 8) {
    ((float4 *)s_A1)[threadIdx.x] = p;
    s_B1[(((x_b << 2) + 0) << 7) + y_b] = q.x;
    s_B1[(((x_b << 2) + 1) << 7) + y_b] = q.y;
    s_B1[(((x_b << 2) + 2) << 7) + y_b] = q.z;
    s_B1[(((x_b << 2) + 3) << 7) + y_b] = q.w;
    __syncthreads();

    if (i + 8 < n) {
      p_A += (lda << 3);
      p_B += 8;
      p = ((float4 *)p_A)[y_a * (lda >> 2) + x_a];
      q = ((float4 *)p_B)[y_b * (ldb >> 2) + x_b];
    }

    for (int j = 0; j < 8; j++) {
      float *p_s_A = s_A1 + (j << 7) + (x_c << 2);
      float *p_s_B = s_B1 + (j << 7) + (y_c << 2);

      f8 a, b;
      d_load8(p_s_A, a);
      d_load8(p_s_B, b);

      d_mult8v(l_C.a, a, b.a.x);
      d_mult8v(l_C.b, a, b.a.y);
      d_mult8v(l_C.c, a, b.a.z);
      d_mult8v(l_C.d, a, b.a.w);
      d_mult8v(l_C.e, a, b.b.x);
      d_mult8v(l_C.f, a, b.b.y);
      d_mult8v(l_C.g, a, b.b.z);
      d_mult8v(l_C.h, a, b.b.w);
    }

    Swap(s_A1, s_A2);
    Swap(s_B1, s_B2);
  }

  float *p_C = d_C + ((blockIdx.x << 7) + (x_c << 2)) +
               ((blockIdx.y << 7) + (y_c << 2)) * ldc;
  d_store8(p_C, l_C.a);
  p_C += ldc;
  d_store8(p_C, l_C.b);
  p_C += ldc;
  d_store8(p_C, l_C.c);
  p_C += ldc;
  d_store8(p_C, l_C.d);
  p_C += (ldc * 61);
  d_store8(p_C, l_C.e);
  p_C += ldc;
  d_store8(p_C, l_C.f);
  p_C += ldc;
  d_store8(p_C, l_C.g);
  p_C += ldc;
  d_store8(p_C, l_C.h);
}
}; // namespace fynv

namespace wuk {
#define IDX2C(i, j, ld) (((j) * (ld)) + (i))

template <typename T>
__global__
__launch_bounds__(1024) void gemm_32x32_v0(int m, int n, int k, T alpha,
                                           const T *A, int lda, const T *B,
                                           int ldb, T beta, T *C, int ldc) {
  const int y_range = ((int)blockIdx.y + 1 << 5) - m,
            x_range = ((int)blockIdx.x + 1 << 5) - n;
  if (y_range > 0) {
    A -= y_range;
    C -= y_range;
  }
  if (x_range > 0) {
    B -= x_range * ldb;
    C -= x_range * ldc;
  }
  A += blockIdx.y << 5;
  B += (blockIdx.x << 5) * ldb;
  C += (blockIdx.y << 5) + (blockIdx.x << 5) * ldc;

  T resC = 0;
  for (int i = 0; i < k; ++i) {
    const T resA = A[IDX2C(threadIdx.y, i, lda)],
            resB = B[IDX2C(i, threadIdx.x, ldb)];
    resC += resA * resB;
  }
  resC = resC * alpha + C[IDX2C(threadIdx.y, threadIdx.x, ldc)] * beta;
  C[IDX2C(threadIdx.y, threadIdx.x, ldc)] = resC;
}

template <typename T>
__global__
__launch_bounds__(1024) void gemm_32x32_v1(int m, int n, int k, T alpha,
                                           const T *A, int lda, const T *B,
                                           int ldb, T beta, T *C, int ldc) {
  const int x_range = ((int)blockIdx.x + 1 << 5) - m,
            y_range = ((int)blockIdx.y + 1 << 5) - n;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }
  A += blockIdx.x << 5;
  B += (blockIdx.y << 5) * ldb;
  C += (blockIdx.x << 5) + (blockIdx.y << 5) * ldc;
  T resC = 0;
  for (int i = 0; i < k; ++i) {
    const T resA = A[IDX2C(threadIdx.x, i, lda)],
            resB = B[IDX2C(i, threadIdx.y, ldb)];
    resC += resA * resB;
  }
  resC = resC * alpha + C[IDX2C(threadIdx.x, threadIdx.y, ldc)] * beta;
  C[IDX2C(threadIdx.x, threadIdx.y, ldc)] = resC;
}

template <typename T>
__global__
__launch_bounds__(1024) void gemm_32x32_v2(int m, int n, int k, T alpha,
                                           const T *A, int lda, const T *B,
                                           int ldb, T beta, T *C, int ldc) {
  const int idx = threadIdx.x, idy = threadIdx.y,
            x_range = ((int)blockIdx.x + 1 << 5) - m,
            y_range = ((int)blockIdx.y + 1 << 5) - n;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }
  A += blockIdx.x << 5;
  B += (blockIdx.y << 5) * ldb;
  C += (blockIdx.x << 5) + (blockIdx.y << 5) * ldc;
  T resC = 0;
  for (int i = 0; i < k; ++i) {
    const T resA = A[IDX2C(idx, i, lda)], resB = B[IDX2C(i, idy, ldb)];
    resC += resA * resB;
  }
  resC = resC * alpha + C[IDX2C(idx, idy, ldc)] * beta;
  C[IDX2C(idx, idy, ldc)] = resC;
}

template <typename T, int TILE_WIDTH>
__global__ __launch_bounds__(TILE_WIDTH *TILE_WIDTH) void gemm_32x32_v3(
    int m, int n, int k, T alpha, const T *A, int lda, const T *B, int ldb,
    T beta, T *C, int ldc) {
  const int idx = threadIdx.x, idy = threadIdx.y,
            x_range = ((int)blockIdx.x + 1 << 5) - m,
            y_range = ((int)blockIdx.y + 1 << 5) - n;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }
  A += blockIdx.x << 5;
  B += (blockIdx.y << 5) * ldb;
  C += (blockIdx.x << 5) + (blockIdx.y << 5) * ldc;
  T resC = 0;
  __shared__ T sA[TILE_WIDTH][TILE_WIDTH];
  __shared__ T sB[TILE_WIDTH][TILE_WIDTH];
  for (int i = 0; i < k; i += TILE_WIDTH) {
    sA[idx][idy] = A[IDX2C(idx, i + idy, lda)];
    sB[idx][idy] = B[IDX2C(i + idx, idy, ldb)];
    __syncthreads();
    for (int j = 0; j < TILE_WIDTH; ++j)
      resC += sA[idx][j] * sB[j][idy];
    __syncthreads();
  }
  resC = resC * alpha + C[IDX2C(idx, idy, ldc)] * beta;
  C[IDX2C(idx, idy, ldc)] = resC;
}

template <typename T, int TILE_WIDTH>
__global__ __launch_bounds__(TILE_WIDTH *TILE_WIDTH) void gemm_32x32_v4(
    int m, int n, int k, T alpha, const T *A, int lda, const T *B, int ldb,
    T beta, T *C, int ldc) {
  const int idx = threadIdx.x, idy = threadIdx.y,
            x_range = ((int)blockIdx.x + 1 << 5) - m,
            y_range = ((int)blockIdx.y + 1 << 5) - n;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }
  A += blockIdx.x << 5;
  B += (blockIdx.y << 5) * ldb;
  C += (blockIdx.x << 5) + (blockIdx.y << 5) * ldc;
  T resC = 0;
  __shared__ T sA[TILE_WIDTH][TILE_WIDTH | 1];
  __shared__ T sB[TILE_WIDTH][TILE_WIDTH | 1];
  for (int i = 0; i < k; i += TILE_WIDTH) {
    sA[idx][idy] = A[IDX2C(idx, i + idy, lda)];
    sB[idx][idy] = B[IDX2C(i + idx, idy, ldb)];
    __syncthreads();
    for (int j = 0; j < TILE_WIDTH; ++j)
      resC += sA[idx][j] * sB[j][idy];
    __syncthreads();
  }
  resC = resC * alpha + C[IDX2C(idx, idy, ldc)] * beta;
  C[IDX2C(idx, idy, ldc)] = resC;
}
#undef IDX2C

template <typename T, typename Tv, int TWO_BUFFER>
__global__
__launch_bounds__(256) void gemm_32x32_v5(int m, int n, int k, T alpha,
                                          const T *A, int lda, const T *B,
                                          int ldb, T beta, T *C, int ldc) {
  const int TvSize = sizeof(Tv) / sizeof(T), idx = threadIdx.x,
            x_range = ((int)blockIdx.x + 1 << 5) - m,
            y_range = ((int)blockIdx.y + 1 << 5) - n;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }
  A += blockIdx.x << 5;
  B += (blockIdx.y << 5) * ldb;
  C += (blockIdx.x << 5) + (blockIdx.y << 5) * ldc;
  Tv ansC;
  memset(&ansC, 0, sizeof(ansC));

  __shared__ T s_buffer[2048];
  T *s_A = s_buffer;
  T *s_B = s_buffer + 1024;

#if TWO_BUFFER
  __shared__ T s_tbuffer[2048];
  T *s_tA = s_tbuffer;
  T *s_tB = s_tbuffer + 1024;
#endif

  Tv resA = *(Tv *)&A[((idx * TvSize) & 31) + ((idx * TvSize) >> 5) * lda],
     resB = *(Tv *)&B[((idx * TvSize) & 31) + ((idx * TvSize) >> 5) * ldb];

  for (int l = 0; l < k; l += 32) {
    ((Tv *)s_A)[idx] = resA;
    for (int i = 0; i < TvSize; ++i)
      s_B[(((idx * TvSize) & 31) << 5) + ((idx * TvSize) >> 5)] =
          ((T *)&resB)[i];
    __syncthreads();

    if (l + 32 < k) {
      A += lda << 5;
      B += 32;
      resA = *(Tv *)&A[((idx * TvSize) & 31) + ((idx * TvSize) >> 5) * lda];
      resB = *(Tv *)&B[((idx * TvSize) & 31) + ((idx * TvSize) >> 5) * ldb];
    }

    for (int j = 0; j < 32; ++j) {
      Tv tmpA = *(Tv *)&s_A[((idx * TvSize) & 31) + (j << 5)];
      T tmpB = s_B[(j << 5) + ((idx * TvSize) >> 5)];
      for (int i = 0; i < TvSize; ++i)
        ((T *)&ansC)[i] += ((T *)&tmpA)[i] * tmpB;
    }

#if TWO_BUFFER
    {
      T *tmp_A = s_A;
      s_A = s_tA;
      s_tA = tmp_A;
    }
    {
      T *tmp_B = s_B;
      s_B = s_tB;
      s_tB = tmp_B;
    }
#else
    __syncthreads();
#endif
  }

  {
    Tv *devC = (Tv *)&C[((idx * TvSize) & 31) + ((idx * TvSize) >> 5) * ldc],
       resC = *devC;
    for (int i = 0; i < TvSize; ++i)
      ((T *)&resC)[i] = alpha * ((T *)&ansC)[i] + beta * ((T *)&resC)[i];
    *devC = resC;
  }
}

template <typename T, typename Tv, int TWO_BUFFER>
__global__ __launch_bounds__(256) void gemm_64x64(int m, int n, int k, T alpha,
                                                  const T *A, int lda,
                                                  const T *B, int ldb, T beta,
                                                  T *C, int ldc) {
  const int TvSize = sizeof(Tv) / sizeof(T), idx = threadIdx.x,
            x_range = ((int)blockIdx.x + 1 << 6) - m,
            y_range = ((int)blockIdx.y + 1 << 6) - n;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }
  A += blockIdx.x << 6;
  B += (blockIdx.y << 6) * ldb;
  C += (blockIdx.x << 6) + (blockIdx.y << 6) * ldc;
  Tv ansC[TvSize];
  memset(ansC, 0, sizeof(ansC));

  __shared__ T s_buffer[2048];
  T *s_A = s_buffer;
  T *s_B = s_buffer + 1024;

#if TWO_BUFFER
  __shared__ T s_tbuffer[2048];
  T *s_tA = s_tbuffer;
  T *s_tB = s_tbuffer + 1024;
#endif

  Tv resA = *(Tv *)&(A[((idx * TvSize) & 63) + ((idx * TvSize) >> 6) * lda]),
     resB = *(Tv *)&(B[((idx * TvSize) & 15) + ((idx * TvSize) >> 4) * ldb]);
  for (int l = 0; l < k; l += 16) {
    ((Tv *)s_A)[idx] = resA;
    for (int i = 0; i < TvSize; ++i)
      s_B[((((idx * TvSize) & 15) + i) << 6) + ((idx * TvSize) >> 4)] =
          ((T *)&resB)[i];
    __syncthreads();

    if (l + 16 < k) {
      A += lda << 4;
      B += 16;
      resA = *(Tv *)&(A[((idx * TvSize) & 63) + ((idx * TvSize) >> 6) * lda]);
      resB = *(Tv *)&(B[((idx * TvSize) & 15) + ((idx * TvSize) >> 4) * ldb]);
    }

    for (int j = 0; j < 16; ++j) {
      const Tv tmpA = *(Tv *)&(s_A[((idx * TvSize) & 63) + (j << 6)]),
               tmpB = *(Tv *)&(s_B[(j << 6) + (idx >> 4) * TvSize]);
      for (int i = 0; i < TvSize; ++i)
        for (int h = 0; h < TvSize; ++h)
          ((T *)&ansC[i])[h] += ((T *)&tmpA)[i] * ((T *)&tmpB)[h];
    }

#if TWO_BUFFER
    {
      T *tmp_A = s_A;
      s_A = s_tA;
      s_tA = tmp_A;
    }
    {
      T *tmp_B = s_B;
      s_B = s_tB;
      s_tB = tmp_B;
    }
#else
    __syncthreads();
#endif
  }
  for (int i = 0; i < TvSize; ++i) {
    Tv *devC = (Tv *)&(C[(idx & 15) * 4 + ((idx >> 4) * 4 + i) * ldc]),
       resC = *devC;
    for (int h = 0; h < TvSize; ++h)
      ((T *)&resC)[h] = alpha * ((T *)&ansC[i])[h] + beta * ((T *)&resC)[h];
    *devC = resC;
  }
}

template <typename T, typename Tv, int TWO_BUFFER>
__global__
__launch_bounds__(256) void gemm_128x128(int m, int n, int k, T alpha,
                                         const T *A, int lda, const T *B,
                                         int ldb, T beta, T *C, int ldc) {
  const int TvSize = sizeof(Tv) / sizeof(T), idx = threadIdx.x,
            x_range = ((int)blockIdx.x + 1 << 7) - m,
            y_range = ((int)blockIdx.y + 1 << 7) - n, x_a = idx & 31,
            y_a = idx >> 5, x_b = idx & 1, y_b = idx >> 1, x_c = idx & 15,
            y_c = idx >> 4;
  if (x_range > 0) {
    A -= x_range;
    C -= x_range;
  }
  if (y_range > 0) {
    B -= y_range * ldb;
    C -= y_range * ldc;
  }

  A += blockIdx.x << 7;
  B += (blockIdx.y << 7) * ldb;
  C +=
      ((blockIdx.x << 7) + (x_c << 2)) + ((blockIdx.y << 7) + (y_c << 2)) * ldc;

  __shared__ T s_buffer[2048];
  T *s_A = s_buffer;
  T *s_B = s_buffer + 1024;

#if TWO_BUFFER
  __shared__ T s_tbuffer[2048];
  T *s_tA = s_tbuffer;
  T *s_tB = s_tbuffer + 1024;
#endif
  Tv ansC[2][2][TvSize];
  memset(ansC, 0, sizeof(ansC));

  Tv resA = ((Tv *)A)[x_a + y_a * (lda >> 2)];
  Tv resB = ((Tv *)B)[x_b + y_b * (ldb >> 2)];

  for (int l = 0; l < k; l += 8) {
    ((Tv *)s_A)[idx] = resA;
    for (int i = 0; i < TvSize; ++i)
      s_B[(((x_b << 2) + i) << 7) + y_b] = ((T *)&resB)[i];
    __syncthreads();

    if (l + 8 < k) {
      A += lda << 3;
      B += 8;
      resA = ((Tv *)A)[x_a + y_a * (lda >> 2)];
      resB = ((Tv *)B)[x_b + y_b * (ldb >> 2)];
    }

#pragma unroll(4)
    for (int j = 0; j < 8; ++j) {
      Tv a[2], b[2];
      for (int p = 0; p < 2; ++p)
        a[p] = ((Tv *)(s_A + (j << 7) + (x_c << 2)))[p << 4];
      for (int p = 0; p < 2; ++p)
        b[p] = ((Tv *)(s_B + (j << 7) + (y_c << 2)))[p << 4];
      for (int p = 0; p < 2; ++p)
        for (int q = 0; q < 2; ++q)
          for (int i = 0; i < TvSize; ++i)
            for (int j = 0; j < TvSize; ++j)
              ((T *)&ansC[p][q][i])[j] += ((T *)&a[p])[j] * ((T *)&b[q])[i];
    }

#if TWO_BUFFER
    {
      T *tmp_A = s_A;
      s_A = s_tA;
      s_tA = tmp_A;
    }
    {
      T *tmp_B = s_B;
      s_B = s_tB;
      s_tB = tmp_B;
    }
#else
    __syncthreads();
#endif
  }

  for (int p = 0; p < 2; ++p)
    for (int q = 0; q < 2; ++q)
      for (int i = 0; i < TvSize; ++i) {
        Tv *devC = ((Tv *)(C + ldc * (q * 64 + i))) + (p << 4), resC = *devC;
        for (int j = 0; j < TvSize; ++j)
          ((T *)&resC)[j] =
              alpha * ((T *)&ansC[p][q][i])[j] + beta * ((T *)&resC)[j];
        *devC = resC;
      }
}
}; // namespace wuk

void WuK_Timer(const char *tag, float flo, const std::function<void()> &kernel,
               int test_time = 9) {
  float min_time = 9e99;
  while (test_time--) {
    cudaEvent_t beg, end;
    cudaEventCreate(&beg);
    cudaEventCreate(&end);
    cudaEventRecord(beg);
    kernel();
    cudaEventRecord(end);
    cudaEventSynchronize(beg);
    cudaEventSynchronize(end);
    float elapsed_time;
    cudaEventElapsedTime(&elapsed_time, beg, end);
    min_time = std::min(min_time, elapsed_time);
  }
  std::printf("%s: %f ms, %e FLOPS.\n", tag, min_time, flo * 1e3 / min_time);
}

struct WuK_cublas {
  cublasHandle_t handle;
  WuK_cublas() { cublasCreate(&handle); }
  ~WuK_cublas() { cublasDestroy(handle); }
} wuk_cublas;

const float alpha = 1, beta = 0;
const cublasOperation_t opA = CUBLAS_OP_N, opB = CUBLAS_OP_N, opC = CUBLAS_OP_N;
const int m = 1 << 13, n = 1 << 13, k = 1 << 13,
          lda = opA == CUBLAS_OP_N ? k : m, ldb = opB == CUBLAS_OP_N ? n : k,
          ldc = opC == CUBLAS_OP_N ? n : m;
cutlass::HostTensor<float, cutlass::layout::ColumnMajor> dA({m, k}), dB({k, n}),
    dC({m, n});

int main() {
  cutlass::reference::device::TensorFill(dA.device_view(),
                                         static_cast<float>(2));
  cutlass::reference::device::TensorFill(dB.device_view(),
                                         static_cast<float>(1));
  cutlass::reference::device::TensorFill(dC.device_view(),
                                         static_cast<float>(0));
  WuK_Timer("cublasSgemm", 2.0 * m * k * n, [&] {
    cublasSgemm(wuk_cublas.handle, opA, opB, m, n, k, &alpha, dA.device_data(),
                lda, dB.device_data(), ldb, &beta, dC.device_data(), ldc);
  });
  WuK_Timer("CutlassSgemmNN", 2.0 * m * k * n, [&] {
    CutlassSgemmNN(m, n, k, alpha, dA.device_data(), lda, dB.device_data(), ldb,
                   beta, dC.device_data(), ldc);
  });
  WuK_Timer("Alcanderian::cuda_kernel_sgemm_1", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    Alcanderian::cuda_kernel_sgemm_1<<<gridDim, blockDim>>>(
        dA.device_data(), dB.device_data(), dC.device_data(), n, m, k, alpha,
        beta);
  });
  WuK_Timer("Alcanderian::cuda_kernel_sgemm_2", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    Alcanderian::cuda_kernel_sgemm_2<<<gridDim, blockDim>>>(
        dA.device_data(), dB.device_data(), dC.device_data(), n, m, k, alpha,
        beta);
  });
  WuK_Timer("Alcanderian::cuda_kernel_sgemm_3", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    Alcanderian::cuda_kernel_sgemm_3<<<gridDim, blockDim>>>(
        dA.device_data(), dB.device_data(), dC.device_data(), n, m, k, alpha,
        beta);
  });
  WuK_Timer("fynv::g_fgemm", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 128;
    const dim3 blockDim(256), gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                                      (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m % TILE_WIDTH == 0);
    assert(n % TILE_WIDTH == 0);
    assert(k % TILE_WIDTH == 0);
    fynv::g_fgemm<<<gridDim, blockDim>>>(dC.device_data(), dA.device_data(),
                                         dB.device_data(), k, lda, ldb, ldc);
  });
  WuK_Timer("wuk::gemm_32x32_v0", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((n + TILE_WIDTH - 1) / TILE_WIDTH,
                (m + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= blockDim.y);
    assert(n >= blockDim.x);
    wuk::gemm_32x32_v0<<<gridDim, blockDim>>>(m, n, k, alpha, dA.device_data(),
                                              lda, dB.device_data(), ldb, beta,
                                              dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_32x32_v1", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= blockDim.x);
    assert(n >= blockDim.y);
    wuk::gemm_32x32_v1<<<gridDim, blockDim>>>(m, n, k, alpha, dA.device_data(),
                                              lda, dB.device_data(), ldb, beta,
                                              dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_32x32_v2", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= blockDim.x);
    assert(n >= blockDim.y);
    wuk::gemm_32x32_v2<<<gridDim, blockDim>>>(m, n, k, alpha, dA.device_data(),
                                              lda, dB.device_data(), ldb, beta,
                                              dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_32x32_v3", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= TILE_WIDTH);
    assert(n >= TILE_WIDTH);
    assert(k >= TILE_WIDTH);
    wuk::gemm_32x32_v3<float, TILE_WIDTH><<<gridDim, blockDim>>>(
        m, n, k, alpha, dA.device_data(), lda, dB.device_data(), ldb, beta,
        dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_32x32_v4", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(TILE_WIDTH, TILE_WIDTH),
        gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= TILE_WIDTH);
    assert(n >= TILE_WIDTH);
    assert(k >= TILE_WIDTH);
    wuk::gemm_32x32_v4<float, TILE_WIDTH><<<gridDim, blockDim>>>(
        m, n, k, alpha, dA.device_data(), lda, dB.device_data(), ldb, beta,
        dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_32x32_v5", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 32;
    const dim3 blockDim(256), gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                                      (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= TILE_WIDTH);
    assert(n >= TILE_WIDTH);
    assert(k % 32 == 0);
    wuk::gemm_32x32_v5<float, float4, 0><<<gridDim, blockDim>>>(
        m, n, k, alpha, dA.device_data(), lda, dB.device_data(), ldb, beta,
        dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_64x64", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 64;
    const dim3 blockDim(256), gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                                      (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= TILE_WIDTH);
    assert(n >= TILE_WIDTH);
    assert(k % 16 == 0);
    wuk::gemm_64x64<float, float4, 0><<<gridDim, blockDim>>>(
        m, n, k, alpha, dA.device_data(), lda, dB.device_data(), ldb, beta,
        dC.device_data(), ldc);
  });
  WuK_Timer("wuk::gemm_128x128", 2.0 * m * k * n, [&] {
    const int TILE_WIDTH = 128;
    const dim3 blockDim(256), gridDim((m + TILE_WIDTH - 1) / TILE_WIDTH,
                                      (n + TILE_WIDTH - 1) / TILE_WIDTH);
    assert(opA == CUBLAS_OP_N);
    assert(opB == CUBLAS_OP_N);
    assert(opC == CUBLAS_OP_N);
    assert(m >= TILE_WIDTH);
    assert(n >= TILE_WIDTH);
    assert(k % 8 == 0);
    wuk::gemm_128x128<float, float4, 1><<<gridDim, blockDim>>>(
        m, n, k, alpha, dA.device_data(), lda, dB.device_data(), ldb, beta,
        dC.device_data(), ldc);
  });
}