Support TheRock ROCm “Python-managed” install paths (non-/opt/rocm)

[akao-amd]

After installing the latest ROCm 7.12 (TheRock / ROCm Core SDK stream) following the official guide , and initializing with: rocm-sdk init, the resulting SDK layout does not follow the traditional /opt/rocm or /opt/rocm-x.y.z convention on a vanilla ubuntu:22.04 base image.

Instead, the SDK is installed / expanded under a Python site-packages managed location, e.g.:

/usr/local/lib/python3.10/site-packages/_rocm_sdk_core
(and other rocm_sdk* directories)

This path change breaks tools and scripts that assume ROCm lives under /opt/rocm* (e.g., hardcoded discovery, PATH/CMAKE_PREFIX_PATH logic, or “rocm root” probing).

Request / ask
As TheRock-derived ROCm becomes a priority target, please consider updating the tooling to discover and support ROCm installations rooted in the rocm_sdk* site-packages paths (in addition to /opt/rocm*).

[amd-zfyu]

Supporting TheRock (ROCm Core SDK / site-packages) — Plan

Summary

We'll extend aiter's ROCm discovery so it works when the SDK is installed under Python site-packages (e.g. .../site-packages/_rocm_sdk_core) after rocm-sdk init, while keeping compatibility with /opt/rocm and /opt/rocm-x.y.z.

---

Affected code today

Location	Issue
aiter/jit/utils/cpp_extension.py	_find_rocm_home() only checks env + which(hipcc) + fallback /opt/rocm; no site-packages. executable_path() only looks under ROCM_HOME/bin, not llvm/bin.
csrc/cpp_itfs/utils.py	Hardcoded /opt/rocm/llvm/bin/amdgpu-arch and /opt/rocm/bin/hipconfig.
op_tests/cpp/mha/build_mha.sh	Hardcoded /opt/rocm/bin/hipcc.
csrc/kernels/solver/Makefile	ROCBLAS_DIR / ROCSOLVER_DIR fixed to /opt/rocm.
.githooks/pre-commit	PATH uses /opt/rocm/llvm/bin, /opt/rocm/hcc/bin.

---

Discovery order (to implement)

1. Env — ROCM_HOME / ROCM_PATH (unchanged).
2. hipcc in PATH — Derive root from which("hipcc") (unchanged).
3. TheRock / site-packages — Scan sys.path for dirs named _rocm_sdk* or rocm_sdk*; treat as ROCm root if they contain bin/hipcc or hip/bin/hipcc (or bin/hipconfig). Use first valid candidate.
4. Fallback — /opt/rocm if it exists (unchanged).

---

Implementation phases

Phase 1 — Central discovery (Python)

• In cpp_extension.py: add _find_rocm_home_in_site_packages() and call it from _find_rocm_home() before the /opt/rocm fallback.
• Extend executable_path() to also check ROCM_HOME/llvm/bin when not found under ROCM_HOME/bin.
• (Optional) Document or add a small helper so scripts can get the root, e.g. python -c "from aiter.jit.utils.cpp_extension import ROCM_HOME; print(ROCM_HOME or '')".

Phase 2 — Remove hardcoding in csrc

• In csrc/cpp_itfs/utils.py: get amdgpu-arch and hipconfig via the shared discovery (e.g. executable_path(...) from cpp_extension, or a thin _rocm_root helper in csrc if import cycles are a problem). No more /opt/rocm paths.

Phase 3 — Shell / Makefile

• build_mha.sh: use HIPCC="${ROCM_HOME:-/opt/rocm}/bin/hipcc" or resolve ROCM_HOME via Python when unset.
• Solver Makefile: ROCM_HOME ?= /opt/rocm, then ROCBLAS_DIR ?= $(ROCM_HOME), ROCSOLVER_DIR ?= $(ROCM_HOME).
• Pre-commit (optional): use ROCM_HOME for PATH instead of hardcoded /opt/rocm.

Phase 4 — Tests & docs

• Add a test (or CI check) that discovery works with a mock site-packages layout; keep a traditional /opt/rocm path test to avoid regressions.
• Document in CONTRIBUTE.md/README: TheRock users get auto-discovery from site-packages; setting ROCM_HOME/ROCM_PATH overrides.

---

Risks / notes

• If TheRock changes directory names under site-packages, we may need to adjust the scan; keying on presence of bin/hipcc (or hip/bin/hipcc) keeps it robust.
• If both /opt/rocm and a site-packages ROCm exist, we take the first valid root in the order above; users can force one via ROCM_HOME.
• If importing cpp_extension from csrc/cpp_itfs/utils.py causes circular deps, use a minimal discovery module in csrc that mirrors the same order.

---

Once these phases are done, aiter will discover and use ROCm from TheRock/site-packages while remaining compatible with /opt/rocm.

[astrelsky]

Please see ROCm/TheRock#3976 for guidance on this.

[zufayu]

Minimal fix PR: #2605

Replaces the two remaining hardcoded /opt/rocm paths in aiter's own code (excluding 3rdparty/composable_kernel) with ROCM_HOME/ROCM_PATH env var lookup, falling back to /opt/rocm when unset.

What changed (2 files, ~10 lines):

• csrc/cpp_itfs/utils.py — amdgpu-arch and hipconfig paths now use ROCM_HOME
• csrc/kernels/solver/Makefile — ROCBLAS_DIR/ROCSOLVER_DIR use ?= with ROCM_HOME

For TheRock users:

export ROCM_HOME=$(rocm-sdk path --root)

Note: aiter/jit/utils/cpp_extension.py already supports ROCM_HOME/ROCM_PATH — no changes needed there.

This follows TheRock's recommended pattern (see ROCm/TheRock#3976) — env vars as user-facing convenience, with tool commands (rocm-sdk path --root) to discover the path.

[astrelsky]

Minimal fix PR: #2605

Replaces the two remaining hardcoded /opt/rocm paths in aiter's own code (excluding 3rdparty/composable_kernel) with ROCM_HOME/ROCM_PATH env var lookup, falling back to /opt/rocm when unset.

What changed (2 files, ~10 lines):

• csrc/cpp_itfs/utils.py — amdgpu-arch and hipconfig paths now use ROCM_HOME
• csrc/kernels/solver/Makefile — ROCBLAS_DIR/ROCSOLVER_DIR use ?= with ROCM_HOME

For TheRock users:

export ROCM_HOME=$(rocm-sdk path --root)

Note: aiter/jit/utils/cpp_extension.py already supports ROCM_HOME/ROCM_PATH — no changes needed there.

This follows TheRock's recommended pattern (see ROCm/TheRock#3976) — env vars as user-facing convenience, with tool commands (rocm-sdk path --root) to discover the path.

NO! This is exactly what the recommendations say not to do.

It should use the command line tool rocm-sdk to determine the path and that should take precedence over the environment variable. The environment variables ROCM_HOME or ROCM_PATH shall not need to be set.