Install rocm-develop packages during torch testing

.github/workflows/test_pytorch_wheels.yml

@@ -133,12 +133,22 @@ jobs:
 
       - name: Set up virtual environment
         run: |
+          TORCH_VERSION="${{ inputs.torch_version }}"
+          ROCM_VERSION="${TORCH_VERSION#*+rocm}"
           python build_tools/setup_venv.py ${VENV_DIR} \
-            --packages torch==${{ inputs.torch_version }} \
+            --packages "torch==${TORCH_VERSION} rocm[devel]==${ROCM_VERSION}" \
             --index-url=${{ inputs.package_index_url }} \
             --index-subdir=${{ inputs.amdgpu_family }} \
             --activate-in-future-github-actions-steps
 
+      # rocm[devel] provides development headers needed by tests that compile
+      # C++ extensions (e.g. dummy_allocator). hipconfig --rocmpath points to
+      # the devel package root so PyTorch's _find_rocm_home() resolves
+      # correctly.
+      - name: Set ROCM_PATH
+        run: |
+          echo "ROCM_PATH=$(hipconfig --rocmpath)" >> $GITHUB_ENV
+
       - name: Install test requirements
         run: |
           python -m pip install -r external-builds/pytorch/requirements-test.txt

external-builds/pytorch/README.md

@@ -232,6 +232,16 @@ python run_pytorch_smoke_tests.py -- \
 
 ### Running full PyTorch tests
 
+> [!WARNING]
+> Our CI currently tests with `rocm[devel]` installed alongside `torch` because
+> some tests require development headers at runtime. If you are running tests
+> locally, install the devel packages and set `ROCM_PATH` accordingly:
+>
+> ```bash
+> pip install "rocm[devel]" --index-url <your-index-url>
+> export ROCM_PATH=$(hipconfig --rocmpath)
+> ```
+
 We have a [`run_pytorch_tests.py`](run_pytorch_tests.py) script
 which runs PyTorch unit tests using pytest with additional test exclusion
 capabilities tailored for AMD ROCm GPUs. See the script for detailed

external-builds/pytorch/skip_tests/generic.py

@@ -24,15 +24,6 @@
             # 'num_host_free': 0, 'reserved_bytes.allocated': 0, 'reserved_bytes.current': 0, 'reserved_bytes.freed': 0,
             # 'reserved_bytes.peak': 0, 'segment.allocated': 0, 'segment.current': 0, 'segment.freed': 0, 'segment.peak': 0})
             "test_host_memory_stats",
-            # THIS IS AN OLD ERROR
-            # In file included from /home/tester/.cache/torch_extensions/py312_cpu/dummy_allocator/main_hip.cpp:5:
-            # /home/tester/TheRock/.venv/lib/python3.12/site-packages/torch/include/ATen/hip/Exceptions.h:4:10: fatal error: hipblas/hipblas.h: No such file or directory
-            #     4 | #include <hipblas/hipblas.h>
-            #     |          ^~~~~~~~~~~~~~~~~~~
-            # compilation terminated.
-            # NEW ERROR
-            # RuntimeError: Error building extension 'dummy_allocator'
-            "test_mempool_with_allocator",
             # Change detector test (Cublaslt vs Cublas depending on gcn_arch and torch version)
             # Always skip as this test is very basic and needs manual intervention for new architectures
             # See

external-builds/pytorch/skip_tests/pytorch_2.10.py

@@ -29,10 +29,6 @@
             "test_memory_snapshot_with_cpp",
             # AssertionError: Scalars are not equal!
             "test_mempool_ctx_multithread",
-            # RuntimeError: Error building extension 'dummy_allocator'
-            "test_mempool_empty_cache_inactive",
-            # RuntimeError: Error building extension 'dummy_allocator_v1'
-            "test_mempool_limited_memory_with_allocator",
             # OSError: libhiprtc.so: cannot open shared object file: No such file or directory
             # File "/home/tester/TheRock/.venv/lib/python3.12/site-packages/torch/cuda/_utils.py", line 57, in _get_hiprtc_library
             # lib = ctypes.CDLL("libhiprtc.so")

external-builds/pytorch/skip_tests/pytorch_2.11.py

@@ -23,15 +23,6 @@
             "test_memory_snapshot_with_cpp",
             # AssertionError: Scalars are not equal!
             "test_mempool_ctx_multithread",
-            # RuntimeError: Error building extension 'dummy_allocator'
-            "test_mempool_empty_cache_inactive",
-            # RuntimeError: Error building extension 'dummy_allocator_v1'
-            "test_mempool_limited_memory_with_allocator",
-            # new for pytorch 2.11
-            # RuntimeError: Error building extension 'dummy_allocator_v3'
-            "test_tensor_delete_after_allocator_delete",
-            # RuntimeError: Error building extension 'dummy_allocator'
-            "test_deleted_mempool_not_used_on_oom",
             # ModuleNotFoundError: No module named 'torchvision'
             "test_resnet",
             # RuntimeError: miopenStatusUnknownError

external-builds/pytorch/skip_tests/pytorch_2.8.py

@@ -23,10 +23,6 @@
             "test_memory_snapshot_script",
             "test_memory_snapshot_with_cpp",
             "test_mempool_ctx_multithread",
-            # RuntimeError: Error building extension 'dummy_allocator'
-            "test_mempool_empty_cache_inactive",
-            # RuntimeError: Error building extension 'dummy_allocator_v1'
-            "test_mempool_limited_memory_with_allocator",
             # This test was fixed in torch 2.9, see
             # https://github.com/ROCm/TheRock/issues/2206
             "test_hip_device_count",

external-builds/pytorch/skip_tests/pytorch_2.9.py

@@ -33,10 +33,6 @@
             "test_memory_snapshot_with_cpp",
             # AssertionError: Scalars are not equal!
             "test_mempool_ctx_multithread",
-            # RuntimeError: Error building extension 'dummy_allocator'
-            "test_mempool_empty_cache_inactive",
-            # RuntimeError: Error building extension 'dummy_allocator_v1'
-            "test_mempool_limited_memory_with_allocator",
             # for whatever reason these are also flaky: if run standalone they pass?
             # AttributeError: module 'torch.backends.cudnn.rnn' has no attribute 'fp32_precision'
             "test_fp32_precision_with_float32_matmul_precision",