src/caffe/util/math_functions.cpp

   1 // Copyright 2013 Yangqing Jia
   2
   3 #include <mkl.h>
   4 #include <cublas_v2.h>
   5 #include "caffe/common.hpp"
   6 #include "caffe/util/math_functions.hpp"
   7
   8 namespace caffe {
   9
  10 template<>
  11 void caffe_cpu_gemm<float>(const CBLAS_TRANSPOSE TransA,
  12     const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
  13     const float alpha, const float* A, const float* B, const float beta,
  14     float* C) {
  15   int lda = (TransA == CblasNoTrans) ? K : M;
  16   int ldb = (TransB == CblasNoTrans) ? N : K;
  17   cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
  18       ldb, beta, C, N);
  19 }
  20
  21 template<>
  22 void caffe_cpu_gemm<double>(const CBLAS_TRANSPOSE TransA,
  23     const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
  24     const double alpha, const double* A, const double* B, const double beta,
  25     double* C) {
  26   int lda = (TransA == CblasNoTrans) ? K : M;
  27   int ldb = (TransB == CblasNoTrans) ? N : K;
  28   cblas_dgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
  29       ldb, beta, C, N);
  30 }
  31
  32 template <>
  33 void caffe_gpu_gemm<float>(const CBLAS_TRANSPOSE TransA,
  34     const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
  35     const float alpha, const float* A, const float* B, const float beta,
  36     float* C) {
  37   // Note that cublas follows fortran order.
  38   int lda = (TransA == CblasNoTrans) ? K : M;
  39   int ldb = (TransB == CblasNoTrans) ? N : K;
  40   cublasOperation_t cuTransA =
  41       (TransA == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
  42   cublasOperation_t cuTransB =
  43       (TransB == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
  44   CUBLAS_CHECK(cublasSgemm(Caffe::cublas_handle(), cuTransB, cuTransA,
  45       N, M, K, &alpha, B, ldb, A, lda, &beta, C, N));
  46 }
  47
  48 template <>
  49 void caffe_gpu_gemm<double>(const CBLAS_TRANSPOSE TransA,
  50     const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
  51     const double alpha, const double* A, const double* B, const double beta,
  52     double* C) {
  53   // Note that cublas follows fortran order.
  54   int lda = (TransA == CblasNoTrans) ? K : M;
  55   int ldb = (TransB == CblasNoTrans) ? N : K;
  56   cublasOperation_t cuTransA =
  57       (TransA == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
  58   cublasOperation_t cuTransB =
  59       (TransB == CblasNoTrans) ? CUBLAS_OP_N : CUBLAS_OP_T;
  60   CUBLAS_CHECK(cublasDgemm(Caffe::cublas_handle(), cuTransB, cuTransA,
  61       N, M, K, &alpha, B, ldb, A, lda, &beta, C, N));
  62 }
  63
  64 template <>
  65 void caffe_cpu_gemv<float>(const CBLAS_TRANSPOSE TransA, const int M,
  66     const int N, const float alpha, const float* A, const float* x,
  67     const float beta, float* y) {
  68   cblas_sgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
  69 }
  70
  71 template <>
  72 void caffe_cpu_gemv<double>(const CBLAS_TRANSPOSE TransA, const int M,
  73     const int N, const double alpha, const double* A, const double* x,
  74     const double beta, double* y) {
  75   cblas_dgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
  76 }
  77
  78 template <>
  79 void caffe_gpu_gemv<float>(const CBLAS_TRANSPOSE TransA, const int M,
  80     const int N, const float alpha, const float* A, const float* x,
  81     const float beta, float* y) {
  82   cublasOperation_t cuTransA =
  83       (TransA == CblasNoTrans) ? CUBLAS_OP_T : CUBLAS_OP_N;
  84   CUBLAS_CHECK(cublasSgemv(Caffe::cublas_handle(), cuTransA, N, M, &alpha,
  85       A, N, x, 1, &beta, y, 1));
  86 }
  87
  88 template <>
  89 void caffe_gpu_gemv<double>(const CBLAS_TRANSPOSE TransA, const int M,
  90     const int N, const double alpha, const double* A, const double* x,
  91     const double beta, double* y) {
  92   cublasOperation_t cuTransA =
  93       (TransA == CblasNoTrans) ? CUBLAS_OP_T : CUBLAS_OP_N;
  94   CUBLAS_CHECK(cublasDgemv(Caffe::cublas_handle(), cuTransA, N, M, &alpha,
  95       A, N, x, 1, &beta, y, 1));
  96 }
  97
  98 template <>
  99 void caffe_axpy<float>(const int N, const float alpha, const float* X,
 100     float* Y) { cblas_saxpy(N, alpha, X, 1, Y, 1); }
 101
 102 template <>
 103 void caffe_axpy<double>(const int N, const double alpha, const double* X,
 104     double* Y) { cblas_daxpy(N, alpha, X, 1, Y, 1); }
 105
 106 template <>
 107 void caffe_copy<float>(const int N, const float* X, float* Y) {
 108   cblas_scopy(N, X, 1, Y, 1);
 109 }
 110
 111 template <>
 112 void caffe_copy<double>(const int N, const double* X, double* Y) {
 113   cblas_dcopy(N, X, 1, Y, 1);
 114 }
 115
 116 template <>
 117 void caffe_scal<float>(const int N, const float alpha, float *X) {
 118   cblas_sscal(N, alpha, X, 1);
 119 }
 120
 121 template <>
 122 void caffe_scal<double>(const int N, const double alpha, double *X) {
 123   cblas_dscal(N, alpha, X, 1);
 124 }
 125
 126 template <>
 127 void caffe_gpu_scal<float>(const int N, const float alpha, float *X) {
 128   CUBLAS_CHECK(cublasSscal(Caffe::cublas_handle(), N, &alpha, X, 1));
 129 }
 130
 131 template <>
 132 void caffe_gpu_scal<double>(const int N, const double alpha, double *X) {
 133   CUBLAS_CHECK(cublasDscal(Caffe::cublas_handle(), N, &alpha, X, 1));
 134 }
 135
 136 template <>
 137 void caffe_sqr<float>(const int n, const float* a, float* y) {
 138   vsSqr(n, a, y);
 139 }
 140
 141 template <>
 142 void caffe_sqr<double>(const int n, const double* a, double* y) {
 143   vdSqr(n, a, y);
 144 }
 145
 146 template <>
 147 void caffe_mul<float>(const int n, const float* a, const float* b,
 148     float* y) { vsMul(n, a, b, y); }
 149
 150 template <>
 151 void caffe_mul<double>(const int n, const double* a, const double* b,
 152     double* y) { vdMul(n, a, b, y); }
 153
 154 template <>
 155 void caffe_div<float>(const int n, const float* a, const float* b,
 156     float* y) { vsDiv(n, a, b, y); }
 157
 158 template <>
 159 void caffe_div<double>(const int n, const double* a, const double* b,
 160     double* y) { vdDiv(n, a, b, y); }
 161
 162 template <>
 163 void caffe_powx<float>(const int n, const float* a, const float b,
 164     float* y) { vsPowx(n, a, b, y); }
 165
 166 template <>
 167 void caffe_powx<double>(const int n, const double* a, const double b,
 168     double* y) { vdPowx(n, a, b, y); }
 169
 170 template <>
 171 void caffe_vRngUniform<float>(const int n, float* r,
 172     const float a, const float b) {
 173   VSL_CHECK(vsRngUniform(VSL_RNG_METHOD_UNIFORM_STD, Caffe::vsl_stream(),
 174       n, r, a, b));
 175 }
 176
 177 template <>
 178 void caffe_vRngUniform<double>(const int n, double* r,
 179     const double a, const double b) {
 180   VSL_CHECK(vdRngUniform(VSL_RNG_METHOD_UNIFORM_STD, Caffe::vsl_stream(),
 181       n, r, a, b));
 182 }
 183
 184 template <>
 185 void caffe_vRngGaussian<float>(const int n, float* r, const float a,
 186     const float sigma) {
 187   VSL_CHECK(vsRngGaussian(VSL_RNG_METHOD_GAUSSIAN_BOXMULLER,
 188       Caffe::vsl_stream(), n, r, a, sigma));
 189 }
 190
 191
 192 template <>
 193 void caffe_vRngGaussian<double>(const int n, double* r, const double a,
 194     const double sigma) {
 195   VSL_CHECK(vdRngGaussian(VSL_RNG_METHOD_GAUSSIAN_BOXMULLER,
 196       Caffe::vsl_stream(), n, r, a, sigma));
 197 }
 198
 199 template <>
 200 void caffe_exp<float>(const int n, const float* a, float* y) {
 201   vsExp(n, a, y);
 202 }
 203
 204 template <>
 205 void caffe_exp<double>(const int n, const double* a, double* y) {
 206   vdExp(n, a, y);
 207 }
 208
 209 template <>
 210 float caffe_cpu_dot<float>(const int n, const float* x, const float* y) {
 211   return cblas_sdot(n, x, 1, y, 1);
 212 }
 213
 214 template <>
 215 double caffe_cpu_dot<double>(const int n, const double* x, const double* y) {
 216   return cblas_ddot(n, x, 1, y, 1);
 217 }
 218
 219 }  // namespace caffe