setup.py: ROCm build fixes

pytorch · May 10, 2023 · f2ab2f7 · f2ab2f7
1 parent 078959f
commit f2ab2f7
Showing 1 changed file with 7 additions and 1 deletion.
diff --git 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(