How to use WMMA efficiently

Accelerated Computing CUDA CUDA Programming and Performance
4

When I run the code you have shown now, on a Tesla V100:

$ cat t18.cu
#include <cuda_fp16.h>
#include <mma.h>
using namespace nvcuda;
using namespace std;
__global__ void test_wmma()
{
        // Warps within a block read from 256 byte aligned strided adresses to avoid
        // bank conflicts (makes no difference).
        __shared__ __half smem[1024 * 8];
        __half* A = smem + threadIdx.y * 1024 + threadIdx.y * 16;
        __half* B = smem + threadIdx.y * 1024 + threadIdx.y * 16 + 256;
        __half* C = smem + threadIdx.y * 1024 + threadIdx.y * 16 + 512;

        // Matrix A is read once, and accumulator is filled once.
        wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> a_frag;
        wmma::fragment<wmma::accumulator, 16, 16, 16, __half> acc_frag;
        wmma::fill_fragment( acc_frag, 0.0f );
        wmma::load_matrix_sync( a_frag, A, 16 );

#pragma unroll
        for ( int i = 0; i < 20; i++ )
        {
                wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::col_major> b_frag;
                wmma::load_matrix_sync( b_frag, B, 16 );
                wmma::mma_sync( acc_frag, a_frag, b_frag, acc_frag );
        }

        wmma::store_matrix_sync( C, acc_frag, 16, wmma::mem_col_major );
}
void TestWMMA()
{
        int threads = 256;
        int blocks = 10000;
        test_wmma<<<blocks, threads>>>();
}
int main(){

        TestWMMA();
        cudaDeviceSynchronize();
        TestWMMA();
        cudaDeviceSynchronize();
}
$ nvcc -o t18 t18.cu -arch=sm_70
$ cuda-memcheck ./t18
========= CUDA-MEMCHECK
========= ERROR SUMMARY: 0 errors
$ nvprof ./t18
==49596== NVPROF is profiling process 49596, command: ./t18
==49596== Profiling application: ./t18
==49596== Profiling result:
            Type  Time(%)      Time     Calls       Avg       Min       Max  Name
 GPU activities:  100.00%  264.31us         2  132.16us  131.87us  132.44us  test_wmma(void)
      API calls:   95.91%  438.94ms         2  219.47ms  14.358us  438.92ms  cudaLaunchKernel
                    2.13%  9.7555ms         8  1.2194ms  1.1207ms  1.3072ms  cuDeviceTotalMem
                    1.73%  7.8957ms       808  9.7710us     263ns  498.41us  cuDeviceGetAttribute
                    0.16%  733.96us         8  91.744us  86.852us  114.31us  cuDeviceGetName
                    0.06%  266.37us         2  133.18us  131.21us  135.16us  cudaDeviceSynchronize
                    0.01%  30.125us         8  3.7650us  1.8050us  10.934us  cuDeviceGetPCIBusId
                    0.00%  8.4330us        16     527ns     361ns  1.0820us  cuDeviceGet
                    0.00%  4.5090us         8     563ns     396ns     808ns  cuDeviceGetUuid
                    0.00%  2.8980us         3     966ns     404ns  2.0260us  cuDeviceGetCount
$

I get a kernel execution time of ~130us.

If I do a flops/s calculation based on that, I believe it should be:

16x16x16x2x20x8x10000/0.000130 = 100,824,615,384,615.3… which is 100TF. That looks pretty good to me.