From 514b44fca8542bafe8640adcbf1cccce1900f74c Mon Sep 17 00:00:00 2001
From: Michael Norris <mnorris@meta.com>
Date: Fri, 15 Aug 2025 12:06:27 -0700
Subject: [PATCH] OSS changes: Enable ROCm to work with Faiss on BUCK (#4485)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4485

hipify.sh currently generates .hip files into a separate directory called `gpu-rocm`. The swigfaiss.swig file expects them there too. But to use internal hipification, it needs to be generated into the same directory `gpu`. That means swigfaiss.swig has to change `gpu-rocm` to just `gpu`, same directory as .cu files. If swigfaiss.swig changes, then the hipify.sh has to change as well, otherwise OSS CI breaks.

OSS changes
- Requires changing `swigfaiss.swig`, because the internal BUCK approach builds in the same directory as `.cu` files, aka `gpu/`.
- Because it builds in same `gpu/` directory, we need to change the OSS side `hipify.sh`. After this diff, it just builds in `gpu/`. The rocm CI still passes.

Reviewed By: junjieqi

Differential Revision: D79145713

fbshipit-source-id: a83b805c8a7bf286a7e6cff50c8eb2ddd72f518c
---
 CMakeLists.txt              |   8 +--
 c_api/CMakeLists.txt        |   6 +-
 faiss/gpu/hipify.sh         |  75 ++++----------------
 faiss/python/CMakeLists.txt |  43 ++++-------
 faiss/python/swigfaiss.swig | 137 ++++++++----------------------------
 5 files changed, 59 insertions(+), 210 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index b80a5f4c9..bcd260c4c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -91,11 +91,7 @@ if(FAISS_ENABLE_CUVS AND NOT TARGET cuvs::cuvs)
 add_subdirectory(faiss)
 
 if(FAISS_ENABLE_GPU)
-  if(FAISS_ENABLE_ROCM)
-    add_subdirectory(faiss/gpu-rocm)
-  else()
-    add_subdirectory(faiss/gpu)
-  endif()
+  add_subdirectory(faiss/gpu)
 endif()
 
 if(FAISS_ENABLE_PYTHON)
@@ -119,7 +115,7 @@ if(BUILD_TESTING)
   add_subdirectory(perf_tests)
   if(FAISS_ENABLE_GPU)
     if(FAISS_ENABLE_ROCM)
-      add_subdirectory(faiss/gpu-rocm/test)
+      add_subdirectory(faiss/gpu/test)
     else()
       add_subdirectory(faiss/gpu/test)
     endif()
diff --git a/c_api/CMakeLists.txt b/c_api/CMakeLists.txt
index 07acf2018..aa0eed4fd 100644
--- a/c_api/CMakeLists.txt
+++ b/c_api/CMakeLists.txt
@@ -160,9 +160,5 @@ add_executable(example_c EXCLUDE_FROM_ALL example_c.c)
 target_link_libraries(example_c PRIVATE faiss_c)
 
 if(FAISS_ENABLE_GPU)
-  if(FAISS_ENABLE_ROCM)
-    add_subdirectory(gpu-rocm)
-  else ()
-    add_subdirectory(gpu)
-  endif()
+  add_subdirectory(gpu)
 endif()
diff --git a/faiss/gpu/hipify.sh b/faiss/gpu/hipify.sh
index 9bc4773bb..46916b15c 100755
--- a/faiss/gpu/hipify.sh
+++ b/faiss/gpu/hipify.sh
@@ -39,21 +39,13 @@ function hipify_dir()
         exit
     fi
 
-    # create all destination directories for hipified files into sibling 'gpu-rocm' directory
-    while IFS= read -r -d '' src
-    do
-        dst="${src//gpu-tmp/gpu-rocm}"
-
-        if [ -d $dst ]; then
-            #Clearing out any leftover files and directories
-            echo "Removing old $dst"
-            rm -rf "$dst"
-        fi
-
-        #Making directories
-        echo "Creating $dst"
-        mkdir -p "$dst"
-    done <   <(find ./gpu-tmp -type d -print0)
+    # Create a backup of the original gpu directory
+    if [ -d ./gpu-backup ]; then
+        echo "Removing old ./gpu-backup"
+        rm -rf ./gpu-backup
+    fi
+    echo "Creating backup of original gpu directory"
+    cp -r ./gpu ./gpu-backup
 
     # run hipify-perl against all *.cu *.cuh *.h *.cpp files, no renaming
     # run all files in parallel to speed up
@@ -61,7 +53,8 @@ function hipify_dir()
     do
         while IFS= read -r -d '' src
         do
-            dst="${src//\.\/gpu-tmp/\.\/gpu-rocm}"
+            # Output to the original gpu directory but with .tmp extension
+            dst="${src//\.\/gpu-tmp/\.\/gpu}"
             hipify-perl -o="$dst.tmp" "$src" &
         done <   <(find ./gpu-tmp -name "*.$ext" -print0)
     done
@@ -72,7 +65,7 @@ function hipify_dir()
     do
         dst=${src%.cu.tmp}.hip.tmp
         mv "$src" "$dst"
-    done <   <(find ./gpu-rocm -name "*.cu.tmp" -print0)
+    done <   <(find ./gpu -name "*.cu.tmp" -print0)
 
     if [ -d ./gpu-tmp ]; then
         #Clearing out any leftover files and directories
@@ -80,7 +73,6 @@ function hipify_dir()
         rm -rf ./gpu-tmp
     fi
 
-    # replace header include statements "<faiss/gpu/" with "<faiss/gpu-rocm"
     # replace thrust::cuda::par with thrust::hip::par
     # adjust header path location for hipblas.h to avoid unnecessary deprecation warnings
     # adjust header path location for hiprand_kernel.h to avoid unnecessary deprecation warnings
@@ -88,11 +80,10 @@ function hipify_dir()
     do
         while IFS= read -r -d '' src
         do
-            sed -i 's@#include <faiss/gpu/@#include <faiss/gpu-rocm/@' "$src"
             sed -i 's@thrust::cuda::par@thrust::hip::par@' "$src"
             sed -i 's@#include <hipblas.h>@#include <hipblas/hipblas.h>@' "$src"
             sed -i 's@#include <hiprand_kernel.h>@#include <hiprand/hiprand_kernel.h>@' "$src"
-        done <   <(find ./gpu-rocm -name "*.$ext.tmp" -print0)
+        done <   <(find ./gpu -name "*.$ext.tmp" -print0)
     done
 
     # hipify was run in parallel above
@@ -116,50 +107,10 @@ function hipify_dir()
                 echo "$dst"
                 mv "$src" "$dst"
             fi
-        done <   <(find ./gpu-rocm -name "*.$ext.tmp" -print0)
+        done <   <(find ./gpu -name "*.$ext.tmp" -print0)
     done
 
-    # copy over CMakeLists.txt
-    while IFS= read -r -d '' src
-    do
-        dst="${src//\.\/gpu/\.\/gpu-rocm}"
-        if test -f "$dst"
-        then
-            if diff -q "$src" "$dst" >& /dev/null
-            then
-                echo "$dst [unchanged]"
-            else
-                echo "$dst"
-                cp "$src" "$dst"
-            fi
-        else
-            echo "$dst"
-            cp "$src" "$dst"
-        fi
-    done <   <(find ./gpu -name "CMakeLists.txt" -print0)
-
-    # Copy over other files
-    other_exts="py"
-    for ext in $other_exts
-    do
-        while IFS= read -r -d '' src
-        do
-            dst="${src//\.\/gpu/\.\/gpu-rocm}"
-            if test -f "$dst"
-            then
-                if diff -q "$src" "$dst" >& /dev/null
-                then
-                    echo "$dst [unchanged]"
-                else
-                    echo "$dst"
-                    cp "$src" "$dst"
-                fi
-            else
-                echo "$dst"
-                cp "$src" "$dst"
-            fi
-        done <   <(find ./gpu -name "*.$ext" -print0)
-    done
+    # No need to copy CMakeLists.txt or other files as we're working directly in the gpu directory
 }
 
 # Convert the faiss/gpu dir
diff --git a/faiss/python/CMakeLists.txt b/faiss/python/CMakeLists.txt
index c1bbf088c..b28723761 100644
--- a/faiss/python/CMakeLists.txt
+++ b/faiss/python/CMakeLists.txt
@@ -81,35 +81,20 @@ if(TARGET faiss)
     list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
       "${faiss_SOURCE_DIR}/faiss/${h}")
   endforeach()
-  if(FAISS_ENABLE_ROCM)
-    foreach(h ${FAISS_GPU_HEADERS})
-      list(APPEND SWIG_MODULE_swigfaiss_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_avx2_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_avx512_spr_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
-      list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
-    endforeach()
-  else()
-    foreach(h ${FAISS_GPU_HEADERS})
-      list(APPEND SWIG_MODULE_swigfaiss_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_avx2_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_avx512_spr_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
-      list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
-      list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
-        "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
-    endforeach()
-  endif()
+  foreach(h ${FAISS_GPU_HEADERS})
+    list(APPEND SWIG_MODULE_swigfaiss_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_avx2_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_avx512_spr_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+    list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+    list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
+      "${faiss_SOURCE_DIR}/faiss/gpu/${h}")
+  endforeach()
 else()
   find_package(faiss REQUIRED)
 endif()
diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig
index 7e57ab95d..66d95cb16 100644
--- a/faiss/python/swigfaiss.swig
+++ b/faiss/python/swigfaiss.swig
@@ -330,93 +330,16 @@ void gpu_profiler_stop();
 void gpu_sync_all_devices();
 
 #ifdef GPU_WRAPPER
-#ifdef FAISS_ENABLE_ROCM
 
 %shared_ptr(faiss::gpu::GpuResources);
 %shared_ptr(faiss::gpu::StandardGpuResourcesImpl);
 
 %{
 
-#include <faiss/gpu-rocm/StandardGpuResources.h>
-#include <faiss/gpu-rocm/GpuIndicesOptions.h>
-#include <faiss/gpu-rocm/GpuClonerOptions.h>
-#include <faiss/gpu-rocm/GpuIndex.h>
-#include <faiss/gpu-rocm/GpuIndexFlat.h>
-#include <faiss/gpu-rocm/GpuIndexIVF.h>
-#include <faiss/gpu-rocm/GpuIndexIVFPQ.h>
-#include <faiss/gpu-rocm/GpuIndexIVFFlat.h>
-#include <faiss/gpu-rocm/GpuIndexIVFScalarQuantizer.h>
-#include <faiss/gpu-rocm/GpuIndexBinaryFlat.h>
-#include <faiss/gpu-rocm/GpuAutoTune.h>
-#include <faiss/gpu-rocm/GpuCloner.h>
-#include <faiss/gpu-rocm/GpuDistance.h>
-#include <faiss/gpu-rocm/GpuIcmEncoder.h>
-
-int get_num_gpus()
-{
-    return faiss::gpu::getNumDevices();
-}
-
-void gpu_profiler_start()
-{
-    return faiss::gpu::profilerStart();
-}
-
-void gpu_profiler_stop()
-{
-    return faiss::gpu::profilerStop();
-}
-
-void gpu_sync_all_devices()
-{
-    return faiss::gpu::synchronizeAllDevices();
-}
-
-%}
-
-%template() std::pair<int, uint64_t>;
-%template() std::map<std::string, std::pair<int, uint64_t> >;
-%template() std::map<int, std::map<std::string, std::pair<int, uint64_t> > >;
-
-// causes weird wrapper bug
-%ignore *::allocMemoryHandle;
-%ignore faiss::gpu::GpuMemoryReservation;
-%ignore faiss::gpu::GpuMemoryReservation::operator=(GpuMemoryReservation&&);
-%ignore faiss::gpu::AllocType;
-
-%include  <faiss/gpu-rocm/GpuResources.h>
-%include  <faiss/gpu-rocm/StandardGpuResources.h>
-
-typedef ihipStream_t* hipStream_t;
-
-%inline %{
-
-// interop between pytorch exposed hipStream_t and faiss
-hipStream_t cast_integer_to_cudastream_t(int64_t x) {
-  return (hipStream_t) x;
-}
-
-int64_t cast_cudastream_t_to_integer(hipStream_t x) {
-  return (int64_t) x;
-}
-
-%}
-
-#else // FAISS_ENABLE_ROCM
-
-%shared_ptr(faiss::gpu::GpuResources);
-%shared_ptr(faiss::gpu::StandardGpuResourcesImpl);
-%shared_ptr(faiss::gpu::IVFPQBuildCagraConfig);
-%shared_ptr(faiss::gpu::IVFPQSearchCagraConfig);
-
-%{
-
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/GpuIndicesOptions.h>
 #include <faiss/gpu/GpuClonerOptions.h>
 #include <faiss/gpu/GpuIndex.h>
-#include <faiss/gpu/GpuIndexCagra.h>
-#include <faiss/gpu/GpuIndexBinaryCagra.h>
 #include <faiss/gpu/GpuIndexFlat.h>
 #include <faiss/gpu/GpuIndexIVF.h>
 #include <faiss/gpu/GpuIndexIVFPQ.h>
@@ -463,6 +386,35 @@ void gpu_sync_all_devices()
 %include  <faiss/gpu/GpuResources.h>
 %include  <faiss/gpu/StandardGpuResources.h>
 
+#ifdef FAISS_ENABLE_ROCM
+
+typedef ihipStream_t* hipStream_t;
+
+%inline %{
+
+// interop between pytorch exposed hipStream_t and faiss
+hipStream_t cast_integer_to_cudastream_t(int64_t x) {
+  return (hipStream_t) x;
+}
+
+int64_t cast_cudastream_t_to_integer(hipStream_t x) {
+  return (int64_t) x;
+}
+
+%}
+
+#else // FAISS_ENABLE_ROCM
+
+%shared_ptr(faiss::gpu::IVFPQBuildCagraConfig);
+%shared_ptr(faiss::gpu::IVFPQSearchCagraConfig);
+
+%{
+
+#include <faiss/gpu/GpuIndexCagra.h>
+#include <faiss/gpu/GpuIndexBinaryCagra.h>
+
+%}
+
 typedef CUstream_st* cudaStream_t;
 
 %inline %{
@@ -682,25 +634,6 @@ struct faiss::simd16uint16 {};
 
 #ifdef GPU_WRAPPER
 
-#ifdef FAISS_ENABLE_ROCM
-
-// quiet SWIG warnings
-%ignore faiss::gpu::GpuIndexIVF::GpuIndexIVF;
-
-%include  <faiss/gpu-rocm/GpuIndicesOptions.h>
-%include  <faiss/gpu-rocm/GpuClonerOptions.h>
-%include  <faiss/gpu-rocm/GpuIndex.h>
-%include  <faiss/gpu-rocm/GpuIndexFlat.h>
-%include  <faiss/gpu-rocm/GpuIndexIVF.h>
-%include  <faiss/gpu-rocm/GpuIndexIVFPQ.h>
-%include  <faiss/gpu-rocm/GpuIndexIVFFlat.h>
-%include  <faiss/gpu-rocm/GpuIndexIVFScalarQuantizer.h>
-%include  <faiss/gpu-rocm/GpuIndexBinaryFlat.h>
-%include  <faiss/gpu-rocm/GpuDistance.h>
-%include  <faiss/gpu-rocm/GpuIcmEncoder.h>
-
-#else // FAISS_ENABLE_ROCM
-
 // quiet SWIG warnings
 %ignore faiss::gpu::GpuIndexIVF::GpuIndexIVF;
 
@@ -720,7 +653,6 @@ struct faiss::simd16uint16 {};
 %include  <faiss/gpu/GpuDistance.h>
 %include  <faiss/gpu/GpuIcmEncoder.h>
 
-#endif // FAISS_ENABLE_ROCM
 #endif
 
 
@@ -960,16 +892,6 @@ faiss::Quantizer * downcast_Quantizer (faiss::Quantizer *aq)
 
 #ifdef GPU_WRAPPER
 
-#ifdef FAISS_ENABLE_ROCM
-%include  <faiss/gpu-rocm/GpuAutoTune.h>
-
-%newobject index_gpu_to_cpu;
-%newobject index_cpu_to_gpu;
-%newobject index_cpu_to_gpu_multiple;
-
-%include  <faiss/gpu-rocm/GpuCloner.h>
-
-#else // FAISS_ENABLE_ROCM
 %include  <faiss/gpu/GpuAutoTune.h>
 
 %newobject index_gpu_to_cpu;
@@ -978,7 +900,6 @@ faiss::Quantizer * downcast_Quantizer (faiss::Quantizer *aq)
 
 %include  <faiss/gpu/GpuCloner.h>
 
-#endif // FAISS_ENABLE_ROCM
 #endif