From e389cc5d7202dc8b6ad0d1def30defd7d4e9a6a3 Mon Sep 17 00:00:00 2001 From: "Peter Y. Yeh" Date: Thu, 1 May 2025 14:08:40 -0700 Subject: [PATCH 1/3] Enable ROCm support in build workflow and specify runner configuration for MI300 GPU --- .github/workflows/build_wheels_linux.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build_wheels_linux.yml b/.github/workflows/build_wheels_linux.yml index 7c05cbe8bd..71746c2853 100644 --- a/.github/workflows/build_wheels_linux.yml +++ b/.github/workflows/build_wheels_linux.yml @@ -28,10 +28,11 @@ 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"]' + runner: '{"rocm": "linux.rocm.gpu.mi300.2"}' build: needs: generate-matrix From 31c3c4cc37ac47eba96204ee70e3c9e248dec954 Mon Sep 17 00:00:00 2001 From: "Peter Y. Yeh" Date: Thu, 8 May 2025 10:54:37 -0700 Subject: [PATCH 2/3] Refactor HIP source directory handling in setup.py and remove deprecated runner configuration from build workflow --- .github/workflows/build_wheels_linux.yml | 1 - setup.py | 23 +++++++++++++---------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/.github/workflows/build_wheels_linux.yml b/.github/workflows/build_wheels_linux.yml index 71746c2853..a8d96abc8a 100644 --- a/.github/workflows/build_wheels_linux.yml +++ b/.github/workflows/build_wheels_linux.yml @@ -32,7 +32,6 @@ jobs: with-xpu: enable # Note: if free-threaded python is required add py3.13t here python-versions: '["3.9"]' - runner: '{"rocm": "linux.rocm.gpu.mi300.2"}' build: needs: generate-matrix diff --git a/setup.py b/setup.py index 7e60acbfa8..72fc4a6810 100644 --- a/setup.py +++ b/setup.py @@ -311,16 +311,19 @@ 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: From f9f2d97f589f43740c4fe1ac5c8b15a86b42d918 Mon Sep 17 00:00:00 2001 From: "Peter Y. Yeh" Date: Thu, 8 May 2025 10:56:24 -0700 Subject: [PATCH 3/3] Refactor HIP source collection in setup.py for improved readability --- setup.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index 72fc4a6810..a269d4410d 100644 --- a/setup.py +++ b/setup.py @@ -315,15 +315,13 @@ def get_extensions(): 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") + # 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) - ) + 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: