Repository navigation
Julia JIT support for device APIs #32
Description
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
Activity
ZzEeKkAa commentedon Jun 24, 2025
Hi @ggkountouras! I'm not very familiar with Julia, but this is how we bring mathdx libraries to python using libmathdx:
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 commentedon Jul 14, 2025
It would work, with changes to the Julia compiler to support libNVVM.