BlasExtra.h

Go to the documentation of this file.

/*-------------------------------------------------------------------
Copyright 2011 Ravishankar Sundararaman

This file is part of JDFTx.

JDFTx is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

JDFTx is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with JDFTx.  If not, see <http://www.gnu.org/licenses/>.
-------------------------------------------------------------------*/

#ifndef JDFTX_CORE_BLASEXTRA_H
#define JDFTX_CORE_BLASEXTRA_H

#include <gsl/gsl_cblas.h>
#include <cstdlib>
#include <cstdio>
#include <cfloat>
#include <core/scalar.h>
#include <core/Thread.h>

#ifdef GPU_ENABLED
#include <cublas.h>
#include <cuda_runtime.h>
#endif

template<typename Ty, typename Tx> void eblas_mul(const int N, const Tx* X, const int incX, Ty* Y, const int incY);

inline void eblas_dmul(const int N, const double* X, const int incX, double* Y, const int incY) { eblas_mul(N,X,incX,Y,incY); }
inline void eblas_zmul(const int N, const complex* X, const int incX, complex* Y, const int incY) { eblas_mul(N,X,incX,Y,incY); }
inline void eblas_zmuld(const int N, const double* X, const int incX, complex* Y, const int incY) { eblas_mul(N,X,incX,Y,incY); }
#ifdef GPU_ENABLED
//GPU versions of the above functions implemented in BlasExtra.cu
void eblas_dmul_gpu(const int N, const double* X, const int incX, double* Y, const int incY);
void eblas_zmul_gpu(const int N, const complex* X, const int incX, complex* Y, const int incY);
void eblas_zmuld_gpu(const int N, const double* X, const int incX, complex* Y, const int incY);
#endif

template<typename Ty, typename Tx> void eblas_div(const int N, const Tx* X, const int incX, Ty* Y, const int incY);

inline void eblas_ddiv(const int N, const double* X, const int incX, double* Y, const int incY) { eblas_div(N,X,incX,Y,incY); }
inline void eblas_zdiv(const int N, const complex* X, const int incX, complex* Y, const int incY) { eblas_div(N,X,incX,Y,incY); }
inline void eblas_zdivd(const int N, const double* X, const int incX, complex* Y, const int incY) { eblas_div(N,X,incX,Y,incY); }
#ifdef GPU_ENABLED
//GPU versions of the above functions implemented in BlasExtra.cu
void eblas_ddiv_gpu(const int N, const double* X, const int incX, double* Y, const int incY);
void eblas_zdiv_gpu(const int N, const complex* X, const int incX, complex* Y, const int incY);
void eblas_zdivd_gpu(const int N, const double* X, const int incX, complex* Y, const int incY);
#endif

void eblas_lincomb(const int N,
        const complex& sX, const complex* X, const int incX,
        const complex& sY, const complex* Y, const int incY,
        complex* Z, const int incZ);

#ifdef GPU_ENABLED
void eblas_lincomb_gpu(const int N,
        const complex& sX, const complex* X, const int incX,
        const complex& sY, const complex* Y, const int incY,
        complex* Z, const int incZ);
#endif

void eblas_zgemm(CBLAS_TRANSPOSE TransA, CBLAS_TRANSPOSE TransB, int M, int N, int K,
        const complex& alpha, const complex *A, const int lda, const complex *B, const int ldb,
        const complex& beta, complex *C, const int ldc);
#ifdef GPU_ENABLED
void eblas_zgemm_gpu(CBLAS_TRANSPOSE TransA, CBLAS_TRANSPOSE TransB, int M, int N, int K,
        const complex& alpha, const complex *A, const int lda, const complex *B, const int ldb,
        const complex& beta, complex *C, const int ldc);
#endif

//Sparse<->dense vector operations:
void eblas_scatter_zdaxpy(const int Nindex, double a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_scatter_zaxpy(const int Nindex, complex a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_scatter_daxpy(const int Nindex, double a, const int* index, const double* x, double* y);

void eblas_gather_zdaxpy(const int Nindex, double a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_gather_zaxpy(const int Nindex, complex a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_gather_daxpy(const int Nindex, double a, const int* index, const double* x, double* y);

#ifdef GPU_ENABLED
void eblas_scatter_zdaxpy_gpu(const int Nindex, double a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_scatter_zaxpy_gpu(const int Nindex, complex a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_scatter_daxpy_gpu(const int Nindex, double a, const int* index, const double* x, double* y);
void eblas_gather_zdaxpy_gpu(const int Nindex, double a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_gather_zaxpy_gpu(const int Nindex, complex a, const int* index, const complex* x, complex* y, bool conjugate=false);
void eblas_gather_daxpy_gpu(const int Nindex, double a, const int* index, const double* x, double* y);
#endif

void eblas_accumNorm(int N, const double& a, const complex* x, double* y);
void eblas_accumProd(int N, const double& a, const complex* xU, const complex* xC, double* yRe, double* yIm);
#ifdef GPU_ENABLED
void eblas_accumNorm_gpu(int N, const double& a, const complex* x, double* y);
void eblas_accumProd_gpu(int N, const double& a, const complex* xU, const complex* xC, double* yRe, double* yIm);
#endif

void eblas_symmetrize(int N, int n, const int* symmIndex, double* x);
void eblas_symmetrize(int N, int n, const int* symmIndex, complex* x);
#ifdef GPU_ENABLED
void eblas_symmetrize_gpu(int N, int n, const int* symmIndex, double* x);
void eblas_symmetrize_gpu(int N, int n, const int* symmIndex, complex* x);
#endif

//Threaded-wrappers for BLAS1 functions (Cblas)
template<typename T> void eblas_copy(T* dest, const T* src, int N) { memcpy(dest, src, N*sizeof(T)); }
void eblas_zero(int N, double* x);
void eblas_zero(int N, complex* x);
void eblas_dscal(int N, double a, double* x, int incx);
void eblas_zdscal(int N, double a, complex* x, int incx);
void eblas_zscal(int N, const complex& a, complex* x, int incx);
void eblas_daxpy(int N, double a, const double* x, int incx, double* y, int incy);
void eblas_zaxpy(int N, const complex& a, const complex* x, int incx, complex* y, int incy);
complex eblas_zdotc(int N, const complex* x, int incx, const complex* y, int incy);
double eblas_ddot(int N, const double* x, int incx, const double* y, int ncy);
double eblas_dznrm2(int N, const complex* x, int incx);
double eblas_dnrm2(int N, const double* x, int incx);

//Wrappers/alternate names for CUBLAS functions (auto-selectable with Cblas ones above using callPref)
#ifdef GPU_ENABLED
template<typename T> void eblas_copy_gpu(T* dest, const T* src, int N) { cudaMemcpy(dest, src, N*sizeof(T), cudaMemcpyDeviceToDevice); }
void eblas_zero_gpu(int N, double* x);
void eblas_zero_gpu(int N, complex* x);
#define eblas_dscal_gpu cublasDscal
void eblas_zdscal_gpu(int N, double a, complex* x, int incx);
void eblas_zscal_gpu(int N, const complex& a, complex* x, int incx);
#define eblas_daxpy_gpu cublasDaxpy
void eblas_zaxpy_gpu(int N, const complex& a, const complex* x, int incx, complex* y, int incy);
complex eblas_zdotc_gpu(int N, const complex* x, int incx, const complex* y, int incy);
#define eblas_ddot_gpu cublasDdot
double eblas_dznrm2_gpu(int N, const complex* x, int incx);
#define eblas_dnrm2_gpu cublasDnrm2

#endif

#ifdef GPU_ENABLED
        #define callPref(functionName) functionName##_gpu //gpu versions available, so use them preferentially
#else
        #define callPref(functionName) functionName //only cpu version available
#endif


void eblas_capMinMax(const int N, double* x, double& xMin, double& xMax, double capLo=-DBL_MAX, double capHi=+DBL_MAX);
#ifdef GPU_ENABLED
void eblas_capMinMax_gpu(const int N, double* x, double& xMin, double& xMax, double capLo=-DBL_MAX, double capHi=+DBL_MAX);
#endif


//###################################################################################################
//####  Implementation  ####
//##########################


//Elementwise multiply implementation:
template<typename Ty, typename Tx>
void eblas_mul_sub(size_t iMin, size_t iMax, const Tx* X, const int incX, Ty* Y, const int incY)
{       for(size_t i=iMin; i<iMax; i++) Y[incY*i] *= X[incX*i];
}

template<typename Ty, typename Tx>
void eblas_mul(const int N, const Tx* X, const int incX, Ty* Y, const int incY)
{       if(incY==0) die("incY cannot be = 0")
        threadLaunch((N<100000) ? 1 : 0, //force single threaded for small problem sizes
                eblas_mul_sub<Ty,Tx>, N, X, incX, Y, incY);
}

//Elementwise divide implementation:
template<typename Ty, typename Tx>
void eblas_div_sub(size_t iMin, size_t iMax, const Tx* X, const int incX, Ty* Y, const int incY)
{       for(size_t i=iMin; i<iMax; i++) Y[incY*i] /= X[incX*i];
}
template<typename Ty, typename Tx>
void eblas_div(const int N, const Tx* X, const int incX, Ty* Y, const int incY)
{       if(incY==0) die("incY cannot be = 0")
        threadLaunch((N<100000) ? 1 : 0, //force single threaded for small problem sizes
                eblas_div_sub<Ty,Tx>, N, X, incX, Y, incY);
}

#endif // JDFTX_CORE_BLASEXTRA_H