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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
|
{ lib
, config
, fetchFromGitHub
, symlinkJoin
, stdenv
, cmake
, cudaPackages ? { }
, cudaSupport ? config.cudaSupport
, pythonSupport ? true
, pythonPackages
, llvmPackages
, boost
, blas
, swig
, addOpenGLRunpath
, optLevel ? let
optLevels =
lib.optionals stdenv.hostPlatform.avx2Support [ "avx2" ]
++ lib.optionals stdenv.hostPlatform.sse4_1Support [ "sse4" ]
++ [ "generic" ];
in
# Choose the maximum available optimization level
builtins.head optLevels
, faiss # To run demos in the tests
, runCommand
}@inputs:
let
pname = "faiss";
version = "1.7.4";
inherit (cudaPackages) cudaFlags backendStdenv;
inherit (cudaFlags) cudaCapabilities dropDot;
stdenv = if cudaSupport then backendStdenv else inputs.stdenv;
cudaJoined = symlinkJoin {
name = "cuda-packages-unsplit";
paths = with cudaPackages; [
cuda_cudart # cuda_runtime.h
libcublas
libcurand
cuda_cccl
] ++ lib.optionals (cudaPackages ? cuda_profiler_api) [
cuda_profiler_api # cuda_profiler_api.h
] ++ lib.optionals (!(cudaPackages ? cuda_profiler_api)) [
cuda_nvprof # cuda_profiler_api.h
];
};
in
stdenv.mkDerivation {
inherit pname version;
outputs = [ "out" "demos" ];
src = fetchFromGitHub {
owner = "facebookresearch";
repo = pname;
rev = "v${version}";
hash = "sha256-WSce9X6sLZmGM5F0ZkK54VqpIy8u1VB0e9/l78co29M=";
};
buildInputs = [
blas
swig
] ++ lib.optionals pythonSupport [
pythonPackages.setuptools
pythonPackages.pip
pythonPackages.wheel
] ++ lib.optionals stdenv.cc.isClang [
llvmPackages.openmp
] ++ lib.optionals cudaSupport [
cudaJoined
];
propagatedBuildInputs = lib.optionals pythonSupport [
pythonPackages.numpy
];
nativeBuildInputs = [ cmake ] ++ lib.optionals cudaSupport [
cudaPackages.cuda_nvcc
addOpenGLRunpath
] ++ lib.optionals 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=${builtins.concatStringsSep ";" (map dropDot cudaCapabilities)}"
"-DCUDAToolkit_INCLUDE_DIR=${cudaJoined}/include"
];
# 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/*
'';
# Need buildPythonPackage for this one
# pythonImportsCheck = [
# "faiss"
# ];
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 ];
};
}
|