[KMCompiler][TLERaw] Support nvshmem

.github/workflows/nv3.6-build-and-test.yml

@@ -134,7 +134,11 @@ jobs:
           python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle
           python3 python/tutorials/tle/raw/mlir/06-test-vassert.py
           ## flagtree tle cuda
-          # TODO: These tests are currently skipped because the CLANG environment variable cannot be set.
-          # python3 python/tutorials/tle/raw/cuda/01-vector-add.py
-          # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py
-          # python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py
+          export CLANG=/usr/local/lib/python3.12/dist-packages/mlir/llvm_artifact/bin/clang-22
+          python3 python/tutorials/tle/raw/cuda/01-vector-add.py
+          python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py
+          python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py
+          python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication-smem.py
+          ## flagtree tle nvshmem
+          torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py
+          torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py

python/requirements-nvshmem-cu12.txt

@@ -0,0 +1,2 @@
+#flagtree tle raw nvshmem
+nvidia-nvshmem-cu12

python/requirements-nvshmem-cu13.txt

@@ -0,0 +1,2 @@
+#flagtree tle raw nvshmem
+nvidia-nvshmem-cu13

python/triton/experimental/tle/language/raw/core.py

@@ -23,9 +23,9 @@
 from triton.experimental.tle.language.gpu import buffered_tensor
 
 
-def _resolve_alias_indices(func, llvm, handles, output_indices, _semantic):
+def _resolve_alias_indices(func, llvm, handles, output_indices, extern_func_name, _semantic):
     if output_indices is None:
-        return _semantic.builder.compute_alias_operand_indices(llvm, handles)
+        return _semantic.builder.compute_alias_operand_indices(llvm, handles, extern_func_name)
     return output_indices
 
 
@@ -72,8 +72,10 @@ def _tle_raw_call(func, args, *, output_indices, hint, smem, _semantic):
     else:
         context = _semantic.builder.get_context()
         llvm = func.make_llvm(context)
-        alias_indices = _resolve_alias_indices(func, llvm, handles, output_indices, _semantic)
-        dsl_region_op = func.create_region_by_llvm(_semantic.builder, llvm, handles, alias_indices, hint)
+        extern_func_name = getattr(func, "extern_func_name", None) or ""
+        alias_indices = _resolve_alias_indices(func, llvm, handles, output_indices, extern_func_name, _semantic)
+        dsl_region_op = func.create_region_by_llvm(_semantic.builder, llvm, handles, alias_indices, hint,
+                                                   extern_func_name)
     return _wrap_results(args, alias_indices, dsl_region_op, smem=smem)
 
 

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

@@ -114,6 +114,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

@@ -19,23 +19,136 @@
 # SOFTWARE.
 
 from __future__ import annotations
+
+import ctypes
+import functools
 import os
 import re
 import shlex
+import shutil
 import struct
-from pathlib import Path
 import subprocess
+from pathlib import Path
 from typing import Any, Final
 
 import torch
-
+from triton import knobs
 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.runtime import RawJITFunction
 from triton.experimental.tle.raw.source_store import register_source
 
-# 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", ""))
+# TODO: Temporarily shell out to clang; replace with LLVM Python bindings later.
+_MIN_CLANG_MAJOR = 20
+
+# ---------------------------------------------------------------------------
+# Clang toolchain: knobs override -> discover, always version-check (>= 20)
+# ---------------------------------------------------------------------------
+
+
+def _parse_clang_major(clang: str) -> int | None:
+    try:
+        out = subprocess.check_output([clang, "--version"], text=True, stderr=subprocess.STDOUT)
+    except (OSError, subprocess.CalledProcessError):
+        return None
+    match = re.search(r"clang version (\d+)\.", out)
+    if match is None:
+        match = re.search(r"version (\d+)\.", out)
+    return int(match.group(1)) if match else None
+
+
+def _clang_meets_min_version(clang: str) -> bool:
+    major = _parse_clang_major(clang)
+    return major is not None and major >= _MIN_CLANG_MAJOR
+
+
+def _discover_clang_binaries() -> list[str]:
+    """Prefer newer versioned binaries, then plain ``clang``."""
+    found: list[str] = []
+    seen: set[str] = set()
+    for name in ("clang-22", "clang-21", "clang-20", "clang"):
+        path = shutil.which(name)
+        if path is None:
+            continue
+        resolved = str(Path(path).resolve())
+        if resolved in seen:
+            continue
+        seen.add(resolved)
+        found.append(path)
+    return found
+
+
+@functools.lru_cache()
+def _resolve_clang() -> str:
+    """Use user CLANG if usable; otherwise discover. Every candidate is version-checked."""
+    tried: list[str] = []
+    user_clang = knobs.nvidia.tle_raw_clang
+    if user_clang:
+        tried.append(user_clang)
+        if _clang_meets_min_version(user_clang):
+            return user_clang
+
+    for candidate in _discover_clang_binaries():
+        if candidate in tried:
+            continue
+        tried.append(candidate)
+        if _clang_meets_min_version(candidate):
+            return candidate
+
+    detail = ", ".join(tried) if tried else "<none>"
+    raise RuntimeError(f"TLE raw CUDA requires clang >= {_MIN_CLANG_MAJOR}. "
+                       f"Tried: {detail}. Install clang-20+ or set CLANG to a suitable binary.")
+
+
+# ---------------------------------------------------------------------------
+# Clang compile flags (--cuda-path, includes, optional CLANG_FLAGS)
+# ---------------------------------------------------------------------------
+
+
+def _cuda_home() -> Path:
+    return Path(os.getenv("CUDA_HOME", "/usr/local/cuda"))
+
+
+def _nvidia_backend_include() -> Path | None:
+    for parent in Path(__file__).resolve().parents:
+        candidate = parent / "third_party" / "nvidia" / "backend" / "include"
+        if candidate.is_dir():
+            return candidate
+    return None
+
+
+def _default_clang_flags() -> list[str]:
+    cuda_home = _cuda_home()
+    flags = [f"--cuda-path={cuda_home}", f"-I{cuda_home / 'include'}"]
+    backend_include = _nvidia_backend_include()
+    if backend_include is not None:
+        flags.append(f"-I{backend_include}")
+    try:
+        from triton.experimental.tle.raw.nvshmem.utils import try_get_nvshmem_home
+        nvshmem_home = try_get_nvshmem_home()
+        if nvshmem_home is not None:
+            flags.append(f"-I{nvshmem_home / 'include'}")
+    except Exception:
+        pass
+    return flags
+
+
+def _clang_flags() -> list[str]:
+    extra = knobs.nvidia.tle_raw_clang_flags or ""
+    return [*_default_clang_flags(), *shlex.split(extra)]
+
+
+def _get_cuda_gpu_arch() -> str:
+    arch = os.getenv("TLE_CUDA_ARCH")
+    if arch:
+        return f"--cuda-gpu-arch={arch}"
+    major, minor = torch.cuda.get_device_capability()
+    return f"--cuda-gpu-arch=sm_{major}{minor}"
+
+
+# ---------------------------------------------------------------------------
+# Sanitize clang LLVM IR for this Triton's parser
+# ---------------------------------------------------------------------------
 
 
 def _sanitize_clang_ir(ir: str) -> str:
@@ -58,27 +171,70 @@ def _replace_hex_float(match: re.Match[str]) -> str:
     return re.sub(r"f0x([0-9A-Fa-f]+)", _replace_hex_float, ir)
 
 
-def _get_cuda_gpu_arch() -> str:
-    arch = os.getenv("TLE_CUDA_ARCH")
-    if arch:
-        return f"--cuda-gpu-arch={arch}"
-    major, minor = torch.cuda.get_device_capability()
-    return f"--cuda-gpu-arch=sm_{major}{minor}"
+# ---------------------------------------------------------------------------
+# NVSHMEM: post-compile cumodule init hook
+# ---------------------------------------------------------------------------
+
+_cumodule_hook_installed = False
+_nvshmemx_cumodule_init = None
 
 
-class CUDAJITFunction(object):
+def _get_nvshmemx_cumodule_init():
+    global _nvshmemx_cumodule_init
+    if _nvshmemx_cumodule_init is not None:
+        return _nvshmemx_cumodule_init
+
+    from triton.experimental.tle.raw.nvshmem.utils import (
+        get_nvshmem_home,
+        resolve_nvshmem_host_library,
+    )
+    library = ctypes.CDLL(str(resolve_nvshmem_host_library(get_nvshmem_home())))
+    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
+
+
+# ---------------------------------------------------------------------------
+# Dialect runtime
+# ---------------------------------------------------------------------------
+
+
+class CUDAJITFunction(RawJITFunction):
 
     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")})
-        self.fn: Final[Any] = fn
+        super().__init__(fn, **kwargs)
         self.code: Final[str] = file.read_text()
         self.region_dialect: Final[str] = "cuda"
         self.lowered_region_dialect: Final[str] = "llvm"
         self.arg_dialect: Final[str] = "llvm"
         self.source_file: Final[str] = str(file)
-        self.extern_func_name = kwargs.get("extern_func_name", None)
-        self.deferred: Final[bool] = kwargs.get("deferred", False)
-        self.__triton_builtin__: Final[bool] = True
+
+        if self.library == "nvshmem":
+            from triton.experimental.tle.raw.nvshmem.utils import enable_nvshmem_device_bc
+            enable_nvshmem_device_bc(True)
+        if self.library == "nvshmem" or "nvshmem" in self.code:
+            _install_cumodule_hook()
 
     def register_pending_source(self, *, hint: str = "") -> str:
         if not self.extern_func_name:
@@ -92,15 +248,9 @@ def register_pending_source(self, *, hint: str = "") -> str:
             extra={"source_file": self.source_file},
         )
 
-    def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = ""):
-        return builder.create_tle_raw_region_by_llvm_func(
-            llvm,
-            self.region_dialect,
-            self.arg_dialect,
-            handles,
-            alias_indices,
-            hint,
-        )
+    def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "",
+                              extern_func_name: str = ""):
+        return super().create_region_by_llvm(builder, llvm, handles, alias_indices, hint, extern_func_name)
 
     def create_region_deferred(self, builder, source_id: str, handles, alias_indices, hint: str = ""):
         return builder.create_tle_raw_region_deferred(
@@ -115,7 +265,7 @@ def create_region_deferred(self, builder, source_id: str, handles, alias_indices
     def make_llvm(self, mlir_context) -> str:
         build = subprocess.run(
             [
-                CLANG,
+                _resolve_clang(),
                 "-x",
                 "cuda",
                 "--cuda-device-only",
@@ -126,7 +276,7 @@ def make_llvm(self, mlir_context) -> str:
                 "-",
                 "-o",
                 "-",
-                *CLANG_FLAGS,
+                *_clang_flags(),
             ],
             input=self.code.encode(),
             capture_output=True,

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

@@ -30,16 +30,16 @@
 
 from .codegen import MLIRCodeGenerator
 from triton.experimental.tle.raw.source_store import register_source
+from triton.experimental.tle.raw.runtime import RawJITFunction
 
 _pending_jit_fn_key = "mlir_jit_fn"
 
 
-class MLIRJITFunction(object):
+class MLIRJITFunction(RawJITFunction):
 
     def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optional[ir.Context] = None, *args,
                  **kwargs) -> None:
-        super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")})
-        self.fn: Final[Any] = fn
+        super().__init__(fn, **kwargs)
         self.pipeline: Final[List[str]] = ([*pipeline] if pipeline is not None else [
             "convert-scf-to-cf",
             "finalize-memref-to-llvm",
@@ -53,9 +53,6 @@ def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optio
         self.context: Final[ir.Context] = ir.Context() if context is None else context
         self.region_dialect: Final[str] = "mlir"
         self.arg_dialect: Final[str] = "llvm"
-        self.extern_func_name: Final[Optional[str]] = kwargs.get("extern_func_name")
-        self.deferred: Final[bool] = kwargs.get("deferred", False)
-        self.__triton_builtin__: Final[bool] = True
 
     def __deepcopy__(self, memo: Dict[int, Any]) -> MLIRJITFunction:
         return self.__class__(copy.deepcopy(self.fn, memo), copy.deepcopy(self.pipeline, memo), self.context)
@@ -119,15 +116,9 @@ def create_region_deferred(self, builder, source_id: str, handles, alias_indices
             hint,
         )
 
-    def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = ""):
-        return builder.create_tle_raw_region_by_llvm_func(
-            llvm,
-            self.region_dialect,
-            self.arg_dialect,
-            handles,
-            alias_indices,
-            hint,
-        )
+    def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "",
+                              extern_func_name: str = ""):
+        return super().create_region_by_llvm(builder, llvm, handles, alias_indices, hint, extern_func_name)
 
     @cached_property
     def src(self) -> str: