mirror of
https://github.com/NVIDIA/cuda-samples.git
synced 2026-10-25 22:28:17 +08:00
Apply consistent code formatting across the repo. Add clang-format and pre-commit hooks.
This commit is contained in:
Rob Armstrong
@@ -58,14 +58,14 @@
|
||||
// but carefully enough to avoid local memory use.
|
||||
|
||||
#include <assert.h>
|
||||
#include <stdio.h>
|
||||
#include <cuda.h>
|
||||
#include <mma.h>
|
||||
#include <cuda/pipeline>
|
||||
#include <mma.h>
|
||||
#include <stdio.h>
|
||||
|
||||
// helper functions and utilities to work with CUDA
|
||||
#include <helper_functions.h>
|
||||
#include <helper_cuda.h>
|
||||
#include <helper_functions.h>
|
||||
|
||||
// Externally configurable parameters.
|
||||
|
||||
@@ -106,7 +106,7 @@
|
||||
|
||||
// Implementation constants.
|
||||
|
||||
#define WARPS_PER_BLOCK 8
|
||||
#define WARPS_PER_BLOCK 8
|
||||
#define THREADS_PER_BLOCK (WARP_SIZE * WARPS_PER_BLOCK)
|
||||
|
||||
#if SHARED_MEMORY_LIMIT_64K
|
||||
@@ -124,10 +124,10 @@
|
||||
#define CHUNK_K 8
|
||||
#endif
|
||||
|
||||
#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(float))
|
||||
#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))
|
||||
#define CHUNK_LINE_BYTES (CHUNK_K * K * sizeof(float))
|
||||
#define WARP_COPY_BYTES (WARP_SIZE * sizeof(int4))
|
||||
#define CHUNK_COPY_LINES_PER_WARP (WARP_COPY_BYTES / CHUNK_LINE_BYTES)
|
||||
#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)
|
||||
#define CHUNK_COPY_LINE_LANES (WARP_SIZE / CHUNK_COPY_LINES_PER_WARP)
|
||||
|
||||
#define BLOCK_ROW_WARPS 2
|
||||
#define BLOCK_COL_WARPS 4
|
||||
@@ -157,25 +157,24 @@
|
||||
// we must keep each row and column 256-bit aligned, as required by nvcuda::wmma::load_matrix_sync.
|
||||
#define SKEW_FLOAT 8
|
||||
|
||||
#define checkKernelErrors(expr) do { \
|
||||
expr; \
|
||||
\
|
||||
cudaError_t __err = cudaGetLastError(); \
|
||||
if (__err != cudaSuccess) { \
|
||||
printf("Line %d: '%s' failed: %s\n", __LINE__, # expr, cudaGetErrorString(__err)); \
|
||||
abort(); \
|
||||
} \
|
||||
} while(0)
|
||||
#define checkKernelErrors(expr) \
|
||||
do { \
|
||||
expr; \
|
||||
\
|
||||
cudaError_t __err = cudaGetLastError(); \
|
||||
if (__err != cudaSuccess) { \
|
||||
printf("Line %d: '%s' failed: %s\n", __LINE__, #expr, cudaGetErrorString(__err)); \
|
||||
abort(); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
enum kernels
|
||||
{
|
||||
tf32mma_shmem_gemm_async_copy = 0, // tf32 MMA shmem using kernel with async_copy
|
||||
tf32mma_shmem_gemm = 1, // tf32 MMA shmem using kernel normal copy (without async_copy).
|
||||
simple_tf32mma_gemm = 2 // tf32 MMA non-shmem using simple kernel.
|
||||
enum kernels {
|
||||
tf32mma_shmem_gemm_async_copy = 0, // tf32 MMA shmem using kernel with async_copy
|
||||
tf32mma_shmem_gemm = 1, // tf32 MMA shmem using kernel normal copy (without async_copy).
|
||||
simple_tf32mma_gemm = 2 // tf32 MMA non-shmem using simple kernel.
|
||||
};
|
||||
|
||||
const char* kernelNames[] = {"compute_tf32gemm_async_copy", "compute_tf32gemm",
|
||||
"simple_wmma_tf32gemm"};
|
||||
const char *kernelNames[] = {"compute_tf32gemm_async_copy", "compute_tf32gemm", "simple_wmma_tf32gemm"};
|
||||
|
||||
using namespace nvcuda;
|
||||
|
||||
@@ -183,18 +182,18 @@ __host__ void init_host_matrices(float *a, float *b, float *c)
|
||||
{
|
||||
for (int i = 0; i < M_GLOBAL; i++) {
|
||||
for (int j = 0; j < K_GLOBAL; j++) {
|
||||
a[i*K_GLOBAL+j] = (float)(rand() % 3);
|
||||
a[i * K_GLOBAL + j] = (float)(rand() % 3);
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < N_GLOBAL; i++) {
|
||||
for (int j = 0; j < K_GLOBAL; j++) {
|
||||
b[i*K_GLOBAL+j] = (float)(rand() % 3);
|
||||
b[i * K_GLOBAL + j] = (float)(rand() % 3);
|
||||
}
|
||||
}
|
||||
|
||||
for (int t = 0; t < M_GLOBAL * N_GLOBAL; t++) {
|
||||
c[t] = (float)(rand() % 3);
|
||||
c[t] = (float)(rand() % 3);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -211,10 +210,11 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
const size_t shmem_idx_b_off = BLOCK_COL_TILES * M;
|
||||
|
||||
// This pointer is used to access the C and D matrix tiles this warp computes.
|
||||
float *shmem_warp_tile_ptr = (float*)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;
|
||||
float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS
|
||||
+ (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;
|
||||
|
||||
// This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.
|
||||
float *shmem_warp_stream_ptr = (float*)&shmem[0][0] + warpId * SHMEM_STRIDE * N;
|
||||
float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;
|
||||
|
||||
// Adjust the beta scaler, as it'll be multiplied by alpha at the end of
|
||||
// each tile computation. Technically this is not generally correct (may result
|
||||
@@ -224,7 +224,7 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
// Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the
|
||||
// right and down, and selects the next tile to compute. Once there's no such tile,
|
||||
// all warps in this CTA exit.
|
||||
for(unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {
|
||||
for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {
|
||||
const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);
|
||||
const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;
|
||||
|
||||
@@ -234,14 +234,14 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
}
|
||||
|
||||
// This warp's pointer to the C matrix data to copy memory from to shared memory.
|
||||
const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;
|
||||
const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;
|
||||
const float *src_gmem_warp_stream_ptr = &C[gmem_idx];
|
||||
|
||||
// Stream multiple C tiles to shared memory.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < N; i++) {
|
||||
*((int4*)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =
|
||||
*((int4*)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);
|
||||
*((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId) =
|
||||
*((int4 *)(src_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
@@ -265,7 +265,7 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
|
||||
// Scale the C matrix.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < WARP_COL_TILES; i++) {
|
||||
for (int i = 0; i < WARP_COL_TILES; i++) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < WARP_ROW_TILES; j++) {
|
||||
#pragma unroll
|
||||
@@ -277,16 +277,19 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
|
||||
// Select what warp copies what matrix to shared memory.
|
||||
// Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.
|
||||
const float *warp_ptr = (warpId < (WARPS_PER_BLOCK/2)) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK/2)) * 2) :
|
||||
(&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK/2)) * 2);
|
||||
const float *warp_ptr =
|
||||
(warpId < (WARPS_PER_BLOCK / 2))
|
||||
? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)
|
||||
: (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);
|
||||
|
||||
// Go through the global K dimension by a fixed step at a time.
|
||||
#pragma unroll
|
||||
for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {
|
||||
// Copy slices of the A and B matrices to shared memory.
|
||||
// The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.
|
||||
size_t shmem_idx = warpId < (WARPS_PER_BLOCK/2) ? (M * (warpId % (WARPS_PER_BLOCK/2)) * 2) :
|
||||
(N * (warpId % (WARPS_PER_BLOCK/2)) * 2 + shmem_idx_b_off);
|
||||
size_t shmem_idx = warpId < (WARPS_PER_BLOCK / 2)
|
||||
? (M * (warpId % (WARPS_PER_BLOCK / 2)) * 2)
|
||||
: (N * (warpId % (WARPS_PER_BLOCK / 2)) * 2 + shmem_idx_b_off);
|
||||
|
||||
// First half of the warp copies the first row / column of the matrix,
|
||||
// the second half of the warp copies the next.
|
||||
@@ -296,9 +299,10 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
shmem_idx += laneId / CHUNK_COPY_LINE_LANES;
|
||||
|
||||
#pragma unroll
|
||||
for(int i = 0; i < ((WARP_SIZE/2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {
|
||||
for (int i = 0; i < ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2; i++) {
|
||||
// Copy 16 bytes at once in each lane.
|
||||
*((int4*)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) = *((int4*)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));
|
||||
*((int4 *)&shmem[shmem_idx][0] + (laneId % CHUNK_COPY_LINE_LANES)) =
|
||||
*((int4 *)lane_ptr + (laneId % CHUNK_COPY_LINE_LANES));
|
||||
|
||||
// Advance the global memory pointer and the shared memory index.
|
||||
lane_ptr = lane_ptr + K_GLOBAL * CHUNK_COPY_LINES_PER_WARP;
|
||||
@@ -315,8 +319,8 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < WARP_COL_TILES; i++) {
|
||||
size_t shmem_idx_a = (warpId/BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);
|
||||
const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];
|
||||
size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);
|
||||
const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];
|
||||
|
||||
wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_FLOAT);
|
||||
#pragma unroll
|
||||
@@ -328,8 +332,8 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
if (i == 0) {
|
||||
// Load the B matrix fragment once, because it is going to be reused
|
||||
// against the other A matrix fragments.
|
||||
size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId%2) + (j * N);
|
||||
const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];
|
||||
size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);
|
||||
const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];
|
||||
|
||||
wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_FLOAT);
|
||||
#pragma unroll
|
||||
@@ -370,8 +374,8 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < N; i++) {
|
||||
*((int4*)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =
|
||||
*((int4*)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);
|
||||
*((int4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =
|
||||
*((int4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
@@ -379,7 +383,8 @@ __global__ void compute_tf32gemm(const float *A, const float *B, const float *C,
|
||||
#endif
|
||||
}
|
||||
|
||||
__global__ void compute_tf32gemm_async_copy(const float *A, const float *B, const float *C, float *D, const float alpha, float beta)
|
||||
__global__ void
|
||||
compute_tf32gemm_async_copy(const float *A, const float *B, const float *C, float *D, const float alpha, float beta)
|
||||
{
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
extern __shared__ float shmem[][CHUNK_K * K + SKEW_FLOAT];
|
||||
@@ -389,10 +394,11 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
const unsigned int laneId = threadIdx.x % WARP_SIZE;
|
||||
|
||||
// This pointer is used to access the C and D matrix tiles this warp computes.
|
||||
float *shmem_warp_tile_ptr = (float*)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS + (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;
|
||||
float *shmem_warp_tile_ptr = (float *)&shmem[0][0] + (warpId / BLOCK_ROW_WARPS) * SHMEM_STRIDE * N * BLOCK_ROW_WARPS
|
||||
+ (warpId % BLOCK_ROW_WARPS) * SHMEM_OFFSET;
|
||||
|
||||
// This pointer is used to stream the C and D matrices block-wide tile to and from shared memory.
|
||||
float *shmem_warp_stream_ptr = (float*)&shmem[0][0] + warpId * SHMEM_STRIDE * N;
|
||||
float *shmem_warp_stream_ptr = (float *)&shmem[0][0] + warpId * SHMEM_STRIDE * N;
|
||||
|
||||
// Offset in shared memory from which the B matrix is stored.
|
||||
constexpr size_t shmem_idx_b_off = BLOCK_COL_TILES * M;
|
||||
@@ -402,14 +408,14 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
// in a loss of precision). Zero still needs to be specially handled though.
|
||||
beta /= alpha;
|
||||
|
||||
cuda::pipeline<cuda::thread_scope_thread> pipe = cuda::make_pipeline();
|
||||
const auto shape4 = cuda::aligned_size_t<alignof(float4)>(sizeof(float4));
|
||||
constexpr int loadStride = 2; // load 4 floats, so left-shift by 2.
|
||||
cuda::pipeline<cuda::thread_scope_thread> pipe = cuda::make_pipeline();
|
||||
const auto shape4 = cuda::aligned_size_t<alignof(float4)>(sizeof(float4));
|
||||
constexpr int loadStride = 2; // load 4 floats, so left-shift by 2.
|
||||
|
||||
// Each CTA slides along the 128 x 128 tiles from the top left corner of the matrix to the
|
||||
// right and down, and selects the next tile to compute. Once there's no such tile,
|
||||
// all warps in this CTA exit.
|
||||
for(unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {
|
||||
for (unsigned int block_pos = blockIdx.x;; block_pos += gridDim.x) {
|
||||
const unsigned int block_tile_i = ((block_pos * BLOCK_ROW_TILES) / N_TILES) * (BLOCK_COL_TILES);
|
||||
const unsigned int block_tile_j = (block_pos * BLOCK_COL_TILES) % N_TILES;
|
||||
|
||||
@@ -419,7 +425,7 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
}
|
||||
|
||||
// This warp's pointer to the C matrix data to copy memory from to shared memory.
|
||||
const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;
|
||||
const size_t gmem_idx = (block_tile_i + warpId) * M * GLOBAL_MEM_STRIDE + block_tile_j * N;
|
||||
const float *src_gmem_warp_stream_ptr = &C[gmem_idx];
|
||||
|
||||
// Stream multiple C tiles to shared memory.
|
||||
@@ -427,8 +433,9 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
for (int i = 0; i < N; i++) {
|
||||
pipe.producer_acquire();
|
||||
cuda::memcpy_async(&shmem_warp_stream_ptr[(SHMEM_STRIDE * i) + (laneId << loadStride)],
|
||||
&src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],
|
||||
shape4, pipe);
|
||||
&src_gmem_warp_stream_ptr[(GLOBAL_MEM_STRIDE * i) + (laneId << loadStride)],
|
||||
shape4,
|
||||
pipe);
|
||||
pipe.producer_commit();
|
||||
}
|
||||
// Now wait for all the above issued 8 batches to complete.
|
||||
@@ -461,19 +468,22 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
|
||||
// Select what warp copies what matrix to shared memory.
|
||||
// Warps 0-3 copy the A matrix, warps 4-7 copy the B matrix.
|
||||
const float *warp_ptr = (warpId < (WARPS_PER_BLOCK/2)) ? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK/2)) * 2) :
|
||||
(&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK/2)) * 2);
|
||||
const float *warp_ptr =
|
||||
(warpId < (WARPS_PER_BLOCK / 2))
|
||||
? (&A[block_tile_i * M * K_GLOBAL] + M * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2)
|
||||
: (&B[block_tile_j * N * K_GLOBAL] + N * K_GLOBAL * (warpId % (WARPS_PER_BLOCK / 2)) * 2);
|
||||
|
||||
constexpr int chunksPerLane = ((WARP_SIZE/2) / CHUNK_COPY_LINES_PER_WARP) * 2;
|
||||
const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStride;
|
||||
const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);
|
||||
constexpr int chunksPerLane = ((WARP_SIZE / 2) / CHUNK_COPY_LINES_PER_WARP) * 2;
|
||||
const int laneLoadElem = (laneId % CHUNK_COPY_LINE_LANES) << loadStride;
|
||||
const int stridePerLaneCopy = (laneId / CHUNK_COPY_LINE_LANES);
|
||||
// Go through the global K dimension by a fixed step at a time.
|
||||
#pragma unroll
|
||||
for (int tile_k = 0; tile_k < K_TILES; tile_k += CHUNK_K) {
|
||||
// Copy slices of the A and B matrices to shared memory.
|
||||
// The first half of the warps in the CTA copy the A matrix, the rest copy the B matrix.
|
||||
// As for tf32 MMA M == N we use M for warp 4-7 + shmem_idx_b_off.
|
||||
size_t shmem_idx = (M * (warpId % (WARPS_PER_BLOCK/2)) * 2) + ((warpId / (WARPS_PER_BLOCK/2)) * shmem_idx_b_off);
|
||||
size_t shmem_idx =
|
||||
(M * (warpId % (WARPS_PER_BLOCK / 2)) * 2) + ((warpId / (WARPS_PER_BLOCK / 2)) * shmem_idx_b_off);
|
||||
// First half of the warp copies the first row / column of the matrix,
|
||||
// the second half of the warp copies the next.
|
||||
const float *lane_ptr = (warp_ptr + tile_k * K + stridePerLaneCopy * K_GLOBAL + laneLoadElem);
|
||||
@@ -482,7 +492,7 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
shmem_idx += stridePerLaneCopy;
|
||||
|
||||
#pragma unroll
|
||||
for(int i = 0; i < chunksPerLane; i++) {
|
||||
for (int i = 0; i < chunksPerLane; i++) {
|
||||
// Copy 16 bytes at once in each lane.
|
||||
pipe.producer_acquire();
|
||||
cuda::memcpy_async(&shmem[shmem_idx][laneLoadElem], lane_ptr, shape4, pipe);
|
||||
@@ -504,8 +514,8 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < WARP_COL_TILES; i++) {
|
||||
size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);
|
||||
const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];
|
||||
size_t shmem_idx_a = (warpId / BLOCK_ROW_WARPS) * M * BLOCK_ROW_WARPS + (i * M);
|
||||
const float *tile_ptr = &shmem[shmem_idx_a][k_step * K];
|
||||
|
||||
wmma::load_matrix_sync(a[i], tile_ptr, K * CHUNK_K + SKEW_FLOAT);
|
||||
|
||||
@@ -518,13 +528,13 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
if (i == 0) {
|
||||
// Load the B matrix fragment once, because it is going to be reused
|
||||
// against the other A matrix fragments.
|
||||
size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId%2) + (j * N);
|
||||
const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];
|
||||
size_t shmem_idx_b = shmem_idx_b_off + (WARP_ROW_TILES * N) * (warpId % 2) + (j * N);
|
||||
const float *tile_ptr = &shmem[shmem_idx_b][k_step * K];
|
||||
|
||||
wmma::load_matrix_sync(b[j], tile_ptr, K * CHUNK_K + SKEW_FLOAT);
|
||||
#pragma unroll
|
||||
for (int t = 0; t < b[j].num_elements; t++) {
|
||||
b[j].x[t] = wmma::__float_to_tf32(b[j].x[t]);
|
||||
b[j].x[t] = wmma::__float_to_tf32(b[j].x[t]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -560,8 +570,8 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < N; i++) {
|
||||
*((float4*)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =
|
||||
*((float4*)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);
|
||||
*((float4 *)(dst_gmem_warp_stream_ptr + GLOBAL_MEM_STRIDE * i) + laneId) =
|
||||
*((float4 *)(shmem_warp_stream_ptr + SHMEM_STRIDE * i) + laneId);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
@@ -571,84 +581,90 @@ __global__ void compute_tf32gemm_async_copy(const float *A, const float *B, cons
|
||||
|
||||
// Performs an MxNxK tf32 GEMM (C=alpha*A*B + beta*C) assuming:
|
||||
// 1) Matrices are packed in memory.
|
||||
// 2) M, N and K are multiples of 16, 16 and 8 respectively.
|
||||
// 2) M, N and K are multiples of 16, 16 and 8 respectively.
|
||||
// 3) A is row major, B is column major matrix.
|
||||
// Note: This is a less performant version of the compute_tf32gemm kernel. It is designed for
|
||||
// demonstration purposes only to show the CUDA WMMA API use without relying on
|
||||
// availability of the shared memory.
|
||||
__global__ void simple_wmma_tf32gemm(float *a, float *b, float *c, float *d, int m_ld, int n_ld, int k_ld, float alpha, float beta)
|
||||
__global__ void
|
||||
simple_wmma_tf32gemm(float *a, float *b, float *c, float *d, int m_ld, int n_ld, int k_ld, float alpha, float beta)
|
||||
{
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
// Leading dimensions. Packed with no transpositions.
|
||||
// Leading dimensions. Packed with no transpositions.
|
||||
int lda = k_ld;
|
||||
int ldb = k_ld;
|
||||
int ldc = n_ld;
|
||||
|
||||
// Tile using a 2D grid
|
||||
int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
|
||||
int warpN = (blockIdx.y * blockDim.y + threadIdx.y);
|
||||
|
||||
// Declare the fragments
|
||||
wmma::fragment<wmma::matrix_a, M, N, K, wmma::precision::tf32, wmma::row_major> a_frag;
|
||||
wmma::fragment<wmma::matrix_b, M, N, K, wmma::precision::tf32, wmma::col_major> b_frag;
|
||||
wmma::fragment<wmma::accumulator, M, N, K, float> acc_frag;
|
||||
wmma::fragment<wmma::accumulator, M, N, K, float> c_frag;
|
||||
// Tile using a 2D grid
|
||||
int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
|
||||
int warpN = (blockIdx.y * blockDim.y + threadIdx.y);
|
||||
|
||||
wmma::fill_fragment(acc_frag, 0.0f);
|
||||
// Declare the fragments
|
||||
wmma::fragment<wmma::matrix_a, M, N, K, wmma::precision::tf32, wmma::row_major> a_frag;
|
||||
wmma::fragment<wmma::matrix_b, M, N, K, wmma::precision::tf32, wmma::col_major> b_frag;
|
||||
wmma::fragment<wmma::accumulator, M, N, K, float> acc_frag;
|
||||
wmma::fragment<wmma::accumulator, M, N, K, float> c_frag;
|
||||
|
||||
// Loop over k
|
||||
for (int i = 0; i < k_ld; i += K) {
|
||||
int aCol = i;
|
||||
int aRow = warpM * M;
|
||||
wmma::fill_fragment(acc_frag, 0.0f);
|
||||
|
||||
//int bCol = i;
|
||||
//int bRow = warpN * N;
|
||||
int bCol = warpN * N;
|
||||
int bRow = i;
|
||||
// Loop over k
|
||||
for (int i = 0; i < k_ld; i += K) {
|
||||
int aCol = i;
|
||||
int aRow = warpM * M;
|
||||
|
||||
// Bounds checking
|
||||
if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {
|
||||
// Load the inputs
|
||||
wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);
|
||||
wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);
|
||||
|
||||
#pragma unroll
|
||||
for (int t = 0; t < a_frag.num_elements; t++) {
|
||||
a_frag.x[t] = wmma::__float_to_tf32(a_frag.x[t]);
|
||||
// int bCol = i;
|
||||
// int bRow = warpN * N;
|
||||
int bCol = warpN * N;
|
||||
int bRow = i;
|
||||
|
||||
// Bounds checking
|
||||
if (aRow < m_ld && aCol < k_ld && bRow < k_ld && bCol < n_ld) {
|
||||
// Load the inputs
|
||||
wmma::load_matrix_sync(a_frag, a + aCol + aRow * lda, lda);
|
||||
wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);
|
||||
|
||||
#pragma unroll
|
||||
for (int t = 0; t < a_frag.num_elements; t++) {
|
||||
a_frag.x[t] = wmma::__float_to_tf32(a_frag.x[t]);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int t = 0; t < b_frag.num_elements; t++) {
|
||||
b_frag.x[t] = wmma::__float_to_tf32(b_frag.x[t]);
|
||||
}
|
||||
// Perform the matrix multiplication
|
||||
wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);
|
||||
}
|
||||
}
|
||||
|
||||
// Load in the current value of c, scale it by beta, and add this our result scaled by alpha
|
||||
int cCol = warpN * N;
|
||||
int cRow = warpM * M;
|
||||
|
||||
if (cRow < m_ld && cCol < n_ld) {
|
||||
wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);
|
||||
|
||||
for (int i = 0; i < c_frag.num_elements; i++) {
|
||||
c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int t = 0; t < b_frag.num_elements; t++) {
|
||||
b_frag.x[t] = wmma::__float_to_tf32(b_frag.x[t]);
|
||||
}
|
||||
// Perform the matrix multiplication
|
||||
wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
// Load in the current value of c, scale it by beta, and add this our result scaled by alpha
|
||||
int cCol = warpN * N;
|
||||
int cRow = warpM * M;
|
||||
|
||||
if (cRow < m_ld && cCol < n_ld) {
|
||||
wmma::load_matrix_sync(c_frag, c + cCol + cRow * ldc, ldc, wmma::mem_row_major);
|
||||
|
||||
for(int i=0; i < c_frag.num_elements; i++) {
|
||||
c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];
|
||||
}
|
||||
|
||||
// Store the output
|
||||
wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);
|
||||
}
|
||||
// Store the output
|
||||
wmma::store_matrix_sync(d + cCol + cRow * ldc, c_frag, ldc, wmma::mem_row_major);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__host__ void matMultiplyOnHost(float *A, float *B, float *C,
|
||||
float alpha, float beta,
|
||||
int numARows, int numAColumns,
|
||||
int numBRows, int numBColumns,
|
||||
int numCRows, int numCColumns)
|
||||
__host__ void matMultiplyOnHost(float *A,
|
||||
float *B,
|
||||
float *C,
|
||||
float alpha,
|
||||
float beta,
|
||||
int numARows,
|
||||
int numAColumns,
|
||||
int numBRows,
|
||||
int numBColumns,
|
||||
int numCRows,
|
||||
int numCColumns)
|
||||
{
|
||||
for (int i = 0; i < numCRows; i++) {
|
||||
for (int j = 0; j < numCColumns; j++) {
|
||||
@@ -658,7 +674,7 @@ __host__ void matMultiplyOnHost(float *A, float *B, float *C,
|
||||
temp += A[i * numAColumns + k] * B[j * numBRows + k];
|
||||
}
|
||||
|
||||
C[i*numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];
|
||||
C[i * numCColumns + j] = temp * alpha + beta * C[i * numCColumns + j];
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -686,16 +702,16 @@ int main(int argc, char **argv)
|
||||
float *B_h = NULL;
|
||||
float *C_h = NULL;
|
||||
#if CPU_DEBUG
|
||||
float *result_hD = NULL;
|
||||
float *result_hD = NULL;
|
||||
float *result_host = NULL;
|
||||
#endif
|
||||
|
||||
A_h = (float*) malloc(sizeof(float) * M_GLOBAL * K_GLOBAL);
|
||||
B_h = (float*) malloc(sizeof(float) * K_GLOBAL * N_GLOBAL);
|
||||
C_h = (float*) malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
A_h = (float *)malloc(sizeof(float) * M_GLOBAL * K_GLOBAL);
|
||||
B_h = (float *)malloc(sizeof(float) * K_GLOBAL * N_GLOBAL);
|
||||
C_h = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
#if CPU_DEBUG
|
||||
result_hD = (float*) malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
result_host = (float*) malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
result_hD = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
result_host = (float *)malloc(sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
#endif
|
||||
|
||||
float *A = NULL;
|
||||
@@ -703,10 +719,10 @@ int main(int argc, char **argv)
|
||||
float *C = NULL;
|
||||
float *D = NULL;
|
||||
|
||||
checkCudaErrors(cudaMalloc((void**)&A, sizeof(float) * M_GLOBAL * K_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void**)&B, sizeof(float) * N_GLOBAL * K_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void**)&C, sizeof(float) * M_GLOBAL * N_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void**)&D, sizeof(float) * M_GLOBAL * N_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void **)&A, sizeof(float) * M_GLOBAL * K_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void **)&B, sizeof(float) * N_GLOBAL * K_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void **)&C, sizeof(float) * M_GLOBAL * N_GLOBAL));
|
||||
checkCudaErrors(cudaMalloc((void **)&D, sizeof(float) * M_GLOBAL * N_GLOBAL));
|
||||
|
||||
assert(((unsigned long long)A) % 128 == 0);
|
||||
assert(((unsigned long long)B) % 128 == 0);
|
||||
@@ -734,11 +750,11 @@ int main(int argc, char **argv)
|
||||
printf("Required shared memory size: %lu Kb\n", SHMEM_SZ / 1024UL);
|
||||
|
||||
const float alpha = 1.1f;
|
||||
const float beta = 1.2f;
|
||||
const float beta = 1.2f;
|
||||
|
||||
cudaEvent_t start, stop;
|
||||
|
||||
checkCudaErrors(cudaEventCreate(&start));
|
||||
checkCudaErrors(cudaEventCreate(&start));
|
||||
checkCudaErrors(cudaEventCreate(&stop));
|
||||
checkCudaErrors(cudaEventRecord(start));
|
||||
|
||||
@@ -760,26 +776,30 @@ int main(int argc, char **argv)
|
||||
if ((deviceProp.sharedMemPerMultiprocessor >= SHMEM_SZ) && (selected_kernel != simple_tf32mma_gemm)) {
|
||||
printf("Computing using high performance kernel = %d - %s\n", selected_kernel, kernelNames[selected_kernel]);
|
||||
|
||||
switch (selected_kernel)
|
||||
{
|
||||
case tf32mma_shmem_gemm_async_copy :
|
||||
default:
|
||||
checkCudaErrors(cudaFuncSetAttribute(compute_tf32gemm_async_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));
|
||||
checkKernelErrors((compute_tf32gemm_async_copy<<<deviceProp.multiProcessorCount*2, THREADS_PER_BLOCK, SHMEM_SZ>>>(A, B, C, D, alpha, beta)));
|
||||
break;
|
||||
case tf32mma_shmem_gemm :
|
||||
checkCudaErrors(cudaFuncSetAttribute(compute_tf32gemm, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));
|
||||
checkKernelErrors((compute_tf32gemm<<<deviceProp.multiProcessorCount*2, THREADS_PER_BLOCK, SHMEM_SZ>>>(A, B, C, D, alpha, beta)));
|
||||
break;
|
||||
switch (selected_kernel) {
|
||||
case tf32mma_shmem_gemm_async_copy:
|
||||
default:
|
||||
checkCudaErrors(cudaFuncSetAttribute(
|
||||
compute_tf32gemm_async_copy, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));
|
||||
checkKernelErrors(
|
||||
(compute_tf32gemm_async_copy<<<deviceProp.multiProcessorCount * 2, THREADS_PER_BLOCK, SHMEM_SZ>>>(
|
||||
A, B, C, D, alpha, beta)));
|
||||
break;
|
||||
case tf32mma_shmem_gemm:
|
||||
checkCudaErrors(
|
||||
cudaFuncSetAttribute(compute_tf32gemm, cudaFuncAttributeMaxDynamicSharedMemorySize, SHMEM_SZ));
|
||||
checkKernelErrors((compute_tf32gemm<<<deviceProp.multiProcessorCount * 2, THREADS_PER_BLOCK, SHMEM_SZ>>>(
|
||||
A, B, C, D, alpha, beta)));
|
||||
break;
|
||||
}
|
||||
#if CPU_DEBUG
|
||||
checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float)*M_GLOBAL*N_GLOBAL, cudaMemcpyDeviceToHost));
|
||||
checkCudaErrors(cudaMemcpy(result_hD, D, sizeof(float) * M_GLOBAL * N_GLOBAL, cudaMemcpyDeviceToHost));
|
||||
#endif
|
||||
}
|
||||
else {
|
||||
dim3 gridDim;
|
||||
dim3 blockDim;
|
||||
|
||||
|
||||
// blockDim.x must be a multple of warpSize
|
||||
// 128x4 means we have 16 warps and a block computes a 64x64 output tile
|
||||
blockDim.x = 128;
|
||||
@@ -803,11 +823,7 @@ int main(int argc, char **argv)
|
||||
|
||||
memcpy(result_host, C_h, sizeof(float) * M_GLOBAL * N_GLOBAL);
|
||||
|
||||
matMultiplyOnHost(A_h, B_h, result_host,
|
||||
alpha, beta,
|
||||
M_GLOBAL, K_GLOBAL,
|
||||
K_GLOBAL, N_GLOBAL,
|
||||
M_GLOBAL, N_GLOBAL);
|
||||
matMultiplyOnHost(A_h, B_h, result_host, alpha, beta, M_GLOBAL, K_GLOBAL, K_GLOBAL, N_GLOBAL, M_GLOBAL, N_GLOBAL);
|
||||
|
||||
for (int i = 0; i < N_GLOBAL * M_GLOBAL; i++) {
|
||||
if (fabs(result_hD[i] - result_host[i]) > 0.1f) {
|
||||
@@ -823,15 +839,15 @@ int main(int argc, char **argv)
|
||||
checkCudaErrors(cudaEventElapsedTime(&milliseconds, start, stop));
|
||||
|
||||
printf("Time: %f ms\n", milliseconds);
|
||||
printf("TFLOPS: %.2f\n", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2)/(milliseconds/1000.)) / 1e12);
|
||||
printf("TFLOPS: %.2f\n", (((double)M_GLOBAL * N_GLOBAL * K_GLOBAL * 2) / (milliseconds / 1000.)) / 1e12);
|
||||
|
||||
free(A_h);
|
||||
free(B_h);
|
||||
free(C_h);
|
||||
checkCudaErrors(cudaFree((void*)A));
|
||||
checkCudaErrors(cudaFree((void*)B));
|
||||
checkCudaErrors(cudaFree((void*)C));
|
||||
checkCudaErrors(cudaFree((void*)D));
|
||||
checkCudaErrors(cudaFree((void *)A));
|
||||
checkCudaErrors(cudaFree((void *)B));
|
||||
checkCudaErrors(cudaFree((void *)C));
|
||||
checkCudaErrors(cudaFree((void *)D));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user