about summary refs log tree commit diff
diff options
context:
space:
mode:
authorSomeone Serge <sergei.kozlukov@aalto.fi>2023-03-31 17:36:54 +0300
committerSomeone Serge <sergei.kozlukov@aalto.fi>2023-04-01 05:47:12 +0300
commit36ba5681fab9d18481dae80be0393eedc1f134d1 (patch)
treea461f87e51bcd9e73de9992a5098ee6dd90c04c4
parenta3694989165105fdb0058d7a4a3799c1d180c269 (diff)
faiss: respect config.cudaCapabilities
-rw-r--r--pkgs/development/libraries/science/math/faiss/default.nix13
1 files changed, 9 insertions, 4 deletions
diff --git a/pkgs/development/libraries/science/math/faiss/default.nix b/pkgs/development/libraries/science/math/faiss/default.nix
index 3e4e227f83e2d..9bc7eda5ab74d 100644
--- a/pkgs/development/libraries/science/math/faiss/default.nix
+++ b/pkgs/development/libraries/science/math/faiss/default.nix
@@ -3,9 +3,8 @@
 , fetchFromGitHub
 , stdenv
 , cmake
-, cudaPackages
+, cudaPackages ? { }
 , cudaSupport ? config.cudaSupport or false
-, cudaCapabilities ? [ "60" "70" "80" "86" ]
 , pythonSupport ? true
 , pythonPackages
 , llvmPackages
@@ -27,7 +26,8 @@
 let
   pname = "faiss";
   version = "1.7.2";
-  inherit (cudaPackages) cudatoolkit;
+  inherit (cudaPackages) cudatoolkit cudaFlags;
+  inherit (cudaFlags) cudaCapabilities dropDot;
 in
 stdenv.mkDerivation {
   inherit pname version;
@@ -72,7 +72,7 @@ stdenv.mkDerivation {
     "-DFAISS_ENABLE_PYTHON=${if pythonSupport then "ON" else "OFF"}"
     "-DFAISS_OPT_LEVEL=${optLevel}"
   ] ++ lib.optionals cudaSupport [
-    "-DCMAKE_CUDA_ARCHITECTURES=${lib.concatStringsSep ";" cudaCapabilities}"
+    "-DCMAKE_CUDA_ARCHITECTURES=${builtins.concatStringsSep ";" (map dropDot cudaCapabilities)}"
   ];
 
 
@@ -101,6 +101,11 @@ stdenv.mkDerivation {
     addOpenGLRunpath $demos/bin/*
   '';
 
+  # Need buildPythonPackage for this one
+  # pythonCheckImports = [
+  #   "faiss"
+  # ];
+
   passthru = {
     inherit cudaSupport cudaPackages pythonSupport;