dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

upstream: add GEMM kernel references from 4 repos for BI-V100 porting

Browse Source 
Sources (all CUDA 10.2 compatible, no CUTLASS/Triton dependency):
- leimao/CUDA-GEMM-Optimization: v00-v07, fp16 WMMA variant, double buffered
- siboehm/SGEMM_CUDA: kernel 1-12, warp tiling + double buffering
- wangzyon/NVIDIA_SGEMM_PRACTICE: kernel 1-7
- edtallison/sgemm-cuda: kernel 1-12 (reimplementation with notes)

Key porting issue: ALL kernels hardcode WARPSIZE=32.
BI-V100 has warp_size=64. Need to:
1. Replace all 32U / WARPSIZE constants with 64
2. Adjust warp subtile decomposition (WMITER, WNITER, WSUBM, WSUBN)
3. Adjust shared memory bank conflict avoidance (may have different bank count)
4. Test __shfl_down_sync with mask=0xFFFFFFFFFFFFFFFF (64-bit)
This commit is contained in:
Claude
2026-08-14 15:11:57 +00:00
parent 29ecc2e602
commit 9ca33cf4d5
59 changed files with 8802 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

229
upstream_ref/sgemm_cuda/12_kernel_double_buffering.cuh Normal file
View File

@@ -0,0 +1,229 @@
#pragma once
#include <algorithm>
#include <cassert>
#include <cooperative_groups.h>
#include <cstdio>
#include <cstdlib>
#include <cublas_v2.h>
#include <cuda/barrier>
#include <cuda_runtime.h>
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
namespace {
template <const int BM, const int BN, const int BK, const int rowStrideA,
const int rowStrideB, typename T>
__device__ void loadFromGmem(int N, int K, float *A, float *B, float *As,
float *Bs, int innerRowA, int innerColA,
int innerRowB, int innerColB, T &barrier) {
for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
cuda::memcpy_async(&As[(innerColA * 4 + 0) * BM + innerRowA + offset],
&A[(innerRowA + offset) * K + innerColA * 4],
cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
barrier);
cuda::memcpy_async(&As[(innerColA * 4 + 1) * BM + innerRowA + offset],
&A[(innerRowA + offset) * K + innerColA * 4 + 1],
cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
barrier);
cuda::memcpy_async(&As[(innerColA * 4 + 2) * BM + innerRowA + offset],
&A[(innerRowA + offset) * K + innerColA * 4 + 2],
cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
barrier);
cuda::memcpy_async(&As[(innerColA * 4 + 3) * BM + innerRowA + offset],
&A[(innerRowA + offset) * K + innerColA * 4 + 3],
cuda::aligned_size_t<sizeof(float)>(sizeof(float)),
barrier);
}
for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
cuda::memcpy_async(&Bs[(innerRowB + offset) * BN + innerColB * 4],
&B[(innerRowB + offset) * N + innerColB * 4],
cuda::aligned_size_t<sizeof(float4)>(sizeof(float4)),
barrier);
}
}
template <const int BM, const int BN, const int BK, const int WM, const int WN,
const int WMITER, const int WNITER, const int WSUBM, const int WSUBN,
const int TM, const int TN>
__device__ void
processFromSmem(float *regM, float *regN, float *threadResults, const float *As,
const float *Bs, const uint warpRow, const uint warpCol,
const uint threadRowInWarp, const uint threadColInWarp) {
for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
// populate registers for whole warptile
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
for (uint i = 0; i < TM; ++i) {
regM[wSubRowIdx * TM + i] =
As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
threadRowInWarp * TM + i];
}
}
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
for (uint i = 0; i < TN; ++i) {
regN[wSubColIdx * TN + i] =
Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
threadColInWarp * TN + i];
}
}
// execute warptile matmul
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
// calculate per-thread results
for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
(wSubColIdx * TN) + resIdxN] +=
regM[wSubRowIdx * TM + resIdxM] *
regN[wSubColIdx * TN + resIdxN];
}
}
}
}
}
}
} // namespace
/*
* @tparam BM The threadblock size for M dimension SMEM caching.
* @tparam BN The threadblock size for N dimension SMEM caching.
* @tparam BK The threadblock size for K dimension SMEM caching.
* @tparam WM M dim of continuous tile computed by each warp
* @tparam WN N dim of continuous tile computed by each warp
* @tparam WMITER The number of subwarp tiling steps in M dimension.
* @tparam WNITER The number of subwarp tiling steps in N dimension.
* @tparam TM The per-thread tile size for M dimension.
* @tparam TN The per-thread tile size for N dimension.
*/
template <const int BM, const int BN, const int BK, const int WM, const int WN,
const int WNITER, const int TM, const int TN, const int NUM_THREADS>
__global__ void __launch_bounds__(NUM_THREADS)
runSgemmDoubleBuffering2(int M, int N, int K, float alpha, float *A,
float *B, float beta, float *C) {
auto block = cooperative_groups::this_thread_block();
__shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> frontBarrier;
__shared__ cuda::barrier<cuda::thread_scope::thread_scope_block> backBarrier;
auto frontBarrierPtr = &frontBarrier;
auto backBarrierPtr = &backBarrier;
if (block.thread_rank() == 0) {
init(&frontBarrier, block.size());
init(&backBarrier, block.size());
}
__syncthreads();
const uint cRow = blockIdx.y;
const uint cCol = blockIdx.x;
// Placement of the warp in the threadblock tile
const uint warpIdx = threadIdx.x / WARPSIZE; // the warp this thread is in
const uint warpCol = warpIdx % (BN / WN);
const uint warpRow = warpIdx / (BN / WN);
// size of the warp subtile
constexpr uint WMITER = (WM * WN) / (WARPSIZE * TM * TN * WNITER);
constexpr uint WSUBM = WM / WMITER; // 64/2=32
constexpr uint WSUBN = WN / WNITER; // 32/2=16
// Placement of the thread in the warp subtile
const uint threadIdxInWarp = threadIdx.x % WARPSIZE; // [0, 31]
const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4)
const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4
// allocate space for the current blocktile in SMEM
__shared__ float As[2 * BM * BK];
__shared__ float Bs[2 * BK * BN];
// Move blocktile to beginning of A's row and B's column
A += cRow * BM * K;
B += cCol * BN;
// Move C_ptr to warp's output tile
C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN;
// calculating the indices that this thread will load into SMEM
// we'll load 128bit / 32bit = 4 elements per thread at each step
const uint innerRowA = threadIdx.x / (BK / 4);
const uint innerColA = threadIdx.x % (BK / 4);
constexpr uint rowStrideA = (NUM_THREADS * 4) / BK;
const uint innerRowB = threadIdx.x / (BN / 4);
const uint innerColB = threadIdx.x % (BN / 4);
constexpr uint rowStrideB = NUM_THREADS / (BN / 4);
// allocate thread-local cache for results in registerfile
float threadResults[WMITER * TM * WNITER * TN] = {0.0};
// we cache into registers on the warptile level
float regM[WMITER * TM] = {0.0};
float regN[WNITER * TN] = {0.0};
int As_offset = 0;
int Bs_offset = 0;
// double-buffering: load first blocktile into SMEM
loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
N, K, A, B, As + As_offset * BM * BK, Bs + Bs_offset * BK * BN, innerRowA,
innerColA, innerRowB, innerColB, (*frontBarrierPtr));
// outer-most loop over block tiles
for (uint bkIdx = 0; bkIdx < K - BK; bkIdx += BK) {
// double-buffering: load next blocktile into SMEM
loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
N, K, A + BK, B + BK * N, As + (1 - As_offset) * BM * BK,
Bs + (1 - Bs_offset) * BK * BN, innerRowA, innerColA, innerRowB,
innerColB, (*backBarrierPtr));
// compute the current blocktile
(*frontBarrierPtr).arrive_and_wait();
processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM, TN>(
regM, regN, threadResults, As + As_offset * BM * BK,
Bs + Bs_offset * BK * BN, warpRow, warpCol, threadRowInWarp,
threadColInWarp);
A += BK; // move BK columns to right
B += BK * N; // move BK rows down
As_offset = 1 - As_offset;
Bs_offset = 1 - Bs_offset;
// swap the front and back barriers
auto tmp = frontBarrierPtr;
frontBarrierPtr = backBarrierPtr;
backBarrierPtr = tmp;
__syncthreads();
}
// compute the last blocktile
(*frontBarrierPtr).arrive_and_wait();
processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM, TN>(
regM, regN, threadResults, As + As_offset * BM * BK,
Bs + Bs_offset * BK * BN, warpRow, warpCol, threadRowInWarp,
threadColInWarp);
// write out the results
for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
// move C pointer to current warp subtile
float *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN;
for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) {
for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) {
// load C vector into registers
float4 tmp = reinterpret_cast<float4 *>(
&C_interim[(threadRowInWarp * TM + resIdxM) * N +
threadColInWarp * TN + resIdxN])[0];
// perform GEMM update in reg
const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
wSubColIdx * TN + resIdxN;
tmp.x = alpha * threadResults[i + 0] + beta * tmp.x;
tmp.y = alpha * threadResults[i + 1] + beta * tmp.y;
tmp.z = alpha * threadResults[i + 2] + beta * tmp.z;
tmp.w = alpha * threadResults[i + 3] + beta * tmp.w;
// write back
reinterpret_cast<float4 *>(
&C_interim[(threadRowInWarp * TM + resIdxM) * N +
threadColInWarp * TN + resIdxN])[0] = tmp;
}
}
}
}
}