Julia JIT support for device APIs

[ggkountouras]

As header-only libraries based on CUTLASS, these only work in C++.

It would be nice to have e.g. the new DGEMM via IMMA (https://github.com/NVIDIA/CUDALibrarySamples/tree/master/MathDx/cuBLASDx/16_dgemm_emulation) inside Julia kernels.

Unfortunately, Warp/numba can only JIT compile a subset of python types, and not the (more general) Julia types that we require (e.g. differential equation solvers).

A potential way forward would be to compile libmathdx to PTX, link it in via LLVM, and get full performance via LTO. However, these often contain NVVM IR, which LLVM cannot handle.

Related discussion: https://discourse.julialang.org/t/using-cublasdx-in-julia/125527

[ZzEeKkAa]

Hi @ggkountouras! I'm not very familiar with Julia, but this is how we bring mathdx libraries to python using libmathdx:

1. compile separately mathdx types and functions as an LTO(s) through libmathdx
2. compile numba-cuda kernel with function references to LTO from previous step.
3. Link kernel and and libmathdx produced LTO using libnvvm
4. Dispatch the kernel

Please let me know if the same approach would work at Julia.

BTW, we have DGEMM via IMMA implementation already available at nvmath: https://github.com/NVIDIA/nvmath-python/blob/main/examples/device/cublasdx_fp64_emulation.py

[ggkountouras]

It would work, with changes to the Julia compiler to support libNVVM.