[KMCompiler][TLERaw] Support nvshmem

Author: lizhangyu258

python/triton/experimental/tle/raw/cache_key.py

@@ -94,6 +94,16 @@ def compute_tle_raw_source_cache_key(
     return hasher.hexdigest()
 
 
+def compute_tle_raw_host_cache_key(source: Union[str, Path], arch: str) -> str:
+    """Hash CUDA host source files."""
+    source_path = Path(source).resolve()
+    hasher = hashlib.sha256()
+    hasher.update(_read_source(source_path).encode())
+    hasher.update(str(arch).encode())
+
+    return hasher.hexdigest()
+
+
 def bind_tle_raw_source_cache_key(edsl: Any, **dialect_kwargs) -> None:
     """Attach __triton_tle_raw_source_cache_key__ to a @dialect edsl object."""
     if getattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR, None) is not None:

python/triton/experimental/tle/raw/cuda/runtime.py

@@ -6,17 +6,23 @@
 from pathlib import Path
 import subprocess
 from typing import Any, Final
+import ctypes
+from triton import knobs
 
 import torch
 
 from triton._C.libtriton import llvm  # pyright: ignore[reportMissingImports]
 from triton._C.libtriton.tle.llvm import parse_llvm_ir  # pyright: ignore[reportMissingImports]
 from triton.experimental.tle.raw.source_store import register_source
+from triton.experimental.tle.raw.nvshmem.utils import get_nvshmem_home
 
 # TODO: We use cli tools to compile CUDA code temporarily, and plan to replace it with LLVM components Python bindings in the future.
 CLANG = os.getenv("CLANG", "clang")
 CLANG_FLAGS = shlex.split(os.getenv("CLANG_FLAGS", ""))
 
+_cumodule_hook_installed = False
+_nvshmemx_cumodule_init = None
+
 
 def _sanitize_clang_ir(ir: str) -> str:
     # Newer clang emits attributes that this Triton branch's LLVM parser does
@@ -46,10 +52,43 @@ def _get_cuda_gpu_arch() -> str:
     return f"--cuda-gpu-arch=sm_{major}{minor}"
 
 
+def _get_nvshmemx_cumodule_init():
+    global _nvshmemx_cumodule_init
+    if _nvshmemx_cumodule_init is not None:
+        return _nvshmemx_cumodule_init
+
+    nvshmem_home = get_nvshmem_home()
+    library = ctypes.CDLL(str(Path(nvshmem_home) / "lib" / "libnvshmem_host.so"))
+    fn = library.nvshmemx_cumodule_init
+    fn.argtypes = [ctypes.c_void_p]
+    fn.restype = ctypes.c_int
+    _nvshmemx_cumodule_init = fn
+    return fn
+
+
+def _install_cumodule_hook():
+    global _cumodule_hook_installed
+    if _cumodule_hook_installed:
+        return
+
+    def hook(*args, **kwargs):
+        key = kwargs["key"]
+        function = kwargs["fn"].jit_function
+        device = kwargs["compile"]["device"]
+        kernel = function.device_caches[device][0].get(key)
+        assert kernel is not None
+        kernel._init_handles()
+        result = _get_nvshmemx_cumodule_init()(ctypes.c_void_p(kernel.module))
+        assert result == 0, f"nvshmemx_cumodule_init failed: {result}"
+
+    knobs.runtime.jit_post_compile_hook = hook
+    _cumodule_hook_installed = True
+
+
 class CUDAJITFunction(object):
 
     def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None:
-        super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")})
+        super().__init__()
         self.fn: Final[Any] = fn
         self.code: Final[str] = file.read_text()
         self.region_dialect: Final[str] = "cuda"
@@ -60,6 +99,9 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None:
         self.deferred: Final[bool] = kwargs.get("deferred", False)
         self.__triton_builtin__: Final[bool] = True
 
+        if "nvshmem" in self.code:
+            _install_cumodule_hook()
+
     def register_pending_source(self, *, hint: str = "") -> str:
         if not self.extern_func_name:
             raise RuntimeError("deferred tle_raw CUDA source requires extern_func_name= "

python/triton/experimental/tle/raw/nvshmem/common-host.cu

@@ -0,0 +1,56 @@
+#include <cstring>
+#include <cuda_runtime.h>
+#include <nvshmem.h>
+#include <nvshmemx.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#undef CUDA_CHECK
+#define CUDA_CHECK(stmt)                                                       \
+  do {                                                                         \
+    cudaError_t result = (stmt);                                               \
+    if (cudaSuccess != result) {                                               \
+      fprintf(stderr, "[%s:%d] cuda failed with %s \n", __FILE__, __LINE__,    \
+              cudaGetErrorString(result));                                     \
+      exit(-1);                                                                \
+    }                                                                          \
+  } while (0)
+
+extern "C" int nvshmem_get_unique_id_bytes(void *uid_buffer,
+                                           size_t uid_buffer_size) {
+  if (uid_buffer_size < sizeof(nvshmemx_uniqueid_t)) {
+    return -1;
+  }
+
+  nvshmemx_uniqueid_t uid;
+  nvshmemx_get_uniqueid(&uid);
+  memcpy(uid_buffer, &uid, sizeof(uid));
+  return 0;
+}
+
+extern "C" int nvshmem_init_from_torch_distributed(int rank, int nranks,
+                                                   int cuda_device,
+                                                   void *uid_buffer,
+                                                   size_t uid_buffer_size) {
+  if (uid_buffer_size < sizeof(nvshmemx_uniqueid_t)) {
+    return -1;
+  }
+
+  CUDA_CHECK(cudaSetDevice(cuda_device));
+
+  nvshmemx_uniqueid_t uid;
+  memcpy(&uid, uid_buffer, sizeof(uid));
+
+  nvshmemx_init_attr_t attr;
+  memset(&attr, 0, sizeof(attr));
+  nvshmemx_set_attr_uniqueid_args(rank, nranks, &uid, &attr);
+  nvshmemx_init_attr(NVSHMEMX_INIT_WITH_UNIQUEID, &attr);
+
+  return 0;
+}
+
+extern "C" int nvshmem_finalize_from_torch_distributed() {
+  nvshmem_finalize();
+
+  return 0;
+}