wehub-resource-sync
368 lines
16 KiB
C++
368 lines
16 KiB
C++
/* ******************************************************************************
|
|
*
|
|
*
|
|
* This program and the accompanying materials are made available under the
|
|
* terms of the Apache License, Version 2.0 which is available at
|
|
* https://www.apache.org/licenses/LICENSE-2.0.
|
|
*
|
|
* See the NOTICE file distributed with this work for additional
|
|
* information regarding copyright ownership.
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
* License for the specific language governing permissions and limitations
|
|
* under the License.
|
|
*
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
******************************************************************************/
|
|
|
|
//
|
|
// @author raver119@gmail.com
|
|
//
|
|
|
|
#ifndef LIBND4J_BLAS_HELPER_H
|
|
#define LIBND4J_BLAS_HELPER_H
|
|
|
|
// work around conflict with OpenBLAS
|
|
struct bfloat16;
|
|
#define BFLOAT16 BFLOAT16
|
|
#if !SD_CUDA
|
|
#include <blas/cblas.h>
|
|
#endif
|
|
#include <helpers/logger.h>
|
|
#include <types/float16.h>
|
|
#include <mutex>
|
|
#include <atomic>
|
|
|
|
#ifdef _WIN32
|
|
#define CUBLASWINAPI __stdcall
|
|
#define CUSOLVERAPI __stdcall
|
|
#else
|
|
#define CUBLASWINAPI
|
|
#define CUSOLVERAPI
|
|
#endif
|
|
|
|
namespace sd {
|
|
typedef enum {
|
|
CUBLAS_STATUS_SUCCESS = 0,
|
|
CUBLAS_STATUS_NOT_INITIALIZED = 1,
|
|
CUBLAS_STATUS_ALLOC_FAILED = 3,
|
|
CUBLAS_STATUS_INVALID_VALUE = 7,
|
|
CUBLAS_STATUS_ARCH_MISMATCH = 8,
|
|
CUBLAS_STATUS_MAPPING_ERROR = 11,
|
|
CUBLAS_STATUS_EXECUTION_FAILED = 13,
|
|
CUBLAS_STATUS_INTERNAL_ERROR = 14,
|
|
CUBLAS_STATUS_NOT_SUPPORTED = 15,
|
|
CUBLAS_STATUS_LICENSE_ERROR = 16
|
|
} cublasStatus_t;
|
|
|
|
typedef enum { CUBLAS_OP_N = 0, CUBLAS_OP_T = 1, CUBLAS_OP_C = 2 } cublasOperation_t;
|
|
|
|
struct cublasContext;
|
|
typedef struct cublasContext *cublasHandle_t;
|
|
|
|
typedef enum {
|
|
CUDA_R_16F = 2, /* real as a half */
|
|
CUDA_C_16F = 6, /* complex as a pair of half numbers */
|
|
CUDA_R_32F = 0, /* real as a float */
|
|
CUDA_C_32F = 4, /* complex as a pair of float numbers */
|
|
CUDA_R_64F = 1, /* real as a double */
|
|
CUDA_C_64F = 5, /* complex as a pair of double numbers */
|
|
CUDA_R_8I = 3, /* real as a signed char */
|
|
CUDA_C_8I = 7, /* complex as a pair of signed char numbers */
|
|
CUDA_R_8U = 8, /* real as a unsigned char */
|
|
CUDA_C_8U = 9, /* complex as a pair of unsigned char numbers */
|
|
CUDA_R_32I = 10, /* real as a signed int */
|
|
CUDA_C_32I = 11, /* complex as a pair of signed int numbers */
|
|
CUDA_R_32U = 12, /* real as a unsigned int */
|
|
CUDA_C_32U = 13 /* complex as a pair of unsigned int numbers */
|
|
} cublasDataType_t;
|
|
|
|
#if !defined(SD_CUDA)
|
|
typedef void (*CblasSgemv)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, int M, int N, float alpha, float *A, int lda,
|
|
float *X, int incX, float beta, float *Y, int incY);
|
|
|
|
typedef void (*CblasDgemv)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, int M, int N, double alpha, double *A, int lda,
|
|
double *X, int incX, double beta, double *Y, int incY);
|
|
|
|
typedef void (*CblasSgemm)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, CBLAS_TRANSPOSE TransB, int M, int N, int K,
|
|
float alpha, float *A, int lda, float *B, int ldb, float beta, float *C, int ldc);
|
|
|
|
typedef void (*CblasDgemm)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, CBLAS_TRANSPOSE TransB, int M, int N, int K,
|
|
double alpha, double *A, int lda, double *B, int ldb, double beta, double *C, int ldc);
|
|
|
|
typedef void (*CblasSgemmBatch)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE *TransA_Array, CBLAS_TRANSPOSE *TransB_Array,
|
|
int *M_Array, int *N_Array, int *K_Array, float *alpha_Array, float **A_Array,
|
|
int *lda_Array, float **B_Array, int *ldb_Array, float *beta_Array, float **C_Array,
|
|
int *ldc_Array, int group_count, int *group_size);
|
|
|
|
typedef void (*CblasDgemmBatch)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE *TransA_Array, CBLAS_TRANSPOSE *TransB_Array,
|
|
int *M_Array, int *N_Array, int *K_Array, double *alpha_Array, double **A_Array,
|
|
int *lda_Array, double **B_Array, int *ldb_Array, double *beta_Array, double **C_Array,
|
|
int *ldc_Array, int group_count, int *group_size);
|
|
#endif
|
|
#ifdef LAPACK_ROW_MAJOR
|
|
#undef LAPACK_ROW_MAJOR
|
|
#endif
|
|
|
|
#ifdef LAPACK_COL_MAJOR
|
|
#undef LAPACK_COL_MAJOR
|
|
#endif
|
|
enum LAPACK_LAYOUT { LAPACK_ROW_MAJOR = 101, LAPACK_COL_MAJOR = 102 };
|
|
|
|
typedef int (*LapackeSgesvd)(LAPACK_LAYOUT matrix_layout, char jobu, char jobvt, int m, int n, float *a, int lda,
|
|
float *s, float *u, int ldu, float *vt, int ldvt, float *superb);
|
|
|
|
typedef int (*LapackeDgesvd)(LAPACK_LAYOUT matrix_layout, char jobu, char jobvt, int m, int n, double *a, int lda,
|
|
double *s, double *u, int ldu, double *vt, int ldvt, double *superb);
|
|
|
|
typedef int (*LapackeSgesdd)(LAPACK_LAYOUT matrix_layout, char jobz, int m, int n, float *a, int lda, float *s,
|
|
float *u, int ldu, float *vt, int ldvt);
|
|
typedef int (*LapackeDgesdd)(LAPACK_LAYOUT matrix_layout, char jobz, int m, int n, double *a, int lda, double *s,
|
|
double *u, int ldu, double *vt, int ldvt);
|
|
|
|
#if defined(SD_CUDA)
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasSgemv)(cublasHandle_t handle, cublasOperation_t trans, int m, int n,
|
|
float *alpha, /* host or device pointer */
|
|
float *A, int lda, float *x, int incx,
|
|
float *beta, /* host or device pointer */
|
|
float *y, int incy);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasDgemv)(cublasHandle_t handle, cublasOperation_t trans, int m, int n,
|
|
double *alpha, /* host or device pointer */
|
|
double *A, int lda, double *x, int incx,
|
|
double *beta, /* host or device pointer */
|
|
double *y, int incy);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasHgemm)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
__half *alpha, /* host or device pointer */
|
|
__half *A, int lda, __half *B, int ldb,
|
|
__half *beta, /* host or device pointer */
|
|
__half *C, int ldc);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasSgemm)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
float *alpha, /* host or device pointer */
|
|
float *A, int lda, float *B, int ldb,
|
|
float *beta, /* host or device pointer */
|
|
float *C, int ldc);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasDgemm)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
double *alpha, /* host or device pointer */
|
|
double *A, int lda, double *B, int ldb,
|
|
double *beta, /* host or device pointer */
|
|
double *C, int ldc);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasSgemmEx)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
float *alpha, /* host or device pointer */
|
|
void *A, cublasDataType_t Atype, int lda, void *B,
|
|
cublasDataType_t Btype, int ldb,
|
|
float *beta, /* host or device pointer */
|
|
void *C, cublasDataType_t Ctype, int ldc);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasHgemmBatched)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
__half *alpha, /* host or device pointer */
|
|
__half *Aarray[], int lda, __half *Barray[], int ldb,
|
|
__half *beta, /* host or device pointer */
|
|
__half *Carray[], int ldc, int batchCount);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasSgemmBatched)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
float *alpha, /* host or device pointer */
|
|
float *Aarray[], int lda, float *Barray[], int ldb,
|
|
float *beta, /* host or device pointer */
|
|
float *Carray[], int ldc, int batchCount);
|
|
|
|
typedef cublasStatus_t(CUBLASWINAPI *CublasDgemmBatched)(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k,
|
|
double *alpha, /* host or device pointer */
|
|
double *Aarray[], int lda, double *Barray[], int ldb,
|
|
double *beta, /* host or device pointer */
|
|
double *Carray[], int ldc, int batchCount);
|
|
|
|
typedef enum {
|
|
CUSOLVER_STATUS_SUCCESS = 0,
|
|
CUSOLVER_STATUS_NOT_INITIALIZED = 1,
|
|
CUSOLVER_STATUS_ALLOC_FAILED = 2,
|
|
CUSOLVER_STATUS_INVALID_VALUE = 3,
|
|
CUSOLVER_STATUS_ARCH_MISMATCH = 4,
|
|
CUSOLVER_STATUS_MAPPING_ERROR = 5,
|
|
CUSOLVER_STATUS_EXECUTION_FAILED = 6,
|
|
CUSOLVER_STATUS_INTERNAL_ERROR = 7,
|
|
CUSOLVER_STATUS_MATRIX_TYPE_NOT_SUPPORTED = 8,
|
|
CUSOLVER_STATUS_NOT_SUPPORTED = 9,
|
|
CUSOLVER_STATUS_ZERO_PIVOT = 10,
|
|
CUSOLVER_STATUS_INVALID_LICENSE = 11
|
|
} cusolverStatus_t;
|
|
|
|
typedef enum { CUSOLVER_EIG_TYPE_1 = 1, CUSOLVER_EIG_TYPE_2 = 2, CUSOLVER_EIG_TYPE_3 = 3 } cusolverEigType_t;
|
|
|
|
typedef enum { CUSOLVER_EIG_MODE_NOVECTOR = 0, CUSOLVER_EIG_MODE_VECTOR = 1 } cusolverEigMode_t;
|
|
|
|
struct cusolverDnContext;
|
|
typedef struct cusolverDnContext *cusolverDnHandle_t;
|
|
|
|
typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnSgesvdBufferSize)(cusolverDnHandle_t handle, int m, int n, int *lwork);
|
|
|
|
typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnDgesvdBufferSize)(cusolverDnHandle_t handle, int m, int n, int *lwork);
|
|
|
|
typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnSgesvd)(cusolverDnHandle_t handle, signed char jobu, signed char jobvt,
|
|
int m, int n, float *A, int lda, float *S, float *U, int ldu,
|
|
float *VT, int ldvt, float *work, int lwork, float *rwork,
|
|
int *info);
|
|
|
|
typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnDgesvd)(cusolverDnHandle_t handle, signed char jobu, signed char jobvt,
|
|
int m, int n, double *A, int lda, double *S, double *U, int ldu,
|
|
double *VT, int ldvt, double *work, int lwork, double *rwork,
|
|
int *info);
|
|
#endif
|
|
enum BlasFunctions {
|
|
GEMV = 0,
|
|
GEMM = 1,
|
|
};
|
|
|
|
class BlasHelper {
|
|
private:
|
|
// Private constructor for singleton pattern
|
|
BlasHelper();
|
|
|
|
bool _hasHgemv = false;
|
|
bool _hasHgemm = false;
|
|
bool _hasHgemmBatch = false;
|
|
|
|
bool _hasSgemv = false;
|
|
bool _hasSgemm = false;
|
|
bool _hasSgemmBatch = false;
|
|
|
|
bool _hasDgemv = false;
|
|
bool _hasDgemm = false;
|
|
bool _hasDgemmBatch = false;
|
|
|
|
// Mutex for serializing BLAS calls to prevent OpenBLAS TLS corruption
|
|
// and race conditions when called from multiple threads concurrently
|
|
mutable std::mutex _blasMutex;
|
|
|
|
// Controls whether BLAS calls are serialized (default: true for safety)
|
|
// Can be disabled via SD_BLAS_SERIALIZE=0 for testing or if using a thread-safe BLAS
|
|
std::atomic<bool> _serializeBlasCalls{true};
|
|
|
|
// Number of threads OpenBLAS should use (0 = use OpenBLAS default)
|
|
std::atomic<int> _openblasThreads{0};
|
|
#if !defined(SD_CUDA)
|
|
CblasSgemv cblasSgemv;
|
|
CblasDgemv cblasDgemv;
|
|
CblasSgemm cblasSgemm;
|
|
CblasDgemm cblasDgemm;
|
|
CblasSgemmBatch cblasSgemmBatch;
|
|
CblasDgemmBatch cblasDgemmBatch;
|
|
LapackeSgesvd lapackeSgesvd;
|
|
LapackeDgesvd lapackeDgesvd;
|
|
LapackeSgesdd lapackeSgesdd;
|
|
LapackeDgesdd lapackeDgesdd;
|
|
#endif
|
|
|
|
#if defined(SD_CUDA)
|
|
CublasSgemv cublasSgemv;
|
|
CublasDgemv cublasDgemv;
|
|
CublasHgemm cublasHgemm;
|
|
CublasSgemm cublasSgemm;
|
|
CublasDgemm cublasDgemm;
|
|
CublasSgemmEx cublasSgemmEx;
|
|
CublasHgemmBatched cublasHgemmBatched;
|
|
CublasSgemmBatched cublasSgemmBatched;
|
|
CublasDgemmBatched cublasDgemmBatched;
|
|
CusolverDnSgesvdBufferSize cusolverDnSgesvdBufferSize;
|
|
CusolverDnDgesvdBufferSize cusolverDnDgesvdBufferSize;
|
|
CusolverDnSgesvd cusolverDnSgesvd;
|
|
CusolverDnDgesvd cusolverDnDgesvd;
|
|
#endif
|
|
|
|
public:
|
|
static BlasHelper &getInstance();
|
|
|
|
void initializeFunctions(Pointer *functions);
|
|
void initializeDeviceFunctions(Pointer *functions);
|
|
|
|
template <typename T>
|
|
bool hasGEMV();
|
|
|
|
template <typename T>
|
|
bool hasGEMM();
|
|
|
|
bool hasGEMM(const DataType dtype);
|
|
bool hasGEMV(const DataType dtype);
|
|
|
|
template <typename T>
|
|
bool hasBatchedGEMM();
|
|
#if !defined(SD_CUDA)
|
|
CblasSgemv sgemv();
|
|
CblasDgemv dgemv();
|
|
|
|
CblasSgemm sgemm();
|
|
CblasDgemm dgemm();
|
|
|
|
CblasSgemmBatch sgemmBatched();
|
|
CblasDgemmBatch dgemmBatched();
|
|
|
|
LapackeSgesvd sgesvd();
|
|
LapackeDgesvd dgesvd();
|
|
|
|
LapackeSgesdd sgesdd();
|
|
LapackeDgesdd dgesdd();
|
|
#endif
|
|
|
|
// BLAS call serialization methods to prevent OpenBLAS TLS corruption
|
|
// and race conditions in multi-threaded environments
|
|
|
|
/**
|
|
* Returns a lock guard that serializes BLAS calls.
|
|
* Use this before making any BLAS call to prevent concurrent access.
|
|
* Example:
|
|
* auto lock = BlasHelper::getInstance().lockBlas();
|
|
* sgemm()(...);
|
|
*/
|
|
std::unique_lock<std::mutex> lockBlas() const;
|
|
|
|
/**
|
|
* Check if BLAS call serialization is enabled.
|
|
* Default is true for safety with OpenBLAS.
|
|
*/
|
|
bool isSerializeBlasCalls() const;
|
|
|
|
/**
|
|
* Enable or disable BLAS call serialization.
|
|
* Disable only if using a thread-safe BLAS implementation (e.g., MKL).
|
|
*/
|
|
void setSerializeBlasCalls(bool serialize);
|
|
|
|
/**
|
|
* Get the configured number of OpenBLAS threads.
|
|
* Returns 0 if using OpenBLAS default.
|
|
*/
|
|
int getOpenblasThreads() const;
|
|
|
|
/**
|
|
* Set the number of threads OpenBLAS should use.
|
|
* Set to 0 to use OpenBLAS default, or a specific number for control.
|
|
* This is applied when BLAS is initialized.
|
|
*/
|
|
void setOpenblasThreads(int threads);
|
|
|
|
/**
|
|
* Initialize OpenBLAS thread settings from environment variables.
|
|
* Checks SD_OPENBLAS_THREADS and SD_BLAS_SERIALIZE.
|
|
* Called automatically during getInstance() initialization.
|
|
*/
|
|
void initializeBlasThreading();
|
|
|
|
// destructor
|
|
~BlasHelper() noexcept;
|
|
};
|
|
} // namespace sd
|
|
|
|
#endif
|