向量內積的GPU實現
背景
搜索場景下通常有少量的查詢向量和大量的索引向量求距離,不管是L2距離還是cos距離都和這兩者的內積有關係, 所以本文就來看看Faiss中內積的GPU實現,內積計算本質上是矩陣相乘,Faiss並沒有自己寫kernel,而是調用了大名鼎鼎的cublas,下面來看些細節,本文不會講cublas的實現,而是重點講在faiss中是如何調用cublas的
cublass的奇怪之處
輸入
我們來看一個矩陣,以一個4*2的舉證為例, 在程序裡面,矩陣存儲在一個一維數組中,總共8個元素
[x0,x1,x2,x3,x4,x5,x6,x7]
邏輯上我們認為是這樣的一個矩陣
[
x0 x1
x2 x3
x4,x5
x6,x7
]
但如果放入cublas, 就會成為這樣的形式
[
x0 x2 x4 x6
x1 x3 x5 x7
]
導致這個差別的主要原因是cublas在填充矩陣時，是按照列的順序來填充的,而我們是通常是按行
如果這樣放入,就和我們的意願違背了
輸出
當我們計算完一個數組,並拷貝回主機的時候,也會發現也這樣
假設輸出同樣邏輯上是4*2矩陣
[
x0 x1
x2 x3
x4,x5
x6,x7
]
那麼輸出是[x0,x2,x4,x6,x1,x3,x5,x7]
那麼我們會認為輸出是
[
x0 x2
x4 x6
x1 x3
x5 x7
]
輸出也是按列填充,實際上這也不是我麼想要的
怎麼解
看看我們的目標
查詢向量矩陣(Q): 3個2維的向量
索引向量矩陣(D): 4個2維的向量
結果(R): 3*4的結果矩陣
即 ? ? [式1]
正確的輸出
我們想要正確的輸出R(3*4), 那麼cublas中就不能是3 *4 的R,否則就會像上面的例子那樣輸出有問題,我們應該在cublas中輸出的是啥?
比如我們想輸出
[
R00 R01 R02 R03
R10 R11 R12 R13
R20 R21 R22 R23
]
如果想要這個那麼反推在cublas中的結果矩陣應該是
[
R00 R10 R20
R01 R11 R21
R02 R12 R22
R03 R13 R23
]
這實際上就是?
也就是說cublas中的? , 拷貝到主機上就變成了
正確的輸入
既然我們要得到? ?
將[式1]全部全部轉置,就得到
[式2]
那麼我們如何正確的輸入? 和?
以?? 為例
正如上文提到,如果直接輸入? ?,那麼結果是不對的
也正如上面的輸出例子所示,如果想在cublas中的正確的? ,那麼實際的輸入應該是?
................????????????????????有點暈了,因為cublas這個奇怪的特性,導致繞來繞去
還是先把這個? 的例子搞清楚
我們希望在cublas中得到? (4*2)
[
D00 D01
D10 D11
D20 D21
D30 D31
]
那麼我們的輸入應該
[D00 D10 D20 D30 D01 D11 D21 D31] (如上提到,cublas是按列填充的)
那麼我們可以認為這實際就是? (2*4) 即
[
D00 D10 D20 D30
D01 D11 D21 D31
]
同理可以認為要得到? ,直接輸入
就可以了
休息下,我們稍微總結下
- 我們一定要分清,cublas中的矩陣(GPU設備上的參與計算的矩陣)和cublas外的矩陣(即我們主機端的輸入和輸出矩陣)
- 我們必須看到公式中的矩陣實際上屬於前者
- 由於cublas特殊的按列填充,會使得cublas內和cublas外的矩陣有天然的轉置關係
如下圖
參考代碼
T* pA = transC ? a.data() : b.data();
T* pB = transC ? b.data() : a.data();
T* pC = c.data();
?
int m = c.getSize(1); // stride 1 size
int n = c.getSize(0); // other size
int k = transA ? a.getSize(0) : a.getSize(1);
?
int lda = transC ? a.getStride(0) : b.getStride(0);
int ldb = transC ? b.getStride(0) : a.getStride(0);
int ldc = c.getStride(0);
?
auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
?
......
?
auto err = CublasGemm<T>::gemm(handle,
gemmTrA, gemmTrB,
m, n, k, alpha,
pA, lda, pB, ldb, beta,
pC, ldc, useHgemm);
基本上的原因都在上文中解釋過了,不過還是再稍微解釋下
- pA 指向了b, 即上文提到D, 且需要轉置gemmTrA 為
CUBLAS_OP_T - pB 指向了a, 即上文中提到的Q, 無需轉置gemmTrB為
CUBLAS_OP_N - m,n,k 分別是4,3,2 即正在進行的矩陣計算為(4,2) * (2,3) = (4,3) 和上文中的式2對應起來了
- lda,ldb,ldc 分別為2,2,4 代表了各自矩陣的行或者列數(因為pA轉置了,按列數計算,pB,pC未轉置按行計算,這和實際的輸入對應起來的即D,Q,R的實際列數)
- 最後的就是cublas的矩陣計算函數
CublasGemm<T>::gemm
最後附上完整測試代碼
#include <faiss/gpu/StandardGpuResources.h>
?
#include <faiss/gpu/utils/DeviceTensor.cuh>
#include <faiss/gpu/utils/MatrixMult.cuh>
#include <faiss/gpu/utils/CopyUtils.cuh>
?
#include <iostream>
#include <random>
#include <algorithm>
?
?
void calc(float *d_h,float *q_h,int d_n,int q_n, int dim,float *y_h){
//init gpu resource(memory && stream)
int dev = faiss::gpu::getCurrentDevice();
faiss::gpu::StandardGpuResources gpuResources;
?
auto stream = gpuResources.getDefaultStream(dev);
?
//copy from cpu to gpu
auto d_d_tmp = faiss::gpu::toDevice<float, 2>(&gpuResources, dev, d_h, stream, {d_n, dim});
faiss::gpu::Tensor<float, 2, true> d_d(d_d_tmp.data(),d_d_tmp.sizes());
?
auto q_d_tmp = faiss::gpu::toDevice<float, 2>(&gpuResources, dev, q_h, stream, {q_n, dim});
faiss::gpu::Tensor<float, 2, true> q_d(q_d_tmp.data(),q_d_tmp.sizes());
?
?
//init gpu output
faiss::gpu::DeviceTensor<float, 2, true> y_d_tmp({q_n,d_n});
faiss::gpu::Tensor<float, 2, true> y_d(y_d_tmp.data(), y_d_tmp.sizes());
?
//compute
faiss::gpu::runMatrixMult(y_d, false,
q_d, false,
d_d,true,
1.0f, 0.0f, false,
gpuResources.getBlasHandleCurrentDevice(),
stream);
?
//copy from gpu to cpu
faiss::gpu::fromDevice<float,2>(y_d, y_h, stream);
?
}
?
void run(int d_n,int q_n,int dim){
if(d_n<1 || q_n <1){
return;
}
?
std::random_device rd; // 將用於獲得隨機數引擎的種子
std::mt19937 gen(rd()); // 以 rd() 播種的標準 mersenne_twister_engine
std::uniform_real_distribution<float> dis(0, 1);
?
// init input data
float *d_h = new float[d_n * dim];
for(int i = 0; i < d_n; i++) {
for(int j = 0; j < dim; j++)
d_h[dim * i + j] = dis(gen);
}
?
float *q_h = new float[q_n * dim];
for(int i = 0; i < q_n; i++) {
for(int j = 0; j < dim; j++)
q_h[dim * i + j] = dis(gen);
}
?
//init output data
float *y_h = new float[q_n*d_n];
?
calc(d_h,q_h,d_n,q_n,dim,y_h);
?
int d_n_display=std::min(d_n,10);
int q_n_display=std::min(q_n,10);
?
//display result
for(int i = 0; i <q_n_display; i++) {
for(int j = 0; j < dim; j++){
std::cout <<q_h[dim * i + j]<<",";
}
std::cout <<std::endl;
?
}
?
std::cout <<"====="<<std::endl;
?
for(int i = 0; i <d_n_display; i++) {
for(int j = 0; j < dim; j++){
std::cout <<d_h[dim * i + j]<<",";
}
std::cout <<std::endl;
?
}
?
std::cout <<"====="<<std::endl;
for(int i = 0; i <d_n_display*q_n_display; i++) {
std::cout <<y_h[i]<<",";
}
std::cout <<std::endl;
}
?
int main() {
?
int d_n = 4;
int q_n = 3;
int dim = 2;
?
run(d_n,q_n,dim);
?
}
小結
- 向量的內積使用cublas的矩陣相乘函數
- 注意cublas的矩陣按列填充特性
- 注意分清cublas中的矩陣和之外的矩陣的差異性以及仔細設計輸入和輸出
推薦閱讀：
