facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

OSS changes: Enable ROCm to work with Faiss on BUCK (#4485)

Browse Source 
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
This commit is contained in:
Michael Norris
2025-08-15 12:06:27 -07:00
committed by Facebook GitHub Bot
parent dd637c98d6
commit 514b44fca8
5 changed files with 59 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+2 -6
View File
@@ -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()
c_api/CMakeLists.txt
+1 -5
View File
@@ -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()
faiss/gpu/hipify.sh
+13 -62
View File
@@ -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
faiss/python/CMakeLists.txt
+14 -29
View File
@@ -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()
faiss/python/swigfaiss.swig
+29 -108
View File
@@ -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