ROCm CMake configuration cleanup

CMakeLists.txt

@@ -24,7 +24,6 @@ set(FAISS_LANGUAGES CXX)
 
 if(FAISS_ENABLE_GPU)
   if (FAISS_ENABLE_ROCM)
-    set(USE_ROCM TRUE)
     list(APPEND FAISS_LANGUAGES HIP)
     list(PREPEND CMAKE_MODULE_PATH "/opt/rocm/lib/cmake")
     list(PREPEND CMAKE_PREFIX_PATH "/opt/rocm")
@@ -66,9 +65,9 @@ option(FAISS_ENABLE_PYTHON "Build Python extension." ON)
 option(FAISS_ENABLE_C_API "Build C API." OFF)
 
 if(FAISS_ENABLE_GPU)
-  if(USE_ROCM)
+  if(FAISS_ENABLE_ROCM)
     enable_language(HIP)
-    add_definitions(-DUSE_ROCM)
+    add_definitions(-DUSE_AMD_ROCM)
     find_package(HIP REQUIRED)
     find_package(hipBLAS REQUIRED)
     set(GPU_EXT_PREFIX "hip")
@@ -83,15 +82,10 @@ if(FAISS_ENABLE_RAFT AND NOT TARGET raft::raft)
   find_package(raft COMPONENTS compiled distributed)
 endif()
 
-if(USE_ROCM)
-  find_package(HIP REQUIRED)
-  find_package(hipBLAS REQUIRED)
-endif()
-
 add_subdirectory(faiss)
 
 if(FAISS_ENABLE_GPU)
-  if(USE_ROCM)
+  if(FAISS_ENABLE_ROCM)
     add_subdirectory(faiss/gpu-rocm)
   else()
     add_subdirectory(faiss/gpu)
@@ -116,7 +110,7 @@ if(BUILD_TESTING)
   add_subdirectory(tests)
 
   if(FAISS_ENABLE_GPU)
-    if(USE_ROCM)
+    if(FAISS_ENABLE_ROCM)
       add_subdirectory(faiss/gpu-rocm/test)
     else()
       add_subdirectory(faiss/gpu/test)

c_api/CMakeLists.txt

@@ -56,7 +56,7 @@ add_executable(example_c EXCLUDE_FROM_ALL example_c.c)
 target_link_libraries(example_c PRIVATE faiss_c)
 
 if(FAISS_ENABLE_GPU)
-  if(USE_ROCM)
+  if(FAISS_ENABLE_ROCM)
     add_subdirectory(gpu-rocm)
   else ()
     add_subdirectory(gpu)

c_api/gpu/CMakeLists.txt

@@ -15,13 +15,13 @@ 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)
+if (FAISS_ENABLE_ROCM)
+  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>)
+  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)

faiss/gpu/CMakeLists.txt

@@ -197,7 +197,7 @@ function(generate_ivf_interleaved_code)
     "64|2048|8"
   )
 
-  if (USE_ROCM)
+  if (FAISS_ENABLE_ROCM)
      list(TRANSFORM FAISS_GPU_SRC REPLACE cu$ hip)
   endif()
 
@@ -294,7 +294,7 @@ if(FAISS_ENABLE_RAFT)
   target_compile_definitions(faiss_gpu PUBLIC USE_NVIDIA_RAFT=1)
 endif()
 
-if (USE_ROCM)
+if (FAISS_ENABLE_ROCM)
   list(TRANSFORM FAISS_GPU_SRC REPLACE cu$ hip)
 endif()
 
@@ -313,8 +313,8 @@ foreach(header ${FAISS_GPU_HEADERS})
   )
 endforeach()
 
-if (USE_ROCM)
-  target_link_libraries(faiss_gpu PRIVATE $<$<BOOL:${USE_ROCM}>:hip::host> $<$<BOOL:${USE_ROCM}>:roc::hipblas>)
+if (FAISS_ENABLE_ROCM)
+  target_link_libraries(faiss_gpu PRIVATE hip::host roc::hipblas)
   target_compile_options(faiss_gpu PRIVATE)
 else()
   # Prepares a host linker script and enables host linker to support
@@ -333,6 +333,13 @@ else()
   target_link_options(faiss_gpu PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")
 
   find_package(CUDAToolkit REQUIRED)
-  target_link_libraries(faiss_gpu PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::compiled> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass> $<$<BOOL:${FAISS_ENABLE_RAFT}>:OpenMP::OpenMP_CXX>)
-  target_compile_options(faiss_gpu PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr $<$<BOOL:${FAISS_ENABLE_RAFT}>:-Xcompiler=${OpenMP_CXX_FLAGS}>>)
+  target_link_libraries(faiss_gpu PRIVATE CUDA::cudart CUDA::cublas
+    $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft>
+    $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::compiled>
+    $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>
+    $<$<BOOL:${FAISS_ENABLE_RAFT}>:OpenMP::OpenMP_CXX>)
+  target_compile_options(faiss_gpu PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all
+    --expt-extended-lambda --expt-relaxed-constexpr
+    $<$<BOOL:${FAISS_ENABLE_RAFT}>:-Xcompiler=${OpenMP_CXX_FLAGS}>>)
 endif()

faiss/gpu/GpuFaissAssert.h

@@ -15,7 +15,7 @@
 /// Assertions
 ///
 
-#if defined(__CUDA_ARCH__) || defined(USE_ROCM)
+#if defined(__CUDA_ARCH__) || defined(USE_AMD_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

@@ -363,7 +363,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
             prop.major,
             prop.minor);
 
-#if USE_ROCM
+#if USE_AMD_ROCM
     // Our code is pre-built with and expects warpSize == 32 or 64, validate
     // that
     FAISS_ASSERT_FMT(