wehub-resource-sync
231 lines
9.8 KiB
C++
231 lines
9.8 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 "cublasLtWrapper.h"
|
|
#include "common/checkMacrosPlugin.h"
|
|
#include "cudaDriverWrapper.h"
|
|
|
|
#if defined(_WIN32)
|
|
#if !defined(WIN32_LEAN_AND_MEAN)
|
|
#define WIN32_LEAN_AND_MEAN
|
|
#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 kCUBLASLT_PLUGIN_LIBNAME
|
|
= std::string{"cublasLt64_"} + 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 kCUBLASLT_PLUGIN_LIBNAME = std::string{"libcublasLt.so."} + std::to_string(nvinfer1::getCudaLibVersionMaj());
|
|
#endif // defined(_WIN32)
|
|
|
|
namespace nvinfer1::pluginInternal
|
|
{
|
|
using namespace nvinfer1;
|
|
|
|
// If tryLoadingCublasLt failed, the CublasLtWrapper object won't be created.
|
|
CublasLtWrapper::CublasLtWrapper(bool initHandle)
|
|
: mLibrary(tryLoadingCublasLt())
|
|
{
|
|
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 cublasLt library.";
|
|
PLUGIN_VALIDATE(ret != nullptr, loadError.c_str());
|
|
return ret;
|
|
};
|
|
*(void**) (&_cublasLtCreate) = load_sym(mLibrary, "cublasLtCreate");
|
|
*(void**) (&_cublasLtDestroy) = load_sym(mLibrary, "cublasLtDestroy");
|
|
*(void**) (&_cublasLtMatmul) = load_sym(mLibrary, "cublasLtMatmul");
|
|
*(void**) (&_cublasLtMatmulDescCreate) = load_sym(mLibrary, "cublasLtMatmulDescCreate");
|
|
*(void**) (&_cublasLtMatmulDescDestroy) = load_sym(mLibrary, "cublasLtMatmulDescDestroy");
|
|
*(void**) (&_cublasLtMatmulPreferenceCreate) = load_sym(mLibrary, "cublasLtMatmulPreferenceCreate");
|
|
*(void**) (&_cublasLtMatmulPreferenceDestroy) = load_sym(mLibrary, "cublasLtMatmulPreferenceDestroy");
|
|
*(void**) (&_cublasLtMatmulPreferenceSetAttribute) = load_sym(mLibrary, "cublasLtMatmulPreferenceSetAttribute");
|
|
*(void**) (&_cublasLtMatmulAlgoInit) = load_sym(mLibrary, "cublasLtMatmulAlgoInit");
|
|
*(void**) (&_cublasLtMatmulAlgoCheck) = load_sym(mLibrary, "cublasLtMatmulAlgoCheck");
|
|
*(void**) (&_cublasLtMatmulAlgoGetIds) = load_sym(mLibrary, "cublasLtMatmulAlgoGetIds");
|
|
*(void**) (&_cublasLtMatrixLayoutCreate) = load_sym(mLibrary, "cublasLtMatrixLayoutCreate");
|
|
*(void**) (&_cublasLtMatrixLayoutDestroy) = load_sym(mLibrary, "cublasLtMatrixLayoutDestroy");
|
|
*(void**) (&_cublasLtMatrixLayoutSetAttribute) = load_sym(mLibrary, "cublasLtMatrixLayoutSetAttribute");
|
|
*(void**) (&_cublasLtMatmulAlgoConfigSetAttribute) = load_sym(mLibrary, "cublasLtMatmulAlgoConfigSetAttribute");
|
|
*(void**) (&_cublasLtMatmulAlgoConfigGetAttribute) = load_sym(mLibrary, "cublasLtMatmulAlgoConfigGetAttribute");
|
|
*(void**) (&_cublasLtMatmulAlgoCapGetAttribute) = load_sym(mLibrary, "cublasLtMatmulAlgoCapGetAttribute");
|
|
*(void**) (&_cublasLtMatmulDescSetAttribute) = load_sym(mLibrary, "cublasLtMatmulDescSetAttribute");
|
|
|
|
if (initHandle)
|
|
{
|
|
PLUGIN_VALIDATE(cublasLtCreate(&mHandle) == CUBLAS_STATUS_SUCCESS, "Could not create cublasLt handle.");
|
|
PLUGIN_VALIDATE(mHandle != nullptr);
|
|
}
|
|
}
|
|
|
|
CublasLtWrapper::~CublasLtWrapper()
|
|
{
|
|
if (mHandle != nullptr)
|
|
{
|
|
PLUGIN_VALIDATE(cublasLtDestroy(mHandle) == CUBLAS_STATUS_SUCCESS, "Could not destroy cublas handle.");
|
|
mHandle = nullptr;
|
|
}
|
|
|
|
if (mLibrary != nullptr)
|
|
{
|
|
dllClose(mLibrary);
|
|
}
|
|
}
|
|
|
|
void* CublasLtWrapper::tryLoadingCublasLt()
|
|
{
|
|
void* cublasLtLib = dllOpen(kCUBLASLT_PLUGIN_LIBNAME.c_str());
|
|
std::string errorMsg = "Failed to load " + kCUBLASLT_PLUGIN_LIBNAME + ".";
|
|
PLUGIN_VALIDATE(cublasLtLib != nullptr, errorMsg.c_str());
|
|
return cublasLtLib;
|
|
}
|
|
|
|
cublasLtContext* CublasLtWrapper::getCublasLtHandle()
|
|
{
|
|
return mHandle;
|
|
}
|
|
|
|
bool CublasLtWrapper::isValid() const
|
|
{
|
|
return mHandle != nullptr;
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtCreate(cublasLtHandle_t* handle)
|
|
{
|
|
return (*_cublasLtCreate)(handle);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtDestroy(cublasLtHandle_t handle)
|
|
{
|
|
return (*_cublasLtDestroy)(handle);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmul(cublasLtHandle_t lightHandle, cublasLtMatmulDesc_t computeDesc,
|
|
void const* alpha, void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc,
|
|
void const* beta, void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
|
|
cublasLtMatmulAlgo_t const* algo, void* workspace, size_t workspaceSizeInBytes, cudaStream_t stream)
|
|
{
|
|
return (*_cublasLtMatmul)(lightHandle, computeDesc, alpha, A, Adesc, B, Bdesc, beta, C, Cdesc, D, Ddesc, algo,
|
|
workspace, workspaceSizeInBytes, stream);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulDescCreate(
|
|
cublasLtMatmulDesc_t* matmulDesc, cublasComputeType_t computeType, cudaDataType_t scaleType)
|
|
{
|
|
return (*_cublasLtMatmulDescCreate)(matmulDesc, computeType, scaleType);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulDescDestroy(cublasLtMatmulDesc_t matmulDesc)
|
|
{
|
|
return (*_cublasLtMatmulDescDestroy)(matmulDesc);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulPreferenceCreate(cublasLtMatmulPreference_t* pref)
|
|
{
|
|
return (*_cublasLtMatmulPreferenceCreate)(pref);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulPreferenceDestroy(cublasLtMatmulPreference_t pref)
|
|
{
|
|
return (*_cublasLtMatmulPreferenceDestroy)(pref);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulPreferenceSetAttribute(
|
|
cublasLtMatmulPreference_t pref, cublasLtMatmulPreferenceAttributes_t attr, void const* buf, size_t sizeInBytes)
|
|
{
|
|
return (*_cublasLtMatmulPreferenceSetAttribute)(pref, attr, buf, sizeInBytes);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoInit(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
|
|
cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype, cudaDataType_t Dtype,
|
|
int algoId, cublasLtMatmulAlgo_t* algo)
|
|
{
|
|
return (*_cublasLtMatmulAlgoInit)(lightHandle, computeType, scaleType, Atype, Btype, Ctype, Dtype, algoId, algo);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoCheck(cublasLtHandle_t lightHandle,
|
|
cublasLtMatmulDesc_t operationDesc, cublasLtMatrixLayout_t Adesc, cublasLtMatrixLayout_t Bdesc,
|
|
cublasLtMatrixLayout_t Cdesc, cublasLtMatrixLayout_t Ddesc, cublasLtMatmulAlgo_t const* algo,
|
|
cublasLtMatmulHeuristicResult_t* result)
|
|
{
|
|
return (*_cublasLtMatmulAlgoCheck)(lightHandle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, algo, result);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoGetIds(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
|
|
cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype, cudaDataType_t Dtype,
|
|
int requestedAlgoCount, int algoIdsArray[], int* returnAlgoCount)
|
|
{
|
|
return (*_cublasLtMatmulAlgoGetIds)(lightHandle, computeType, scaleType, Atype, Btype, Ctype, Dtype,
|
|
requestedAlgoCount, algoIdsArray, returnAlgoCount);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatrixLayoutCreate(
|
|
cublasLtMatrixLayout_t* matLayout, cudaDataType type, uint64_t rows, uint64_t cols, int64_t ld)
|
|
{
|
|
return (*_cublasLtMatrixLayoutCreate)(matLayout, type, rows, cols, ld);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatrixLayoutDestroy(cublasLtMatrixLayout_t matLayout)
|
|
{
|
|
return (*_cublasLtMatrixLayoutDestroy)(matLayout);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatrixLayoutSetAttribute(
|
|
cublasLtMatrixLayout_t matLayout, cublasLtMatrixLayoutAttribute_t attr, void const* buf, size_t sizeInBytes)
|
|
{
|
|
return (*_cublasLtMatrixLayoutSetAttribute)(matLayout, attr, buf, sizeInBytes);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoConfigGetAttribute(cublasLtMatmulAlgo_t const* algo,
|
|
cublasLtMatmulAlgoConfigAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten)
|
|
{
|
|
return (*_cublasLtMatmulAlgoConfigGetAttribute)(algo, attr, buf, sizeInBytes, sizeWritten);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoConfigSetAttribute(
|
|
cublasLtMatmulAlgo_t* algo, cublasLtMatmulAlgoConfigAttributes_t attr, void const* buf, size_t sizeInBytes)
|
|
{
|
|
return (*_cublasLtMatmulAlgoConfigSetAttribute)(algo, attr, buf, sizeInBytes);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoCapGetAttribute(cublasLtMatmulAlgo_t const* algo,
|
|
cublasLtMatmulAlgoCapAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten)
|
|
{
|
|
return (*_cublasLtMatmulAlgoCapGetAttribute)(algo, attr, buf, sizeInBytes, sizeWritten);
|
|
}
|
|
|
|
cublasStatus_t CublasLtWrapper::cublasLtMatmulDescSetAttribute(
|
|
cublasLtMatmulDesc_t matmulDesc, cublasLtMatmulDescAttributes_t attr, void const* buf, size_t sizeInBytes)
|
|
{
|
|
return (*_cublasLtMatmulDescSetAttribute)(matmulDesc, attr, buf, sizeInBytes);
|
|
}
|
|
|
|
CublasLtWrapper& getCublasLtWrapper()
|
|
{
|
|
// Initialize a global cublasLtWrapper instance to be used to call cublasLt functions.
|
|
static CublasLtWrapper sGCublasLtWrapper;
|
|
return sGCublasLtWrapper;
|
|
}
|
|
|
|
} // namespace nvinfer1::pluginInternal
|