Triaging ROCm wheel build (#2161)

* Enable ROCm support in build workflow and specify runner configuration for MI300 GPU

* Refactor HIP source directory handling in setup.py and remove deprecated runner configuration from build workflow

* Refactor HIP source collection in setup.py for improved readability

Author: petrex

.github/workflows/build_wheels_linux.yml

@@ -28,7 +28,7 @@ jobs:
       os: linux
       with-cpu: enable
       with-cuda: enable
-      with-rocm: disable
+      with-rocm: enable
       with-xpu: enable
       # Note: if free-threaded python is required add py3.13t here
       python-versions: '["3.9"]'

setup.py

@@ -311,16 +311,17 @@ def get_extensions():
         glob.glob(os.path.join(extensions_cuda_dir, "**/*.cu"), recursive=True)
     )
 
-    extensions_hip_dir = os.path.join(
-        extensions_dir, "cuda", "tensor_core_tiled_layout"
-    )
-    hip_sources = list(
-        glob.glob(os.path.join(extensions_hip_dir, "*.cu"), recursive=True)
-    )
-    extensions_hip_dir = os.path.join(extensions_dir, "cuda", "sparse_marlin")
-    hip_sources += list(
-        glob.glob(os.path.join(extensions_hip_dir, "*.cu"), recursive=True)
-    )
+    # Define HIP source directories
+    hip_source_dirs = [
+        os.path.join(extensions_dir, "cuda", "tensor_core_tiled_layout"),
+        # TODO: Add sparse_marlin back in once we have a ROCm build for it
+        # os.path.join(extensions_dir, "cuda", "sparse_marlin")
+    ]
+
+    # Collect all HIP sources from the defined directories
+    hip_sources = []
+    for hip_dir in hip_source_dirs:
+        hip_sources.extend(glob.glob(os.path.join(hip_dir, "*.cu"), recursive=True))
 
     # Collect CUDA source files if needed
     if not IS_ROCM and use_cuda: