setup.py: fix ROCm build (#7573)

20d90dfc · Paul Mulders · GitHub · 078959f5 · 20d90dfc
Unverified Commit 20d90dfc authored May 11, 2023 by Paul Mulders Committed by GitHub May 10, 2023
Show whitespace changes
Inline Side-by-side

Showing with 7 additions and 1 deletion

setup.py setup.py +7 -1

No files found.
--- a/setup.py
+++ b/setup.py
@@ -328,9 +328,15 @@ def get_extensions():
    image_src = (
        glob.glob(os.path.join(image_path, "*.cpp"))
        + glob.glob(os.path.join(image_path, "cpu", "*.cpp"))
-        + glob.glob(os.path.join(image_path, "cuda", "*.cpp"))
    )

+    if is_rocm_pytorch:
+        image_src += glob.glob(os.path.join(image_path, "hip", "*.cpp"))
+        # we need to exclude this in favor of the hipified source
+        image_src.remove(os.path.join(image_path, "image.cpp"))
+    else:
+        image_src += glob.glob(os.path.join(image_path, "cuda", "*.cpp"))
+
    if use_png or use_jpeg:
        ext_modules.append(
            extension(