1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
{ lib
, config
, fetchFromGitHub
, stdenv
, cmake
, cudaPackages
, cudaSupport ? config.cudaSupport or false
, cudaCapabilities ? [ "60" "70" "80" "86" ]
, pythonSupport ? true
, pythonPackages
, blas
, swig
, addOpenGLRunpath
, optLevel ? let
optLevels =
lib.optional stdenv.hostPlatform.avx2Support "avx2"
++ lib.optional stdenv.hostPlatform.sse4_1Support "sse4"
++ [ "generic" ];
in
# Choose the maximum available optimization level
builtins.head optLevels
, faiss # To run demos in the tests
, runCommand
}:
let
pname = "faiss";
version = "1.7.2";
inherit (cudaPackages) cudatoolkit;
in
stdenv.mkDerivation {
inherit pname version;
outputs = [ "out" "demos" ];
src = fetchFromGitHub {
owner = "facebookresearch";
repo = pname;
rev = "v${version}";
hash = "sha256-Tklf5AaqJbOs9qtBZVcxXPLAp+K54EViZLSOvEhmswg=";
};
buildInputs = [
blas
swig
] ++ lib.optionals pythonSupport [
pythonPackages.setuptools
pythonPackages.pip
pythonPackages.wheel
];
propagatedBuildInputs = lib.optionals pythonSupport [
pythonPackages.numpy
];
nativeBuildInputs = [ cmake ] ++ lib.optionals cudaSupport [
cudatoolkit
addOpenGLRunpath
] ++ lib.optional pythonSupport [
pythonPackages.python
];
passthru.extra-requires.all = [
pythonPackages.numpy
];
cmakeFlags = [
"-DFAISS_ENABLE_GPU=${if cudaSupport then "ON" else "OFF"}"
"-DFAISS_ENABLE_PYTHON=${if pythonSupport then "ON" else "OFF"}"
"-DFAISS_OPT_LEVEL=${optLevel}"
] ++ lib.optionals cudaSupport [
"-DCMAKE_CUDA_ARCHITECTURES=${lib.concatStringsSep ";" cudaCapabilities}"
];
# pip wheel->pip install commands copied over from opencv4
buildPhase = ''
make -j faiss
make demo_ivfpq_indexing
'' + lib.optionalString pythonSupport ''
make -j swigfaiss
(cd faiss/python &&
python -m pip wheel --verbose --no-index --no-deps --no-clean --no-build-isolation --wheel-dir dist .)
'';
installPhase = ''
make install
mkdir -p $demos/bin
cp ./demos/demo_ivfpq_indexing $demos/bin/
'' + lib.optionalString pythonSupport ''
mkdir -p $out/${pythonPackages.python.sitePackages}
(cd faiss/python && python -m pip install dist/*.whl --no-index --no-warn-script-location --prefix="$out" --no-cache)
'';
fixupPhase = lib.optionalString (pythonSupport && cudaSupport) ''
addOpenGLRunpath $out/${pythonPackages.python.sitePackages}/faiss/*.so
addOpenGLRunpath $demos/bin/*
'';
passthru = {
inherit cudaSupport cudaPackages pythonSupport;
tests = {
runDemos = runCommand "${pname}-run-demos"
{ buildInputs = [ faiss.demos ]; }
# There are more demos, we run just the one that documentation mentions
''
demo_ivfpq_indexing && touch $out
'';
} // lib.optionalAttrs pythonSupport {
pytest = pythonPackages.callPackage ./tests.nix { };
};
};
meta = with lib; {
description = "A library for efficient similarity search and clustering of dense vectors by Facebook Research";
homepage = "https://github.com/facebookresearch/faiss";
license = licenses.mit;
platforms = platforms.unix;
maintainers = with maintainers; [ SomeoneSerge ];
};
}
|
