Resolve accelerator-vendor dependency variants from pyproject.toml

Dockerfile.rocm

@@ -79,33 +79,34 @@ COPY LICENSE CHANGELOG.md CODE_OF_CONDUCT.md CONTRIBUTING.md README.md versionee
 COPY tests ./tests
 COPY monai ./monai
 
-# Use print_dependencies.py rather than -e .[all,testing] to filter CUDA-only packages:
-#   cucim-cu* pulls in cuda-toolkit (~1.2 GB); nvidia-ml-py fails at import on ROCm; nni depends on it.
+# Set HIPCIM_INDEX_URL="" to build without WSI/cucim support.
+ARG HIPCIM_INDEX_URL="https://pypi.amd.com/rocm-${ROCM_SERIES}.0/simple/"
+
+# print_dependencies.py resolves the ROCm variants of the GPU imaging stack (amd-hipcim, amd-cupy)
+# and drops the NVIDIA-only packages, see monai/config/vendor_deps.py.
 # BUILD_MONAI=1 builds the C++/HIP extensions ahead of time; FORCE_CUDA=1 is required because the
 # build host has no GPU, so setup.py's `torch.cuda.is_available()` check would otherwise skip them.
 # Compilation itself needs only the toolkit and PYTORCH_ROCM_ARCH (set above), not a device.
 # pytest is not in the "testing" extra; it is added explicitly for running tests by hand.
 RUN python monai/config/print_dependencies.py build-system \
       | xargs -d '\n' pip install --no-cache-dir --no-build-isolation \
-  && python monai/config/print_dependencies.py all testing \
-      | grep -vE '^cucim-cu|^nvidia-ml-py|^nni' > /tmp/rocm-requirements-$$.txt \
+  && python monai/config/print_dependencies.py all testing > /tmp/rocm-requirements-$$.txt \
+  && if [ -z "${HIPCIM_INDEX_URL}" ]; then sed -i '/^amd-hipcim/d' /tmp/rocm-requirements-$$.txt; fi \
   && BUILD_MONAI=1 FORCE_CUDA=1 pip install --no-cache-dir --no-build-isolation \
+      ${HIPCIM_INDEX_URL:+--extra-index-url "${HIPCIM_INDEX_URL}"} \
       -r /tmp/rocm-requirements-$$.txt pytest -e . \
   && rm -f /tmp/rocm-requirements-$$.txt
 
 # Required at runtime too: monai.config.deviceconfig gates USE_COMPILED on this variable, so
 # without it the extensions built above would be present but never used.
 ENV BUILD_MONAI=1
 
-# Set HIPCIM_INDEX_URL="" to build without WSI/cucim support.
 # CuImage is imported (not just cucim) because cucim uses lazy_loader -- a bare import
 # succeeds even when the native library is unresolvable. Failing here is deliberate: if
 # hipCIM was requested, an image where the cucim backends silently do not work is worse
 # than no image at all.
-ARG HIPCIM_INDEX_URL="https://pypi.amd.com/rocm-${ROCM_SERIES}.0/simple/"
 RUN if [ -n "${HIPCIM_INDEX_URL}" ]; then \
-      pip install --no-cache-dir --extra-index-url "${HIPCIM_INDEX_URL}" "amd-hipcim" \
-      && python -c "from cucim import CuImage"; \
+      python -c "from cucim import CuImage"; \
     else \
       echo "hipCIM not installed; whole-slide-image (cucim) backends are unavailable."; \
     fi

docs/source/installation.md

@@ -89,6 +89,62 @@ The `nvimgcodec` extra installs GPU-accelerated DICOM decoding for `NvImgCodecPy
 (`pip install 'monai[nvimgcodec]'`). It is Linux-only in the extra definition; CUDA 13 is the
 default. CUDA 12 users should install matching `cupy-cuda12x` and `nvidia-nvimgcodec-cu12` wheels.
 
+On AMD ROCm the extras above keep their names. When MONAI is built or installed against a ROCm
+build of PyTorch, the CUDA-only distributions they pull in are replaced by their AMD counterparts:
+`cucim-cu12`/`cucim-cu13` become [`amd-hipcim`](https://rocm.docs.amd.com/projects/hipCIM/en/latest/)
+(hipCIM), `cupy-cuda*` becomes `amd-cupy`, and `nvidia-ml-py`, `nni` and `nvidia-nvimgcodec-cu*` are
+dropped. `amd-hipcim` ships the `cucim` Python namespace and `amd-cupy` ships `cupy`, so all existing
+MONAI code importing from either (e.g. `WSIReader` with `backend="cucim"`, or `convert_to_cupy`)
+works on ROCm without any code changes. The required dependencies additionally gain the AMD GPU
+device extras (`torch[device-gfx942,device-gfx950]`) and a `rocm[libraries,devel,device-gfx*]`
+requirement, which is what steers pip towards the ROCm build of PyTorch. `devel` is included because
+`monai/_extensions` JIT-compiles its HIP sources on first use and the link step needs the ROCm
+development tree.
+
+This is driven by a general accelerator-vendor mechanism: `monai/config/vendor_deps.py` holds the
+generic rewriting and a registry of vendors, and each vendor contributes a `vendor_<name>.py` plugin
+listing its package substitutions. A plugin is imported only once its registry probe matches, so a
+build for one vendor never runs another vendor's code, and a build with no vendor detected is left
+exactly as declared in `pyproject.toml`. Set `MONAI_VENDOR` to a registered vendor name to force one,
+or to `none` to disable rewriting. For ROCm, the GPU architectures come from `GPU_TARGETS` or
+`AMDGPU_TARGETS` and the ROCm series from `MONAI_ROCM_SERIES`.
+
+Two AMD package indexes are involved, and neither is on PyPI. The ROCm SDK (`rocm`) and the ROCm
+build of PyTorch come from the ROCm wheel index. `amd-hipcim` and `amd-cupy` come from the AMD
+extensions index, which is published per ROCm series — at the time of writing, `rocm-10.0.0`. Those
+wheels are supported on the series they were built for and the one after it, which is the range the
+generated `rocm` requirement allows.
+
+Install the ROCm runtime and PyTorch first, selecting the architecture of your GPU (`gfx942` for
+MI300X/MI325X, `gfx950` for MI350X/MI355X), then expand the development tree:
+
+```bash
+pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ \
+    "rocm[libraries,devel,device-gfx942]" "torch[device-gfx942]"
+rocm-sdk init
+```
+
+`rocm-sdk init` is required: the `rocm-sdk-devel` wheel ships its contents as an archive that this
+command expands. Without it the ROCm development tree is absent, and `monai/_extensions` fails to
+JIT-compile on first use because the link step cannot find `libamdhip64.so`.
+
+Then install MONAI with both indexes available, so the `rocm` requirement and the AMD imaging
+packages can both resolve:
+
+```bash
+pip install \
+    --extra-index-url https://stable.repo.amd.com/rocm/whl-next/ \
+    --extra-index-url https://pypi.amd.com/rocm-10.0.0/simple/ \
+    'monai[cucim]'
+```
+
+AMD currently publishes the imaging wheels for CPython 3.12 on x86_64 Linux only; on other
+interpreters or architectures the install fails with `No matching distribution found for
+amd-hipcim`.
+
+`Dockerfile.rocm` in the repo root is the authoritative reference for the complete environment setup
+(library paths, compiler flags, `rocm-sdk init`, and the full container recipe).
+
 The `hyena` extra pulls in [`nvsubquadratic`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic),
 required by `HyenaNDUNETR` / `HyenaMixer` / `HyenaTransformerBlock` (subquadratic
 O(N log N) alternatives to windowed self-attention). Install with
@@ -333,7 +389,7 @@ MONAI itself:
 ```bash
 git clone https://github.com/Project-MONAI/MONAI.git
 cd MONAI/
-python monai/config/print_dependencies.py \* > requirements.txt
+python monai/config/print_dependencies.py all testing > requirements.txt
 pip install -r requirements.txt
 ```
 

monai/__init__.py

@@ -94,6 +94,9 @@ def filter(self, record):
         "(.*(__main__)$)",
         "(.*(video_dataset)$)",
         "(.*(nnunet).*$)",
+        # packaging-time only, and a vendor plugin must not be imported on another vendor's hardware
+        "(^(monai.config.vendor_))",
+        "(^(monai.config.print_dependencies)$)",
     ]
 )
 

monai/config/check_env.py

@@ -123,7 +123,11 @@ def check_torch_cuda():
     """
     import torch
 
-    fprint("CUDA version:", torch.version.cuda)
+    # `torch.version.cuda` is None on a ROCm build, where the toolkit version is `torch.version.hip`.
+    if torch.version.hip:
+        fprint("HIP version:", torch.version.hip)
+    else:
+        fprint("CUDA version:", torch.version.cuda)
 
     try:
         dcount = torch.cuda.device_count()

monai/config/deviceconfig.py

@@ -214,7 +214,12 @@ def get_gpu_info() -> OrderedDict:
     _dict_append(output, "Has CUDA", lambda: bool(torch.cuda.is_available()))
 
     if output["Has CUDA"]:
-        _dict_append(output, "CUDA version", lambda: torch.version.cuda)
+        # On a ROCm build `torch.version.cuda` is None and `torch.version.hip` carries the toolkit
+        # version. Reporting "CUDA version: None" is misleading, so the key name switches with the build.
+        if torch.version.hip:
+            _dict_append(output, "HIP version", lambda: torch.version.hip)
+        else:
+            _dict_append(output, "CUDA version", lambda: torch.version.cuda)
     cudnn_ver = torch.backends.cudnn.version()
     _dict_append(output, "cuDNN enabled", lambda: bool(cudnn_ver))
     _dict_append(output, "NVIDIA_TF32_OVERRIDE", os.environ.get("NVIDIA_TF32_OVERRIDE"))

monai/config/print_dependencies.py

@@ -14,13 +14,21 @@
 be piped to a requirements file to work with pip. All required dependencies are always printed, those for builing are
 included in "build-system" is given as an argument, and all optional requirements are included if "*" is given. This
 assumes the pyproject.toml file is in the current working directory.
+
+On an accelerator-vendor PyTorch the NVIDIA-only distributions are swapped for that vendor's
+equivalents, see ``monai.config.vendor_deps``.
 """
 
 from __future__ import annotations
 
 import sys
 from collections.abc import Collection
 
+try:
+    from .vendor_deps import active_vendor
+except ImportError:  # run as a script rather than imported from the package
+    from vendor_deps import active_vendor  # type: ignore[no-redef]
+
 BUILD_SYSTEM_KEY = "build-system"
 PROJ_KEY = "project"
 OPTS_KEY = "optional-dependencies"
@@ -57,10 +65,16 @@ def parse_dependencies(filename: str | None = None, sections: Collection[str] |
     opts = proj[OPTS_KEY]
     dependencies = list(proj[DEP_KEY])
     sections = set(sections or [])
+    vendor = active_vendor()
+
+    if vendor is not None:
+        dependencies = vendor.apply_to_dependencies(dependencies)
+        opts = vendor.apply_to_optional_dependencies(opts)
 
     if BUILD_SYSTEM_KEY in sections:
         sections.remove(BUILD_SYSTEM_KEY)
-        dependencies += data[BUILD_SYSTEM_KEY][REQ_KEY]
+        build_requires = data[BUILD_SYSTEM_KEY][REQ_KEY]
+        dependencies += vendor.apply_to_dependencies(build_requires) if vendor else build_requires
 
     if "*" in sections:
         dependencies += sum(opts.values(), [])