From 439aee24b8ee87f9477207a25d76bf40dd763e60 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 14 Apr 2023 15:45:41 +0200
Subject: [PATCH] Add pylibraft

Signed-off-by: Mickael Ide <mide@nvidia.com>
---
 CMakeLists.txt | 9 +++++++++
 README.md      | 4 ++--
 conanfile.py   | 6 ++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 64c0ffec4..4aa9acd92 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -29,6 +29,7 @@ knowhere_option(WITH_BENCHMARK "Build with benchmark" OFF)
 knowhere_option(WITH_COVERAGE "Build with coverage" OFF)
 knowhere_option(WITH_CCACHE "Build with ccache" ON)
 knowhere_option(WITH_PROFILER "Build with profiler" OFF)
+knowhere_option(HINT_LIBRAFT "Hint directory for libraft" "")
 
 if(KNOWHERE_VERSION)
   message(STATUS "Building KNOWHERE version: ${KNOWHERE_VERSION}")
@@ -147,6 +148,14 @@ add_library(knowhere SHARED ${KNOWHERE_SRCS})
 add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS})
 if(WITH_RAFT)
   list(APPEND KNOWHERE_LINKER_LIBS raft::raft raft::compiled)
+  find_library(LIBRAFT_FOUND raft HINTS ${HINT_LIBRAFT})
+  if (NOT LIBRAFT_FOUND)
+    message(WARNING "libraft not found")
+  else()
+    message(STATUS "libraft found")
+    list(APPEND KNOWHERE_LINKER_LIBS ${LIBRAFT_FOUND})
+    add_definitions(-DRAFT_COMPILED)
+  endif()
 endif()
 target_link_libraries(knowhere PUBLIC ${KNOWHERE_LINKER_LIBS})
 target_include_directories(
diff --git a/README.md b/README.md
index 51992f292..8b93486ef 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ Here's a list of verified OS types where Knowhere can successfully build and run
 
 ```bash
 $ sudo apt install build-essential libopenblas-dev libaio-dev python3-dev python3-pip
-$ pip3 install conan==1.59.0 --user
+$ pip3 install conan==1.59.0 pylibraft-cu11 --user --extra-index=https://pypi.nvidia.com
 $ export PATH=$PATH:$HOME/.local/bin
 ```
 
@@ -48,7 +48,7 @@ $ conan install .. --build=missing -o with_ut=True -o with_raft=True -s compiler
 #DISKANN SUPPORT
 $ conan install .. --build=missing -o with_ut=True -o with_diskann=True -s compiler.libcxx=libstdc++11 -s build_type=Debug/Release
 #build with conan
-$conan build ..
+$ conan build ..
 #verbose
 export VERBOSE=1
 ```
diff --git a/conanfile.py b/conanfile.py
index d66ecac11..d927df28a 100644
--- a/conanfile.py
+++ b/conanfile.py
@@ -132,6 +132,12 @@ def generate(self):
             if cxx_std_flag
             else "c++{}".format(self._minimum_cpp_standard)
         )
+        if self.options.with_raft:
+            try:
+                import pylibraft
+                tc.variables["HINT_LIBRAFT"] = pylibraft.__path__[0]
+            except:
+                pass
         tc.variables["CXX_STD"] = cxx_std_value
         if is_msvc(self):
             tc.variables["MSVC_LANGUAGE_VERSION"] = cxx_std_value