Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuSelect.cu
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9
10 #include "../utils/DeviceUtils.h"
11 #include "../utils/BlockSelectKernel.cuh"
12 #include "../utils/WarpSelectKernel.cuh"
13 #include "../utils/HostTensor.cuh"
14 #include "../utils/DeviceTensor.cuh"
15 #include "../test/TestUtils.h"
16 #include <algorithm>
17 #include <gtest/gtest.h>
18 #include <sstream>
19 #include <unordered_map>
20 #include <vector>
21
22 void testForSize(int rows, int cols, int k, bool dir, bool warp) {
23 std::vector<float> v = faiss::gpu::randVecs(rows, cols);
24 faiss::gpu::HostTensor<float, 2, true> hostVal({rows, cols});
25
26 for (int r = 0; r < rows; ++r) {
27 for (int c = 0; c < cols; ++c) {
28 hostVal[r][c] = v[r * cols + c];
29 }
30 }
31
32 // row -> (val -> idx)
33 std::unordered_map<int, std::vector<std::pair<int, float>>> hostOutValAndInd;
34 for (int r = 0; r < rows; ++r) {
35 std::vector<std::pair<int, float>> closest;
36
37 for (int c = 0; c < cols; ++c) {
38 closest.emplace_back(c, (float) hostVal[r][c]);
39 }
40
41 auto dirFalseFn =
42 [](std::pair<int, float>& a, std::pair<int, float>& b) {
43 return a.second < b.second;
44 };
45 auto dirTrueFn =
46 [](std::pair<int, float>& a, std::pair<int, float>& b) {
47 return a.second > b.second;
48 };
49
50 std::sort(closest.begin(), closest.end(), dir ? dirTrueFn : dirFalseFn);
51 hostOutValAndInd.emplace(r, closest);
52 }
53
54 // Select top-k on GPU
55 faiss::gpu::DeviceTensor<float, 2, true> gpuVal(hostVal, 0);
56 faiss::gpu::DeviceTensor<float, 2, true> gpuOutVal({rows, k});
57 faiss::gpu::DeviceTensor<int, 2, true> gpuOutInd({rows, k});
58
59 if (warp) {
60 faiss::gpu::runWarpSelect(gpuVal, gpuOutVal, gpuOutInd, dir, k, 0);
61 } else {
62 faiss::gpu::runBlockSelect(gpuVal, gpuOutVal, gpuOutInd, dir, k, 0);
63 }
64
65 // Copy back to CPU
66 faiss::gpu::HostTensor<float, 2, true> outVal(gpuOutVal, 0);
67 faiss::gpu::HostTensor<int, 2, true> outInd(gpuOutInd, 0);
68
69 for (int r = 0; r < rows; ++r) {
70 std::unordered_map<int, int> seenIndices;
71
72 for (int i = 0; i < k; ++i) {
73 float gpuV = outVal[r][i];
74 float cpuV = hostOutValAndInd[r][i].second;
75
76 EXPECT_EQ(gpuV, cpuV) <<
77 "rows " << rows << " cols " << cols << " k " << k << " dir " << dir
78 << " row " << r << " ind " << i;
79
80 // If there are identical elements in a row that should be
81 // within the top-k, then it is possible that the index can
82 // differ, because the order in which the GPU will see the
83 // equivalent values is different than the CPU (and will remain
84 // unspecified, since this is affected by the choice of
85 // k-selection algorithm that we use)
86 int gpuInd = outInd[r][i];
87 int cpuInd = hostOutValAndInd[r][i].first;
88
89 // We should never see duplicate indices, however
90 auto itSeenIndex = seenIndices.find(gpuInd);
91
92 EXPECT_EQ(itSeenIndex, seenIndices.end()) <<
93 "Row " << r << " user index " << gpuInd << " was seen at both " <<
94 itSeenIndex->second << " and " << i;
95
96 seenIndices[gpuInd] = i;
97
98 if (gpuInd != cpuInd) {
99 // Gather the values from the original data via index; the
100 // values should be the same
101 float gpuGatherV = hostVal[r][gpuInd];
102 float cpuGatherV = hostVal[r][cpuInd];
103
104 EXPECT_EQ(gpuGatherV, cpuGatherV) <<
105 "rows " << rows << " cols " << cols << " k " << k << " dir " << dir
106 << " row " << r << " ind " << i << " source ind "
107 << gpuInd << " " << cpuInd;
108 }
109 }
110 }
111 }
112
113 // General test
114 TEST(TestGpuSelect, test) {
115 for (int i = 0; i < 10; ++i) {
116 int rows = faiss::gpu::randVal(10, 100);
117 int cols = faiss::gpu::randVal(1, 30000);
118 int k = std::min(cols, faiss::gpu::randVal(1, 1024));
119 bool dir = faiss::gpu::randBool();
120
121 testForSize(rows, cols, k, dir, false);
122 }
123 }
124
125 // Test for k = 1
126 TEST(TestGpuSelect, test1) {
127 for (int i = 0; i < 5; ++i) {
128 int rows = faiss::gpu::randVal(10, 100);
129 int cols = faiss::gpu::randVal(1, 30000);
130 bool dir = faiss::gpu::randBool();
131
132 testForSize(rows, cols, 1, dir, false);
133 }
134 }
135
136 // Test for where k = #cols exactly (we are returning all the values,
137 // just sorted)
138 TEST(TestGpuSelect, testExact) {
139 for (int i = 0; i < 5; ++i) {
140 int rows = faiss::gpu::randVal(10, 100);
141 int cols = faiss::gpu::randVal(1, 1024);
142 bool dir = faiss::gpu::randBool();
143
144 testForSize(rows, cols, cols, dir, false);
145 }
146 }
147
148 // General test
149 TEST(TestGpuSelect, testWarp) {
150 for (int i = 0; i < 10; ++i) {
151 int rows = faiss::gpu::randVal(10, 100);
152 int cols = faiss::gpu::randVal(1, 30000);
153 int k = std::min(cols, faiss::gpu::randVal(1, 1024));
154 bool dir = faiss::gpu::randBool();
155
156 testForSize(rows, cols, k, dir, true);
157 }
158 }
159
160 // Test for k = 1
161 TEST(TestGpuSelect, test1Warp) {
162 for (int i = 0; i < 5; ++i) {
163 int rows = faiss::gpu::randVal(10, 100);
164 int cols = faiss::gpu::randVal(1, 30000);
165 bool dir = faiss::gpu::randBool();
166
167 testForSize(rows, cols, 1, dir, true);
168 }
169 }
170
171 // Test for where k = #cols exactly (we are returning all the values,
172 // just sorted)
173 TEST(TestGpuSelect, testExactWarp) {
174 for (int i = 0; i < 5; ++i) {
175 int rows = faiss::gpu::randVal(10, 100);
176 int cols = faiss::gpu::randVal(1, 1024);
177 bool dir = faiss::gpu::randBool();
178
179 testForSize(rows, cols, cols, dir, true);
180 }
181 }
182
183 int main(int argc, char** argv) {
184 testing::InitGoogleTest(&argc, argv);
185
186 // just run with a fixed test seed
187 faiss::gpu::setTestSeed(100);
188
189 return RUN_ALL_TESTS();
190 }
faiss::gpu::HostTensor
Definition: HostTensor.cuh:21
faiss::gpu::DeviceTensor< float, 2, true >