deepseek-ai/DeepJIT

DeepJIT – JIT‑compilation for CUDA and Ascend GPUs/NPUs

What it is – DeepJIT is a header‑only C++20 library that lets you compile CUDA or Huawei Ascend kernel code at runtime and launch it from C++ or Python extensions. It handles the whole JIT workflow: source hashing, invoking the device compiler, caching the resulting binaries (both in memory and on disk), loading them onto the device, and providing launch helpers.

Why it matters – Writing custom kernels for deep‑learning frameworks often requires a separate build step. DeepJIT removes that friction, enabling just‑in‑time compilation that can be cached across runs and even shared between users or nodes. This speeds up development cycles and reduces the need for pre‑compiled binaries, especially when kernels depend on version‑specific headers or flags.

Key features

• Dual‑backend support – Same API works for NVIDIA CUDA (NVCC → CUBIN) and Huawei Ascend (Bisheng/ld.lld → Ascend assembly). Choose the backend with deep_jit::Runtime<deep_jit::CUDA> or ...<deep_jit::Ascend>.
• Kernel caching – Cache keys incorporate source, includes, compiler version, options, and a user‑provided extra signature. Cached artifacts are stored on disk and can be shared via a common directory (DJ_JIT_CACHE_DIR).
• Distributed/shared cache – Works on networked filesystems that support atomic rename and fsync, allowing multiple processes or nodes to reuse compiled kernels.
• Lazy initialization – Device and compiler discovery are deferred until the first compile/launch call.
• PyTorch integration – Uses the current PyTorch CUDA or torch_npu stream by default and exposes a get_jit() function through pybind11.
• Diagnostics & hooks – Runtime‑wide or per‑kernel compiler options, ability to dump PTX/SASS, and a Python post‑compilation hook that can modify the generated binary before it is cached.

How you use it

1. Add the headers to your project (include path).
2. Select a backend and create a runtime, e.g.:

   using JIT = deep_jit::Runtime<deep_jit::CUDA>;
   auto jit = deep_jit::create_lazy_jit<deep_jit::CUDA>(
       deep_jit::Config("/abs/path/to/lib", "MYLIB"));
   

3. Compile a kernel (source string) and launch it:

   auto kernel = jit->compile("scale", R"(
   extern "C" __global__ void scale(float* out, const float* in, int n){
       int i = blockIdx.x*blockDim.x + threadIdx.x;
       if(i<n) out[i] = in[i]*2.0f; }
   )");
   jit->launch(kernel, {.grid_dim={...}, .block_dim={256,1,1}}, out, in, n);
   

4. From Python you can obtain the same runtime via:

   import my_lib._C as lib
   jit = lib.get_jit()
   

Configuration – deep_jit::Config lets you set:

• Cache directory (env var MYLIB_JIT_CACHE_DIR or global DJ_JIT_CACHE_DIR).
• Include directories and which angle‑bracket prefixes to track.
• An extra_signature string to capture external dependencies (e.g., a specific CUTLASS version).
• Per‑kernel compiler and launch defaults.

Cache layout – Each cached entry lives under <cache‑root>/cache/<tag>.<digest>/ and contains the source, compiled binary (.cubin or Ascend artifact), optional PTX/SASS dumps, and a meta.json with the compilation command.

Current status – Core functionality (CUDA & Ascend compilation, caching, Python API) is stable. Two WIP items are mentioned: a cache‑warmup feature and a direct‑from‑Python compilation API.

Who might need it – Developers of custom PyTorch extensions, high‑performance scientific libraries, or any C++/Python code that wants to generate device kernels on the fly without managing a separate build system.

---

All details above are taken directly from the repository’s README.