[ROCm] initial port

CMakeLists.txt

@@ -23,7 +23,13 @@ cmake_minimum_required(VERSION 3.23.1 FATAL_ERROR)
 set(FAISS_LANGUAGES CXX)
 
 if(FAISS_ENABLE_GPU)
-  list(APPEND FAISS_LANGUAGES CUDA)
+  # if ROCm install detected, assume ROCm/HIP is GPU device
+  if (EXISTS /opt/rocm)
+    set(USE_ROCM TRUE)
+    list(APPEND FAISS_LANGUAGES HIP)
+  else()
+    list(APPEND FAISS_LANGUAGES CUDA)
+  endif()
 endif()
 
 if(FAISS_ENABLE_RAFT)
@@ -58,8 +64,13 @@ option(FAISS_ENABLE_PYTHON "Build Python extension." ON)
 option(FAISS_ENABLE_C_API "Build C API." OFF)
 
 if(FAISS_ENABLE_GPU)
-  set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER})
-  enable_language(CUDA)
+  if(USE_ROCM)
+    enable_language(HIP)
+    add_definitions(-DUSE_ROCM)
+  else ()
+    set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER})
+    enable_language(CUDA)
+  endif()
 endif()
 
 if(FAISS_ENABLE_RAFT)
@@ -69,7 +80,11 @@ endif()
 add_subdirectory(faiss)
 
 if(FAISS_ENABLE_GPU)
-  add_subdirectory(faiss/gpu)
+  if(USE_ROCM)
+    add_subdirectory(faiss/gpu-rocm)
+  else()
+    add_subdirectory(faiss/gpu)
+  endif()
 endif()
 
 if(FAISS_ENABLE_PYTHON)
@@ -90,6 +105,10 @@ if(BUILD_TESTING)
   add_subdirectory(tests)
 
   if(FAISS_ENABLE_GPU)
-    add_subdirectory(faiss/gpu/test)
+    if(USE_ROCM)
+      add_subdirectory(faiss/gpu-rocm/test)
+    else()
+      add_subdirectory(faiss/gpu/test)
+    endif()
   endif()
 endif()

c_api/gpu/CMakeLists.txt

@@ -15,8 +15,14 @@ target_sources(faiss_c PRIVATE
 file(GLOB FAISS_C_API_GPU_HEADERS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.h")
 faiss_install_headers("${FAISS_C_API_GPU_HEADERS}" c_api/gpu)
 
+if (USE_ROCM)
+find_package(HIP REQUIRED)
+find_package(hipBLAS REQUIRED)
+target_link_libraries(faiss_c PUBLIC hip::host roc::hipblas)
+else()
 find_package(CUDAToolkit REQUIRED)
 target_link_libraries(faiss_c PUBLIC CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
+endif()
 
 add_executable(example_gpu_c EXCLUDE_FROM_ALL example_gpu_c.c)
 target_link_libraries(example_gpu_c PRIVATE faiss_c)

faiss/gpu/CMakeLists.txt

@@ -197,6 +197,12 @@ function(generate_ivf_interleaved_code)
     "64|2048|8"
   )
 
+  if(USE_ROCM)
+    set(CU_OR_HIP "hip")
+  else()
+    set(CU_OR_HIP "cu")
+  endif()
+
   # Traverse through the Cartesian product of X and Y
   foreach(sub_codec ${SUB_CODEC_TYPE})
   foreach(metric_type ${SUB_METRIC_TYPE})
@@ -210,10 +216,10 @@ function(generate_ivf_interleaved_code)
     set(filename "template_${sub_codec}_${metric_type}_${sub_threads}_${sub_num_warp_q}_${sub_num_thread_q}")
     # Remove illegal characters from filename
     string(REGEX REPLACE "[^A-Za-z0-9_]" "" filename ${filename})
-    set(output_file "${CMAKE_CURRENT_BINARY_DIR}/${filename}.cu")
+    set(output_file "${CMAKE_CURRENT_BINARY_DIR}/${filename}.${CU_OR_HIP}")
 
     # Read the template file
-    file(READ "${CMAKE_CURRENT_SOURCE_DIR}/impl/scan/IVFInterleavedScanKernelTemplate.cu" template_content)
+    file(READ "${CMAKE_CURRENT_SOURCE_DIR}/impl/scan/IVFInterleavedScanKernelTemplate.${CU_OR_HIP}" template_content)
 
     # Replace the placeholders
     string(REPLACE "SUB_CODEC_TYPE" "${sub_codec}" template_content "${template_content}")
@@ -248,6 +254,10 @@ if(FAISS_ENABLE_RAFT)
   target_compile_definitions(faiss_avx2 PUBLIC USE_NVIDIA_RAFT=1)
 endif()
 
+if (USE_ROCM)
+  list(TRANSFORM FAISS_GPU_SRC REPLACE cu$ hip)
+endif()
+
 # Export FAISS_GPU_HEADERS variable to parent scope.
 set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE)
 
@@ -261,8 +271,15 @@ foreach(header ${FAISS_GPU_HEADERS})
   )
 endforeach()
 
+if (USE_ROCM)
+find_package(HIP REQUIRED)
+find_package(hipBLAS REQUIRED)
+target_link_libraries(faiss PRIVATE hip::host roc::hipblas)
+target_link_libraries(faiss_avx2 PRIVATE hip::host roc::hipblas)
+else()
 find_package(CUDAToolkit REQUIRED)
 target_link_libraries(faiss PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::compiled> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
 target_link_libraries(faiss_avx2 PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::compiled>  $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
 target_compile_options(faiss PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)
 target_compile_options(faiss_avx2 PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)
+endif()

faiss/gpu/GpuFaissAssert.h

@@ -15,7 +15,7 @@
 /// Assertions
 ///
 
-#ifdef __CUDA_ARCH__
+#if defined(__CUDA_ARCH__) || defined(USE_ROCM)
 #define GPU_FAISS_ASSERT(X) assert(X)
 #define GPU_FAISS_ASSERT_MSG(X, MSG) assert(X)
 #define GPU_FAISS_ASSERT_FMT(X, FMT, ...) assert(X)

faiss/gpu/StandardGpuResources.cpp

@@ -351,8 +351,8 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
 
     // Our code is pre-built with and expects warpSize == 32, validate that
     FAISS_ASSERT_FMT(
-            prop.warpSize == 32,
-            "Device id %d does not have expected warpSize of 32",
+            prop.warpSize == 32 || prop.warpSize == 64,
+            "Device id %d does not have expected warpSize of 32 or 64",
             device);
 
     // Create streams

faiss/gpu/hipify.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+
+# go one level up from faiss/gpu
+top=$(dirname "${BASH_SOURCE[0]}")/..
+echo "top=$top"
+cd $top
+echo "pwd=`pwd`"
+
+# create all destination directories for hipified files into sibling 'gpu-rocm' directory
+for src in $(find ./gpu -type d)
+do
+    dst=$(echo $src | sed 's/gpu/gpu-rocm/')
+    echo "Creating $dst"
+    mkdir -p $dst
+done
+
+# run hipify-perl against all *.cu *.cuh *.h *.cpp files, no renaming
+# run all files in parallel to speed up
+for ext in cu cuh h cpp
+do
+    for src in $(find ./gpu -name "*.$ext")
+    do
+        dst=$(echo $src | sed 's@./gpu@./gpu-rocm@')
+        hipify-perl -o=$dst.tmp $src &
+    done
+done
+wait
+
+# rename all hipified *.cu files to *.hip
+for src in $(find ./gpu-rocm -name "*.cu.tmp")
+do
+    dst=${src%.cu.tmp}.hip.tmp
+    mv $src $dst
+done
+
+# 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
+for ext in hip cuh h cpp
+do
+    for src in $(find ./gpu-rocm -name "*.$ext.tmp")
+    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
+done
+
+# hipify was run in parallel above
+# don't copy the tmp file if it is unchanged
+for ext in hip cuh h cpp
+do
+    for src in $(find ./gpu-rocm -name "*.$ext.tmp")
+    do
+        dst=${src%.tmp}
+        if test -f $dst
+        then
+            if diff -q $src $dst >& /dev/null
+            then
+                echo "$dst [unchanged]"
+                rm $src
+            else
+                echo "$dst"
+                mv $src $dst
+            fi
+        else
+            echo "$dst"
+            mv $src $dst
+        fi
+    done
+done
+
+# copy over CMakeLists.txt
+for src in $(find ./gpu -name "CMakeLists.txt")
+do
+    dst=$(echo $src | sed 's@./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
+
+# run hipify-perl against python/swigfaiss.swig
+# replace header include statements "<faiss/gpu/" with "<faiss/gpu-rocm" in python
+for src in ./python/swigfaiss.swig
+do
+    hipify-perl -inplace $src
+    sed -i 's@#include <faiss/gpu/@#include <faiss/gpu-rocm/@' $src
+    sed -i 's@thrust::cuda::par@thrust::hip::par@' $src
+done