diff --git a/.github/workflows/nv3.6-build-and-test.yml b/.github/workflows/nv3.6-build-and-test.yml index 73f660c85..cf725bc53 100644 --- a/.github/workflows/nv3.6-build-and-test.yml +++ b/.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 diff --git a/python/requirements-nvshmem-cu12.txt b/python/requirements-nvshmem-cu12.txt new file mode 100644 index 000000000..720391de9 --- /dev/null +++ b/python/requirements-nvshmem-cu12.txt @@ -0,0 +1,2 @@ +#flagtree tle raw nvshmem +nvidia-nvshmem-cu12 diff --git a/python/requirements-nvshmem-cu13.txt b/python/requirements-nvshmem-cu13.txt new file mode 100644 index 000000000..147178646 --- /dev/null +++ b/python/requirements-nvshmem-cu13.txt @@ -0,0 +1,2 @@ +#flagtree tle raw nvshmem +nvidia-nvshmem-cu13 diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index 8340e3115..100b23eb3 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/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) diff --git a/python/triton/experimental/tle/raw/cache_key.py b/python/triton/experimental/tle/raw/cache_key.py index fec85cac7..d9957194c 100644 --- a/python/triton/experimental/tle/raw/cache_key.py +++ b/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: diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 536b4604b..7f3e6e5fc 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/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 "" + 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, diff --git a/python/triton/experimental/tle/raw/mlir/runtime.py b/python/triton/experimental/tle/raw/mlir/runtime.py index 186405510..4179567a1 100644 --- a/python/triton/experimental/tle/raw/mlir/runtime.py +++ b/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: diff --git a/python/triton/experimental/tle/raw/nvshmem/common-host.cu b/python/triton/experimental/tle/raw/nvshmem/common-host.cu new file mode 100644 index 000000000..5f1382b7f --- /dev/null +++ b/python/triton/experimental/tle/raw/nvshmem/common-host.cu @@ -0,0 +1,56 @@ +#include +#include +#include +#include +#include +#include + +#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; +} diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py new file mode 100644 index 000000000..c6e76fba2 --- /dev/null +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -0,0 +1,350 @@ +import ctypes +import datetime +import os +import warnings +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.cache_key import compute_tle_raw_host_cache_key + +_cuda = None +_cudart = None + + +def _get_cuda_modules(): + global _cuda, _cudart + if _cuda is not None and _cudart is not None: + return _cuda, _cudart + try: + from cuda.bindings import driver as cuda + from cuda.bindings import runtime as cudart + except ImportError: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"The cuda\.(cuda|cudart) module is deprecated.*", + category=FutureWarning, + ) + from cuda import cuda as cuda + from cuda import cudart as cudart + _cuda = cuda + _cudart = cudart + return _cuda, _cudart + + +def __getattr__(name): + # Preserve `from ...utils import cuda/cudart` after lazy-loading. + if name in ("cuda", "cudart"): + cuda, cudart = _get_cuda_modules() + return cuda if name == "cuda" else cudart + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def CUDA_CHECK(err): + cuda, cudart = _get_cuda_modules() + if isinstance(err, cuda.CUresult): + if err != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"Cuda Error: {err}: {cuda.cuGetErrorName(err)}") + elif isinstance(err, cudart.cudaError_t): + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"Cuda Error: {err}: {cudart.cudaGetErrorString(err)}") + else: + raise RuntimeError(f"Unknown error type: {err}") + + +@functools.lru_cache() +def get_nvshmem_home() -> Path: + from triton import knobs + if knobs.nvidia.nvshmem_home is not None: + return Path(knobs.nvidia.nvshmem_home) + + try: + import nvidia.nvshmem + return Path(nvidia.nvshmem.__path__[0]) + except Exception: + raise RuntimeError("Cannot resolve NVSHMEM_HOME: environment variable not set and nvidia.nvshmem not found") + + +def try_get_nvshmem_home() -> Path | None: + try: + return get_nvshmem_home() + except RuntimeError: + return None + + +def resolve_nvshmem_host_library(nvshmem_home: Path | None = None) -> Path: + """Prefer unversioned .so; fall back to .so.N (pip wheels often omit the symlink).""" + home = Path(nvshmem_home) if nvshmem_home is not None else get_nvshmem_home() + lib_dir = home / "lib" + for name in ("libnvshmem_host.so", "libnvshmem_host.so.3"): + path = lib_dir / name + if path.exists(): + return path + raise RuntimeError(f"Cannot find libnvshmem_host.so[.3] under {lib_dir}") + + +# Set by @dialect(..., library="nvshmem"); CUDA backend links device .bc only when True. +_nvshmem_device_bc_enabled: bool = False + + +def enable_nvshmem_device_bc(enabled: bool = True) -> None: + global _nvshmem_device_bc_enabled + _nvshmem_device_bc_enabled = bool(enabled) + + +def is_nvshmem_device_bc_enabled() -> bool: + return _nvshmem_device_bc_enabled + + +def get_nvshmem_extern_libs(arch: str | int | None = None) -> dict[str, str]: + """Return {libnvshmem_device: path} when enabled; else {}.""" + if not is_nvshmem_device_bc_enabled(): + return {} + bc = resolve_nvshmem_device_bitcode(arch=arch) + return {"libnvshmem_device": str(bc)} if bc is not None else {} + + +def resolve_nvshmem_device_bitcode(nvshmem_home: Path | None = None, arch: str | int | None = None) -> Path | None: + """Resolve device bitcode: unified .bc, else per-SM .bc (NVSHMEM >= ~3.7).""" + home = Path(nvshmem_home) if nvshmem_home is not None else try_get_nvshmem_home() + if home is None: + return None + lib_dir = home / "lib" + unified = lib_dir / "libnvshmem_device.bc" + if unified.is_file(): + return unified + # arch: sm90 / sm_90a / 90 -> try exact then next-lower known ships + if isinstance(arch, int): + sm = arch + else: + s = str(arch or "").lower().replace("sm_", "").replace("sm", "").rstrip("a") + sm = int(s) if s.isdigit() else None + sms = (90, 89, 80, 75, 70) + for n in ([sm] + [x for x in sms if sm is not None and x < sm]) if sm is not None else sms: + path = lib_dir / f"libnvshmem_device_sm_{n}.bc" + if path.is_file(): + return path + return None + + +@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", "")] + + 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: + from triton.experimental.tle.raw.cuda.runtime import _get_cuda_gpu_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() + lib_dir = nvshmem_home / "lib" + host_lib = resolve_nvshmem_host_library(nvshmem_home) + device_lib = lib_dir / "libnvshmem_device.a" + if not device_lib.is_file(): + raise RuntimeError(f"Cannot find libnvshmem_device.a under {lib_dir}") + # nvcc cannot take .so as a positional input; use -L/-l (-l: for versioned SONAME). + command = [ + nvcc, + "-shared", + "-Xcompiler", + "-fPIC", + "-rdc=true", + f"-arch={arch}", + f"-I{nvshmem_home / 'include'}", + f"-L{lib_dir}", + f"-l:{host_lib.name}", + "-lnvshmem_device", + "-Xlinker", + "-rpath", + "-Xlinker", + str(lib_dir), + "-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, library_path): + self.path = Path(library_path).expanduser().resolve() + self.library = ctypes.CDLL(str(self.path)) + + def __getattr__(self, name): + return getattr(self.library, name) + + +def get_common_host_source() -> Path: + return Path(__file__).resolve().parents[0] / "common-host.cu" + + +def load_host(source, nvshmem_home=None, arch=None, force: bool = False) -> CudaHostLibrary: + path, _, _ = _compile_cuda_host_to_cache(source, nvshmem_home, arch, force=force) + return CudaHostLibrary(path) + + +def load_common_host(source=None, nvshmem_home=None, arch=None, force: bool = False) -> CudaHostLibrary: + return load_host(source or get_common_host_source(), nvshmem_home, arch, force) + + +def init_torch_distributed(): + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + torch.distributed.init_process_group( + backend="cpu:gloo,cuda:nccl", + rank=rank, + world_size=world_size, + device_id=torch.device("cuda", local_rank), + timeout=datetime.timedelta(seconds=1800), + ) + group = torch.distributed.new_group(ranks=list(range(world_size)), backend="nccl") + torch.distributed.barrier(group=group) + return group + + +def init_nvshmem_by_torch_pg(common, group): + rank = group.rank() + world_size = group.size() + uid_size = 1024 + + if rank == 0: + temp_buffer = ctypes.create_string_buffer(uid_size) + result = common.nvshmem_get_unique_id_bytes(temp_buffer, uid_size) + assert result == 0, f"nvshmemx_get_uniqueid failed: {result}" + uid = bytes(temp_buffer.raw) + else: + uid = bytes(uid_size) + + objects = [uid] + torch.distributed.broadcast_object_list( + objects, + src=torch.distributed.get_global_rank(group, 0), + group=group, + ) + uid_buffer = ctypes.create_string_buffer(objects[0], uid_size) + result = common.nvshmem_init_from_torch_distributed( + rank, + world_size, + int(os.environ["LOCAL_RANK"]), + uid_buffer, + uid_size, + ) + assert result == 0, f"NVSHMEM init failed: {result}" + torch.distributed.barrier(group=group) + + +def tensor_from_pointer(pointer, shape, dtype, device): + elements = 1 + for extent in shape: + elements *= extent + storage = torch._C._construct_storage_from_data_pointer( + pointer.value, + device, + elements * dtype.itemsize, + ) + return torch.empty(0, dtype=dtype, device=device).set_(storage).view(shape) + + +def set_signal_cuda_ptr(signal_ptr, signal, stream): + cuda, _ = _get_cuda_modules() + (err, ) = cuda.cuStreamWriteValue64( + stream.cuda_stream, + signal_ptr, + signal, + cuda.CUstreamWriteValue_flags.CU_STREAM_WRITE_VALUE_DEFAULT, + ) + CUDA_CHECK(err) + + +def print_perf( + name: str, + value: float, + group, + rank: int, + world_size: int, + unit: str = "ms", +): + for index in range(world_size): + torch.distributed.barrier(group=group) + if rank == index: + print(f"{name} #{rank}: {value:.3f} {unit}", flush=True) + + +def print_perf_mean( + name: str, + value: float, + group, + rank: int, + world_size: int, + unit: str = "ms", +): + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + value_tensor = torch.tensor(float(value), device=device, dtype=torch.float64) + torch.distributed.all_reduce( + value_tensor, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + mean_value = (value_tensor / world_size).item() + if rank == 0: + print(f"{name} mean: {mean_value:.3f} {unit}", flush=True) diff --git a/python/triton/experimental/tle/raw/runtime.py b/python/triton/experimental/tle/raw/runtime.py index a41c2be38..e1f5cc7d9 100644 --- a/python/triton/experimental/tle/raw/runtime.py +++ b/python/triton/experimental/tle/raw/runtime.py @@ -1,27 +1,40 @@ -# Copyright 2026 FlagOS Contributors -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. +from __future__ import annotations + +from typing import Any from .cache_key import bind_tle_raw_source_cache_key -from .cuda import CUDAJITFunction -registry = {"cuda": CUDAJITFunction} + +class RawJITFunction: + """Shared @dialect state and default LLVM region materialization.""" + + def __init__(self, fn: Any, **kwargs) -> None: + self.fn = fn + self.extern_func_name = kwargs.get("extern_func_name", "") + self.deferred = kwargs.get("deferred", False) + self.library = kwargs.get("library", "") or "" + self.__triton_builtin__ = True + + def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", + extern_func_name: str = ""): + return builder.create_tle_raw_region_by_llvm_func( + llvm, + self.region_dialect, + self.arg_dialect, + handles, + alias_indices, + hint, + extern_func_name, + ) + + +registry = {} + +try: + from .cuda import CUDAJITFunction + registry["cuda"] = CUDAJITFunction +except ImportError: + pass try: from .mlir import MLIRJITFunction @@ -45,9 +58,13 @@ def dialect( ): def decorator(fn): - if name == "mlir" and name not in registry: - from .mlir import MLIRJITFunction - registry[name] = MLIRJITFunction + if name not in registry: + if name == "cuda": + from .cuda import CUDAJITFunction + registry[name] = CUDAJITFunction + elif name == "mlir": + from .mlir import MLIRJITFunction + registry[name] = MLIRJITFunction edsl = registry[name](fn, **kwargs) bind_tle_raw_source_cache_key(edsl, name=name, **kwargs) return edsl diff --git a/python/triton/experimental/tle/raw/tops/mlir_runtime.py b/python/triton/experimental/tle/raw/tops/mlir_runtime.py index 86b63eba3..50c04a605 100644 --- a/python/triton/experimental/tle/raw/tops/mlir_runtime.py +++ b/python/triton/experimental/tle/raw/tops/mlir_runtime.py @@ -33,11 +33,12 @@ from mlir.passmanager import PassManager from ..mlir.codegen import MLIRCodeGenerator +from triton.experimental.tle.raw.runtime import RawJITFunction _GCU_COMPILER_OPT = "/opt/triton_gcu/bin/gcu-compiler-opt" -class TOPSMLIRJITFunction(object): +class TOPSMLIRJITFunction(RawJITFunction): """TLE-Raw dialect for TOPS MLIR EDSL: writes MLIR that lowers to GCU-compatible LLVM IR. Usage: @@ -52,8 +53,7 @@ def edsl(x: Input[L["!llvm.ptr<1>"]], ...): def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optional[ir.Context] = None, arch: str = "gcu400", use_gcu_opt: bool = True, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.fn: Final[Any] = fn + super().__init__(fn, **kwargs) self.arch: Final[str] = arch self.use_gcu_opt: Final[bool] = use_gcu_opt self.region_dialect: Final[str] = "tops" @@ -73,7 +73,6 @@ def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optio ctx = ir.Context() if context is None else context ctx.allow_unregistered_dialects = True self.context: Final[ir.Context] = ctx - self.__triton_builtin__: Final[bool] = True def __deepcopy__(self, memo: Dict[int, Any]) -> TOPSMLIRJITFunction: return self.__class__( @@ -365,15 +364,9 @@ def _unwrap_gpu_module(text: str) -> str: body = "\n".join(lines[start:end]) return f"module {{\n{body}\n}}\n" - 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 make_llvm(self, context=None) -> str: if self.use_gcu_opt and os.path.isfile(_GCU_COMPILER_OPT): diff --git a/python/triton/experimental/tle/raw/tops/runtime.py b/python/triton/experimental/tle/raw/tops/runtime.py index 20e81fa71..4f0c2b2bc 100644 --- a/python/triton/experimental/tle/raw/tops/runtime.py +++ b/python/triton/experimental/tle/raw/tops/runtime.py @@ -28,6 +28,7 @@ from triton._C.libtriton import llvm from triton._C.libtriton.tle.llvm import parse_llvm_ir from triton.backends.enflame.gcu_intrinsics import rewrite_intrinsics_to_placeholders +from triton.experimental.tle.raw.runtime import RawJITFunction def _find_tops_include_dir() -> str: @@ -61,7 +62,7 @@ def _get_gcu_arch() -> str: return os.getenv("GCU_ARCH", "gcu400") -class TOPSJITFunction(object): +class TOPSJITFunction(RawJITFunction): """TLE-Raw dialect for TOPS C++ (.tops) files compiled via topscc. Usage: @@ -72,13 +73,11 @@ def edsl(*args, **kwargs): def __init__(self, fn: Any, file: Optional[Path] = None, arch: Optional[str] = None, extra_flags: Optional[List[str]] = None, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.fn: Final[Any] = fn + super().__init__(fn, **kwargs) self.arch: Final[str] = arch or _get_gcu_arch() self.extra_flags: Final[List[str]] = extra_flags or [] self.region_dialect: Final[str] = "tops" self.arg_dialect: Final[str] = "llvm" - self.__triton_builtin__: Final[bool] = True if file is not None: self.code: Final[str] = Path(file).read_text() @@ -193,15 +192,9 @@ def _rewrite_gcu_intrinsics(llvm_ir: str) -> str: print("// ---- end ----") return result - 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 make_llvm(self, mlir_context) -> str: llvm_ir_text = self._compile_tops_to_llvm_ir() diff --git a/python/triton/knobs.py b/python/triton/knobs.py index 82bb529cd..02baa5c6e 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -588,6 +588,11 @@ class nvidia_knobs(base_knobs): libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH") libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH") + # flagtree tle-raw cuda nvshmem + nvshmem_home: env_opt_str = env_opt_str("NVSHMEM_HOME") + tle_raw_clang: env_opt_str = env_opt_str("CLANG") + tle_raw_clang_flags: env_opt_str = env_opt_str("CLANG_FLAGS") + class amd_knobs(base_knobs): use_buffer_ops: env_bool = env_bool("AMDGCN_USE_BUFFER_OPS", True) diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu new file mode 100644 index 000000000..f2e4d780c --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu @@ -0,0 +1,11 @@ +extern "C" __device__ int nvshmem_my_pe(); +extern "C" __device__ int nvshmem_n_pes(); +extern "C" __device__ void nvshmem_int_p(int *dest, int value, int pe); + +extern "C" __device__ __attribute__((always_inline)) void +simple_shift(__attribute__((address_space(1))) int *destination) { + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + nvshmem_int_p(destination, mype, peer); +} diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu new file mode 100644 index 000000000..da769f89c --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu @@ -0,0 +1,44 @@ +#include +#include +#include +#include + +#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" void simple_shift_before_launch(int *mype, int *npes, + int *mype_in_node, int *npes_in_node, + cudaStream_t *stream, int **dst, + int **data_h) { + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *npes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + CUDA_CHECK(cudaStreamCreate(stream)); + + *dst = (int *)nvshmem_malloc(sizeof(int)); + *data_h = (int *)malloc(sizeof(int)); +} + +extern "C" void simple_shift_after_launch(cudaStream_t stream, void *dst, + void *data_h, int mype, int npes) { + nvshmemx_barrier_all_on_stream(stream); + cudaMemcpyAsync(data_h, dst, sizeof(int), cudaMemcpyDeviceToHost, stream); + cudaStreamSynchronize(stream); + + int *data = (int *)data_h; + printf("%d: received message %d\n", mype, data[0]); + + nvshmem_free(dst); + free(data_h); +} diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py new file mode 100644 index 000000000..925b02807 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py @@ -0,0 +1,86 @@ +import os +import ctypes +import torch +import triton +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect + +from triton.experimental.tle.raw.nvshmem.utils import ( + load_common_host, + load_host, + init_torch_distributed, + init_nvshmem_by_torch_pg, + tensor_from_pointer, +) + + +@dialect( + name="cuda", + library="nvshmem", + compiler="clang", + file=(Path(__file__).parent / "simple-shift-device.cu"), + extern_func_name="simple_shift", +) +def simple_shift(*args, **kwargs): + ... + + +@triton.jit +def simple_shift_kernel(destination_ptr, ): + tle_raw.call(simple_shift, [destination_ptr]) + + +def simpe_shift(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device("cuda", local_rank) + + group = init_torch_distributed() + host_source = Path(__file__).with_name("simple-shift-host.cu") + host_lib = load_host(source=host_source) + common = load_common_host() + init_nvshmem_by_torch_pg(common, group) + + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + stream_ptr = ctypes.c_void_p() + destination_ptr = ctypes.c_void_p() + host_data_ptr = ctypes.c_void_p() + host_lib.simple_shift_before_launch( + ctypes.byref(mype), + ctypes.byref(npes), + ctypes.byref(mype_in_node), + ctypes.byref(npes_in_node), + ctypes.byref(stream_ptr), + ctypes.byref(destination_ptr), + ctypes.byref(host_data_ptr), + ) + + destination = tensor_from_pointer( + destination_ptr, + (1, ), + torch.int32, + device, + ) + stream = torch.cuda.ExternalStream(stream_ptr.value, device=device) + + with torch.cuda.stream(stream): + simple_shift_kernel[(1, )](destination) + + host_lib.simple_shift_after_launch( + stream_ptr, + destination_ptr, + host_data_ptr, + mype_in_node.value, + npes_in_node.value, + ) + + common.nvshmem_finalize_from_torch_distributed() + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + simpe_shift() diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu new file mode 100644 index 000000000..15eafbb33 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu @@ -0,0 +1,31 @@ +#include +#include +#include + +extern "C" __device__ uint64_t nvshmem_signal_wait_until(uint64_t *sig_addr, + int cmp, + uint64_t cmp_value); + +enum { + NVSHMEM_CMP_GE = 5, + NVSHMEM_SIGNAL_SET = 9, +}; + +extern "C" __device__ __attribute__((always_inline)) void +ag_mark_local_ready(__attribute__((address_space(1))) uint64_t *ready, + int rank) { + if (threadIdx.x == 0) { + __threadfence_system(); + ready[(size_t)rank] = 1; + } + __syncthreads(); +} + +extern "C" __device__ __attribute__((always_inline)) void +ag_wait_ready(__attribute__((address_space(1))) uint64_t *ready, + int signal_index) { + if (threadIdx.x == 0) { + nvshmem_signal_wait_until(ready + signal_index, NVSHMEM_CMP_GE, 1); + } + __syncthreads(); +} diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu new file mode 100644 index 000000000..72d0898f9 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu @@ -0,0 +1,63 @@ +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(stmt) \ + do { \ + cudaError_t result = (stmt); \ + if (result != cudaSuccess) { \ + fprintf(stderr, "[%s:%d] CUDA failed: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(result)); \ + return -1; \ + } \ + } while (0) + +extern "C" int ag_gemm_workspace_create(int elements_per_rank, void **workspace, + uint64_t **ready, int *mype, int *npes, + int *mype_in_node, int *npes_in_node) { + if (elements_per_rank <= 0 || workspace == nullptr || ready == nullptr) { + return -1; + } + + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *npes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + + size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * sizeof(__half); + *workspace = nvshmem_malloc(workspace_bytes); + *ready = (uint64_t *)nvshmem_calloc((size_t)(*npes), sizeof(uint64_t)); + if (*workspace == nullptr || *ready == nullptr) { + if (*ready != nullptr) { + nvshmem_free(*ready); + } + if (*workspace != nullptr) { + nvshmem_free(*workspace); + } + return -2; + } + + nvshmem_barrier_all(); + return 0; +} + +extern "C" int ag_gemm_workspace_destroy(void *workspace, void *ready) { + CUDA_CHECK(cudaDeviceSynchronize()); + nvshmem_barrier_all(); + nvshmem_free(ready); + nvshmem_free(workspace); + return 0; +} + +extern "C" void *ag_gemm_peer_workspace_ptr(void *workspace, int peer) { + return nvshmem_ptr(workspace, peer); +} + +extern "C" uint64_t *ag_gemm_peer_ready_ptr(uint64_t *ready, int peer) { + return (uint64_t *)nvshmem_ptr(ready, peer); +} diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py new file mode 100644 index 000000000..88a792b21 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py @@ -0,0 +1,556 @@ +import argparse +import ctypes +import os +from pathlib import Path +from typing import Optional + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language.raw as tle_raw +from triton.experimental.tle.raw import dialect + +from triton.experimental.tle.raw.nvshmem.utils import ( + print_perf_mean, + load_common_host, + load_host, + init_torch_distributed, + init_nvshmem_by_torch_pg, + tensor_from_pointer, + set_signal_cuda_ptr, + cudart, +) + + +def _device_dialect(function_name): + return dialect( + name="cuda", + library="nvshmem", + compiler="clang", + file=Path(__file__).parent / "ag-gemm-device.cu", + extern_func_name=function_name, + ) + + +@_device_dialect("ag_mark_local_ready") +def mark_local_ready(*args, **kwargs): + ... + + +@_device_dialect("ag_wait_ready") +def wait_ready(*args, **kwargs): + ... + + +@triton.jit(do_not_specialize=["rank"]) +def set_local_ready(ready, rank): + tle_raw.call(mark_local_ready, [ready, rank]) + + +def cp_engine_producer_all_gather_put(host, local_tensor, ag_buffer, signal_buffer, M_per_rank, N, signal_target, rank, + local_world_size, world_size, intranode_ag_stream): + local_rank = rank % local_world_size + nbytes = M_per_rank * N * local_tensor.element_size() + for i in range(1, local_world_size): + segment = rank * M_per_rank * N + local_dst_rank = (local_rank + local_world_size - i) % local_world_size + peer_workspace = host.ag_gemm_peer_workspace_ptr( + ctypes.c_void_p(ag_buffer.data_ptr()), + local_dst_rank, + ) + peer_ready = host.ag_gemm_peer_ready_ptr( + ctypes.c_void_p(signal_buffer.data_ptr()), + local_dst_rank, + ) + assert peer_workspace, f"failed to get peer workspace pointer for PE {local_dst_rank}" + assert peer_ready, f"failed to get peer ready pointer for PE {local_dst_rank}" + + src_ptr = ag_buffer.data_ptr() + segment * local_tensor.element_size() + dst_ptr = peer_workspace + segment * local_tensor.element_size() + (err, ) = cudart.cudaMemcpyAsync( + dst_ptr, + src_ptr, + nbytes, + cudart.cudaMemcpyKind.cudaMemcpyDefault, + intranode_ag_stream.cuda_stream, + ) + set_signal_cuda_ptr( + peer_ready + rank * ctypes.sizeof(ctypes.c_uint64), + signal_target, + intranode_ag_stream, + ) + + +@triton.jit +def ag_gemm_consumer( + a_ptr, + b_ptr, + c_ptr, + ready, + M, + N, + K, + RANK: tl.constexpr, + NUM_RANKS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + LOCAL_WORLD_SIZE: tl.constexpr, + NUM_SMS: tl.constexpr, + EPILOGUE_SUBTILE: tl.constexpr, +): + dtype = c_ptr.dtype.element_ty + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + node_id = RANK // LOCAL_WORLD_SIZE + nnodes = NUM_RANKS // LOCAL_WORLD_SIZE + + a_desc = tl.make_tensor_descriptor( + a_ptr, + shape=[M, K], + strides=[K, 1], + block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], + ) + b_desc = tl.make_tensor_descriptor( + b_ptr, + shape=[N, K], + strides=[K, 1], + block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K], + ) + c_desc = tl.make_tensor_descriptor( + c_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[ + BLOCK_SIZE_M, + BLOCK_SIZE_N if not EPILOGUE_SUBTILE else BLOCK_SIZE_N // 2, + ], + ) + + tiles_per_SM = num_tiles // NUM_SMS + if start_pid < num_tiles % NUM_SMS: + tiles_per_SM += 1 + + tile_id = start_pid - NUM_SMS + ki = -1 + + pid_m = 0 + pid_n = 0 + offs_am = 0 + offs_bn = 0 + + M_per_rank = M // NUM_RANKS + pid_ms_per_rank = tl.cdiv(M_per_rank, BLOCK_SIZE_M) + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for _ in range(0, k_tiles * tiles_per_SM): + ki = tl.where(ki == k_tiles - 1, 0, ki + 1) + if ki == 0: + tile_id += NUM_SMS + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + + # swizzle m + if nnodes == 1: + alpha = 0 + beta = 0 + pid_m = (pid_m + ((((RANK ^ alpha) + beta) % NUM_RANKS) * pid_ms_per_rank)) % num_pid_m + else: + m_rank = pid_m // pid_ms_per_rank + pid_m_intra_rank = pid_m - m_rank * pid_ms_per_rank + m_node_id = m_rank // LOCAL_WORLD_SIZE + m_local_rank = m_rank % LOCAL_WORLD_SIZE + swizzle_m_node_id = (m_node_id + node_id) % nnodes + swizzle_m_local_rank = (m_local_rank + RANK) % LOCAL_WORLD_SIZE + swizzle_m_rank = swizzle_m_node_id * LOCAL_WORLD_SIZE + swizzle_m_local_rank + + pid_m = swizzle_m_rank * pid_ms_per_rank + pid_m_intra_rank + + offs_am = pid_m * BLOCK_SIZE_M + offs_bn = pid_n * BLOCK_SIZE_N + + source_rank = offs_am // M_per_rank + tle_raw.call(wait_ready, [ready, source_rank]) + + offs_k = ki * BLOCK_SIZE_K + a = a_desc.load([offs_am, offs_k]) + b = b_desc.load([offs_bn, offs_k]) + accumulator = tl.dot(a, b.T, accumulator) + + if ki == k_tiles - 1: + if EPILOGUE_SUBTILE: + acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2)) + acc = tl.permute(acc, (0, 2, 1)) + acc0, acc1 = tl.split(acc) + c0 = acc0.to(dtype) + c_desc.store([offs_am, offs_bn], c0) + c1 = acc1.to(dtype) + c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c1) + else: + c = accumulator.to(dtype) + c_desc.store([offs_am, offs_bn], c) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + +def ag_gemm_op( + a, + b, + c, + rank, + num_ranks, + workspace, + ready, + comm_stream, + compute_stream, + host, + local_world_size, + block_m=128, + block_n=256, + block_k=64, + group_m=8, + num_warps=8, + num_stages=3, +): + assert a.shape[1] == b.shape[1], "incompatible GEMM dimensions" + assert a.dtype == b.dtype == c.dtype, "incompatible GEMM dtypes" + + m, k = workspace.shape + m_per_rank = m // num_ranks + n_per_rank = b.shape[0] + total_tiles = triton.cdiv(m, block_m) * triton.cdiv(n_per_rank, block_n) + + num_ag_sms = 0 # The copy engine does not occupy SM resources. + num_gemm_sms = torch.cuda.get_device_properties("cuda").multi_processor_count - num_ag_sms + + grid = (min(total_tiles, num_gemm_sms), ) + + current_stream = torch.cuda.current_stream(b.device) + comm_stream.wait_stream(current_stream) + cp_engine_producer_all_gather_put(host=host, local_tensor=a, ag_buffer=workspace, signal_buffer=ready, + M_per_rank=m_per_rank, N=k, signal_target=1, rank=rank, + local_world_size=local_world_size, world_size=num_ranks, + intranode_ag_stream=comm_stream) + + def alloc_fn(size: int, alignment: int, stream: Optional[int]): + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(alloc_fn) + + compiled = None + with torch.cuda.stream(compute_stream): + compute_stream.wait_stream(current_stream) + compiled = ag_gemm_consumer[grid]( + workspace, + b, + c, + ready, + m, + n_per_rank, + k, + rank, + num_ranks, + block_m, + block_n, + block_k, + group_m, + local_world_size, + NUM_SMS=num_gemm_sms, + EPILOGUE_SUBTILE=False, + num_warps=num_warps, + num_stages=num_stages, + ) + + current_stream.wait_stream(comm_stream) + current_stream.wait_stream(compute_stream) + + return compiled + + +def triton_prepare(a_local, workspace, ready, rank): + ready.zero_() + m_per_rank = a_local.shape[0] + workspace[rank * m_per_rank:(rank + 1) * m_per_rank].copy_(a_local) + set_local_ready[(1, )](ready, rank, num_warps=1) + + +def torch_ag_gemm(group, a_local, b, gathered): + torch.distributed.all_gather_into_tensor(gathered, a_local, group=group) + return torch.matmul(gathered, b.T) + + +def configure_host_library(host): + host.ag_gemm_workspace_create.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ] + host.ag_gemm_workspace_create.restype = ctypes.c_int + host.ag_gemm_workspace_destroy.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + host.ag_gemm_workspace_destroy.restype = ctypes.c_int + + # The return value is a 64-bit address. + # Without declaring it as c_void_p, it will be truncated to 32 bits. + host.ag_gemm_peer_workspace_ptr.argtypes = [ctypes.c_void_p, ctypes.c_int] + host.ag_gemm_peer_workspace_ptr.restype = ctypes.c_void_p + host.ag_gemm_peer_ready_ptr.argtypes = [ctypes.c_void_p, ctypes.c_int] + host.ag_gemm_peer_ready_ptr.restype = ctypes.c_void_p + + +def create_workspace(host, world_size, m_per_rank, k, device): + workspace_ptr = ctypes.c_void_p() + ready_ptr = ctypes.c_void_p() + mype = ctypes.c_int() + npes = ctypes.c_int() + local_pe = ctypes.c_int() + local_npes = ctypes.c_int() + result = host.ag_gemm_workspace_create( + m_per_rank * k, + ctypes.byref(workspace_ptr), + ctypes.byref(ready_ptr), + ctypes.byref(mype), + ctypes.byref(npes), + ctypes.byref(local_pe), + ctypes.byref(local_npes), + ) + assert result == 0, f"workspace allocation failed: {result}" + assert npes.value == world_size + workspace = tensor_from_pointer( + workspace_ptr, + (world_size * m_per_rank, k), + torch.float16, + device, + ) + ready = tensor_from_pointer( + ready_ptr, + (world_size, ), + torch.uint64, + device, + ) + return workspace_ptr, ready_ptr, workspace, ready, mype, local_pe, local_npes + + +def perf(fn, group, warmup, iters, prepare_fn=None): + assert warmup >= 0 and iters > 0 + + def prepare(): + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + if prepare_fn is not None: + prepare_fn() + + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + for _ in range(warmup): + prepare() + fn() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + for _ in range(iters): + prepare() + fn() + end.record() + torch.cuda.synchronize() + + profile = {"total": start.elapsed_time(end) / iters} + return profile + + +def parse_args(): + parser = argparse.ArgumentParser(description="NVSHMEM allgather GEMM") + parser.add_argument("--m-per-rank", type=int, default=1024) + parser.add_argument("--n-per-rank", type=int, default=1024) + parser.add_argument("--k", type=int, default=1024) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--case", choices=("check", "perf"), default="check") + + return parser.parse_args() + + +def main(): + args = parse_args() + rank = int(os.environ.get("RANK", 0)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + + assert world_size >= 2, "AG-GEMM requires at least two GPUs" + assert world_size == local_world_size, ("This example is designed for single-node testing: " + "WORLD_SIZE must equal LOCAL_WORLD_SIZE") + + if torch.cuda.get_device_capability()[0] < 9: + print("Skip the test because the device is not sm90 or higher") + import sys + sys.exit() + + group = init_torch_distributed() + host_source = Path(__file__).with_name("ag-gemm-host.cu") + host = load_host(host_source) + common = load_common_host() + configure_host_library(host) + init_nvshmem_by_torch_pg(common, group) + + dtype = torch.float16 + a_local = torch.randn((args.m_per_rank, args.k), dtype=dtype, device=device) + b = torch.randn((args.n_per_rank, args.k), dtype=dtype, device=device) + c = torch.empty( + (world_size * args.m_per_rank, args.n_per_rank), + dtype=a_local.dtype, + device=device, + ) + gathered = torch.empty( + (world_size * args.m_per_rank, args.k), + dtype=a_local.dtype, + device=device, + ) + + ( + workspace_ptr, + ready_ptr, + workspace, + ready, + mype, + local_pe, + local_npes, + ) = create_workspace( + host, + world_size, + args.m_per_rank, + args.k, + device, + ) + assert mype.value == rank + assert local_pe.value == local_rank + + comm_stream = torch.cuda.Stream(device=device, priority=-1) + compute_stream = torch.cuda.Stream(device=device, priority=-1) + + def triton_prapare_func(): + triton_prepare(a_local, workspace, ready, rank) + + def triton_func(): + return ag_gemm_op(a_local, b, c, rank, world_size, workspace, ready, comm_stream, compute_stream, host, + local_world_size) + + def torch_func(): + return torch_ag_gemm(group, a_local, b, gathered) + + def run_correctness(): + if rank == 0: + print("[check] start correctness validation", flush=True) + + triton_prapare_func() + triton_func() + result = host.ag_gemm_workspace_destroy(workspace_ptr, ready_ptr) + assert result == 0 + result = common.nvshmem_finalize_from_torch_distributed() + assert result == 0 + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + golden = torch_func() + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + torch.distributed.destroy_process_group() + + torch.testing.assert_close(c, golden, atol=1e-3, rtol=1e-3) + if rank == 0: + print("[check] Pass!", flush=True) + print(f"configuration: GPUs={world_size}, " + f"A_local=({args.m_per_rank}, {args.k}), " + f"B_local=({args.n_per_rank}, {args.k}), ") + + # For performance testing + # It's best to test them separately + # For example, when testing Triton, comment out the Torch part. + def run_benchmark(): + if rank == 0: + print( + f"[bench] start benchmark: warmup={args.warmup}, " + f"iters={args.iters}", + flush=True, + ) + + triton_profile = perf( + triton_func, + group, + args.warmup, + args.iters, + prepare_fn=triton_prapare_func, + ) + result = host.ag_gemm_workspace_destroy(workspace_ptr, ready_ptr) + assert result == 0 + result = common.nvshmem_finalize_from_torch_distributed() + assert result == 0 + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + torch_profile = perf( + torch_func, + group, + args.warmup, + args.iters, + ) + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + print_perf_mean( + "triton ag-gemm", + triton_profile["total"], + group, + rank, + world_size, + ) + print_perf_mean( + "torch ag-gemm", + torch_profile["total"], + group, + rank, + world_size, + ) + print_perf_mean( + "speedup", + torch_profile["total"] / triton_profile["total"], + group, + rank, + world_size, + unit="x", + ) + + if rank == 0: + print(f"configuration: GPUs={world_size}, " + f"A_local=({args.m_per_rank}, {args.k}), " + f"B_local=({args.n_per_rank}, {args.k}).") + + torch.distributed.destroy_process_group() + + if args.case == "check": + run_correctness() + if args.case == "perf": + run_benchmark() + + +if __name__ == "__main__": + main() diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index 8fee79a6e..fa27dcc23 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -139,6 +139,19 @@ def __post_init__(self): if not extern_libs.get('libdevice', None): extern_libs['libdevice'] = knobs.nvidia.libdevice_path or str(default_libdir / 'libdevice.10.bc') + # flagtree tle raw: libnvshmem_device when @dialect(library="nvshmem") enabled in utils. + try: + from triton.experimental.tle.raw.nvshmem.utils import ( + is_nvshmem_device_bc_enabled, + resolve_nvshmem_device_bitcode, + ) + if is_nvshmem_device_bc_enabled(): + nvshmem_bc = resolve_nvshmem_device_bitcode(arch=self.arch) + if nvshmem_bc is not None: + extern_libs.update({"libnvshmem_device": str(nvshmem_bc)}) + except Exception: + pass + # flagtree tle distributed: Add distributed bitcode library(libflagcx_device.bc) if distributed features are enabled. extern_libs.update(Distributed().get_extern_libs()) object.__setattr__(self, 'extern_libs', tuple(extern_libs.items())) @@ -428,7 +441,7 @@ def make_llir(self, src, metadata, options, capability): passes.ttgpuir.add_concurrency_sanitizer(pm) passes.ttgpuir.add_allocate_global_scratch_memory(pm) nvidia.passes.ttnvgpuir.add_proxy_fence_insertion(pm, capability) - # Materialize deferred tle_raw sources before inlining DSL regions. + # flagtree tle raw: Materialize deferred tle_raw sources before inlining DSL regions. from .deferred_raw import ( finish_deferred_raw_materialize, deferred_raw_materialize, diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 7709e90b2..3aefb71ea 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -62,13 +62,16 @@ namespace tle = triton::tle; extern std::vector computeAliasOperandIndices(TritonOpBuilder &self, std::string_view text, - const std::vector &args); + const std::vector &args, + std::string_view funcName); -extern tle::DSLRegionOp createTLERawRegionByLLVMFunc( - TritonOpBuilder &self, std::string_view text, - std::string_view regionDialect, std::string_view argDialect, - const std::vector &args, - const std::vector &aliasOperandIndices, std::string_view hint); +extern tle::DSLRegionOp +createTLERawRegionByLLVMFunc(TritonOpBuilder &self, std::string_view text, + std::string_view regionDialect, + std::string_view argDialect, + const std::vector &args, + const std::vector &aliasOperandIndices, + std::string_view hint, std::string_view funcName); extern tle::DSLRegionOp createTLERawRegionDeferred( TritonOpBuilder &self, std::string_view sourceId, @@ -697,12 +700,13 @@ void init_tle_raw_ir(py::module &&m) { .def("dump", &tle::YieldOp::dump); auto *builder_cls = ir::getBuilderClass(); - builder_cls->def("compute_alias_operand_indices", - &computeAliasOperandIndices); - builder_cls->def( - "create_tle_raw_region_by_llvm_func", &createTLERawRegionByLLVMFunc, - py::arg("text"), py::arg("region_dialect"), py::arg("arg_dialect"), - py::arg("args"), py::arg("output_operand_indices"), py::arg("hint") = ""); + builder_cls->def("compute_alias_operand_indices", &computeAliasOperandIndices, + py::arg("text"), py::arg("args"), py::arg("func_name") = ""); + builder_cls->def("create_tle_raw_region_by_llvm_func", + &createTLERawRegionByLLVMFunc, py::arg("text"), + py::arg("region_dialect"), py::arg("arg_dialect"), + py::arg("args"), py::arg("output_operand_indices"), + py::arg("hint") = "", py::arg("func_name") = ""); builder_cls->def( "create_tle_raw_region_deferred", &createTLERawRegionDeferred, py::arg("source_id"), py::arg("region_dialect"), py::arg("arg_dialect"), diff --git a/third_party/tle/triton_tle_raw.cc b/third_party/tle/triton_tle_raw.cc index f790e9ab3..6409b3893 100644 --- a/third_party/tle/triton_tle_raw.cc +++ b/third_party/tle/triton_tle_raw.cc @@ -7,6 +7,7 @@ #include "tle/utils/include/AnalyzeReturnType.h" #include "tle/utils/include/TleRawMaterialize.h" #include "llvm/ADT/STLExtras.h" +#include using namespace mlir; namespace tle = triton::tle; @@ -18,6 +19,12 @@ StringAttr getOptionalStringAttr(OpBuilder &builder, std::string_view value) { return builder.getStringAttr(value); } +std::optional getOptionalFuncName(std::string_view value) { + if (value.empty()) + return std::nullopt; + return llvm::StringRef(value.data(), value.size()); +} + void setDeferredMetadataAttrs(tle::DSLRegionOp op, OpBuilder &builder, std::string_view sourceId) { if (!sourceId.empty()) @@ -37,13 +44,15 @@ tle::DSLRegionOp createDSLRegionOp( } } // namespace -std::vector -computeAliasOperandIndices(TritonOpBuilder &self, std::string_view text, - const std::vector &args) { +std::vector computeAliasOperandIndices(TritonOpBuilder &self, + std::string_view text, + const std::vector &args, + std::string_view funcName) { OwningOpRef module = tle::raw::parseLLVMModule(self.getContext(), text); assert(module && "Failed to parse LLVM IR text"); - LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc(module.get(), {}); + LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc( + module.get(), getOptionalFuncName(funcName)); assert(func && "No function found in LLVM IR text"); SmallVector funcArgToDslArg = @@ -61,15 +70,18 @@ computeAliasOperandIndices(TritonOpBuilder &self, std::string_view text, return std::vector(result.begin(), result.end()); } -tle::DSLRegionOp createTLERawRegionByLLVMFunc( - TritonOpBuilder &self, std::string_view text, - std::string_view regionDialect, std::string_view argDialect, - const std::vector &args, - const std::vector &aliasOperandIndices, std::string_view hint) { +tle::DSLRegionOp +createTLERawRegionByLLVMFunc(TritonOpBuilder &self, std::string_view text, + std::string_view regionDialect, + std::string_view argDialect, + const std::vector &args, + const std::vector &aliasOperandIndices, + std::string_view hint, std::string_view funcName) { OwningOpRef module = tle::raw::parseLLVMModule(self.getContext(), text); assert(module && "Failed to parse LLVM IR text"); - LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc(module.get(), {}); + LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc( + module.get(), getOptionalFuncName(funcName)); assert(func && "No function found in LLVM IR text"); OpBuilder &builder = self.getBuilder(); @@ -79,8 +91,8 @@ tle::DSLRegionOp createTLERawRegionByLLVMFunc( } ModuleOp curModule = cast(curOp); - auto funcOpOrErr = - tle::raw::cloneLLVMSymbolsAndLookupFunc(curModule, module.get(), {}); + auto funcOpOrErr = tle::raw::cloneLLVMSymbolsAndLookupFunc( + curModule, module.get(), getOptionalFuncName(funcName)); assert(succeeded(funcOpOrErr)); LLVM::LLVMFuncOp funcOp = *funcOpOrErr;