File size: 3,074 Bytes
ef6eb55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | // PAX Pipeline GEMM — 3-stage async cp.async with overlap guarantee
// Ahmad Ali Parr · PAX Architecture · sm_86
// Proof obligations: PO4 (happens-before), PO6 (barrier conservation), PO7 (data-race freedom)
#include <cuda_fp16.h>
#include <mma.h>
using namespace nvcuda;
#define STAGES 3
#define TILE_M 64
#define TILE_N 64
#define TILE_K 16
__shared__ __half smem_a[STAGES][TILE_K][TILE_M];
__shared__ __half smem_b[STAGES][TILE_K][TILE_N];
// cp.async descriptor for one tile
inline __device__ void async_load_tile(
__half* dst, const __half* src, int bytes
) {
asm volatile(
"cp.async.ca.shared.global [%0], [%1], %2;\n"
: : "r"((unsigned)__cvta_generic_to_shared(dst)),
"l"(src), "n"(32)
);
}
extern "C" __global__ void pax_gemm_pipeline_sm86(
const __half* A, const __half* B, float* C,
int M, int N, int K
) {
int lane = threadIdx.x % 32;
int warp = threadIdx.x / 32;
wmma::fragment<wmma::accumulator, 16, 8, 8, float> acc;
wmma::fill_fragment(acc, 0.0f);
// Prologue: fill pipeline (stages 0..STAGES-2)
for (int s = 0; s < STAGES - 1 && s * TILE_K < K; s++) {
int kOff = s * TILE_K;
async_load_tile(&smem_a[s][0][0], A + kOff * M + blockIdx.y * TILE_M, 32);
async_load_tile(&smem_b[s][0][0], B + kOff * N + blockIdx.x * TILE_N, 32);
asm volatile("cp.async.commit_group;\n");
}
// Steady state: compute stage s while loading stage s+STAGES-1
for (int k = 0; k < K; k += TILE_K) {
int cur = (k / TILE_K) % STAGES;
int pre = (k / TILE_K + STAGES - 1) % STAGES;
// PO4: wait_group 1 = wait for all but 1 outstanding group
// HB(copy[k], compute[k]) enforced here
asm volatile("cp.async.wait_group 1;\n");
__syncthreads(); // PO6: barrier consumes all permissions from prior cp.async
wmma::fragment<wmma::matrix_a, 16, 8, 8, __half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 8, 8, __half, wmma::col_major> b_frag;
// PO3: whole warp executes mma.sync
wmma::load_matrix_sync(a_frag, &smem_a[cur][0][0], TILE_M);
wmma::load_matrix_sync(b_frag, &smem_b[cur][0][0], TILE_N);
wmma::mma_sync(acc, a_frag, b_frag, acc);
// Prefetch next tile — PO4: HB(compute[k], copy[k+STAGES-1])
int nextK = k + (STAGES - 1) * TILE_K;
if (nextK < K) {
async_load_tile(&smem_a[pre][0][0], A + nextK * M + blockIdx.y * TILE_M, 32);
async_load_tile(&smem_b[pre][0][0], B + nextK * N + blockIdx.x * TILE_N, 32);
asm volatile("cp.async.commit_group;\n");
}
}
// Epilogue: drain pipeline
asm volatile("cp.async.wait_all;\n");
__syncthreads();
// PO5: store — write permission is warp-local
int row = blockIdx.y * TILE_M;
int col = blockIdx.x * TILE_N;
if (row < M && col < N)
wmma::store_matrix_sync(C + row * N + col, acc, N, wmma::mem_row_major);
}
|