refactor(rocke): Split the torch-optional path from a torch-free hip-only core#9302
Open
shumway wants to merge 6 commits into
Open
refactor(rocke): Split the torch-optional path from a torch-free hip-only core#9302shumway wants to merge 6 commits into
shumway wants to merge 6 commits into
Conversation
…e reorg Pin the behaviors an upcoming rocke.runtime module split must preserve, so the moves are provably behavior-preserving rather than a leap. All host-only (no GPU), run in the torch-free platform venv. - TestRuntimeLaunchKeepAlive: async launch parks the extra-path args buffer until drain (parked extra ptr == addressof(args_buf)); launch_blocking parks nothing and syncs exactly once. This args-buffer lifetime is a non-torch invariant that protects numpy callers too. - TestResolveStream: passthrough of a nonzero handle without touching torch; fall back to the null stream when torch is absent; substitute torch's current stream when present (the tripwire against a silent collapse to 0 inside a torch process). - TestPackArgsKernargABI: AMDGPU kernarg natural-alignment (padding before a trailing pointer after three i32s), pack_args/pack_args_kernelparams value agreement, and the missing-arg / unknown-type guards. - TestLibDiscoveryOrder: newest-ROCm-first version ordering, env override first in candidate resolution, and that the torch-bundled probe never imports torch for the loader side effect. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Split the shared library-resolution and lazy-ctypes plumbing out of hip_module.py so the HIP and comgr wrappers become peers over a shared foundation instead of comgr reaching up into hip_module. - runtime_coexistence.py: "which ROCm runtime do we bind to, and whose" -- torch-bundled vs system-install discovery (_torch_bundled_lib, _rocm_sdk_dll, _rocm_root_libdirs, _version_key, _candidate_lib_paths, _add_dll_dir, _IS_WINDOWS). Framed as sharing one loaded HIP/comgr runtime instance (not an HSA context), the distinction that avoids the hipError(500) named-symbol-not-found from a double-loaded runtime. - _ctypes_bind.py: the _LazyFn lazy-binding helper (pure ctypes mechanism). - hip_module.py and comgr.py now import both from these modules; the inverted comgr->hip_module import is removed. _ensure_hip_init and get_device_arch stay in hip_module with the HIP function table (they bind the primary context / read gcnArchName). - Bodies moved verbatim; only the _version_key `import re` was hoisted to module scope. Doc cross-refs updated (comgr, lower_llvm, AGENTS, BUILD). Behavior-preserving: full host suite 207 passed / 17 skipped, unchanged. The characterization tests from the prior commit now import the moved helpers from runtime_coexistence. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
pack_args / pack_args_kernelparams / _as_ptr are torch-agnostic (they duck-type device pointers on .data_ptr() and accept plain ints), yet lived in torch_module.py -- so a torch-free caller (heuristics/_measure_moe, whose own docstring says "No torch dependency") had to import from a module named torch_module. Move them to a new torch-free packing.py (imports only ctypes/struct/typing) that the numpy/manifest and torch paths both use. Re-pointed launcher, heuristics/_measure_moe (+ its docstring), the runtime __init__ re-export, and the pack_args characterization tests. torch_module shed its now-unused ctypes/struct/List imports. Public rocke.runtime.pack_args path is unchanged (re-exported via __init__). Behavior-preserving: host suite 207 passed / 17 skipped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
torch_module was a false cognate beside hip_module (which genuinely wraps the hipModule API) -- there is no "torch module" being wrapped. After the packing extraction, what remains is torch-tensor launch glue: resolve_stream, empty_workspace, launch_torch_kernel. Rename to torch_interop to name the role. resolve_stream stays here (it is torch-aware, just torch-optional). Re-pointed the runtime __init__ import + layer map, launcher, and the resolve_stream characterization tests. Updated the documented public path and the dsl_docs surface (api_index split into packing/torch_interop sections; file_index rows for the new runtime_coexistence / _ctypes_bind / packing / torch_interop modules; BUILD/AGENTS/mental_model/limitations/ runbook/manifest/comgr doc refs), routing pack_args mentions to packing and module/stream mentions to torch_interop. Behavior-preserving: host suite 207 passed / 17 skipped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…stic The _pending_args FIFO reads as torch machinery, but it is a stream-scoped lifetime manager: the bucket only ever holds ctypes buffers + HIP events, retain_for_stream takes *objects: Any and filters ints/None, and Runtime never imports torch. It serves two lifetimes -- the ctypes args buffer (non-torch; parked on every async launch, numpy callers included) and any caller-retained object (torch tensors being the motivating case, retained by the torch-aware launcher). Reframe the class comment to say so, so a future reader doesn't "tidy up" by moving retention into a torch module -- which would split one FIFO across two owners and drop the args-buffer keep-alive for non-torch callers. Docstring-only; no behavior change. Host suite 207 passed / 17 skipped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Holistic review follow-ups (no logic change): - torch_interop.py module docstring: drop the stale "packed from torch.Tensor.data_ptr()" claim (packing moved to packing.py) and the torch-tensor-only charter; name its torch-optional resolve_stream tenant. - packing.py: error text "unsupported torch arg type" -> "kernel arg type" so the torch-agnostic module doesn't blame torch on the numpy path. - __init__ layer map: name the torch-free-core / torch-aware-edge seam and give _ctypes_bind its own line. - file_index intro: "primary public exports" -> "primary contents", since whole-module-private files (runtime_coexistence, _ctypes_bind) list their key internal helpers. Host suite 207 passed / 17 skipped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
shumway
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue ID : AICK-1518
Summary
Reorganize the rocke platform runtime (
dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/runtime/) so the hip-only path is torch-free by construction while torch stays optional. No behavior change; the publicrocke.runtimeAPI (__all__) is unchanged. This is a step toward removing the python torch dependency from rocKE.The
hip_module.pyhad picked up a lot of torch logic. With this change,hip_module.pyis now consistent with its docstring: a minimal ctypes wrapper overlibamdhip64.sofor the hipModule API.The extra bloat from
hip_module.pyis now split out into three modules, and we renamed one module:runtime_coexistance.pyto logically bind using either system or torch-bundled discovery of ROCm._ctypes_bind.pyfor_LaxyFnneeded by bothhip_module.pyandcomgr.py.packing.pyfor args packing.torch_interop.pywastorch_module.py, renamed to clarify that it is glue for torch inter operability.The changes are:
runtime_coexistence.py(which ROCm runtime we bind to, and whose: torch-bundled vs system discovery)._ctypes_bind.py(_LazyFn); bothcomgrandhip_modulenow depend on them — removing the invertedcomgr → hip_moduleimport.packing.py(torch-agnosticpack_args/pack_args_kernelparams), so the numpy/manifest path no longer imports from a torch-named module.torch_module.py→torch_interop.py(it holds torch-tensor launch glue, not a wrapped "torch module");resolve_streamstays there (torch-aware but torch-optional).Runtime._pending_argskeep-alive queue as a backend-agnostic, stream-scoped lifetime manager.Risk Assessment
Risk 3 (core runtime subsystem, multi-file), low residual. Behavior-preserving, host-only Python (no kernel/device/default-selection change), public import surface unchanged. Characterization tests were landed first to lock the moved
behaviors, and the host suite stayed green (207 passed / 17 skipped) at every commit. Verified: with
import torchpoisoned, the hip-only core imports, packs args, and constructsRuntime; no core module imports torch at module scope(
comgr.prefer_bundled_libkeeps one deliberate, documented function-localimport torch).Device / Architecture Coverage
Host-only Python plumbing; no kernel, dispatch, or default-selection change. Passing PR CI is sufficient — no device sweep required. The byte-identity parity gate does not apply (the runtime layer is downstream of
.llemission; there isno C++ runtime twin).
Testing Summary
Testing Checklist
python -m pytest tests/test_rocke.py(platform/.venv) - Status: Passed (207 passed, 17 skipped)python tests/run_all.py- Status: Not runFlags / Guardrails
None. Behavior-preserving refactor; no new runtime behavior or feature flag.
Adjacent Tests Considered
The 13 new characterization tests cover the moved seams (args-buffer keep-alive on async vs blocking launch;
resolve_streampresent/absent/passthrough;pack_argskernarg-alignment ABI + kernelparams agreement; lib-discoveryordering). The full
test_rocke.py(224 host tests) exercises the runtime, launcher, and instance-build paths that import these modules and stayed green.Technical Changes
_ctypes_bindextraction →packingextraction →torch_module→torch_interoprename →Runtimedocstring reframe → doc alignment. Each behavior-preserving.