diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/AGENTS.md b/dnn-providers/hip-kernel-provider/rocke/platform/AGENTS.md index 0dcf13ee6b88..fc63ac581479 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/AGENTS.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/AGENTS.md @@ -114,7 +114,7 @@ circular. torch is therefore **optional** and the only hard dep is `numpy` 1. the `torch.fx` fusion frontend (`torch_backend.py`, `helpers/fuse.py` graph capture); -2. torch-tensor launch integration (`runtime/torch_module.py`); +2. torch-tensor launch integration (`runtime/torch_interop.py`); 3. on-GPU numeric verification against torch eager. Everything else runs torch-free: IR build, lower, `comgr` compile, numpy launch, @@ -125,8 +125,8 @@ with neither torch nor a GPU installed. ### ROCm library discovery (libamd_comgr / libamdhip64) The runtime resolves the ROCm shared libs WITHOUT importing torch -(`runtime/hip_module._candidate_lib_paths` / `_rocm_root_libdirs`), in priority -order: +(`runtime/runtime_coexistence._candidate_lib_paths` / `_rocm_root_libdirs`), in +priority order: 1. explicit full-path override env var (`ROCKE_COMGR_LIB`, `ROCKE_HIP_LIB`); 2. torch-bundled `/lib/lib*.so` — opportunistic fast-path **only if torch diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/BUILD.md b/dnn-providers/hip-kernel-provider/rocke/platform/BUILD.md index 51ed4a1dc416..b5974f3bedc0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/BUILD.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/BUILD.md @@ -74,8 +74,8 @@ Optional sanitizer build for diagnostics (not for shipping): `-DROCKE_SANITIZE=O In-process compile + launch needs the ROCm shared libs at runtime. The Python runtime resolves them WITHOUT importing torch -(`Python/rocke/runtime/hip_module._candidate_lib_paths` / `_rocm_root_libdirs`), -in priority order: +(`Python/rocke/runtime/runtime_coexistence._candidate_lib_paths` / +`_rocm_root_libdirs`), in priority order: 1. explicit full-path override env var: `ROCKE_COMGR_LIB`, `ROCKE_HIP_LIB`; 2. torch-bundled `/lib/lib*.so` — only if torch is already imported (the @@ -87,7 +87,7 @@ in priority order: 4. bare `lib.so` on the dynamic linker's search path (last resort). torch is therefore **optional** — required only for the `torch.fx` fusion -frontend, torch-tensor launch (`runtime/torch_module.py`), and on-GPU torch-eager +frontend, torch-tensor launch (`runtime/torch_interop.py`), and on-GPU torch-eager numeric checks. Building the engine, lowering, `comgr` compile, numpy launch, and the byte-identity gate need no torch. If a torch-less process reports `cannot load libamd_comgr.so`, set `ROCM_PATH` (or `ROCKE_COMGR_LIB`) to your ROCm install. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/README.md index fd2abd242f5c..63aee0b1fbd2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/README.md @@ -157,7 +157,7 @@ Conventional anchors: - COMGR: `python/rocke/runtime/comgr.py`. - HIP runtime: `python/rocke/runtime/hip_module.py`. - Launcher / workspace / timing: `python/rocke/runtime/launcher.py`. -- Torch arg packing: `python/rocke/runtime/torch_module.py`. +- Kernel arg packing: `python/rocke/runtime/packing.py`. - High-level compile: `python/rocke/helpers/compile.py`. - Manifest schema: `python/rocke/helpers/manifest.py`. - Optimization runbook: `gpu-op-optimization-runbook` Cursor skill. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/mental_model.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/mental_model.md index a5f4085ee258..ca04ec73982c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/mental_model.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/mental_model.md @@ -39,7 +39,7 @@ core/lower_llvm.py runtime/comgr.py ctypes over libamd_comgr; LLVM IR -> bitcode -> relocatable -> HSACO. -runtime/hip_module.py, runtime/launcher.py, runtime/torch_module.py +runtime/hip_module.py, runtime/launcher.py, runtime/torch_interop.py Load HSACO, pack args, resolve streams, launch, time, manage workspace and per-stream pending-args lifetime. ``` diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 642dc8d8f251..4619dcd9609d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1988,8 +1988,8 @@ moved measured throughput by under 1 % on direct-conv kernels. | `no_fence` context manager | `runtime/launcher.py` | Skip per-call sync inside an event-timed loop (graph-style) | | `time_launches(fn, warmup, iters, stream)` | `runtime/launcher.py` | The canonical HIP-event timer | | `StreamConfig` | `runtime/launcher.py` | Mirror of CK Tile `stream_config` | -| `resolve_stream(stream=0)` | `runtime/torch_module.py` | Substitute torch's current stream to keep allocator coherent | -| `pack_args` vs `pack_args_kernelparams` | `runtime/torch_module.py` | AMDGPU kernarg buffer vs the safer `kernelParams` path | +| `resolve_stream(stream=0)` | `runtime/torch_interop.py` | Substitute torch's current stream to keep allocator coherent | +| `pack_args` vs `pack_args_kernelparams` | `runtime/packing.py` | AMDGPU kernarg buffer vs the safer `kernelParams` path | | HIP graph capture | torch | Amortizes launch overhead — pair with `no_fence` and many iters | | `rocm-smi --setperflevel high && --setsclk 7` | shell | Lock clocks to avoid thermal / DVFS noise during measurement | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md index c097400faf36..a2af7ef0ddda 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/api_index.md @@ -203,10 +203,15 @@ wait_stream_and_release release_retained_for_stream ``` -## `from rocke.runtime.torch_module import ...` +## `from rocke.runtime.packing import ...` + +```text +pack_args, pack_args_kernelparams # torch-agnostic kernarg packing +``` + +## `from rocke.runtime.torch_interop import ...` ```text -pack_args, pack_args_kernelparams resolve_stream empty_workspace launch_torch_kernel # back-compat shim diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md index 3a0631d2bc49..e527c9f32b30 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md @@ -1,6 +1,6 @@ # File Index -A by-file map of the `rocke` package. Symbols listed are the primary public exports — see each file for the rest. +A by-file map of the `rocke` package. Symbols listed are the primary contents (public exports, or the key internal helpers for whole-module-private files) — see each file for the rest. ## Top Level @@ -32,9 +32,12 @@ A by-file map of the `rocke` package. Symbols listed are the primary public expo | Path | Primary contents | |-------------------------------|--------------------------------------------------------------------------------------------------------| +| `runtime/runtime_coexistence.py` | `_candidate_lib_paths`, `_rocm_root_libdirs`, `_torch_bundled_lib`, ... . Which ROCm runtime we bind to, and whose (torch-bundled vs system). Shared by `comgr` + `hip_module`. | +| `runtime/_ctypes_bind.py` | `_LazyFn`. Lazy ctypes function binder shared by `comgr` + `hip_module`. | | `runtime/comgr.py` | `build_hsaco_from_llvm_ir(ir_text, isa=..., options=...) -> (bytes, ComgrTimings)`. Ctypes over `libamd_comgr`. | | `runtime/hip_module.py` | `Runtime`, `Module`, `Event`, `HipError`. Ctypes over `libamdhip64`. Per-stream pending-args queue. | -| `runtime/torch_module.py` | `pack_args`, `pack_args_kernelparams`, `resolve_stream`, `empty_workspace`, `launch_torch_kernel`. | +| `runtime/packing.py` | `pack_args`, `pack_args_kernelparams`, `_as_ptr`. Torch-agnostic AMDGPU kernarg packing. | +| `runtime/torch_interop.py` | `resolve_stream`, `empty_workspace`, `launch_torch_kernel`, `TorchLaunchSummary`. Torch-tensor launch glue. | | `runtime/launcher.py` | `KernelLauncher`, `PipelineLauncher`, `LaunchConfig`, `LaunchSummary`, `WorkspaceSpec`, `WorkspacePool`, `DeviceMem`, `time_launches`, `no_fence`, `synchronize_and_release`, `wait_stream_and_release`, `release_retained_for_stream`. | ## `helpers/` — Authoring Layer diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/comgr_and_hipmodule.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/comgr_and_hipmodule.md index b69ed7112e2a..32a97a1f3a28 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/comgr_and_hipmodule.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/comgr_and_hipmodule.md @@ -166,7 +166,7 @@ extra[] = { } ``` -`args_bytes` is the packed arg buffer (built by `runtime/torch_module.py::pack_args` from the signature dict list). For an `(A, B, C, M, N, K)` GEMM: +`args_bytes` is the packed arg buffer (built by `runtime/packing.py::pack_args` from the signature dict list). For an `(A, B, C, M, N, K)` GEMM: ```text struct.pack(" bytes # packs kernel args in declaration order diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/limitations.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/limitations.md index 6142a52f5b67..bcefa3dbf46f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/limitations.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/limitations.md @@ -207,7 +207,7 @@ What does work for multi-GPU usage of the DSL: This pattern is verified end-to-end by `dsl_docs/development/verify_dsl_docs.py` on the GPU available to this box; the per-device launcher constructed under `torch.cuda.device(d)` correctly launches on that device. -- **Torch-stream-aware launches**: `runtime/torch_module.py::resolve_stream(stream, device=None)` honors `torch.cuda.current_stream(device).cuda_stream`, so the standard pattern of allocating tensors and launching all under `torch.cuda.device(d):` interoperates with torch's caching allocator on the right device. +- **Torch-stream-aware launches**: `runtime/torch_interop.py::resolve_stream(stream, device=None)` honors `torch.cuda.current_stream(device).cuda_stream`, so the standard pattern of allocating tensors and launching all under `torch.cuda.device(d):` interoperates with torch's caching allocator on the right device. - **Composing with `torch.distributed` (RCCL on ROCm)**: the user runs `torch.distributed.init_process_group(backend="nccl")` (PyTorch's "nccl" backend on ROCm is implemented by RCCL — `librccl.so`). DSL kernels then execute on each rank's local device; collectives (`torch.distributed.all_reduce`, `all_gather`, etc.) execute outside the DSL on a separate communication stream. The DSL never invokes RCCL itself. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/manifest_schema.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/manifest_schema.md index 4dbbc2ba5c0c..f5dad1fca6fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/manifest_schema.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/runtime/manifest_schema.md @@ -147,7 +147,7 @@ The runner uses `default_shape` to allocate `X`, `Gamma`/`Beta` (norms), and `Y` "f32" # 4-byte scalar ``` -`size_bytes` is for the host arg-packer (`runtime/torch_module.py::pack_args`) and the runner. Pointers are always 8 bytes; scalars match the canonical width. +`size_bytes` is for the host arg-packer (`runtime/packing.py::pack_args`) and the runner. Pointers are always 8 bytes; scalars match the canonical width. ## Standard Signatures diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py index e2c0a003f944..6c2585510610 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/lower_llvm.py @@ -130,7 +130,7 @@ def _torch_hip_version() -> Optional[Tuple[int, int]]: bundle their own ``libamd_comgr.so`` whose LLVM version follows the wheel's ROCm release, not the system ``/opt/rocm`` one. When rocke is paired with a torch-bundled comgr (see - :func:`runtime.hip_module._torch_bundled_lib`), the flavor must + :func:`runtime.runtime_coexistence._torch_bundled_lib`), the flavor must match torch's ROCm vintage or comgr will reject the IR or silently auto-upgrade declares the lowerer didn't intend. """ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/heuristics/_measure_moe.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/heuristics/_measure_moe.py index 74ddbcf7d8fa..4a00187493f7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/heuristics/_measure_moe.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/heuristics/_measure_moe.py @@ -11,7 +11,7 @@ proxy, since the streaming trio has no GEMM FLOP) plus ``is_valid``. No torch dependency: device buffers come from ``Runtime.alloc`` (hipMalloc), -args are packed with ``runtime.torch_module.pack_args`` (which accepts integer +args are packed with ``runtime.packing.pack_args`` (which accepts integer device pointers directly). """ @@ -29,7 +29,7 @@ from rocke.core.lower_llvm import lower_kernel_to_llvm from rocke.runtime.comgr import build_hsaco_from_llvm_ir from rocke.runtime.hip_module import Runtime -from rocke.runtime.torch_module import pack_args +from rocke.runtime.packing import pack_args _DTYPE_BYTES = {"f16": 2, "fp16": 2, "bf16": 2} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py index 746e2bbbee7d..c67817542127 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/__init__.py @@ -7,7 +7,17 @@ in-process pipeline that turns AMDGPU LLVM IR text into a running kernel. -Layered modules (bottom-up): +Layered modules (bottom-up). The first five run torch-free -- the +hip-only core; ``torch_interop`` and ``launcher`` are the torch-aware +edge: + + - ``runtime_coexistence`` : which ROCm runtime we bind to, and whose + (torch-bundled vs a system ROCm install). Shared + library resolution used by both ``comgr`` and + ``hip_module``. + + - ``_ctypes_bind``: the lazy ctypes function binder (`_LazyFn`) shared + by ``comgr`` and ``hip_module``. - ``comgr`` : ctypes wrapper over `libamd_comgr.so`. Implements `LLVM IR (text) -> BC -> relocatable ELF -> HSA @@ -23,10 +33,14 @@ `HIP_LAUNCH_PARAM_BUFFER_POINTER` arg-buffer lifetime race. - - ``torch_module``: torch-tensor arg packing + `resolve_stream` - (which collapses ``stream=0`` to torch's current - stream so the caching allocator sees our - launches). + - ``packing`` : torch-agnostic kernel-arg packing (`pack_args`, + `pack_args_kernelparams`) for the AMDGPU kernarg + ABI. Used by both the numpy and torch paths. + + - ``torch_interop``: torch-tensor launch glue -- `resolve_stream` + (collapses ``stream=0`` to torch's current stream + so the caching allocator sees our launches), + `empty_workspace`, `launch_torch_kernel`. - ``launcher`` : long-lived launch abstractions (CK Tile / FlyDSL / Triton inspired). The recommended entry point @@ -84,11 +98,11 @@ time_launches, wait_stream_and_release, ) -from .torch_module import ( +from .packing import pack_args +from .torch_interop import ( TorchLaunchSummary, empty_workspace, launch_torch_kernel, - pack_args, resolve_stream, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/_ctypes_bind.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/_ctypes_bind.py new file mode 100644 index 000000000000..18e595a88664 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/_ctypes_bind.py @@ -0,0 +1,58 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Lazy ctypes function binding shared by the HIP / comgr wrappers. + +A tiny ergonomics helper: defer ``getattr`` on the underlying ``CDLL`` +until the first call so the runtime wrappers (`hip_module`, `comgr`) can +be imported before their shared library is resolved -- which lets rocke +and torch be imported in any order and still bind the same loaded HIP / +comgr runtime instance (see :mod:`rocke.runtime.runtime_coexistence`). +""" + +from __future__ import annotations + +import ctypes +from typing import Any, Callable, List, Optional + + +class _LazyFn: + """Lazy ctypes function wrapper. + + Defers ``getattr`` on the underlying CDLL until the first call so + that rocke and torch can be imported in any order without ending + up with two HIP runtimes (see + ``runtime_coexistence._torch_bundled_lib``). Resolved on + first use; subsequent calls dispatch directly through the cached + function pointer. + + ``lib_resolver`` returns the shared ``ctypes.CDLL`` for this lib + family (HIP runtime / comgr / ...). It is invoked exactly once per + function on first call. + """ + + __slots__ = ("_name", "_argtypes", "_restype", "_lib_resolver", "_fn") + + def __init__( + self, + name: str, + argtypes: List[Any], + restype: Any, + lib_resolver: "Callable[[], ctypes.CDLL]", + ) -> None: + self._name = name + self._argtypes = argtypes + self._restype = restype + self._lib_resolver = lib_resolver + self._fn: Optional[Any] = None + + def _resolve(self) -> Any: + fn = getattr(self._lib_resolver(), self._name) + fn.argtypes = self._argtypes + fn.restype = self._restype + self._fn = fn + return fn + + def __call__(self, *args: Any) -> Any: + fn = self._fn or self._resolve() + return fn(*args) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/comgr.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/comgr.py index f0d72e473a46..662361b8e449 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/comgr.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/comgr.py @@ -12,7 +12,7 @@ -> AMD_COMGR_ACTION_LINK_RELOCATABLE_TO_EXECUTABLE -> HSA code object The resulting HSACO bytes are returned to Python and can be handed -straight to `hipModuleLoadData` (see `_hip_module.py`). No subprocesses, +straight to `hipModuleLoadData` (see `hip_module.py`). No subprocesses, no `` parsing, no clang spawn. The library is loaded from the default ROCm library locations or the dynamic @@ -28,7 +28,8 @@ from dataclasses import dataclass from typing import List, Optional, Tuple -from .hip_module import _IS_WINDOWS, _LazyFn, _add_dll_dir, _candidate_lib_paths +from ._ctypes_bind import _LazyFn +from .runtime_coexistence import _IS_WINDOWS, _add_dll_dir, _candidate_lib_paths # Status codes. @@ -58,9 +59,9 @@ class ComgrError(RuntimeError): def _load_lib() -> ctypes.CDLL: # Pair the loader with ``hip_module._load_lib`` so the two halves of # the process always share a single HIP/comgr runtime instance. See - # ``_torch_bundled_lib`` in ``hip_module`` for why a torch-shipped - # libamd_comgr is preferred over /opt/rocm when torch is in the - # process. + # ``_torch_bundled_lib`` in ``runtime_coexistence`` for why a + # torch-shipped libamd_comgr is preferred over /opt/rocm when torch is + # in the process. err = None for p in _candidate_lib_paths("amd_comgr", "ROCKE_COMGR_LIB", ["3"]): try: @@ -73,7 +74,7 @@ def _load_lib() -> ctypes.CDLL: # Lazy: resolved on first call so that rocke and torch can be imported -# in any order. See ``hip_module._torch_bundled_lib`` for context. +# in any order. See ``runtime_coexistence._torch_bundled_lib`` for context. _lib: Optional[ctypes.CDLL] = None @@ -86,7 +87,7 @@ def _resolve_lib() -> ctypes.CDLL: def resolved_lib_path() -> Optional[str]: """Path of the ``libamd_comgr`` this module will load (torch-bundled - preferred over ``/opt/rocm``; see :func:`hip_module._torch_bundled_lib`). + preferred over ``/opt/rocm``; see :func:`runtime_coexistence._torch_bundled_lib`). Returns the already-loaded lib's path once :func:`_resolve_lib` has run, else the first existing candidate. Pure lookup -- does NOT ``dlopen``, so it @@ -196,7 +197,7 @@ def prefer_bundled_lib() -> Optional[Tuple[int, int]]: The resolver never imports torch as a side effect (a library must not), so it only prefers torch's bundled (newest) ``libamd_comgr`` when torch is ALREADY - in the process -- see :func:`hip_module._torch_bundled_lib`. A CLI / runner + in the process -- see :func:`runtime_coexistence._torch_bundled_lib`. A CLI / runner that lowers IR should call this ONCE at startup, BEFORE the first lowering, so the bundled comgr (e.g. ROCm 7.2 / llvm22) is in the process and the LLVM flavor cannot be locked to a stale ``/opt/rocm`` by import order. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py index 7e6832c5978e..36fdf241760b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/hip_module.py @@ -3,7 +3,7 @@ """Minimal ctypes wrapper over `libamdhip64.so` for the hipModule API. -This is the runtime twin of `_comgr.py`: it takes the HSACO bytes that +This is the runtime twin of `comgr.py`: it takes the HSACO bytes that comgr produced from our LLVM IR and runs the kernel via `hipModuleLoadData` + `hipModuleLaunchKernel`. No host compilation, no `` parsing — the same code-object path AMDGPU @@ -20,15 +20,11 @@ from __future__ import annotations import ctypes -import glob -import os -import re -import sys from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple - -_IS_WINDOWS = sys.platform == "win32" +from ._ctypes_bind import _LazyFn +from .runtime_coexistence import _IS_WINDOWS, _add_dll_dir, _candidate_lib_paths HIP_LAUNCH_PARAM_BUFFER_POINTER = ctypes.c_void_p(1) @@ -55,204 +51,6 @@ class _HipEventHandle(ctypes.Structure): _fields_ = [("p", ctypes.c_void_p)] -def _torch_bundled_lib(stem: str) -> Optional[str]: - """Return path to ``/lib/lib.so`` if torch is in this process. - - Newer PyTorch+ROCm wheels (e.g. torch>=2.12 / ROCm 7.2) ship their - own ``libamdhip64.so`` and ``libamd_comgr.so`` inside the wheel's - ``torch/lib/`` directory. When torch is imported, those bundled - libraries get loaded into the process and own torch's HIP context. - A second copy of HIP loaded by rocke from ``/opt/rocm/lib`` is a - *different* runtime instance with disjoint state — modules loaded - via one are invisible to ``hipModuleGetFunction`` from the other, - surfacing as ``hipError(500) named symbol not found`` even when the - HSACO is well-formed and the symbol is present in its ELF. - - To keep both halves of the process talking to the same HIP/comgr - runtime, prefer torch's bundled lib when torch is already imported. - Avoids importing torch as a side effect: only honors a torch that - is *already* in :data:`sys.modules`. - """ - torch_mod = sys.modules.get("torch") - if torch_mod is None: - return None - torch_file = getattr(torch_mod, "__file__", None) - if not torch_file: - return None - libdir = os.path.join(os.path.dirname(torch_file), "lib") - if _IS_WINDOWS: - # ROCm-for-Windows torch wheels (TheRock / AMD nightlies) bundle - # ``amdhip64.dll`` and a version-stamped ``amd_comgr*.dll`` (no - # ``lib`` prefix). Prefer an exact match, else glob the versioned - # comgr name. - direct = os.path.join(libdir, f"{stem}.dll") - if os.path.exists(direct): - return direct - matches = sorted(glob.glob(os.path.join(libdir, f"{stem}*.dll"))) - return matches[0] if matches else None - candidate = os.path.join(libdir, f"lib{stem}.so") - return candidate if os.path.exists(candidate) else None - - -def _rocm_sdk_dll(stem: str) -> Optional[str]: - """Locate a ROCm runtime DLL shipped by the ``rocm-sdk-core`` wheel. - - ROCm-for-Windows torch nightlies (AMD's gfx1151 index / TheRock) put - the HIP runtime and comgr in ``_rocm_sdk_core/bin`` with a version - suffix (e.g. ``amdhip64_7.dll``, ``amd_comgr0702.dll``) rather than in - ``torch/lib``. Returns the first match for ``*.dll`` or None. - """ - if not _IS_WINDOWS: - return None - try: - import importlib.util - - spec = importlib.util.find_spec("_rocm_sdk_core") - except Exception: - return None - if spec is None or not spec.submodule_search_locations: - return None - for loc in spec.submodule_search_locations: - bindir = os.path.join(loc, "bin") - direct = os.path.join(bindir, f"{stem}.dll") - if os.path.exists(direct): - return direct - matches = sorted(glob.glob(os.path.join(bindir, f"{stem}*.dll"))) - if matches: - return matches[0] - return None - - -def _version_key(path: str) -> Any: - """Sort key that orders ROCm install dirs newest-first. - - A plain string sort puts ``rocm-7.10`` *before* ``rocm-7.2`` (because - ``'1' < '2'`` lexically) -- wrong for picking the newest toolkit. Extract - every run of digits from the path and compare them as an integer tuple so - ``7.10`` > ``7.2``. Non-numeric paths sort last. Callers reverse the result - to get descending (newest-first) order. - """ - nums = tuple(int(n) for n in re.findall(r"\d+", path)) - return (len(nums) > 0, nums) - - -def _rocm_root_libdirs() -> List[str]: - """Existing ``/lib`` directories discovered WITHOUT importing torch. - - This is the crux of removing rocke's accidental torch dependency. The ROCm - torch wheel bundles ``libamdhip64.so`` / ``libamd_comgr.so`` inside - ``torch/lib`` and, as a side effect of ``import torch``, drops that - directory onto the process's loader search path -- which is the *only* - reason a bare ``ctypes.CDLL("libamd_comgr.so")`` used to succeed here. A - library must never ``import torch`` to obtain that side effect (it would - invert the dependency and drag a multi-hundred-MB wheel into a pure-IR - process), so we discover a real ROCm install directly instead. - - Priority, newest-version-first within each tier: - 1. ``$ROCM_PATH`` / ``$ROCM_HOME`` -> ``/lib`` (operator override). - 2. Globbed real install layouts. On a packaged ROCm 7.2 there is often no - ``/opt/rocm/lib`` with the runtime in it; the libs live under a - versioned ``core-/lib`` subdir (e.g. - ``/opt/rocm-7.2.0/core-7.13/lib``). Cover both ``/opt/rocm*/lib`` and - ``/opt/rocm*/core-*/lib``. - - Returns directories that exist, de-duplicated, in resolution order. - """ - roots: List[str] = [] - seen: set = set() - - def _add(d: str) -> None: - # De-dupe on the resolved real path so a symlinked root and its target - # don't both get probed; keep the original string for readable candidate - # paths. - if not d: - return - rp = os.path.realpath(d) - if rp and rp not in seen and os.path.isdir(rp): - seen.add(rp) - roots.append(d) - - # Tier 1: explicit env roots win over any globbed install. - for env in ("ROCM_PATH", "ROCM_HOME"): - root = os.environ.get(env) - if root: - _add(os.path.join(root, "lib")) - - # Tier 2: glob real install trees, newest version first. ``core-*/lib`` is - # listed ahead of plain ``lib`` because that is where a packaged install - # actually keeps the runtime .so's. - for pattern in ("/opt/rocm*/core-*/lib", "/opt/rocm*/lib"): - for d in sorted(glob.glob(pattern), key=_version_key, reverse=True): - _add(d) - return roots - - -def _candidate_lib_paths(stem: str, env_var: str, sonames: List[str]) -> List[str]: - """Resolution order for the HIP runtime / COMGR shared libraries. - - Order: - 1. ``$ROCKE_HIP_LIB`` / ``$ROCKE_COMGR_LIB`` (explicit override, full path). - 2. ``/lib/lib.so`` if torch is *already* imported -- an - opportunistic fast-path only (see :func:`_torch_bundled_lib`); we never - import torch to populate it. - 3. A real ROCm install discovered without torch (see - :func:`_rocm_root_libdirs`): ``$ROCM_PATH``/``$ROCM_HOME`` then globbed - ``/opt/rocm*`` trees, newest version first, each with the bare ``.so`` - and the requested SONAME variants. - 4. Bare ``lib.so`` for the dynamic linker's search path -- last - resort. Historically this was the *only* non-torch candidate, which is - why a torch-less process failed with ``cannot load libamd_comgr.so``: - nothing had put the lib on the loader path. Tier 3 fixes that. - """ - paths: List[str] = [] - override = os.environ.get(env_var) - if override: - paths.append(override) - bundled = _torch_bundled_lib(stem) - if bundled is not None: - paths.append(bundled) - sdk = _rocm_sdk_dll(stem) - if sdk is not None: - paths.append(sdk) - if _IS_WINDOWS: - # ROCm-for-Windows / HIP SDK install: ``%HIP_PATH%\bin`` / - # ``%ROCM_PATH%\bin`` / ``%ROCM_HOME%\bin`` then the bare DLL name - # (resolved via the default DLL search path). The comgr DLL carries a - # version suffix, so glob it. - for root_env in ("HIP_PATH", "ROCM_PATH", "ROCM_HOME"): - root = os.environ.get(root_env) - if not root: - continue - bindir = os.path.join(root, "bin") - paths.append(os.path.join(bindir, f"{stem}.dll")) - paths.extend(sorted(glob.glob(os.path.join(bindir, f"{stem}*.dll")))) - paths.append(f"{stem}.dll") - return paths - # POSIX: each discovered ROCm ``/lib`` contributes the bare .so plus - # SONAME-suffixed variants, newest install first. - for libdir in _rocm_root_libdirs(): - paths.append(os.path.join(libdir, f"lib{stem}.so")) - for soname in sonames: - paths.append(os.path.join(libdir, f"lib{stem}.so.{soname}")) - paths.append(f"lib{stem}.so") - return paths - - -def _add_dll_dir(path: str) -> None: - """On Windows, register a resolved DLL's own directory so its - dependent DLLs (bundled alongside it in ``torch/lib`` or the HIP SDK - ``bin``) are found by the loader. No-op off Windows or for bare names. - """ - if not _IS_WINDOWS: - return - d = os.path.dirname(path) - if d and os.path.isdir(d): - try: - os.add_dll_directory(d) - except (OSError, AttributeError): - pass - - def _load_lib() -> ctypes.CDLL: err = None for p in _candidate_lib_paths("amdhip64", "ROCKE_HIP_LIB", ["7"]): @@ -280,47 +78,6 @@ def _resolve_hip() -> ctypes.CDLL: return _hip -class _LazyFn: - """Lazy ctypes function wrapper. - - Defers ``getattr`` on the underlying CDLL until the first call so - that rocke and torch can be imported in any order without ending - up with two HIP runtimes (see ``_torch_bundled_lib``). Resolved on - first use; subsequent calls dispatch directly through the cached - function pointer. - - ``lib_resolver`` returns the shared ``ctypes.CDLL`` for this lib - family (HIP runtime / comgr / ...). It is invoked exactly once per - function on first call. - """ - - __slots__ = ("_name", "_argtypes", "_restype", "_lib_resolver", "_fn") - - def __init__( - self, - name: str, - argtypes: List[Any], - restype: Any, - lib_resolver: "Callable[[], ctypes.CDLL]", - ) -> None: - self._name = name - self._argtypes = argtypes - self._restype = restype - self._lib_resolver = lib_resolver - self._fn: Optional[Any] = None - - def _resolve(self) -> Any: - fn = getattr(self._lib_resolver(), self._name) - fn.argtypes = self._argtypes - fn.restype = self._restype - self._fn = fn - return fn - - def __call__(self, *args: Any) -> Any: - fn = self._fn or self._resolve() - return fn(*args) - - def _b(name: str, *argtypes, restype=ctypes.c_int) -> _LazyFn: return _LazyFn(name, list(argtypes), restype, _resolve_hip) @@ -510,29 +267,47 @@ def destroy(self) -> None: class Runtime: # Per-stream FIFO of ``(refs_tuple, completion_event_or_None)`` - # entries. Every launch appends exactly one entry; tensor lifetimes - # (set up via :meth:`retain_for_stream`) merge into the most-recent - # entry so they share the launch's completion event. + # entries: a backend-agnostic, stream-scoped **lifetime manager**. + # Every async launch appends exactly one entry holding the ctypes + # objects that launch enqueued; a caller may merge additional objects + # to keep alive into the most-recent entry via + # :meth:`retain_for_stream` so they share that launch's completion + # event. + # + # This is NOT torch-specific, and must not migrate into a torch + # module: the bucket only ever holds ctypes buffers and HIP + # ``Event``s, :meth:`retain_for_stream` takes ``*objects: Any`` and + # filters out ints/None, and ``Runtime`` never imports torch. Torch's + # caching allocator (below) is the *motivating caller* of the retain + # path, not something this runtime knows about. # # Why this exists # --------------- # Raw ``hipModuleLaunchKernel`` calls go through ctypes and are - # invisible to torch's stream-aware caching allocator. Two failure - # modes follow: + # invisible to Python's GC (and to torch's stream-aware caching + # allocator, when torch owns the tensors). Two lifetimes must outlive + # an in-flight kernel: # - # 1. The HIP_LAUNCH_PARAM_BUFFER_POINTER ("extra") path does not - # promise to copy the packed-args buffer at enqueue time; - # observation on ROCm 6/7 is that the GPU command processor - # reads it later, when it actually starts the kernel. If the - # Python-owned ctypes buffer has been garbage-collected by then, - # the kernel reads stale memory and writes to whatever pointer - # those bytes now decode as. + # 1. The ctypes args buffer, on the HIP_LAUNCH_PARAM_BUFFER_POINTER + # ("extra") path. That path does not promise to copy the + # packed-args buffer at enqueue time; observation on ROCm 6/7 is + # that the GPU command processor reads it later, when it actually + # starts the kernel. If the Python-owned buffer has been + # garbage-collected by then, the kernel reads stale memory and + # writes to whatever pointer those bytes now decode as. This + # applies to EVERY async launch -- numpy / manifest callers + # included -- so :meth:`launch` parks the buffer here + # unconditionally (torch-independent). # - # 2. Output / workspace tensors built with ``torch.empty(...)`` are - # tracked by torch's caching allocator against torch's - # *current* stream. Once the Python reference drops, the - # allocator can recycle that memory while the raw HIP launch is - # still in flight, mutating the kernel's destination buffer. + # 2. Any caller-owned object backing a kernel argument, retained via + # :meth:`retain_for_stream`. The motivating case: output / + # workspace tensors built with ``torch.empty(...)`` are tracked by + # torch's caching allocator against torch's *current* stream; once + # the Python reference drops, the allocator can recycle that memory + # while the raw HIP launch is still in flight, mutating the + # kernel's destination buffer. The launcher (the torch-aware layer) + # calls ``retain_for_stream`` for that; the mechanism here stays + # agnostic to what it holds. # # The mitigation in both cases is the same: tie the Python # references' lifetime to a HIP completion event recorded on the diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py index dc67ce4e518f..f2caa5c0855a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/launcher.py @@ -95,7 +95,8 @@ from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Sequence, Tuple from .hip_module import Runtime -from .torch_module import pack_args, resolve_stream +from .packing import pack_args +from .torch_interop import resolve_stream __all__ = [ "DeviceMem", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/packing.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/packing.py new file mode 100644 index 000000000000..623aa32db571 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/packing.py @@ -0,0 +1,133 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Kernel-argument packing for the AMDGPU kernarg ABI. + +Torch-agnostic by construction: imports only ``ctypes`` + ``struct`` and +duck-types device pointers on ``.data_ptr()``, so the numpy / manifest +(hip-only) path and the torch-tensor path share one packer. Two forms: + +- :func:`pack_args` -> a single ``bytes`` blob for the + ``HIP_LAUNCH_PARAM_BUFFER_POINTER`` ("extra") launch path. +- :func:`pack_args_kernelparams` -> a list of individual ``ctypes`` + scalars for the ``kernelParams`` launch path (driver-copied at enqueue). + +Both respect the ABI's natural-alignment rule (8-byte ptr/i64, 4-byte +i32/f32) so a mixed ``(ptr ptr ptr i32 i32 i32 ptr)`` signature packs its +trailing pointer at the correct offset. +""" + +from __future__ import annotations + +import ctypes +import struct +from typing import Any, List, Mapping, Sequence, Tuple + + +def pack_args( + signature: Sequence[Mapping[str, Any]], values: Mapping[str, Any] +) -> bytes: + """Pack args from a manifest-style signature, respecting the AMDGPU + kernarg ABI's natural-alignment rule. + + Each AMDGPU kernarg sits at an offset aligned to its own size: + 8-byte alignment for ``ptr`` / ``i64``, 4-byte alignment for + ``i32`` / ``f32``. The flat ``struct.pack`` we used to use packed + fields back-to-back with no inter-field padding, which is fine + when the signature has only ptrs *or* only ints, but fails the + instant a (ptr ptr ptr i32 i32 i32 ptr)-shaped signature shows + up (e.g. a GEMM + bias-fused kernel) — the trailing pointer + lands 4 bytes before its expected kernarg slot and the kernel + reads garbage as the pointer, then segfaults on first access. + + Supported types: ``ptr<..., global>``, ``i32``, ``i64``, ``f32``. + """ + # Map manifest type -> (struct format char, size in bytes, align). + # The Python struct format is built with no implicit padding; + # we insert explicit ``B`` bytes when alignment requires it. + _TY_FMT: Mapping[str, Tuple[str, int, int]] = { + "i32": ("i", 4, 4), + "i64": ("q", 8, 8), + "f32": ("f", 4, 4), + } + fmt_parts: List[str] = ["<"] + packed: List[Any] = [] + offset = 0 + for arg in signature: + name = str(arg["name"]) + ty = str(arg["type"]) + if name not in values: + raise KeyError(f"missing kernel arg {name!r}") + v = values[name] + if ty.startswith("ptr<"): + fmt_char, size, align = "Q", 8, 8 + arg_val: Any = _as_ptr(v) + elif ty in _TY_FMT: + fmt_char, size, align = _TY_FMT[ty] + if ty == "f32": + arg_val = float(v) + else: + arg_val = int(v) + else: + raise ValueError(f"unsupported kernel arg type {ty!r} for {name}") + # Insert padding bytes so this arg lands at its natural alignment. + pad = (-offset) % align + if pad: + fmt_parts.append(f"{pad}x") + offset += pad + fmt_parts.append(fmt_char) + packed.append(arg_val) + offset += size + return struct.pack("".join(fmt_parts), *packed) + + +def pack_args_kernelparams( + signature: Sequence[Mapping[str, Any]], values: Mapping[str, Any] +) -> List[Any]: + """Pack args as a list of individual ``ctypes`` scalars for the + ``kernelParams`` path of ``hipModuleLaunchKernel``. + + Returning one ``ctypes`` object per kernel argument lets the launcher + build a ``void* params[]`` array whose entries point to each scalar. + The CUDA/HIP semantics for ``kernelParams`` guarantee that the + driver copies each parameter into driver-owned memory at enqueue + time -- in contrast to the ``extra`` / + ``HIP_LAUNCH_PARAM_BUFFER_POINTER`` path, whose copy semantics are + underspecified and have been observed to read the host buffer + *after* ``hipModuleLaunchKernel`` returns. The ``extra`` race + produced the parity-harness "max_abs jumps to 2.5 / 512 on the + second CK call" symptom investigated in + ``ck/dsl/unified_attention_creative_results.md``. + + This is the same approach Triton's AMD driver uses (see + ``triton/backends/amd/driver.py:392-402``: ``void *params[] = { ... }; + hipModuleLaunchKernel(..., params, 0);``). + """ + out: List[Any] = [] + for arg in signature: + name = str(arg["name"]) + ty = str(arg["type"]) + if name not in values: + raise KeyError(f"missing kernel arg {name!r}") + v = values[name] + if ty.startswith("ptr<"): + out.append(ctypes.c_uint64(_as_ptr(v))) + elif ty == "i32": + out.append(ctypes.c_int32(int(v))) + elif ty == "i64": + out.append(ctypes.c_int64(int(v))) + elif ty == "f32": + out.append(ctypes.c_float(float(v))) + else: + raise ValueError(f"unsupported kernel arg type {ty!r} for {name}") + return out + + +def _as_ptr(v: Any) -> int: + if isinstance(v, int): + return v + if hasattr(v, "data_ptr"): + return int(v.data_ptr()) + if v is None: + return 0 + raise TypeError(f"cannot convert {type(v).__name__} to device pointer") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/runtime_coexistence.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/runtime_coexistence.py new file mode 100644 index 000000000000..a6051e6a1710 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/runtime_coexistence.py @@ -0,0 +1,244 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Which ROCm runtime do we bind to -- and whose? + +Both the HIP (`hip_module`) and comgr (`comgr`) ctypes wrappers must +resolve the *same* loaded runtime instance. The crux is process-level +coexistence with a host that may already own a ROCm runtime: + + A ROCm PyTorch wheel bundles ``libamdhip64.so`` / ``libamd_comgr.so`` + inside ``torch/lib`` and, as a side effect of ``import torch``, loads + them into the process. A *second* copy of HIP loaded by rocke from + ``/opt/rocm`` is a **different runtime instance** with disjoint state + -- a module loaded via one is invisible to ``hipModuleGetFunction`` + from the other, surfacing as ``hipError(500) named symbol not found`` + even when the HSACO is well-formed. This is a loader / runtime-instance + phenomenon (two separate HSA runtime inits, two handle tables), not a + context-binding one. + +So when torch is already in the process we prefer *its* bundled ``.so``, +sharing one loaded HIP/comgr runtime instance across both halves of the +process. A library must never ``import torch`` merely to obtain that side +effect (it would invert the dependency and drag a multi-hundred-MB wheel +into a pure-IR process): we only honor a torch that is *already* imported +(via :data:`sys.modules`) and otherwise discover a real ROCm install +directly. The sole sanctioned exception is the explicit +:func:`rocke.runtime.comgr.prefer_bundled_lib` entrypoint hook, which +imports torch deliberately to pin the bundled comgr before lowering. + +This module owns only *resolution* -- candidate path discovery + the +Windows DLL-directory registration. The actual ``dlopen`` (per-family +``_load_lib``), the HIP primary-context binding (``_ensure_hip_init``), +and device-property introspection (``get_device_arch``) live with their +respective runtime wrappers. +""" + +from __future__ import annotations + +import glob +import os +import re +import sys +from typing import Any, List, Optional + + +_IS_WINDOWS = sys.platform == "win32" + + +def _torch_bundled_lib(stem: str) -> Optional[str]: + """Return path to ``/lib/lib.so`` if torch is in this process. + + Newer PyTorch+ROCm wheels (e.g. torch>=2.12 / ROCm 7.2) ship their + own ``libamdhip64.so`` and ``libamd_comgr.so`` inside the wheel's + ``torch/lib/`` directory. When torch is imported, those bundled + libraries become torch's loaded HIP runtime instance. + A second copy of HIP loaded by rocke from ``/opt/rocm/lib`` is a + *different* runtime instance with disjoint state — modules loaded + via one are invisible to ``hipModuleGetFunction`` from the other, + surfacing as ``hipError(500) named symbol not found`` even when the + HSACO is well-formed and the symbol is present in its ELF. + + To keep both halves of the process talking to the same HIP/comgr + runtime, prefer torch's bundled lib when torch is already imported. + Avoids importing torch as a side effect: only honors a torch that + is *already* in :data:`sys.modules`. + """ + torch_mod = sys.modules.get("torch") + if torch_mod is None: + return None + torch_file = getattr(torch_mod, "__file__", None) + if not torch_file: + return None + libdir = os.path.join(os.path.dirname(torch_file), "lib") + if _IS_WINDOWS: + # ROCm-for-Windows torch wheels (TheRock / AMD nightlies) bundle + # ``amdhip64.dll`` and a version-stamped ``amd_comgr*.dll`` (no + # ``lib`` prefix). Prefer an exact match, else glob the versioned + # comgr name. + direct = os.path.join(libdir, f"{stem}.dll") + if os.path.exists(direct): + return direct + matches = sorted(glob.glob(os.path.join(libdir, f"{stem}*.dll"))) + return matches[0] if matches else None + candidate = os.path.join(libdir, f"lib{stem}.so") + return candidate if os.path.exists(candidate) else None + + +def _rocm_sdk_dll(stem: str) -> Optional[str]: + """Locate a ROCm runtime DLL shipped by the ``rocm-sdk-core`` wheel. + + ROCm-for-Windows torch nightlies (AMD's gfx1151 index / TheRock) put + the HIP runtime and comgr in ``_rocm_sdk_core/bin`` with a version + suffix (e.g. ``amdhip64_7.dll``, ``amd_comgr0702.dll``) rather than in + ``torch/lib``. Returns the first match for ``*.dll`` or None. + """ + if not _IS_WINDOWS: + return None + try: + import importlib.util + + spec = importlib.util.find_spec("_rocm_sdk_core") + except Exception: + return None + if spec is None or not spec.submodule_search_locations: + return None + for loc in spec.submodule_search_locations: + bindir = os.path.join(loc, "bin") + direct = os.path.join(bindir, f"{stem}.dll") + if os.path.exists(direct): + return direct + matches = sorted(glob.glob(os.path.join(bindir, f"{stem}*.dll"))) + if matches: + return matches[0] + return None + + +def _version_key(path: str) -> Any: + """Sort key that orders ROCm install dirs newest-first. + + A plain string sort puts ``rocm-7.10`` *before* ``rocm-7.2`` (because + ``'1' < '2'`` lexically) -- wrong for picking the newest toolkit. Extract + every run of digits from the path and compare them as an integer tuple so + ``7.10`` > ``7.2``. Non-numeric paths sort last. Callers reverse the result + to get descending (newest-first) order. + """ + nums = tuple(int(n) for n in re.findall(r"\d+", path)) + return (len(nums) > 0, nums) + + +def _rocm_root_libdirs() -> List[str]: + """Existing ``/lib`` directories discovered WITHOUT importing torch. + + This is the crux of removing rocke's accidental torch dependency. The ROCm + torch wheel bundles ``libamdhip64.so`` / ``libamd_comgr.so`` inside + ``torch/lib`` and, as a side effect of ``import torch``, drops that + directory onto the process's loader search path -- which is the *only* + reason a bare ``ctypes.CDLL("libamd_comgr.so")`` used to succeed here. A + library must never ``import torch`` to obtain that side effect (it would + invert the dependency and drag a multi-hundred-MB wheel into a pure-IR + process), so we discover a real ROCm install directly instead. + + Priority, newest-version-first within each tier: + 1. ``$ROCM_PATH`` / ``$ROCM_HOME`` -> ``/lib`` (operator override). + 2. Globbed real install layouts. On a packaged ROCm 7.2 there is often no + ``/opt/rocm/lib`` with the runtime in it; the libs live under a + versioned ``core-/lib`` subdir (e.g. + ``/opt/rocm-7.2.0/core-7.13/lib``). Cover both ``/opt/rocm*/lib`` and + ``/opt/rocm*/core-*/lib``. + + Returns directories that exist, de-duplicated, in resolution order. + """ + roots: List[str] = [] + seen: set = set() + + def _add(d: str) -> None: + # De-dupe on the resolved real path so a symlinked root and its target + # don't both get probed; keep the original string for readable candidate + # paths. + if not d: + return + rp = os.path.realpath(d) + if rp and rp not in seen and os.path.isdir(rp): + seen.add(rp) + roots.append(d) + + # Tier 1: explicit env roots win over any globbed install. + for env in ("ROCM_PATH", "ROCM_HOME"): + root = os.environ.get(env) + if root: + _add(os.path.join(root, "lib")) + + # Tier 2: glob real install trees, newest version first. ``core-*/lib`` is + # listed ahead of plain ``lib`` because that is where a packaged install + # actually keeps the runtime .so's. + for pattern in ("/opt/rocm*/core-*/lib", "/opt/rocm*/lib"): + for d in sorted(glob.glob(pattern), key=_version_key, reverse=True): + _add(d) + return roots + + +def _candidate_lib_paths(stem: str, env_var: str, sonames: List[str]) -> List[str]: + """Resolution order for the HIP runtime / comgr shared libraries. + + Order: + 1. ``$ROCKE_HIP_LIB`` / ``$ROCKE_COMGR_LIB`` (explicit override, full path). + 2. ``/lib/lib.so`` if torch is *already* imported -- an + opportunistic fast-path only (see :func:`_torch_bundled_lib`); we never + import torch to populate it. + 3. A real ROCm install discovered without torch (see + :func:`_rocm_root_libdirs`): ``$ROCM_PATH``/``$ROCM_HOME`` then globbed + ``/opt/rocm*`` trees, newest version first, each with the bare ``.so`` + and the requested SONAME variants. + 4. Bare ``lib.so`` for the dynamic linker's search path -- last + resort. Historically this was the *only* non-torch candidate, which is + why a torch-less process failed with ``cannot load libamd_comgr.so``: + nothing had put the lib on the loader path. Tier 3 fixes that. + """ + paths: List[str] = [] + override = os.environ.get(env_var) + if override: + paths.append(override) + bundled = _torch_bundled_lib(stem) + if bundled is not None: + paths.append(bundled) + sdk = _rocm_sdk_dll(stem) + if sdk is not None: + paths.append(sdk) + if _IS_WINDOWS: + # ROCm-for-Windows / HIP SDK install: ``%HIP_PATH%\bin`` / + # ``%ROCM_PATH%\bin`` / ``%ROCM_HOME%\bin`` then the bare DLL name + # (resolved via the default DLL search path). The comgr DLL carries a + # version suffix, so glob it. + for root_env in ("HIP_PATH", "ROCM_PATH", "ROCM_HOME"): + root = os.environ.get(root_env) + if not root: + continue + bindir = os.path.join(root, "bin") + paths.append(os.path.join(bindir, f"{stem}.dll")) + paths.extend(sorted(glob.glob(os.path.join(bindir, f"{stem}*.dll")))) + paths.append(f"{stem}.dll") + return paths + # POSIX: each discovered ROCm ``/lib`` contributes the bare .so plus + # SONAME-suffixed variants, newest install first. + for libdir in _rocm_root_libdirs(): + paths.append(os.path.join(libdir, f"lib{stem}.so")) + for soname in sonames: + paths.append(os.path.join(libdir, f"lib{stem}.so.{soname}")) + paths.append(f"lib{stem}.so") + return paths + + +def _add_dll_dir(path: str) -> None: + """On Windows, register a resolved DLL's own directory so its + dependent DLLs (bundled alongside it in ``torch/lib`` or the HIP SDK + ``bin``) are found by the loader. No-op off Windows or for bare names. + """ + if not _IS_WINDOWS: + return + d = os.path.dirname(path) + if d and os.path.isdir(d): + try: + os.add_dll_directory(d) + except (OSError, AttributeError): + pass diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/torch_module.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/torch_interop.py similarity index 51% rename from dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/torch_module.py rename to dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/torch_interop.py index 9e32185e6ee3..3823c5e9e85f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/torch_module.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/torch_interop.py @@ -1,20 +1,23 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Torch tensor launcher for CK DSL HSACO kernels. +"""Torch-tensor launch glue, plus the torch-optional stream resolver. -This runtime is for integrations like AITER where tensors already live on the -GPU. It avoids host staging: kernel arguments are packed from -`torch.Tensor.data_ptr()` and Python scalar values, then launched through the -same hipModule path as `run_manifest`. +For integrations like AITER where tensors already live on the GPU: launch +through the same hipModule path as `run_manifest`, without host staging. +This module holds `resolve_stream` (substitute torch's current stream so +its caching allocator stays coherent -- degrades to the HIP null stream +when torch is absent, so the numpy/manifest hip-only path calls it too), +`empty_workspace`, and the `launch_torch_kernel` back-compat shim. + +Kernel-argument packing is torch-agnostic and lives in `packing.py`; this +module builds on it. """ from __future__ import annotations -import ctypes -import struct from dataclasses import dataclass -from typing import Any, List, Mapping, Sequence, Tuple +from typing import Any, Mapping, Sequence, Tuple @dataclass(frozen=True) @@ -27,7 +30,7 @@ def _require_torch(): try: import torch except Exception as e: # pragma: no cover - environment dependent - raise RuntimeError("rocke.runtime.torch_module requires torch") from e + raise RuntimeError("rocke.runtime.torch_interop requires torch") from e return torch @@ -74,115 +77,6 @@ def resolve_stream(stream, device=None) -> int: return 0 -def pack_args( - signature: Sequence[Mapping[str, Any]], values: Mapping[str, Any] -) -> bytes: - """Pack args from a manifest-style signature, respecting the AMDGPU - kernarg ABI's natural-alignment rule. - - Each AMDGPU kernarg sits at an offset aligned to its own size: - 8-byte alignment for ``ptr`` / ``i64``, 4-byte alignment for - ``i32`` / ``f32``. The flat ``struct.pack`` we used to use packed - fields back-to-back with no inter-field padding, which is fine - when the signature has only ptrs *or* only ints, but fails the - instant a (ptr ptr ptr i32 i32 i32 ptr)-shaped signature shows - up (e.g. a GEMM + bias-fused kernel) — the trailing pointer - lands 4 bytes before its expected kernarg slot and the kernel - reads garbage as the pointer, then segfaults on first access. - - Supported types: ``ptr<..., global>``, ``i32``, ``i64``, ``f32``. - """ - # Map manifest type -> (struct format char, size in bytes, align). - # The Python struct format is built with no implicit padding; - # we insert explicit ``B`` bytes when alignment requires it. - _TY_FMT: Mapping[str, Tuple[str, int, int]] = { - "i32": ("i", 4, 4), - "i64": ("q", 8, 8), - "f32": ("f", 4, 4), - } - fmt_parts: List[str] = ["<"] - packed: List[Any] = [] - offset = 0 - for arg in signature: - name = str(arg["name"]) - ty = str(arg["type"]) - if name not in values: - raise KeyError(f"missing kernel arg {name!r}") - v = values[name] - if ty.startswith("ptr<"): - fmt_char, size, align = "Q", 8, 8 - arg_val: Any = _as_ptr(v) - elif ty in _TY_FMT: - fmt_char, size, align = _TY_FMT[ty] - if ty == "f32": - arg_val = float(v) - else: - arg_val = int(v) - else: - raise ValueError(f"unsupported torch arg type {ty!r} for {name}") - # Insert padding bytes so this arg lands at its natural alignment. - pad = (-offset) % align - if pad: - fmt_parts.append(f"{pad}x") - offset += pad - fmt_parts.append(fmt_char) - packed.append(arg_val) - offset += size - return struct.pack("".join(fmt_parts), *packed) - - -def pack_args_kernelparams( - signature: Sequence[Mapping[str, Any]], values: Mapping[str, Any] -) -> List[Any]: - """Pack args as a list of individual ``ctypes`` scalars for the - ``kernelParams`` path of ``hipModuleLaunchKernel``. - - Returning one ``ctypes`` object per kernel argument lets the launcher - build a ``void* params[]`` array whose entries point to each scalar. - The CUDA/HIP semantics for ``kernelParams`` guarantee that the - driver copies each parameter into driver-owned memory at enqueue - time -- in contrast to the ``extra`` / - ``HIP_LAUNCH_PARAM_BUFFER_POINTER`` path, whose copy semantics are - underspecified and have been observed to read the host buffer - *after* ``hipModuleLaunchKernel`` returns. The ``extra`` race - produced the parity-harness "max_abs jumps to 2.5 / 512 on the - second CK call" symptom investigated in - ``ck/dsl/unified_attention_creative_results.md``. - - This is the same approach Triton's AMD driver uses (see - ``triton/backends/amd/driver.py:392-402``: ``void *params[] = { ... }; - hipModuleLaunchKernel(..., params, 0);``). - """ - out: List[Any] = [] - for arg in signature: - name = str(arg["name"]) - ty = str(arg["type"]) - if name not in values: - raise KeyError(f"missing kernel arg {name!r}") - v = values[name] - if ty.startswith("ptr<"): - out.append(ctypes.c_uint64(_as_ptr(v))) - elif ty == "i32": - out.append(ctypes.c_int32(int(v))) - elif ty == "i64": - out.append(ctypes.c_int64(int(v))) - elif ty == "f32": - out.append(ctypes.c_float(float(v))) - else: - raise ValueError(f"unsupported torch arg type {ty!r} for {name}") - return out - - -def _as_ptr(v: Any) -> int: - if isinstance(v, int): - return v - if hasattr(v, "data_ptr"): - return int(v.data_ptr()) - if v is None: - return 0 - raise TypeError(f"cannot convert {type(v).__name__} to device pointer") - - def empty_workspace(shape: Sequence[int], *, dtype: Any, device: Any): """Allocate a fresh torch workspace tensor. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py index a0e951ec97fe..59eed65830c4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/test_rocke.py @@ -4689,6 +4689,272 @@ def destroy(self) -> None: self._restore_pending(prior) +class TestRuntimeLaunchKeepAlive(unittest.TestCase): + """``Runtime.launch`` / ``launch_blocking`` args-buffer lifetime. + + Characterizes the non-torch race the pending-args queue exists to + prevent: the ``extra``-path packed-args buffer must outlive an + async ``launch`` (the GPU command processor reads it after + enqueue), whereas ``launch_blocking`` needs no bucket entry because + its trailing ``hipStreamSynchronize`` is the drain barrier. Runs on + host only -- the two HIP entry points are stubbed to return success. + """ + + def _isolate_pending(self): + from rocke.runtime.hip_module import Runtime + + prior = dict(Runtime._pending_args) + Runtime._pending_args.clear() + return prior + + def _restore_pending(self, prior): + from rocke.runtime.hip_module import Runtime + + Runtime._pending_args.clear() + Runtime._pending_args.update(prior) + + def test_async_launch_parks_args_buffer_until_drain(self): + import ctypes + from unittest import mock + + from rocke.runtime.hip_module import Runtime, _HipFunctionHandle + + rt = Runtime() + prior = self._isolate_pending() + try: + with mock.patch( + "rocke.runtime.hip_module._hipModuleLaunchKernel", return_value=0 + ): + rt.launch( + _HipFunctionHandle(), + (1, 1, 1), + (64, 1, 1), + b"\x01\x02\x03\x04", + stream=7, + record_event=False, + ) + bucket = Runtime._pending_args.get(7) + self.assertIsNotNone(bucket) + self.assertEqual(len(bucket), 1) + refs, evt = bucket[-1] + # No event recorded when record_event=False; refs hold the + # ctypes objects the launch pointed the CP at. + self.assertIsNone(evt) + args_buf, _size_buf, extra = refs + # The parked `extra` array must still point at the *same* + # args buffer the launch enqueued -- proving the buffer the + # CP reads later is the one we kept alive, not a copy that + # was already freed. + self.assertEqual(extra[1], ctypes.addressof(args_buf)) + finally: + self._restore_pending(prior) + + def test_blocking_launch_parks_nothing_and_syncs_once(self): + from unittest import mock + + from rocke.runtime.hip_module import Runtime, _HipFunctionHandle + + rt = Runtime() + prior = self._isolate_pending() + try: + with mock.patch( + "rocke.runtime.hip_module._hipModuleLaunchKernel", return_value=0 + ), mock.patch( + "rocke.runtime.hip_module._hipStreamSynchronize", return_value=0 + ) as sync_stub: + rt.launch_blocking( + _HipFunctionHandle(), + (1, 1, 1), + (64, 1, 1), + b"\x01\x02\x03\x04", + stream=3, + ) + # The sync IS the barrier + args-buffer drain, so no bucket + # bookkeeping is needed for a blocking launch. + self.assertNotIn(3, Runtime._pending_args) + self.assertEqual(sync_stub.call_count, 1) + finally: + self._restore_pending(prior) + + +class TestResolveStream(unittest.TestCase): + """``torch_interop.resolve_stream`` -- the caching-allocator hinge. + + Its whole contract is torch-*optional*: pass through a nonzero + handle, substitute torch's current stream when torch is present, and + fall back to the HIP null stream (0) only when torch is genuinely + absent. The torch-present case is the tripwire that fails loudly if a + future edit ever collapses the torch branch to 0 inside a torch + process. + """ + + def test_passes_through_nonzero_without_touching_torch(self): + import sys + from unittest import mock + + from rocke.runtime.torch_interop import resolve_stream + + # Poison torch so any import attempt would raise -- proves the + # nonzero fast-path never reaches for torch at all. + with mock.patch.dict(sys.modules, {"torch": None}): + self.assertEqual(resolve_stream(1234), 1234) + + def test_returns_zero_when_torch_absent(self): + import sys + from unittest import mock + + from rocke.runtime.torch_interop import resolve_stream + + # sys.modules["torch"] = None makes `import torch` raise + # ImportError, standing in for a torch-free environment. + with mock.patch.dict(sys.modules, {"torch": None}): + self.assertEqual(resolve_stream(0), 0) + + def test_uses_current_stream_when_torch_present(self): + import sys + from unittest import mock + + from rocke.runtime.torch_interop import resolve_stream + + fake = mock.MagicMock() + fake.cuda.current_device.return_value = 0 + fake.cuda.current_stream.return_value.cuda_stream = 99 + with mock.patch.dict(sys.modules, {"torch": fake}): + self.assertEqual(resolve_stream(0), 99) + + +class TestPackArgsKernargABI(unittest.TestCase): + """``pack_args`` / ``pack_args_kernelparams`` -- AMDGPU kernarg ABI. + + Each kernarg sits at an offset aligned to its own size (8 for + ptr/i64, 4 for i32/f32). The regression that motivated the padded + packer is a trailing pointer after three i32s landing 4 bytes early + and being read as garbage. These are pure host-assertable byte + checks -- no GPU, no torch. + """ + + _MIXED_SIG = [ + {"name": "a", "type": "ptr"}, + {"name": "b", "type": "ptr"}, + {"name": "c", "type": "ptr"}, + {"name": "m", "type": "i32"}, + {"name": "n", "type": "i32"}, + {"name": "k", "type": "i32"}, + {"name": "d", "type": "ptr"}, + ] + _MIXED_VALS = { + "a": 0x1000, + "b": 0x2000, + "c": 0x3000, + "m": 4, + "n": 5, + "k": 6, + "d": 0xABCD, + } + + def test_inserts_padding_before_misaligned_pointer(self): + import struct + + from rocke.runtime.packing import pack_args + + packed = pack_args(self._MIXED_SIG, self._MIXED_VALS) + # 3 ptr (24) + 3 i32 (12) = 36, pad 4 to reach 8-alignment, then + # the trailing ptr (8) => 48 bytes, trailing ptr at offset 40. + self.assertEqual(len(packed), 48) + self.assertEqual(struct.unpack_from(""}, + {"name": "i", "type": "i32"}, + {"name": "q", "type": "i64"}, + {"name": "f", "type": "f32"}, + ] + vals = {"p": 0x10, "i": 7, "q": 0x1122334455, "f": 1.5} + packed = pack_args(sig, vals) + # p@0(8) i@8(4) pad4 q@16(8) f@24(4) => 28 bytes. + self.assertEqual(len(packed), 28) + self.assertEqual(struct.unpack_from("