Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
MatrixMult.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 // Copyright 2004-present Facebook. All Rights Reserved.
10 
11 #include "MatrixMult.cuh"
12 #include "DeviceMemory.h"
13 #include "DeviceUtils.h" // CUDA_VERIFY
14 #include "DeviceTensor.cuh"
15 #include "HostTensor.cuh"
16 
17 namespace faiss { namespace gpu {
18 
19 template <typename T>
20 struct CublasGemm {
21 };
22 
23 template <>
24 struct CublasGemm<float> {
25  static cublasStatus_t gemm(cublasHandle_t handle,
26  cublasOperation_t transa,
27  cublasOperation_t transb,
28  int m,
29  int n,
30  int k,
31  float fAlpha,
32  const float *A,
33  int lda,
34  const float *B,
35  int ldb,
36  float fBeta,
37  float *C,
38  int ldc,
39  bool useHgemm) {
40  return cublasSgemm(handle, transa, transb, m, n, k,
41  &fAlpha, A, lda, B, ldb, &fBeta, C, ldc);
42  }
43 };
44 
45 #ifdef FAISS_USE_FLOAT16
46 template <>
47 struct CublasGemm<half> {
48  static cublasStatus_t gemm(cublasHandle_t handle,
49  cublasOperation_t transa,
50  cublasOperation_t transb,
51  int m,
52  int n,
53  int k,
54  const float fAlpha,
55  const half *A,
56  int lda,
57  const half *B,
58  int ldb,
59  const float fBeta,
60  half *C,
61  int ldc,
62  bool useHgemm) {
63  if (getDeviceSupportsFloat16Math(getCurrentDevice()) && useHgemm) {
64  half hAlpha = hostFloat2Half(fAlpha);
65  half hBeta = hostFloat2Half(fBeta);
66 
67  return cublasHgemm(handle, transa, transb, m, n, k,
68  &hAlpha, A, lda, B, ldb, &hBeta, C, ldc);
69  }
70 
71  // CUDA 8.0 changes the half datatype specifier
72 #if CUDA_VERSION == 7050
73  auto halfType = CUBLAS_DATA_HALF;
74 #else
75  auto halfType = CUDA_R_16F;
76 #endif // CUDA_VERSION
77 
78  return cublasSgemmEx(handle, transa, transb, m, n, k,
79  &fAlpha, A, halfType, lda,
80  B, halfType, ldb,
81  &fBeta,
82  C, halfType, ldc);
83  }
84 };
85 #endif // FAISS_USE_FLOAT16
86 
87 
88 template <typename T>
89 void
90 runMatrixMult(Tensor<T, 2, true>& c, bool transC,
91  Tensor<T, 2, true>& a, bool transA,
92  Tensor<T, 2, true>& b, bool transB,
93  float alpha,
94  float beta,
95  bool useHgemm,
96  cublasHandle_t handle,
97  cudaStream_t stream) {
98  cublasSetStream(handle, stream);
99 
100  // Check that we have (m x k) * (k x n) = (m x n)
101  // using the input row-major layout
102  int aM = transA ? a.getSize(1) : a.getSize(0);
103  int aK = transA ? a.getSize(0) : a.getSize(1);
104 
105  int bK = transB ? b.getSize(1) : b.getSize(0);
106  int bN = transB ? b.getSize(0) : b.getSize(1);
107 
108  int cM = transC ? c.getSize(1) : c.getSize(0);
109  int cN = transC ? c.getSize(0) : c.getSize(1);
110 
111  FAISS_ASSERT(aM == cM);
112  FAISS_ASSERT(aK == bK);
113  FAISS_ASSERT(bN == cN);
114 
115  FAISS_ASSERT(a.getStride(1) == 1);
116  FAISS_ASSERT(b.getStride(1) == 1);
117  FAISS_ASSERT(c.getStride(1) == 1);
118 
119  // Now, we have to represent the matrix multiplication in
120  // column-major layout
121  T* pA = transC ? a.data() : b.data();
122  T* pB = transC ? b.data() : a.data();
123  T* pC = c.data();
124 
125  int m = c.getSize(1); // stride 1 size
126  int n = c.getSize(0); // other size
127  int k = transA ? a.getSize(0) : a.getSize(1);
128 
129  int lda = transC ? a.getStride(0) : b.getStride(0);
130  int ldb = transC ? b.getStride(0) : a.getStride(0);
131  int ldc = c.getStride(0);
132 
133  auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
134  auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
135 
136  if (transC) {
137  gemmTrA = transA ? CUBLAS_OP_N : CUBLAS_OP_T;
138  gemmTrB = transB ? CUBLAS_OP_N : CUBLAS_OP_T;
139  }
140 
141  auto err = CublasGemm<T>::gemm(handle,
142  gemmTrA, gemmTrB,
143  m, n, k, alpha,
144  pA, lda, pB, ldb, beta,
145  pC, ldc, useHgemm);
146 
147  FAISS_ASSERT_FMT(err == CUBLAS_STATUS_SUCCESS,
148  "cublas failed (%d): %s "
149  "(%d, %d)%s x (%d, %d)%s = (%d, %d)%s",
150  (int) err,
151  useHgemm ? "Hgemm" : "Sgemm",
152  a.getSize(0), a.getSize(1), transA ? "'" : "",
153  b.getSize(0), b.getSize(1), transB ? "'" : "",
154  c.getSize(0), c.getSize(1), transC ? "'" : "");
155  CUDA_TEST_ERROR();
156 }
157 
158 void runMatrixMult(Tensor<float, 2, true>& c, bool transC,
159  Tensor<float, 2, true>& a, bool transA,
160  Tensor<float, 2, true>& b, bool transB,
161  float alpha,
162  float beta,
163  bool useHgemm,
164  cublasHandle_t handle,
165  cudaStream_t stream) {
166  return runMatrixMult<float>(c, transC, a, transA, b, transB,
167  alpha, beta, useHgemm, handle, stream);
168 }
169 
170 #ifdef FAISS_USE_FLOAT16
171 void runMatrixMult(Tensor<half, 2, true>& c, bool transC,
172  Tensor<half, 2, true>& a, bool transA,
173  Tensor<half, 2, true>& b, bool transB,
174  float alpha,
175  float beta,
176  bool useHgemm,
177  cublasHandle_t handle,
178  cudaStream_t stream) {
179  return runMatrixMult<half>(c, transC, a, transA, b, transB,
180  alpha, beta, useHgemm, handle, stream);
181 }
182 #endif
183 
184 void
185 runIteratedMatrixMult(Tensor<float, 3, true>& c, bool transC,
186  Tensor<float, 3, true>& a, bool transA,
187  Tensor<float, 3, true>& b, bool transB,
188  float alpha,
189  float beta,
190  cublasHandle_t handle,
191  cudaStream_t stream) {
192  FAISS_ASSERT(c.getSize(0) == a.getSize(0));
193  FAISS_ASSERT(a.getSize(0) == b.getSize(0));
194 
195  for (int i = 0; i < a.getSize(0); ++i) {
196  auto cView = c[i].view();
197  auto aView = a[i].view();
198  auto bView = b[i].view();
199 
200  runMatrixMult(cView, transC,
201  aView, transA,
202  bView, transB,
203  alpha, beta, false, handle, stream);
204  }
205 }
206 
207 void
208 runBatchMatrixMult(Tensor<float, 3, true>& c, bool transC,
209  Tensor<float, 3, true>& a, bool transA,
210  Tensor<float, 3, true>& b, bool transB,
211  float alpha,
212  float beta,
213  DeviceMemory& mem,
214  cublasHandle_t handle,
215  cudaStream_t stream) {
216  FAISS_ASSERT(c.getSize(0) == a.getSize(0));
217  FAISS_ASSERT(a.getSize(0) == b.getSize(0));
218  cublasSetStream(handle, stream);
219 
220  // Check that we have (m x k) * (k x n) = (m x n)
221  // using the input row-major layout
222  int aM = transA ? a.getSize(2) : a.getSize(1);
223  int aK = transA ? a.getSize(1) : a.getSize(2);
224 
225  int bK = transB ? b.getSize(2) : b.getSize(1);
226  int bN = transB ? b.getSize(1) : b.getSize(2);
227 
228  int cM = transC ? c.getSize(2) : c.getSize(1);
229  int cN = transC ? c.getSize(1) : c.getSize(2);
230 
231  FAISS_ASSERT(aM == cM);
232  FAISS_ASSERT(aK == bK);
233  FAISS_ASSERT(bN == cN);
234 
235  // Now, we have to represent the matrix multiplication in
236  // column-major layout
237  float* pA = transC ? a.data() : b.data();
238  float* pB = transC ? b.data() : a.data();
239  float* pC = c.data();
240 
241  int m = c.getSize(2); // stride 1 size
242  int n = c.getSize(1); // other size
243  int k = transA ? a.getSize(1) : a.getSize(2);
244 
245  int lda = transC ? a.getStride(1) : b.getStride(1);
246  int ldb = transC ? b.getStride(1) : a.getStride(1);
247  int ldc = c.getStride(1);
248 
249  auto gemmTrA = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
250  auto gemmTrB = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
251 
252  if (transC) {
253  gemmTrA = transA ? CUBLAS_OP_N : CUBLAS_OP_T;
254  gemmTrB = transB ? CUBLAS_OP_N : CUBLAS_OP_T;
255  }
256 
257  HostTensor<float*, 1, true> hostA({a.getSize(0)});
258  HostTensor<float*, 1, true> hostB({b.getSize(0)});
259  HostTensor<float*, 1, true> hostC({c.getSize(0)});
260 
261  size_t aOffset = a.getStride(0);
262  size_t bOffset = b.getStride(0);
263  size_t cOffset = c.getStride(0);
264 
265  for (int i = 0; i < a.getSize(0); ++i) {
266  hostA[i] = transC ? a.data() + i * aOffset : b.data() + i * bOffset;
267  hostB[i] = transC ? b.data() + i * bOffset : a.data() + i * aOffset;
268  hostC[i] = c.data() + i * cOffset;
269  }
270 
271  DeviceTensor<float*, 1, true> deviceA(mem, hostA, stream);
272  DeviceTensor<float*, 1, true> deviceB(mem, hostB, stream);
273  DeviceTensor<float*, 1, true> deviceC(mem, hostC, stream);
274 
275  auto err =
276  cublasSgemmBatched(handle,
277  gemmTrA, gemmTrB,
278  m, n, k, &alpha,
279  (const float**) deviceA.data(), lda,
280  (const float**) deviceB.data(), ldb, &beta,
281  deviceC.data(), ldc, a.getSize(0));
282  FAISS_ASSERT_FMT(err == CUBLAS_STATUS_SUCCESS,
283  "cublasSgemmBatched failed (%d)", (int) err);
284  CUDA_TEST_ERROR();
285 }
286 
287 } } // namespace
Generated by   doxygen 1.8.5