tjh.dev
1{
2 lib,
3 config,
4 fetchFromGitHub,
5 stdenv,
6 cmake,
7 cudaPackages ? { },
8 cudaSupport ? config.cudaSupport,
9 pythonSupport ? true,
10 python3Packages,
11 llvmPackages,
12 blas,
13 swig,
14 autoAddDriverRunpath,
15 optLevel ?
16 let
17 optLevels =
18 lib.optionals stdenv.hostPlatform.avx2Support [ "avx2" ]
19 ++ lib.optionals stdenv.hostPlatform.sse4_1Support [ "sse4" ]
20 ++ [ "generic" ];
21 in
22 # Choose the maximum available optimization level
23 builtins.head optLevels,
24}@inputs:
25
26let
27 pname = "faiss";
28 version = "1.11.0";
29
30 inherit (cudaPackages) flags backendStdenv;
31
32 stdenv = if cudaSupport then backendStdenv else inputs.stdenv;
33
34 cudaComponents = with cudaPackages; [
35 cuda_cudart # cuda_runtime.h
36 libcublas
37 libcurand
38 cuda_cccl
39
40 # cuda_profiler_api.h
41 (cudaPackages.cuda_profiler_api or cudaPackages.cuda_nvprof)
42 ];
43in
44stdenv.mkDerivation {
45 inherit pname version;
46
47 outputs = [ "out" ] ++ lib.optionals pythonSupport [ "dist" ];
48
49 src = fetchFromGitHub {
50 owner = "facebookresearch";
51 repo = "faiss";
52 tag = "v${version}";
53 hash = "sha256-N8UkL+KS9Da6RtaHI9pY5gAzFtTSMJ9R5h4RSX9b1Ro=";
54 };
55
56 nativeBuildInputs = [
57 cmake
58 ]
59 ++ lib.optionals cudaSupport [
60 cudaPackages.cuda_nvcc
61 autoAddDriverRunpath
62 ]
63 ++ lib.optionals pythonSupport [
64 python3Packages.python
65 python3Packages.setuptools
66 python3Packages.pip
67 ];
68
69 buildInputs = [
70 blas
71 swig
72 ]
73 ++ lib.optionals pythonSupport [ python3Packages.numpy ]
74 ++ lib.optionals stdenv.cc.isClang [ llvmPackages.openmp ]
75 ++ lib.optionals cudaSupport cudaComponents;
76
77 cmakeFlags = [
78 (lib.cmakeBool "FAISS_ENABLE_GPU" cudaSupport)
79 (lib.cmakeBool "FAISS_ENABLE_PYTHON" pythonSupport)
80 (lib.cmakeFeature "FAISS_OPT_LEVEL" optLevel)
81 ]
82 ++ lib.optionals cudaSupport [
83 (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)
84 ];
85
86 buildFlags = [ "faiss" ] ++ lib.optionals pythonSupport [ "swigfaiss" ];
87
88 # pip wheel->pip install commands copied over from opencv4
89
90 postBuild = lib.optionalString pythonSupport ''
91 (cd faiss/python &&
92 python -m pip wheel --verbose --no-index --no-deps --no-clean --no-build-isolation --wheel-dir dist .)
93 '';
94
95 postInstall = lib.optionalString pythonSupport ''
96 mkdir "$dist"
97 cp faiss/python/dist/*.whl "$dist/"
98 '';
99
100 passthru = {
101 inherit cudaSupport cudaPackages pythonSupport;
102 };
103
104 meta = {
105 description = "Library for efficient similarity search and clustering of dense vectors by Facebook Research";
106 mainProgram = "demo_ivfpq_indexing";
107 homepage = "https://github.com/facebookresearch/faiss";
108 changelog = "https://github.com/facebookresearch/faiss/blob/v${version}/CHANGELOG.md";
109 license = lib.licenses.mit;
110 platforms = lib.platforms.unix;
111 maintainers = with lib.maintainers; [ SomeoneSerge ];
112 };
113}