Discover the power of Tensor Cores

Environment

CPU

Intel(R) Xeon(R) CPU @ 2.80GHz

GPU

A100

GEFORCE RTX 3090

T4

Preparation

#define	M	512
#define	K	1024
#define	N	2048

Using standard C++

void
MatPro( float* a, float* b, float* c ) {
	for ( auto m = 0; m < M; m++ ) {
		for ( auto n = 0; n < N; n++ ) {
			auto $ = 0.;
			for ( auto k = 0; k < K; k++ ) $ += a[ m * K + k ] * b[ k * N + n ];
			c[ m * N + n ] = $;
		}
	}
}

Using CUDA Cores

__global__ void
MatPro( float* a, float* b, float* c ) {
	auto $ = 0.;
	auto n = blockIdx.x * blockDim.x + threadIdx.x;
	auto m = blockIdx.y * blockDim.y + threadIdx.y;
	for ( auto k = 0; k < K; k++ ) $ += a[ m * K + k ] * b[ k * K + n ];
	c[ m * N + n ] = $;
}

MatPro<<< dim3( N / 32, M / 32 ), dim3( 32, 32 ) >>>( a, b, c );

Using Tensor Cores with half precision data.

template < typename F > __global__ void
MatPro(
	const	half*	_a
,	const	half*	_b
,			F*		_c
) {
	wmma::fragment< wmma::matrix_a, 16, 16, 16, half, wmma::row_major > a;
	wmma::fragment< wmma::matrix_b, 16, 16, 16, half, wmma::row_major > b;
	wmma::fragment< wmma::accumulator, 16, 16, 16, F > c;

	wmma::fill_fragment( c, 0 );

	for ( auto k = 0; k < K; k += 16 ) {
		wmma::load_matrix_sync( a, _a + ( blockIdx.y * K * 16 + k ), K );
		wmma::load_matrix_sync( b, _b + ( k * N + blockIdx.x * 16 ), N );
		wmma::mma_sync( c, a, b, c );
	}

	wmma::store_matrix_sync( _c + ( blockIdx.y * N * 16 + blockIdx.x * 16 ), c, N, wmma::mem_row_major );
}

MatPro< F ><<< dim3( N / 16, M / 16 ), 32 >>>( a, b, c );

Conclusion

Bonus