add tle_raw cuda host cache key

Author: lizhangyu258

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

@@ -9,6 +9,8 @@
 # Protocol attribute checked by triton.runtime.jit.DependenciesFinder.
 # Value may be a str or a zero-arg callable returning the current key fragment.
 TLE_RAW_SOURCE_CACHE_KEY_ATTR = "__triton_tle_raw_source_cache_key__"
+TLE_RAW_COMMON_HOST_CACHE_KEY_ATTR = "__triton_tle_raw_common_host_cache_key__"
+TLE_RAW_HOST_CACHE_KEY_ATTR = "__triton_tle_raw_host_cache_key__"
 
 _DIALECT_KWARG_KEYS = (
     "name",
@@ -107,3 +109,33 @@ def tle_raw_source_cache_key() -> str:
         return compute_tle_raw_source_cache_key(normalized, edsl=edsl)
 
     setattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR, tle_raw_source_cache_key)
+
+
+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(str(source_path).encode())
+    hasher.update(_read_source(source_path).encode())
+    hasher.update(str(arch).encode())
+
+    return hasher.hexdigest()
+
+
+def bind_tle_raw_common_host_cache_key(edsl: Any, **kwargs) -> None:
+    """Attach __triton_tle_raw_source_cache_key__ to a @dialect edsl object."""
+    if getattr(edsl, TLE_RAW_COMMON_HOST_CACHE_KEY_ATTR, None) is not None:
+        return
+
+    setattr(edsl, TLE_RAW_COMMON_HOST_CACHE_KEY_ATTR, kwargs.get("key"))
+
+
+def bind_tle_raw_host_cache_key(edsl: Any, **kwargs) -> None:
+    """Attach __triton_tle_raw_source_cache_key__ to a @dialect edsl object."""
+    if getattr(edsl, TLE_RAW_HOST_CACHE_KEY_ATTR, None) is not None:
+        return
+
+    # def tle_raw_host_cache_key() -> str:
+    #     return compute_tle_raw_host_cache_key(kwargs)
+
+    setattr(edsl, TLE_RAW_HOST_CACHE_KEY_ATTR, kwargs.get("key"))

…ls/tle/raw/nvshmem/common/common-host.cu …experimental/tle/raw/cuda/common-host.cupython/tutorials/tle/raw/nvshmem/common/common-host.cu renamed to python/triton/experimental/tle/raw/cuda/common-host.cu python/tutorials/tle/raw/nvshmem/common/common-host.cu renamed to python/triton/experimental/tle/raw/cuda/common-host.cu

@@ -5,6 +5,15 @@
 from pathlib import Path
 
 import torch
+import functools
+import sysconfig
+import subprocess
+import re
+import tempfile
+import shlex
+from triton.runtime.cache import get_cache_manager
+from triton.experimental.tle.raw.cuda.runtime import _get_cuda_gpu_arch
+from triton.experimental.tle.raw.cache_key import compute_tle_raw_host_cache_key, bind_tle_raw_common_host_cache_key, bind_tle_raw_host_cache_key
 
 try:
     from cuda.bindings import driver as cuda
@@ -20,6 +29,135 @@
         from cuda import cudart as cudart
 
 
+@functools.lru_cache()
+def get_nvshmem_home() -> Path:
+    if (nvshmem_home := os.getenv("NVSHMEM_HOME")) is not None:
+        return Path(nvshmem_home)
+
+    try:
+        import nvidia.nvshmem
+        return Path(nvidia.nvshmem.__path__[0])
+    except Exception:
+        pass
+
+
+@functools.lru_cache()
+def get_nvcc():
+    return _path_to_binary("nvcc")
+
+
+@functools.lru_cache()
+def _path_to_binary(binary: str):
+    binary += sysconfig.get_config_var("EXE")
+    paths = [
+        os.environ.get(f"TRITON_{binary.upper()}_PATH", ""),
+        os.path.join(Path(os.path.dirname(__file__)).parent, "triton/backends/nvidia/bin", binary),
+    ]
+
+    cuda_home = os.getenv("CUDA_HOME", "/usr/local/cuda")
+
+    paths += [f"{cuda_home}/bin/{binary}"]
+
+    for path in paths:
+        if os.path.exists(path) and os.path.isfile(path):
+            result = subprocess.check_output([path, "--version"], stderr=subprocess.STDOUT)
+            if result is not None:
+                version = re.search(r".*release (\d+\.\d+).*", result.decode("utf-8"), flags=re.MULTILINE)
+                if version is not None:
+                    return path, version.group(1)
+    raise RuntimeError(f"Cannot find {binary}")
+
+
+def _compile_cuda_host_to_cache(
+    source,
+    nvshmem_home,
+    arch: str = None,
+    force: bool = False,
+) -> tuple[Path, str, bool]:
+    source_path = Path(source).expanduser().resolve()
+    if not arch:
+        arch = _get_cuda_gpu_arch().split('=')[1]
+    nvshmem_home = get_nvshmem_home()
+    host_cache_key = compute_tle_raw_host_cache_key(source_path, arch)
+    output_name = source_path.with_suffix(".so").name
+    cache = get_cache_manager(host_cache_key)
+
+    cached = None if force else cache.get_file(output_name)
+    if cached is not None:
+        return Path(cached), host_cache_key, True
+
+    temporary = tempfile.NamedTemporaryFile(
+        prefix=f".{output_name}.",
+        suffix=".tmp",
+        delete=False,
+    )
+    temporary_path = Path(temporary.name)
+    temporary.close()
+    nvcc, _ = get_nvcc()
+    command = [
+        nvcc,
+        "-shared",
+        "-Xcompiler",
+        "-fPIC",
+        "-rdc=true",
+        f"-arch={arch}",
+        f"-I{nvshmem_home / 'include'}",
+        f"-L{nvshmem_home / 'lib'}",
+        "-lnvshmem_host",
+        "-lnvshmem_device",
+        "-o",
+        str(temporary_path),
+        str(source_path),
+    ]
+    try:
+        build = subprocess.run(command, capture_output=True)
+        if build.returncode != 0:
+            raise RuntimeError("nvcc failed while compiling CUDA host library\n"
+                               f"command: {shlex.join(command)}\n"
+                               f"stderr:\n{build.stderr.decode()}")
+        cached_path = cache.put(temporary_path.read_bytes(), output_name, binary=True)
+        return Path(cached_path), host_cache_key, False
+    finally:
+        temporary_path.unlink(missing_ok=True)
+
+
+class CudaHostLibrary:
+
+    def __init__(self, bind_edsls, library_path, host_cache_key, type):
+        self.edsls = bind_edsls
+        self.path = Path(library_path).expanduser().resolve()
+        self.host_cache_key = host_cache_key
+        self.library = ctypes.CDLL(str(self.path))
+
+        if type == "common":
+            for edsl in self.edsls:
+                bind_tle_raw_common_host_cache_key(edsl, key=host_cache_key)
+        elif type == "op":
+            for edsl in self.edsls:
+                bind_tle_raw_host_cache_key(edsl, key=host_cache_key)
+        else:
+            raise RuntimeError("Unsupported cuda-host type.")
+
+    def __getattr__(self, name):
+        return getattr(self.library, name)
+
+
+def get_common_host_source() -> Path:
+    return Path(__file__).resolve().parents[1] / "cuda" / "common-host.cu"
+
+
+def load_host(bind_edsls, source, type="op", nvshmem_home=None, arch: str | None = None,
+              force: bool = False) -> CudaHostLibrary:
+    path, host_cache_key, _ = _compile_cuda_host_to_cache(source, nvshmem_home, arch, force=force)
+    return CudaHostLibrary(bind_edsls, path, host_cache_key, type)
+
+
+def load_common_host(bind_edsls, source=None, type="common", nvshmem_home=None, arch: str | None = None,
+                     force: bool = False) -> CudaHostLibrary:
+    return load_host(bind_edsls, source or get_common_host_source(), type=type, nvshmem_home=nvshmem_home, arch=arch,
+                     force=force)
+
+
 def CUDA_CHECK(err):
     if isinstance(err, cuda.CUresult):
         if err != cuda.CUresult.CUDA_SUCCESS:

python/triton/experimental/tle/raw/utils/nvshmem.py

@@ -128,6 +128,19 @@ def record_reference(self, val, var_dict=None, name=None):
             part = (tle_raw_source_cache_key() if callable(tle_raw_source_cache_key) else tle_raw_source_cache_key)
             self.hasher.update(str(part).encode("utf-8"))
 
+        tle_raw_common_host_cache_key = getattr(val, "__triton_tle_raw_common_host_cache_key__", None)
+        if tle_raw_common_host_cache_key is not None:
+            print(f"common host key: {tle_raw_common_host_cache_key}")
+            part = (tle_raw_common_host_cache_key()
+                    if callable(tle_raw_common_host_cache_key) else tle_raw_common_host_cache_key)
+            self.hasher.update(str(part).encode("utf-8"))
+
+        tle_raw_host_cache_key = getattr(val, "__triton_tle_raw_host_cache_key__", None)
+        if tle_raw_host_cache_key is not None:
+            print(f"host key: {tle_raw_host_cache_key}")
+            part = (tle_raw_host_cache_key() if callable(tle_raw_host_cache_key) else tle_raw_host_cache_key)
+            self.hasher.update(str(part).encode("utf-8"))
+
         if getattr(val, "__triton_aggregate__", False):
             for attr in val.hash_attrs:
                 self.record_reference(attr)

python/triton/runtime/jit.py

@@ -7,7 +7,7 @@
 from triton.experimental.tle.raw import dialect
 
 from triton.experimental.tle.raw.utils.nvshmem import (
-    load_library,
+    load_host,
     tensor_from_pointer,
 )
 
@@ -28,8 +28,8 @@ def simple_shift_kernel(destination_ptr, ):
 
 
 def simpe_shift():
-    host_path = Path(__file__).with_name("simple-shift-host.so")
-    host_lib = load_library(host_path)
+    host_source = Path(__file__).with_name("simple-shift-host.cu")
+    host_lib = load_host(bind_edsl=simple_shift, source=host_source)
 
     mype = ctypes.c_int()
     npes = ctypes.c_int()

python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py

@@ -12,7 +12,8 @@
 
 from triton.experimental.tle.raw.utils.nvshmem import (
     print_perf_mean,
-    load_library,
+    load_common_host,
+    load_host,
     init_torch_distributed,
     init_nvshmem_by_torch_pg,
     tensor_from_pointer,
@@ -469,10 +470,9 @@ def main():
         sys.exit()
 
     group = init_torch_distributed()
-    common_path = Path(__file__).parents[1] / "common" / "common-host.so"
-    host_path = Path(__file__).with_name("ag-gemm-host.so")
-    common = load_library(common_path)
-    host = load_library(host_path)
+    host_source = Path(__file__).with_name("ag-gemm-host.cu")
+    host = load_host([mark_local_ready, wait_ready], host_source)
+    common = load_common_host([mark_local_ready, wait_ready])
     configure_host_library(host)
     init_nvshmem_by_torch_pg(common, group)
 

python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py

@@ -205,15 +205,15 @@ def main():
     nvshmem_home = resolve_nvshmem_home(args.nvshmem_home)
     arch = detect_arch(args.arch)
     generated = generate_extern_files(target.parent)
-    common = compile_common_host(nvshmem_home, arch, args.force)
-    libraries = compile_host_files(target.parent, nvshmem_home, arch, args.force)
+    # common = compile_common_host(nvshmem_home, arch, args.force)
+    # libraries = compile_host_files(target.parent, nvshmem_home, arch, args.force)
 
     for path in generated:
         print(f"[prepare] extern file: {path}")
-    if common is not None:
-        print(f"[prepare] common host: {common}")
-    for path in libraries:
-        print(f"[prepare] host: {path}")
+    # if common is not None:
+    #     print(f"[prepare] common host: {common}")
+    # for path in libraries:
+    #     print(f"[prepare] host: {path}")
     print(f"[prepare] NVSHMEM_HOME={nvshmem_home}")
     print(f"[prepare] CUDA architecture={arch}")