wehub-resource-sync
226 lines
9.1 KiB
C++
226 lines
9.1 KiB
C++
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* 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.
|
|
*/
|
|
|
|
#include "cublasWrapper.h"
|
|
#include "common/checkMacrosPlugin.h"
|
|
#include "cudaDriverWrapper.h"
|
|
|
|
#if defined(_WIN32)
|
|
#if !defined(WIN32_LEAN_AND_MEAN)
|
|
#define WIN32_LEAN_AND_MEAN
|
|
// Ensure that macros appearing in multiple files are only defined once.
|
|
#endif // defined(WIN32_LEAN_AND_MEAN)
|
|
#include <windows.h>
|
|
#define dllOpen(name) (void*) LoadLibraryA(name)
|
|
#define dllClose(handle) FreeLibrary(static_cast<HMODULE>(handle))
|
|
#define dllGetSym(handle, name) GetProcAddress(static_cast<HMODULE>(handle), name)
|
|
auto const kCUBLAS_PLUGIN_LIBNAME
|
|
= std::string{"cublas64_"} + std::to_string(nvinfer1::getCudaLibVersionMaj()) + ".dll";
|
|
#else // defined(_WIN32)
|
|
#include <dlfcn.h>
|
|
#define dllOpen(name) dlopen(name, RTLD_LAZY)
|
|
#define dllClose(handle) dlclose(handle)
|
|
#define dllGetSym(handle, name) dlsym(handle, name)
|
|
auto const kCUBLAS_PLUGIN_LIBNAME = std::string{"libcublas.so."} + std::to_string(nvinfer1::getCudaLibVersionMaj());
|
|
#endif // defined(_WIN32)
|
|
|
|
namespace nvinfer1::pluginInternal
|
|
{
|
|
using namespace nvinfer1;
|
|
|
|
// If tryLoadingCublas failed, the CublasWrapper object won't be created.
|
|
CublasWrapper::CublasWrapper(bool initHandle)
|
|
: mLibrary(tryLoadingCublas())
|
|
{
|
|
PLUGIN_VALIDATE(mLibrary != nullptr);
|
|
auto load_sym = [](void* handle, char const* name) {
|
|
void* ret = dllGetSym(handle, name);
|
|
std::string loadError = "Fail to load symbol " + std::string(name) + " from the cublas library.";
|
|
PLUGIN_VALIDATE(ret != nullptr, loadError.c_str());
|
|
return ret;
|
|
};
|
|
*(void**) (&_cublasCreate) = load_sym(mLibrary, "cublasCreate_v2");
|
|
*(void**) (&_cublasDestroy) = load_sym(mLibrary, "cublasDestroy_v2");
|
|
*(void**) (&_cublasSetStream) = load_sym(mLibrary, "cublasSetStream_v2");
|
|
*(void**) (&_cublasGetPointerMode) = load_sym(mLibrary, "cublasGetPointerMode_v2");
|
|
*(void**) (&_cublasSetPointerMode) = load_sym(mLibrary, "cublasSetPointerMode_v2");
|
|
*(void**) (&_cublasGetMathMode) = load_sym(mLibrary, "cublasGetMathMode");
|
|
*(void**) (&_cublasSetMathMode) = load_sym(mLibrary, "cublasSetMathMode");
|
|
*(void**) (&_cublasDscal) = load_sym(mLibrary, "cublasDscal_v2");
|
|
*(void**) (&_cublasSasum) = load_sym(mLibrary, "cublasSasum_v2");
|
|
*(void**) (&_cublasScopy) = load_sym(mLibrary, "cublasScopy_v2");
|
|
*(void**) (&_cublasSscal) = load_sym(mLibrary, "cublasSscal_v2");
|
|
*(void**) (&_cublasSgemm) = load_sym(mLibrary, "cublasSgemm_v2");
|
|
*(void**) (&_cublasHgemm) = load_sym(mLibrary, "cublasHgemm");
|
|
*(void**) (&_cublasHgemmStridedBatched) = load_sym(mLibrary, "cublasHgemmStridedBatched");
|
|
*(void**) (&_cublasSgemmStridedBatched) = load_sym(mLibrary, "cublasSgemmStridedBatched");
|
|
*(void**) (&_cublasGemmEx) = load_sym(mLibrary, "cublasGemmEx");
|
|
*(void**) (&_cublasGemmStridedBatchedEx) = load_sym(mLibrary, "cublasGemmStridedBatchedEx");
|
|
|
|
if (initHandle)
|
|
{
|
|
PLUGIN_VALIDATE(cublasCreate(&mHandle) == CUBLAS_STATUS_SUCCESS, "Could not create cublas handle.");
|
|
PLUGIN_VALIDATE(mHandle != nullptr);
|
|
}
|
|
}
|
|
|
|
CublasWrapper::~CublasWrapper()
|
|
{
|
|
if (mHandle != nullptr)
|
|
{
|
|
PLUGIN_VALIDATE(cublasDestroy(mHandle) == CUBLAS_STATUS_SUCCESS, "Could not destroy cublas handle.");
|
|
mHandle = nullptr;
|
|
}
|
|
|
|
if (mLibrary != nullptr)
|
|
{
|
|
dllClose(mLibrary);
|
|
}
|
|
}
|
|
|
|
void* CublasWrapper::tryLoadingCublas()
|
|
{
|
|
void* cublasLib = dllOpen(kCUBLAS_PLUGIN_LIBNAME.c_str());
|
|
std::string errorMsg = "Failed to load " + kCUBLAS_PLUGIN_LIBNAME + ".";
|
|
PLUGIN_VALIDATE(cublasLib != nullptr, errorMsg.c_str());
|
|
return cublasLib;
|
|
}
|
|
|
|
cublasContext* CublasWrapper::getCublasHandle()
|
|
{
|
|
return mHandle;
|
|
}
|
|
|
|
bool CublasWrapper::isValid() const
|
|
{
|
|
return mHandle != nullptr;
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasCreate(cublasContext** handle)
|
|
{
|
|
return (*_cublasCreate)(handle);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasDestroy(cublasContext* handle)
|
|
{
|
|
return (*_cublasDestroy)(handle);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSetStream(cublasHandle_t handle, cudaStream_t streamId)
|
|
{
|
|
return (*_cublasSetStream)(handle, streamId);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasGetPointerMode(cublasHandle_t handle, cublasPointerMode_t* mode)
|
|
{
|
|
return (*_cublasGetPointerMode)(handle, mode);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSetPointerMode(cublasHandle_t handle, cublasPointerMode_t mode)
|
|
{
|
|
return (*_cublasSetPointerMode)(handle, mode);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasGetMathMode(cublasHandle_t handle, cublasMath_t* mode)
|
|
{
|
|
return (*_cublasGetMathMode)(handle, mode);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSetMathMode(cublasHandle_t handle, cublasMath_t mode)
|
|
{
|
|
return (*_cublasSetMathMode)(handle, mode);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasDscal(cublasHandle_t handle, int n, float const* alpha, float* x, int incx)
|
|
{
|
|
return (*_cublasDscal)(handle, n, alpha, x, incx);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSasum(cublasHandle_t handle, int n, float const* x, int incx, float* result)
|
|
{
|
|
return (*_cublasSasum)(handle, n, x, incx, result);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasScopy(cublasHandle_t handle, int n, float const* x, int incx, float* y, int incy)
|
|
{
|
|
return (*_cublasScopy)(handle, n, x, incx, y, incy);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSscal(cublasHandle_t handle, int n, float const* alpha, float* x, int incx)
|
|
{
|
|
return (*_cublasSscal)(handle, n, alpha, x, incx);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSgemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
|
int m, int n, int k, float const* alpha, float const* A, int lda, float const* B, int ldb, float const* beta,
|
|
float* C, int ldc)
|
|
{
|
|
return (*_cublasSgemm)(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasHgemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
|
int m, int n, int k, __half const* alpha, __half const* A, int lda, __half const* B, int ldb, __half const* beta,
|
|
__half* C, int ldc)
|
|
{
|
|
return (*_cublasHgemm)(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasHgemmStridedBatched(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k, __half const* alpha, __half const* A, int lda, long long int strideA,
|
|
__half const* B, int ldb, long long int strideB, __half const* beta, __half* C, int ldc, long long int strideC,
|
|
int batchCount)
|
|
{
|
|
return (*_cublasHgemmStridedBatched)(
|
|
handle, transa, transb, m, n, k, alpha, A, lda, strideA, B, ldb, strideB, beta, C, ldc, strideC, batchCount);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasSgemmStridedBatched(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k, float const* alpha, float const* A, int lda, long long int strideA,
|
|
float const* B, int ldb, long long int strideB, float const* beta, float* C, int ldc, long long int strideC,
|
|
int batchCount)
|
|
{
|
|
return (*_cublasSgemmStridedBatched)(
|
|
handle, transa, transb, m, n, k, alpha, A, lda, strideA, B, ldb, strideB, beta, C, ldc, strideC, batchCount);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasGemmEx(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
|
|
int m, int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda, void const* B,
|
|
cudaDataType Btype, int ldb, void const* beta, void* C, cudaDataType Ctype, int ldc, cudaDataType computeType,
|
|
cublasGemmAlgo_t algo)
|
|
{
|
|
return (*_cublasGemmEx)(
|
|
handle, transa, transb, m, n, k, alpha, A, Atype, lda, B, Btype, ldb, beta, C, Ctype, ldc, computeType, algo);
|
|
}
|
|
|
|
cublasStatus_t CublasWrapper::cublasGemmStridedBatchedEx(cublasHandle_t handle, cublasOperation_t transa,
|
|
cublasOperation_t transb, int m, int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda,
|
|
long long int strideA, void const* B, cudaDataType Btype, int ldb, long long int strideB, void const* beta, void* C,
|
|
cudaDataType Ctype, int ldc, long long int strideC, int batchCount, cudaDataType computeType, cublasGemmAlgo_t algo)
|
|
{
|
|
return (*_cublasGemmStridedBatchedEx)(handle, transa, transb, m, n, k, alpha, A, Atype, lda, strideA, B, Btype, ldb,
|
|
strideB, beta, C, Ctype, ldc, strideC, batchCount, computeType, algo);
|
|
}
|
|
|
|
CublasWrapper& getCublasWrapper()
|
|
{
|
|
// Initialize a global cublasWrapper instance to be used to call cublas functions.
|
|
static CublasWrapper sGCublasWrapper;
|
|
return sGCublasWrapper;
|
|
}
|
|
|
|
} // namespace nvinfer1::pluginInternal
|