Add support for Tensor cores

[jeanfeydy]

Hi @bcharlier, @joanglaunes,

With the new CUDA 10.2 and 11 releases, Tensor cores are progressively becoming usable for KeOps computations. Upgrading our generic reduction routines to take advantage of these new chips should be a priority for early 2021: we can reasonably expect a x10 speed-up for common use cases such as matrix-vector products with RBF kernel matrices (#98) or nearest neighbor searches with Euclidean/Manhattan/Cosine/Hyperbolic metrics.

I have spent the last few days going through the CUDA documentation to look for code examples. Some useful links:

• The official CUDA API, that defines the "wmma" low-level primitives.

• The reference CUDA samples, with simple "KeOps++-like" implementations of a TensorCore-enabled matrix-vector product. The float32 example is a must-read: feel free to check the float16, bloat16, float64 and integer examples too.

• The CUTLASS collection of optimized kernels for matrix-matrix products at all levels of a CUDA chip. This header-only library is modeled after the reference CUB package for e.g. sorting operations. The introductory slides are interesting and the partial documentation on the GitHub repository is definitely worth reading... But unfortunately, the actual C++ API is still barely documented and I haven't been able to find the "kernel-side matrix multiply" that would actually fit our needs. The PyTorch_block_sparse repository is an example of project that uses CUTLASS, so it could be a useful reference too.

As far as I can tell today, the best way of supporting Tensor cores within KeOps would be to merge it with @joanglaunes 's work on "chunks". Our current reduction schemes (GPU_1D, GPU_2D, etc.) let CUDA threads work independently from each other: with the exception of the memory load of the y_j's, everything is done the thread registers. In order to work efficiently with large variables and/or TensorCores, we have to implement an other strategy: let the B threads (B=CUDA block size) cooperate to compute and store in the Shared memory some important buffers, such as distance and Gram matrices. This enables an optimal management of memory loads, as in Slide 12 of the CUTLASS presentation, and allows us to rely on TensorCores whenever possible.

I believe that @joanglaunes 's strategy is a good one:

1. Identify the costly/high-dimensional sub-formulas (e.g. squared distances and dot products).
2. At the beginning of the tile, precompute their values in large (B,B) buffers and store the results in the Shared Memory or better: in the Registers with one line per thread. This is a standard kernel-side (B,K)@(K,B) -> (B,B) matrix-matrix product (or "gemm"): an operation that is discussed in depth in the links above, and can be done efficiently by looping over the dimension K.
3. Execute the remainder of the reduction "as usual", letting each thread loop over the reduction index j. Crucially, we would only have to load in the Registers and Shared Memory the "small" variables. The only important exception here would be a kernel dot product with a wide right-hand side (i.e. sum_j k(x_i, y_j) * b_j with a large vector b_j) that could be handled using a specific "hack".

As a final note: a quick link to the important numbers on Turing/Ampere architectures:

• The physical piece of hardware that roughly corresponds to a CUDA "block" or "CTA" is the Streaming Multiprocessor (SM).
• There is around 96Kb = 24,000 float32 numbers of Shared Memory available per SM.
• There is up to 255 float32 Registers per thread.
• Threads are grouped by the compiler in Warps of 32 threads. These are executed on 4 physical blocks per SM, with 16 Int32 cores, 16 Float32 cores and 2 Tensor Cores each. All the nice graphics are in the Turing whitepaper.

To enable all of this, we will probably need to reduce the default CUDA block size from 192 to 128 or 64, making it more difficult for the compiler to hide latencies... But I am very optimistic :-)

See you soon!
Jean

N.B.: We could start experimenting with "float16*float16 + float32" computations on our RTX 2080 Ti. @joanglaunes has already done some work with half precision, so it wouldn't be out of reach I suppose.