Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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