pax-coder / src /rtx_gemm_pipeline.cu

chore: push pax-coder from SNAPKITTYWEST GitHub

ef6eb55 verified about 2 months ago

3.07 kB

// 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);
}

	// 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);
	}