[Easy] Add ROCm support to nightly pull tool (#141282)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/141282
Approved by: https://github.com/malfet
ghstack dependencies: #143263

.github/scripts/generate_binary_build_matrix.py

@@ -20,6 +20,7 @@ CUDA_ARCHES = ["11.8", "12.4", "12.6"]
 CUDA_ARCHES_FULL_VERSION = {"11.8": "11.8.0", "12.4": "12.4.1", "12.6": "12.6.3"}
 CUDA_ARCHES_CUDNN_VERSION = {"11.8": "9", "12.4": "9", "12.6": "9"}
 
+# NOTE: Also update the ROCm sources in tools/nightly.py when changing this list
 ROCM_ARCHES = ["6.2.4", "6.3"]
 
 XPU_ARCHES = ["xpu"]

CONTRIBUTING.md

@@ -78,7 +78,9 @@ git clone git@github.com:<USERNAME>/pytorch.git
 cd pytorch
 git remote add upstream git@github.com:pytorch/pytorch.git
 
-make setup-env  # or make setup-env-cuda for pre-built CUDA binaries
+make setup-env
+# Or run `make setup-env-cuda` for pre-built CUDA binaries
+# Or run `make setup-env-rocm` for pre-built ROCm binaries
 source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
 ```
 
@@ -193,6 +195,13 @@ To install the nightly binaries built with CUDA, you can pass in the flag `--cud
 source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
 ```
 
+To install the nightly binaries built with ROCm, you can pass in the flag `--rocm`:
+
+```bash
+./tools/nightly.py checkout -b my-nightly-branch --rocm
+source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
+```
+
 You can also use this tool to pull the nightly commits into the current branch:
 
 ```bash

Makefile

@@ -35,8 +35,12 @@ setup-env: ensure-branch-clean
 setup-env-cuda:
 	$(MAKE) setup-env PYTHON="$(PYTHON)" NIGHTLY_TOOL_OPTS="$(NIGHTLY_TOOL_OPTS) --cuda"
 
+setup-env-rocm:
+	$(MAKE) setup-env PYTHON="$(PYTHON)" NIGHTLY_TOOL_OPTS="$(NIGHTLY_TOOL_OPTS) --rocm"
+
 setup_env: setup-env
 setup_env_cuda: setup-env-cuda
+setup_env_rocm: setup-env-rocm
 
 setup-lint:
 	$(PIP) install lintrunner

tools/nightly.py

@@ -20,6 +20,11 @@ To install the nightly binaries built with CUDA, you can pass in the flag --cuda
     $ ./tools/nightly.py checkout -b my-nightly-branch --cuda
     $ source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
 
+To install the nightly binaries built with ROCm, you can pass in the flag --rocm::
+
+    $ ./tools/nightly.py checkout -b my-nightly-branch --rocm
+    $ source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
+
 You can also use this tool to pull the nightly commits into the current branch as
 well. This can be done with::
 
@@ -134,6 +139,12 @@ PIP_SOURCES = {
         supported_platforms={"Linux", "Windows"},
         accelerator="cuda",
     ),
+    "rocm-6.2.4": PipSource(
+        name="rocm-6.2.4",
+        index_url=f"{PYTORCH_NIGHTLY_PIP_INDEX_URL}/rocm6.2.4",
+        supported_platforms={"Linux"},
+        accelerator="rocm",
+    ),
 }
 
 
@@ -882,6 +893,17 @@ def make_parser() -> argparse.ArgumentParser:
             default=argparse.SUPPRESS,
             metavar="VERSION",
         )
+        subparser.add_argument(
+            "--rocm",
+            help=(
+                "ROCm version to install "
+                "(defaults to the latest version available on the platform)"
+            ),
+            dest="rocm",
+            nargs="?",
+            default=argparse.SUPPRESS,
+            metavar="VERSION",
+        )
     return parser
 
 
@@ -889,6 +911,8 @@ def parse_arguments() -> argparse.Namespace:
     parser = make_parser()
     args = parser.parse_args()
     args.branch = getattr(args, "branch", None)
+    if hasattr(args, "cuda") and hasattr(args, "rocm"):
+        parser.error("Cannot specify both CUDA and ROCm versions.")
     return args
 
 
@@ -901,26 +925,32 @@ def main() -> None:
         sys.exit(status)
 
     pip_source = None
-    if hasattr(args, "cuda"):
+
+    for toolkit in ("CUDA", "ROCm"):
+        accel = toolkit.lower()
+        if hasattr(args, accel):
+            requested = getattr(args, accel)
             available_sources = {
-            src.name[len("cuda-") :]: src
+                src.name[len(f"{accel}-") :]: src
                 for src in PIP_SOURCES.values()
-            if src.name.startswith("cuda-") and PLATFORM in src.supported_platforms
+                if src.name.startswith(f"{accel}-")
+                and PLATFORM in src.supported_platforms
             }
             if not available_sources:
-            print(f"No CUDA versions available on platform {PLATFORM}.")
+                print(f"No {toolkit} versions available on platform {PLATFORM}.")
                 sys.exit(1)
-        if args.cuda is not None:
+            if requested is not None:
-            pip_source = available_sources.get(args.cuda)
+                pip_source = available_sources.get(requested)
                 if pip_source is None:
                     print(
-                    f"CUDA {args.cuda} is not available on platform {PLATFORM}. "
+                        f"{toolkit} {requested} is not available on platform {PLATFORM}. "
                         f"Available version(s): {', '.join(sorted(available_sources, key=Version))}"
                     )
                     sys.exit(1)
             else:
                 pip_source = available_sources[max(available_sources, key=Version)]
-    else:
+
+    if pip_source is None:
         pip_source = PIP_SOURCES["cpu"]  # always available
 
     with logging_manager(debug=args.verbose) as logger: