[RELEASE ONLY CHANGES] Use the packaged ROCm development SDK

Install the development files matching the torch-provided ROCm runtime and point TorchAO and ExecuTorch builds at its expanded SDK root, where the HIP linker library is available.

Assisted by Codex.

Author: JacobSzwejbka

.ci/scripts/install-rocm-dependencies.sh

@@ -18,6 +18,15 @@ export LD_LIBRARY_PATH="${ROCM_PYTHON_PREFIX}/lib:${LD_LIBRARY_PATH:-}"
 TORCH_VERSION="$(python -c 'from torch_pin import TORCH_VERSION; print(TORCH_VERSION)')"
 python -m pip install "torch==${TORCH_VERSION}+rocm${ROCM_VERSION}" \
   --index-url "${PYTORCH_ROCM_INDEX}"
+ROCM_SDK_VERSION="$(python -c 'import rocm_sdk; print(rocm_sdk.__version__)')"
+python -m pip install "rocm[devel]==${ROCM_SDK_VERSION}" \
+  --index-url "${PYTORCH_ROCM_INDEX}"
+ROCM_PATH="$(python -m rocm_sdk path --root)"
+export ROCM_PATH
+export ROCM_HOME="${ROCM_PATH}"
+export HIP_PATH="${ROCM_PATH}"
+export PATH="${ROCM_PATH}/bin:${PATH}"
+export LD_LIBRARY_PATH="${ROCM_PYTHON_PREFIX}/lib:${ROCM_PATH}/lib:${ROCM_PATH}/lib64:${LD_LIBRARY_PATH:-}"
 python -m pip install -r requirements-dev.txt ninja
 # ROCm 7.14 TorchAO wheels are Python-only. Build the pinned extension against
 # the installed torch and HIP libraries instead of reusing a ROCm 7.2 binary.

.ci/scripts/tests/test_rocm_dependencies.py

@@ -19,7 +19,11 @@ def test_compiled_torchao_uses_the_selected_torch_and_propagates_failures(self):
 conda() { echo RUNTIME; }
 python() {
   if [[ "$1" == "-c" ]]; then
-    if [[ "$2" == *torch_pin* ]]; then echo 2.15.0; else echo /env; fi
+    if [[ "$2" == *torch_pin* ]]; then echo 2.15.0
+    elif [[ "$2" == *rocm_sdk* ]]; then echo 7.14.1
+    else echo /env; fi
+  elif [[ "$*" == "-m rocm_sdk path --root" ]]; then
+    echo /sdk
   elif [[ "$1" == "-" ]]; then
     cat >/dev/null
     echo VERIFY
@@ -28,29 +32,34 @@ def test_compiled_torchao_uses_the_selected_torch_and_propagates_failures(self):
     echo "PIP $* CPP=${USE_CPP:-unset} CPU=${USE_CPU_KERNELS:-unset} JOBS=${MAX_JOBS:-unset}"
     if [[ "$*" == *third-party/ao* ]]; then return "$AO_STATUS"; fi
     if [[ "$*" == *torch==* ]]; then return "$TORCH_STATUS"; fi
+    if [[ "$*" == *'rocm[devel]'* ]]; then return "$SDK_STATUS"; fi
   fi
 }
 """
-        for torch_status, ao_status, verify_status in (
-            (0, 0, 0),
-            (1, 0, 0),
-            (0, 1, 0),
-            (0, 0, 1),
+        for torch_status, sdk_status, ao_status, verify_status in (
+            (0, 0, 0, 0),
+            (1, 0, 0, 0),
+            (0, 1, 0, 0),
+            (0, 0, 1, 0),
+            (0, 0, 0, 1),
         ):
-            with self.subTest(torch=torch_status, ao=ao_status, verify=verify_status):
+            with self.subTest(
+                torch=torch_status, sdk=sdk_status, ao=ao_status, verify=verify_status
+            ):
                 result = subprocess.run(
                     [
                         "bash",
                         "-c",
                         stubs
-                        + "source .ci/scripts/install-rocm-dependencies.sh\necho INSTALL_ET",
+                        + "source .ci/scripts/install-rocm-dependencies.sh\necho INSTALL_ET\necho SDK_HOME=$ROCM_HOME HIP_PATH=$HIP_PATH ROCM_PATH=$ROCM_PATH",
                     ],
                     cwd=ROOT,
                     env={
                         **os.environ,
                         "ROCM_VERSION": "7.14",
                         "PYTORCH_ROCM_INDEX": "https://download.pytorch.org/whl/rocm7.14",
                         "TORCH_STATUS": str(torch_status),
+                        "SDK_STATUS": str(sdk_status),
                         "AO_STATUS": str(ao_status),
                         "VERIFY_STATUS": str(verify_status),
                         "MAX_JOBS": "8",
@@ -60,7 +69,7 @@ def test_compiled_torchao_uses_the_selected_torch_and_propagates_failures(self):
                 )
                 self.assertEqual(
                     result.returncode,
-                    max(torch_status, ao_status, verify_status),
+                    max(torch_status, sdk_status, ao_status, verify_status),
                     result.stderr,
                 )
                 self.assertTrue(result.stdout.startswith("RUNTIME\n"))
@@ -69,16 +78,26 @@ def test_compiled_torchao_uses_the_selected_torch_and_propagates_failures(self):
                     result.stdout,
                 )
                 if not torch_status:
+                    self.assertIn(
+                        "rocm[devel]==7.14.1 --index-url https://download.pytorch.org/whl/rocm7.14",
+                        result.stdout,
+                    )
+                if not (torch_status or sdk_status):
                     self.assertIn(
                         "--no-build-isolation --no-deps third-party/ao CPP=1 CPU=1 JOBS=8",
                         result.stdout,
                     )
                 else:
                     self.assertNotIn("third-party/ao", result.stdout)
                 self.assertEqual(
-                    "VERIFY" in result.stdout, not (torch_status or ao_status)
+                    "VERIFY" in result.stdout,
+                    not (torch_status or sdk_status or ao_status),
                 )
                 self.assertEqual("INSTALL_ET" in result.stdout, result.returncode == 0)
+                if result.returncode == 0:
+                    self.assertIn(
+                        "SDK_HOME=/sdk HIP_PATH=/sdk ROCM_PATH=/sdk", result.stdout
+                    )
 
 
 if __name__ == "__main__":