wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 17 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plugin/fcPlugin/CMakeLists.txt
+22
View File
@@ -0,0 +1,22 @@
#
# 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.
#
add_plugin_source(
fcPlugin.cpp
fcPlugin.h
)
plugin/fcPlugin/CustomFCPluginDynamic_PluginConfig.yaml
+92
View File
@@ -0,0 +1,92 @@
#
# SPDX-FileCopyrightText: Copyright (c) 2022-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.
#
---
name: CustomFCPluginDynamic
interface: IPluginV2DynamicExt
versions:
"1":
inputs:
- input
outputs:
- output
input_dims:
input: 5
input_dim_constraints:
- "input_2 == W_0"
input_dim_range:
input:
min: "=1, =1, =1, =1, =1"
max: "=pinf, =pinf, =pinf, =1, =1"
output_dims:
output: "input_0, input_1, out_dims_0, 1, 1"
attributes:
- out_dims
- type_id
- W
supported_input_types:
combination1:
input: float32
combination2:
input: float16
attribute_types:
out_dims: int32
type_id: int32
W: float32
attribute_dims:
out_dims: 1
type_id: 1
W: 2
attribute_dim_range:
out_dims:
min: "=1"
max: "=1"
type_id:
min: "=1"
max: "=1"
W:
min: "=1, =1"
max: "=pinf, =pinf"
attribute_options:
out_dims:
from_shape: "W_1"
type_id:
- 0
- 1
W:
min: "=ninf"
max: "=pinf"
attributes_required:
- out_dims
- type_id
- W
golden_io_path: "plugin/CustomFCPluginDynamic_PluginGoldenIO.json"
golden_reference_script: "plugin/CustomFCPluginDynamic_PluginReference.py"
abs_tol: 1e-5
rel_tol: 1e-5
fp16_atol: 1e-3
fp16_rtol: 1e-3
configs:
config1:
input_types:
input: float16
attribute_options:
"type_id":
value: 1
shape: "1"
output_types:
output: float16
...
plugin/fcPlugin/README.md
+62
View File
@@ -0,0 +1,62 @@
# fcPlugin [DEPRECATED]
**Table Of Contents**
- [Description](#description)
* [Structure](#structure)
- [Parameters](#parameters)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
> NOTE: This plugin is deprecated since TensorRT 10.6. Its functionality has been superseded by the [`IMatrixMultiplyLayer`](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_matrix_multiply_layer.html) (Can be added to the network definition using [`addMatrixMultiply()`](https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_network_definition.html#acf109d93e91c86afbd263f5fea29ffe8))
Performs a matrix multiplication similar to the FullyConnected Layer in TensorRT, but without bias. The main difference is that the weights are not transposed.
Always dispatches to cuBLAS. At engine build time, the plugin runs a search over the parameters of the available algorithms to find the fastest one available.
### Structure
The `fcPlugin` takes one input; `input`.
`input`
input is a tensor with shape `[N, B, K, 1, 1]` where `B` is the batch size.
The `fcPlugin` generates the following output:
`output`
output is a tensor with shape `[N, B, out_dims, 1, 1]` where `B` is the batch size, and `out_dims` is specified as a plugin parameter.
The trailing singleton dimensions in the input and output are added for compatibility with the default TRT FC layer.
## Parameters
`fcPlugin` has plugin creator class `FCPluginDynamicCreator` and plugin class `CustomFCPluginDynamic`.
The parameters are defined below and consists of the following attributes:
| Type | Parameter | Description
|----------|-----------------------------------------|-------------------------------------------------------------------
|`int` |`out_dims` |Integer specifying the length of the third dimension of the output.
|`int` |`type_id` |Integer encoding the DataType (0: FP32, 1: FP16)
|`Weights` |`W` |The weights to multiply with. Shape: `[K, out_dims]`
## License
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
documentation.
## Changelog
- October 2024: Add deprecation note.
- November 2019: This is the first release of this `README.md` file.
## Known issues
This plugin only supports GPUs with compute capability >= 7.0. For more information see the [CUDA GPU Compute Capability Support Matrix](https://developer.nvidia.com/cuda-gpus#compute)
plugin/fcPlugin/fcPlugin.cpp
+825
View File
@@ -0,0 +1,825 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
*/
// cspell:ignore accu CUBLASLT splitk
// cublasLT was introduced in CUDA 10.1
#include <cuda.h>
#if CUDA_VERSION >= 10010
#include "NvInfer.h"
#include "common/serialize.hpp"
#include "fcPlugin.h"
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <cuda_runtime.h>
#include <memory>
#include <string_view>
#include <vector>
using namespace nvinfer1;
using namespace nvinfer1::plugin;
using namespace nvinfer1::plugin::bert;
using namespace nvinfer1::pluginInternal;
// plugin specific constants
namespace
{
using namespace std::string_view_literals;
char const* const kFC_VERSION{"1"};
char const* const kFC_NAME{"CustomFCPluginDynamic"};
constexpr size_t kMAX_WORKSPACE_BYTES = 4 * 1024 * 1024; // 4MiB
} // namespace
REGISTER_TENSORRT_PLUGIN(FCPluginDynamicCreator);
// Utility function to print customMatmulPerf_t structure
static void printPerfStructure(customMatmulPerf_t const& perf, int32_t const m, int32_t const n, int32_t const k)
{
AlgoProps p;
p.populate(perf.algo);
// Calculate GFLOPS
double timeAvg
= perf.time * 1e-3; // Convert to seconds. It has been divided by kNB_KERNEL_REPEATS in customMatmulRun().
double gflop = (2 * static_cast<uint64_t>(m * n) * k) * 1e-9; // Real
gLogVerbose << "Algo=" << p.algoId << " Tile=" << p.tile << " (" << matmulTileName[p.tile] << ") K=" << p.numSplitsK
<< " Red.Sch.=" << p.reductionScheme << " Swiz=" << p.swizzle << " Cust=" << p.customOption
<< " Stat=" << perf.status << " Time=" << perf.time << " WSbytes=" << perf.workspaceSize
<< " math=" << p.numericImpl << " waves=" << perf.wavesCount << "GFlops=" << (gflop / timeAvg)
<< std::endl;
}
static bool timeCompare(customMatmulPerf_t const& perf_a, customMatmulPerf_t const& perf_b)
{
return ((perf_a.status == CUBLAS_STATUS_SUCCESS) && (perf_a.time < perf_b.time));
}
static cublasStatus_t customMatmulRun(cublasLtHandle_t ltHandle, // to get the capabilities (required a GPU)
cublasLtMatmulDesc_t operationDesc, void const* alpha, // host or device pointer
void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc,
void const* beta, // host or device pointer
void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
cublasLtMatmulAlgo_t const& algo, void* workSpace, size_t workSpaceSizeInBytes, customMatmulPerf_t& perfResults,
cudaStream_t stream, cudaEvent_t& startEvent, cudaEvent_t& stopEvent)
{
cublasLtMatmulHeuristicResult_t heurResult;
CublasLtWrapper& cublasLtWrapper = getCublasLtWrapper();
// Looping over the Algo
cublasStatus_t algoStatus = cublasLtWrapper.cublasLtMatmulAlgoCheck(
ltHandle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, &algo, &heurResult);
if (algoStatus == CUBLAS_STATUS_SUCCESS)
{
if (heurResult.workspaceSize <= workSpaceSizeInBytes)
{
if (cudaEventRecord(startEvent, stream) != cudaSuccess)
{
return CUBLAS_STATUS_INTERNAL_ERROR;
}
for (int32_t loop = 0; loop < kNB_KERNEL_REPEATS; loop++)
{
cublasStatus_t oneRunStatus
= cublasLtWrapper.cublasLtMatmul(ltHandle, operationDesc, alpha, // host or device pointer
A, Adesc, B, Bdesc, beta, // host or device pointer
C, Cdesc, D, Ddesc, &algo, workSpace, workSpaceSizeInBytes, stream);
if (oneRunStatus != CUBLAS_STATUS_SUCCESS)
{
algoStatus = oneRunStatus;
break;
}
}
if (cudaEventRecord(stopEvent, stream) != cudaSuccess)
{
return CUBLAS_STATUS_INTERNAL_ERROR;
}
if (cudaEventSynchronize(stopEvent) != cudaSuccess)
{
return CUBLAS_STATUS_INTERNAL_ERROR;
}
float time;
if (cudaEventElapsedTime(&time, startEvent, stopEvent) != cudaSuccess)
{
return CUBLAS_STATUS_INTERNAL_ERROR;
}
// For the moment only add successful findings
perfResults.algo = algo;
perfResults.time = time / kNB_KERNEL_REPEATS; // Average time
perfResults.workspaceSize = heurResult.workspaceSize;
perfResults.wavesCount = heurResult.wavesCount;
}
else
{
algoStatus = CUBLAS_STATUS_NOT_SUPPORTED; // Not enough workspace
}
}
return algoStatus;
}
// Sample wrapper running through multiple algo and config attributes
// combination for single precision gemm using cublasLt low-level API
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
void nvinfer1::plugin::bert::LtGemmSearch(cublasLtHandle_t ltHandle, cublasOperation_t transa, cublasOperation_t transb,
int32_t const& m, int32_t const& n, int32_t const& k, void const* alpha, // host pointer
void const* A, int32_t const& lda, void const* B, int32_t const& ldb, void const* beta, // host pointer
void* C, int32_t const& ldc, void* workSpace, size_t workSpaceSize, cublasComputeType_t computeType,
cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype,
std::vector<customMatmulPerf_t>& perfResults, cudaStream_t stream)
{
cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
cublasLtMatmulDesc_t operationDesc = nullptr;
cublasLtMatrixLayout_t Adesc = nullptr;
cublasLtMatrixLayout_t Bdesc = nullptr;
cublasLtMatrixLayout_t Cdesc = nullptr;
cublasLtMatmulPreference_t preference = nullptr;
cudaEvent_t startEvent = nullptr;
cudaEvent_t stopEvent = nullptr;
CublasLtWrapper& cublasLtWrapper = getCublasLtWrapper();
// SplitK value that we are going to try when SplitK is supported for a given algo.
int32_t const splitKSequenceA[] = {2, 3, 4, 5, 6, 8, 12, 16, 32};
// Let try a fixed number of combinations
int32_t algoCount = 0;
int32_t nbAlgoIds = 0;
int32_t algoIdA[kNB_ALGO_IDS];
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceCreate(&preference));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceSetAttribute(
preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workSpaceSize, sizeof(workSpaceSize)));
uint64_t const numericImplPrefer
= Ctype == CUDA_R_16F ? CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA : CUBLASLT_NUMERICAL_IMPL_FLAGS_FMA;
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceSetAttribute(
preference, CUBLASLT_MATMUL_PREF_IMPL_MASK, &numericImplPrefer, sizeof(numericImplPrefer)));
// Create operation descriptor; see cublasLtMatmulDescAttributes_t for details
// about defaults; here we just need to set the transforms for A and B
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescCreate(&operationDesc, computeType, scaleType));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescSetAttribute(
operationDesc, nvinfer1::pluginInternal::CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescSetAttribute(
operationDesc, nvinfer1::pluginInternal::CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa)));
// Create matrix descriptors. We are good with the details here so no need to
// set any extra attributes
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(
&Adesc, Atype, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(
&Bdesc, Btype, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(&Cdesc, Ctype, m, n, ldc));
// Request the 4 first AlgoId available for SGEMM ( computeType = scaleType =
// Atype = Btype = Ctype = Dtype = CUDA_R_32F)
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoGetIds(
ltHandle, computeType, scaleType, Atype, Btype, Ctype, Ctype, kNB_ALGO_IDS, algoIdA, &nbAlgoIds));
gLogVerbose << "Number of algos" << nbAlgoIds << std::endl;
// Create CUDA event to time the execution time of each algo
PLUGIN_CUASSERT(cudaEventCreate(&startEvent, cudaEventBlockingSync));
PLUGIN_CUASSERT(cudaEventCreate(&stopEvent, cudaEventBlockingSync));
// Loop over the Algo IDs
for (int32_t idx = 0; (idx < nbAlgoIds) && (algoCount < kNB_ALGO_COMBINATIONS); idx++)
{
cublasLtMatmulAlgo_t algo;
size_t nbBytesWritten = 0;
// Initialize algo structure with given Algp ID.
status = cublasLtWrapper.cublasLtMatmulAlgoInit(
ltHandle, computeType, scaleType, Atype, Btype, Ctype, Ctype, algoIdA[idx], &algo);
if (status != CUBLAS_STATUS_SUCCESS)
{
continue;
}
uint64_t numericImpl = -1;
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(
&algo, CUBLASLT_ALGO_CAP_NUMERICAL_IMPL_FLAGS, &numericImpl, sizeof(numericImpl), nullptr));
if (Ctype == CUDA_R_32F && numericImpl == CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA)
{
// skip HMMA-fp32accu kernels
continue;
}
// Query the tiles enums supported by that algo
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(
&algo, CUBLASLT_ALGO_CAP_TILE_IDS, nullptr, 0, &nbBytesWritten));
auto tileA = std::vector<int32_t>(nbBytesWritten / sizeof(int32_t));
if (tileA.empty())
{
tileA = {CUBLASLT_MATMUL_TILE_UNDEFINED};
}
// Retrieve Algo Capabilities attributes to be able to setup loop over the
// different combinations
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(&algo, CUBLASLT_ALGO_CAP_TILE_IDS,
static_cast<void*>(tileA.data()), sizeof(int32_t) * tileA.size(), &nbBytesWritten));
//! Read an int32_t attribute from `cublasLtWrapper` and return the value and update nbBytesWritten.
auto getInt32Attr = [&](cublasLtMatmulAlgoCapAttributes attr, size_t& nbBytesWritten) -> int32_t {
int32_t value;
PLUGIN_CUBLASASSERT(
cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(&algo, attr, &value, sizeof(value), &nbBytesWritten));
return value;
};
auto const splitkSupport = getInt32Attr(CUBLASLT_ALGO_CAP_SPLITK_SUPPORT, nbBytesWritten);
auto const redMask = getInt32Attr(CUBLASLT_ALGO_CAP_REDUCTION_SCHEME_MASK, nbBytesWritten);
auto const swizzlingMax = getInt32Attr(CUBLASLT_ALGO_CAP_CTA_SWIZZLING_SUPPORT, nbBytesWritten);
auto const customOptionMax = getInt32Attr(CUBLASLT_ALGO_CAP_CUSTOM_OPTION_MAX, nbBytesWritten);
std::ignore = getInt32Attr(CUBLASLT_ALGO_CAP_EPILOGUE_MASK, nbBytesWritten);
// Loop over the different tiles
for (uint32_t tileIdx = 0; tileIdx < tileA.size(); tileIdx++)
{
// Loop over the different custom option if any
for (int32_t customOption = 0; customOption <= customOptionMax; customOption++)
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigSetAttribute(
&algo, CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, &customOption, sizeof(customOption)));
// Loop over the CTAs swizzling support
for (int32_t k = 0; k <= swizzlingMax; k++)
{
int32_t const splitkTrial
= splitkSupport ? (sizeof(splitKSequenceA) / sizeof(splitKSequenceA[0])) : 0;
// Loop over the splitK value over a fixed sequence splitKSequenceA in
// addition to the case where splitK is not enabled
for (int32_t l = 0; (l < (1 + splitkTrial)) && (algoCount < kNB_ALGO_COMBINATIONS); l++)
{
//! Set an attribute of the algo to a given value.
auto setAttr = [&](cublasLtMatmulAlgoConfigAttributes attr, auto value) {
static_assert(std::is_same_v<decltype(value), int32_t>, "value must be an int32_t");
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigSetAttribute(
&algo, attr, &value, sizeof(value)));
};
// Setup attribute of the algo to run
setAttr(CUBLASLT_ALGO_CONFIG_TILE_ID, tileA[tileIdx]);
int32_t splitK_val = 0;
setAttr(CUBLASLT_ALGO_CONFIG_SPLITK_NUM, splitK_val);
setAttr(CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, k);
int32_t redScheme = CUBLASLT_REDUCTION_SCHEME_NONE;
setAttr(CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, redScheme);
if (l > 0)
{ // Split-K case
splitK_val = splitKSequenceA[l - 1];
setAttr(CUBLASLT_ALGO_CONFIG_SPLITK_NUM, splitKSequenceA[l - 1]);
// Going over all the reduction scheme
for (redScheme = 1; redScheme < static_cast<int32_t>(CUBLASLT_REDUCTION_SCHEME_MASK)
&& (algoCount < kNB_ALGO_COMBINATIONS);
redScheme <<= 1)
{
if (redScheme & redMask)
{
setAttr(CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, redScheme);
status = customMatmulRun(ltHandle, operationDesc, alpha, // host or device pointer
A, Adesc, B, Bdesc, beta, // host or device pointer
C, Cdesc, C, Cdesc, algo, workSpace, workSpaceSize, perfResults[algoCount],
stream, startEvent, stopEvent);
perfResults[algoCount].status = status;
algoCount += (status == CUBLAS_STATUS_SUCCESS);
}
}
}
else
{ // Non-splitK case
// if user preference is ok with workspace
if (algoCount < kNB_ALGO_COMBINATIONS)
{
status = customMatmulRun(ltHandle, operationDesc, alpha, // host or device pointer
A, Adesc, B, Bdesc, beta, // host or device pointer
C, Cdesc, C, Cdesc, algo, workSpace, workSpaceSize, perfResults[algoCount], stream,
startEvent, stopEvent);
perfResults[algoCount].status = status;
algoCount += (status == CUBLAS_STATUS_SUCCESS);
}
}
} // end l
} // end k
} // end customOption
} // end tileIdx
} // end idx
// Sort the results per run duration
std::sort(perfResults.begin(), perfResults.end(), timeCompare);
// Print timing and perf details of the fastest combinations
for (int32_t i = 0; i < kPRINT_ALGOS && perfResults[i].time != customMatmulPerf_t::kMAX_TIME; i++)
{
printPerfStructure(perfResults[i], m, n, k);
}
// Descriptors are no longer needed as all GPU work was already enqueued
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulPreferenceDestroy(preference));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Cdesc));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Bdesc));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Adesc));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescDestroy(operationDesc));
PLUGIN_CUASSERT(cudaEventDestroy(startEvent));
PLUGIN_CUASSERT(cudaEventDestroy(stopEvent));
}
FCPluginDynamic::FCPluginDynamic(std::string const name, DataType const type, int32_t const outDim, Weights const& W)
: mLayerName(name)
, mType(type)
, mOutDim(outDim)
, mNumParams(W.count)
, mNmax(0)
, mK(0)
, mWdev(nullptr)
{
memset(mAlgo.data, 0, sizeof(mAlgo.data));
mW.convertAndCopy(W, mType);
copyToDevice(mW, getWeightsSize(mW, mType), mWdev);
}
FCPluginDynamic::FCPluginDynamic(std::string const name, void const* data, size_t length)
: mLayerName(name)
, mWdev(nullptr)
{
gLogVerbose << "FCPluginDynamic deserialize\n";
// Deserialize in the same order as serialization
deserialize_value(&data, &length, &mType);
deserialize_value(&data, &length, &mOutDim);
deserialize_value(&data, &length, &mNumParams);
deserialize_value(&data, &length, &mNmax);
deserialize_value(&data, &length, &mK);
deserialize_value(&data, &length, &mAlgo);
char const* d = static_cast<char const*>(data);
mW.convertAndCopy(d, mNumParams, mType);
copyToDevice(mW, getWeightsSize(mW, mType), mWdev);
}
// IPluginV2DynamicExt Methods
IPluginV2DynamicExt* FCPluginDynamic::clone() const noexcept
{
try
{
gLogVerbose << "FCPluginDynamic clone\n";
auto p = std::make_unique<FCPluginDynamic>(mLayerName, mType, mOutDim, mW);
p->mAlgo = mAlgo;
p->setPluginNamespace(mNamespace.c_str());
return p.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
void FCPluginDynamic::attachToContext(
cudnnContext* cudnnContext, cublasContext* cublasContext, nvinfer1::IGpuAllocator* gpuAllocator) noexcept
{
mLtContext.attach();
}
void FCPluginDynamic::detachFromContext() noexcept
{
mLtContext.detach();
}
DimsExprs FCPluginDynamic::getOutputDimensions(
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
{
try
{
PLUGIN_VALIDATE(nbInputs == 1);
PLUGIN_VALIDATE(outputIndex == 0);
PLUGIN_VALIDATE(inputs != nullptr);
DimsExprs ret;
ret.nbDims = 5;
ret.d[0] = inputs[0].d[0];
ret.d[1] = inputs[0].d[1];
ret.d[2] = exprBuilder.constant(mOutDim);
ret.d[3] = exprBuilder.constant(1);
ret.d[4] = exprBuilder.constant(1);
return ret;
}
catch (std::exception const& e)
{
caughtError(e);
}
return DimsExprs{};
}
bool FCPluginDynamic::supportsFormatCombination(
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
PLUGIN_ASSERT(nbInputs == 1);
PLUGIN_ASSERT(nbOutputs == 1);
PLUGIN_ASSERT(inOut != nullptr);
PluginTensorDesc const& in = inOut[pos];
if (pos == 0)
{
return (in.type == mType) && (in.format == TensorFormat::kLINEAR);
}
PluginTensorDesc const& prev = inOut[pos - 1];
// output
return in.type == prev.type && in.format == prev.format;
}
void FCPluginDynamic::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
{
try
{
// Validate input arguments
PLUGIN_VALIDATE(nbOutputs == 1);
PLUGIN_VALIDATE(nbInputs == 1);
PLUGIN_VALIDATE(inputs != nullptr);
PLUGIN_VALIDATE(outputs != nullptr);
PLUGIN_VALIDATE(mType == inputs[0].desc.type);
auto const& inDims0 = inputs[0].desc.dims;
PLUGIN_VALIDATE(inDims0.nbDims == 5);
mK = inDims0.d[HDIM]; // hiddensize
// PLUGIN_ASSERT(hiddenSize * mOutDim == mNumParams);
PLUGIN_VALIDATE(inDims0.d[3] == 1);
PLUGIN_VALIDATE(inDims0.d[4] == 1);
// m and k are mOutDim
// n is B*S
int32_t const S = inputs->max.d[SDIM];
int32_t const B = inputs->max.d[BDIM];
mNmax = S * B;
// Cleanup LtContext descriptors before creating new ones.
mLtContext.destroy();
if (mType == DataType::kFLOAT)
{
Gemm<float> g(mOutDim, mNmax, mK, false, false);
mLtContext.create(g, kMAX_WORKSPACE_BYTES);
}
else if (mType == DataType::kHALF)
{
Gemm<half> g(mOutDim, mNmax, mK, false, false);
mLtContext.create(g, kMAX_WORKSPACE_BYTES);
}
else
{
std::string const msg = "Unsupported type error, expected [kHALF,kFLOAT], but received ";
PLUGIN_VALIDATE(false, (msg + std::to_string(static_cast<int32_t>(mType))).c_str());
}
gLogVerbose << "FCPluginDynamic configurePlugin m=" << mOutDim << ", n=" << mNmax << ", k=" << mK << std::endl;
size_t actualWorkspace = 0;
if (std::all_of(std::begin(mAlgo.data), std::end(mAlgo.data), [](auto v) { return v == 0; }))
{
gLogVerbose << "FCPluginDynamic gemmSearch\n";
if (mSharedStream == nullptr)
{
SharedStream ss{};
mSharedStream = static_cast<SharedStream*>(
getPluginRegistry()->acquirePluginResource(kFCPLUGIN_SHARED_STREAM_KEY, &ss))
->mStream;
}
if (mType == DataType::kFLOAT)
{
mAlgo = gemmSearch<float>(mOutDim, mNmax, mK, kMAX_WORKSPACE_BYTES, actualWorkspace, mSharedStream);
}
else if (mType == DataType::kHALF)
{
mAlgo = gemmSearch<half>(mOutDim, mNmax, mK, kMAX_WORKSPACE_BYTES, actualWorkspace, mSharedStream);
}
}
AlgoProps p;
p.populate(mAlgo);
if (mType == DataType::kFLOAT && p.numericImpl == CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA)
{
gLogWarning << "cuBLAS might use mixed precision instead of FP32" << std::endl;
}
if (mType == DataType::kHALF && p.numericImpl != CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA)
{
gLogWarning << "TensorCore support was not selected" << std::endl;
}
gLogVerbose << "FCPluginDynamic configuration Algo=" << p.algoId << " Tile=" << p.tile << " ("
<< matmulTileName[p.tile] << ") K=" << p.numSplitsK << " Red.Sch.=" << p.reductionScheme
<< " Swiz=" << p.swizzle << " Cust=" << p.customOption << " numericImpl=" << p.numericImpl
<< " ws=" << actualWorkspace << std::endl;
}
catch (std::exception const& e)
{
caughtError(e);
}
}
size_t FCPluginDynamic::getWorkspaceSize(
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
{
return kMAX_WORKSPACE_BYTES;
}
int32_t FCPluginDynamic::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
void const* const* inputs, void* const* outputs, void* workSpace, cudaStream_t stream) noexcept
{
try
{
PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr
&& workSpace != nullptr);
size_t const workspaceSize = getWorkspaceSize(inputDesc, 1, outputDesc, 1);
int32_t const S = inputDesc->dims.d[SDIM];
int32_t const B = inputDesc->dims.d[BDIM];
int32_t const n = S * B;
PLUGIN_VALIDATE(n >= 0);
mLtContext.setN(static_cast<uint64_t>(n));
if (mType == DataType::kFLOAT)
{
auto const* const input = static_cast<float const*>(inputs[0]);
auto* output = static_cast<float*>(outputs[0]);
Gemm<float> g(mOutDim, n, mK, false, false);
if (mWdev == nullptr)
{
return STATUS_FAILURE;
}
g.A = static_cast<float*>(mWdev.get());
g.B = const_cast<float*>(input);
g.C = output;
return cublasLtMatmul(mLtContext, g, mAlgo, workSpace, workspaceSize, stream);
}
if (mType == DataType::kHALF)
{
auto const* const input = static_cast<half const*>(inputs[0]);
auto* output = static_cast<half*>(outputs[0]);
Gemm<half> g(mOutDim, n, mK, false, false);
if (mWdev == nullptr)
{
return STATUS_FAILURE;
}
g.A = static_cast<half*>(mWdev.get());
g.B = const_cast<half*>(input);
g.C = output;
return cublasLtMatmul(mLtContext, g, mAlgo, workSpace, workspaceSize, stream);
}
else
{
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
<< std::endl;
return STATUS_FAILURE;
}
}
catch (std::exception const& e)
{
caughtError(e);
}
return STATUS_FAILURE;
}
// IPluginV2Ext Methods
DataType FCPluginDynamic::getOutputDataType(int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
{
PLUGIN_ASSERT(index == 0);
PLUGIN_ASSERT(nbInputs == 1);
PLUGIN_ASSERT(inputTypes != nullptr);
PLUGIN_ASSERT(inputTypes[0] == DataType::kFLOAT || inputTypes[0] == DataType::kHALF);
return inputTypes[0];
}
// IPluginV2 Methods
char const* FCPluginDynamic::getPluginType() const noexcept
{
return kFC_NAME;
}
char const* FCPluginDynamic::getPluginVersion() const noexcept
{
return kFC_VERSION;
}
int32_t FCPluginDynamic::getNbOutputs() const noexcept
{
return 1;
}
int32_t FCPluginDynamic::initialize() noexcept
{
gLogVerbose << "FCPluginDynamic initialize\n";
return 0;
}
void FCPluginDynamic::terminate() noexcept
{
gLogVerbose << "FCPluginDynamic terminate\n";
if (mSharedStream)
{
TRT_UNUSED(getPluginRegistry()->releasePluginResource(kFCPLUGIN_SHARED_STREAM_KEY));
mSharedStream = nullptr;
}
}
size_t FCPluginDynamic::getSerializationSize() const noexcept
{
size_t wordSize = getElementSize(mType);
return wordSize * mNumParams + sizeof(mType) + sizeof(mOutDim) + sizeof(mNumParams) + sizeof(mAlgo) + sizeof(mNmax)
+ sizeof(mK);
}
void FCPluginDynamic::serialize(void* buffer) const noexcept
{
serialize_value(&buffer, mType);
serialize_value(&buffer, mOutDim);
serialize_value(&buffer, mNumParams);
serialize_value(&buffer, mNmax);
serialize_value(&buffer, mK);
serialize_value(&buffer, mAlgo);
size_t wordSize = getElementSize(mType);
char* d = static_cast<char*>(buffer);
serFromDev(d, static_cast<char*>(mWdev.get()), mNumParams * wordSize);
}
void FCPluginDynamic::destroy() noexcept
{
gLogVerbose << "FCPluginDynamic destroy\n";
// This gets called when the network containing plugin is destroyed
mLtContext.destroy();
mWdev.reset(nullptr);
delete this;
}
void FCPluginDynamic::setPluginNamespace(char const* libNamespace) noexcept
{
try
{
PLUGIN_VALIDATE(libNamespace != nullptr);
mNamespace = libNamespace;
}
catch (std::exception const& e)
{
caughtError(e);
}
}
char const* FCPluginDynamic::getPluginNamespace() const noexcept
{
return mNamespace.c_str();
}
/////////////////////////////////////////////////////////
FCPluginDynamicCreator::FCPluginDynamicCreator()
{
mPluginAttributes.clear();
mPluginAttributes.emplace_back(PluginField("out_dims", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1));
mPluginAttributes.emplace_back(PluginField("W", nullptr, PluginFieldType::kFLOAT32, 1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
char const* FCPluginDynamicCreator::getPluginName() const noexcept
{
return kFC_NAME;
}
char const* FCPluginDynamicCreator::getPluginVersion() const noexcept
{
return kFC_VERSION;
}
PluginFieldCollection const* FCPluginDynamicCreator::getFieldNames() noexcept
{
return &mFC;
}
IPluginV2* FCPluginDynamicCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
{
try
{
gLogVerbose << "Creating FCPluginDynamicCreator...\n";
PLUGIN_VALIDATE(name != nullptr);
PLUGIN_VALIDATE(fc != nullptr);
int32_t outDims = 0;
int32_t typeId = -1;
Weights W{DataType::kFLOAT, nullptr, 0LL};
plugin::validateRequiredAttributesExist({"out_dims", "type_id", "W"}, fc);
for (int32_t i = 0; i < fc->nbFields; i++)
{
std::string_view const fieldName = fc->fields[i].name;
if (fieldName == "out_dims"sv)
{
outDims = static_cast<int32_t const*>(fc->fields[i].data)[0];
gLogVerbose << "Building outDims: " << outDims << std::endl;
}
if (fieldName == "type_id"sv)
{
typeId = static_cast<int32_t const*>(fc->fields[i].data)[0];
gLogVerbose << "Building typeId: " << outDims << std::endl;
}
if (fieldName == "W"sv)
{
gLogVerbose << "Building W...\n";
W.values = fc->fields[i].data;
W.count = fc->fields[i].length;
W.type = fieldTypeToDataType(fc->fields[i].type);
gLogVerbose << "Is W float32: " << (W.type == DataType::kFLOAT) << std::endl;
}
}
if (outDims <= 0)
{
gLogError << "Invalid output dimension" << std::endl;
}
if (typeId < 0 || typeId > 1)
{
gLogError << "Invalid type id" << typeId << std::endl;
}
if (W.count == 0 || W.values == nullptr || W.count < outDims)
{
gLogError << "Invalid weights" << std::endl;
}
DataType type = static_cast<DataType>(typeId);
return new FCPluginDynamic(name, type, outDims, W);
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
IPluginV2* FCPluginDynamicCreator::deserializePlugin(
char const* name, void const* serialData, size_t serialLength) noexcept
{
// This object will be deleted when the network is destroyed, which will
// call FCPluginDynamic::destroy()
try
{
return new FCPluginDynamic(name, serialData, serialLength);
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
void FCPluginDynamicCreator::setPluginNamespace(char const* libNamespace) noexcept
{
try
{
PLUGIN_VALIDATE(libNamespace != nullptr);
mNamespace = libNamespace;
}
catch (std::exception const& e)
{
caughtError(e);
}
}
char const* FCPluginDynamicCreator::getPluginNamespace() const noexcept
{
return mNamespace.c_str();
}
#endif // #if CUDA_VERSION >= 10010
plugin/fcPlugin/fcPlugin.h
+618
View File
@@ -0,0 +1,618 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
*/
// cublasLT was introduced in CUDA 10.1
#include <cuda.h>
#if CUDA_VERSION >= 10010
#ifndef TRT_FC_PLUGIN_H
#define TRT_FC_PLUGIN_H
#include "NvInferPlugin.h"
#include "common/bertCommon.h"
#include "common/cublasLtWrapper.h"
#include <string>
#include <vector>
namespace nvinfer1
{
namespace pluginInternal
{
class SharedStream : public IPluginResource
{
public:
SharedStream(bool init = false)
{
if (init)
{
PLUGIN_CUASSERT(cudaStreamCreate(&mStream));
}
}
void free()
{
if (mStream != nullptr)
{
PLUGIN_CUASSERT(cudaStreamDestroy(mStream));
mStream = nullptr;
}
}
int32_t release() noexcept override
{
try
{
free();
}
catch (std::exception const& e)
{
return -1;
}
return 0;
}
IPluginResource* clone() noexcept override
{
std::unique_ptr<SharedStream> cloned{};
try
{
cloned = std::make_unique<SharedStream>(/* init */ true);
}
catch (std::exception const& e)
{
return nullptr;
}
return cloned.release();
}
~SharedStream() override
{
if (mStream)
{
free();
}
}
cudaStream_t mStream{nullptr};
};
} // namespace pluginInternal
namespace plugin
{
namespace bert
{
template <typename T>
struct GemmTypes
{
};
char const* const kFCPLUGIN_SHARED_STREAM_KEY{"fcPlugin_timing_key"};
template <>
struct GemmTypes<half>
{
static cudaDataType_t const cudaTypeI = CUDA_R_16F;
using dataTypeI = half;
static cudaDataType_t const cudaTypeO = CUDA_R_16F;
using dataTypeO = half;
static cudaDataType_t const cudaTypeS = CUDA_R_16F;
using dataTypeS = half;
static nvinfer1::pluginInternal::cublasComputeType_t const cudaTypeCom
= nvinfer1::pluginInternal::CUBLAS_COMPUTE_16F;
};
template <>
struct GemmTypes<float>
{
static cudaDataType_t const cudaTypeI = CUDA_R_32F;
using dataTypeI = float;
static cudaDataType_t const cudaTypeO = CUDA_R_32F;
using dataTypeO = float;
static cudaDataType_t const cudaTypeS = CUDA_R_32F;
using dataTypeS = float;
static nvinfer1::pluginInternal::cublasComputeType_t const cudaTypeCom
= nvinfer1::pluginInternal::CUBLAS_COMPUTE_32F;
};
template <typename T>
struct Gemm
{
using Types = GemmTypes<T>;
typename Types::dataTypeI* A{nullptr};
typename Types::dataTypeI* B{nullptr};
typename Types::dataTypeO* C{nullptr};
int32_t m, n, k, ldA, ldB, ldC, rA, rB, rC, cA, cB, cC;
size_t bytesA;
size_t bytesB;
size_t bytesC;
size_t elemA;
size_t elemB;
size_t elemC;
bool transA;
bool transB;
nvinfer1::pluginInternal::cublasOperation_t opA;
nvinfer1::pluginInternal::cublasOperation_t opB;
int32_t const word_size{sizeof(T)};
typename Types::dataTypeS alpha;
typename Types::dataTypeS beta;
Gemm() {}
Gemm(int32_t m_, int32_t n_, int32_t k_, bool tA, bool tB)
{
init(m_, n_, k_, tA, tB);
}
void init(int32_t m_, int32_t n_, int32_t k_, bool tA, bool tB) noexcept
{
m = m_;
n = n_;
k = k_;
transA = tA;
transB = tB;
ldA = transA ? k : m;
ldB = transB ? n : k;
ldC = m;
rA = ldA;
rB = ldB;
rC = ldC;
cA = transA ? m : k;
cB = transB ? k : n;
cC = n;
opA = transA ? nvinfer1::pluginInternal::CUBLAS_OP_T : nvinfer1::pluginInternal::CUBLAS_OP_N;
opB = transB ? nvinfer1::pluginInternal::CUBLAS_OP_T : nvinfer1::pluginInternal::CUBLAS_OP_N;
elemA = m * k;
elemB = n * k;
elemC = n * m;
bytesA = word_size * elemA;
bytesB = word_size * elemB;
bytesC = word_size * elemC;
alpha = T(1.f);
beta = T(0.f);
}
};
auto constexpr kNB_ALGO_COMBINATIONS = 6000;
auto constexpr kNB_ALGO_IDS = 40;
auto constexpr kPRINT_ALGOS = 1;
auto constexpr kNB_KERNEL_REPEATS = 10;
auto constexpr kTHREADS_PER_BLOCK = 1024;
// Structure to store information about different run trials
typedef struct customMatMultPerfType_t
{
static constexpr float kMAX_TIME = 1000000.F;
nvinfer1::pluginInternal::cublasLtMatmulAlgo_t algo;
nvinfer1::pluginInternal::cublasStatus_t status;
float time{kMAX_TIME};
size_t workspaceSize; // actual memory workspace needed
nvinfer1::pluginInternal::cublasMath_t mathMode;
nvinfer1::pluginInternal::cublasLtReductionScheme_t reductionScheme;
int32_t customOption;
float wavesCount;
} customMatmulPerf_t;
// clang-format off
void LtGemmSearch(nvinfer1::pluginInternal::cublasLtHandle_t ltHandle,
nvinfer1::pluginInternal::cublasOperation_t transa,
nvinfer1::pluginInternal::cublasOperation_t transb,
int32_t const &m,
int32_t const &n,
int32_t const &k,
void const *alpha,
void const *A,
int32_t const &lda,
void const *B,
int32_t const &ldb,
void const *beta,
void *C,
int32_t const &ldc,
void *workSpace,
size_t workSpaceSize,
nvinfer1::pluginInternal::cublasComputeType_t computeType,
cudaDataType_t scaleType,
cudaDataType_t Atype,
cudaDataType_t Btype,
cudaDataType_t Ctype,
std::vector<customMatmulPerf_t> &perfResults,
cudaStream_t stream);
// clang-format on
template <typename T>
void LtGemmSearch(nvinfer1::pluginInternal::cublasLtHandle_t ltHandle, Gemm<T> const& g, void* workSpace,
size_t workSpaceSize, std::vector<customMatmulPerf_t>& perfResults, cudaStream_t stream)
{
// clang-format off
LtGemmSearch(
ltHandle,
g.opA,
g.opB,
g.m,
g.n,
g.k,
&g.alpha,
g.A,
g.ldA,
g.B,
g.ldB,
&g.beta,
g.C,
g.ldC,
workSpace,
workSpaceSize,
Gemm<T>::Types::cudaTypeCom,
Gemm<T>::Types::cudaTypeS,
Gemm<T>::Types::cudaTypeI,
Gemm<T>::Types::cudaTypeI,
Gemm<T>::Types::cudaTypeO,
perfResults,
stream
);
// clang-format on
}
struct LtContext
{
nvinfer1::pluginInternal::cublasLtHandle_t cublas{nullptr};
nvinfer1::pluginInternal::CublasLtWrapper& cublasLtWrapper = nvinfer1::pluginInternal::getCublasLtWrapper();
cudaDataType_t typeA;
cudaDataType_t typeB;
cudaDataType_t typeC;
nvinfer1::pluginInternal::cublasComputeType_t typeComp;
cudaDataType_t typeS;
nvinfer1::pluginInternal::cublasLtMatmulDesc_t operationDesc{nullptr};
nvinfer1::pluginInternal::cublasLtMatrixLayout_t Adesc{nullptr};
nvinfer1::pluginInternal::cublasLtMatrixLayout_t Bdesc{nullptr};
nvinfer1::pluginInternal::cublasLtMatrixLayout_t Cdesc{nullptr};
nvinfer1::pluginInternal::cublasLtMatmulHeuristicResult_t heuristicResult = {};
void attach()
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtCreate(&cublas));
}
void detach()
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtDestroy(cublas));
}
void destroy()
{
if (operationDesc)
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescDestroy(operationDesc));
operationDesc = nullptr;
}
if (Adesc)
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Adesc));
Adesc = nullptr;
}
if (Bdesc)
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Bdesc));
Bdesc = nullptr;
}
if (Cdesc)
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutDestroy(Cdesc));
Cdesc = nullptr;
}
}
template <typename T>
void create(Gemm<T>& g, size_t workspaceSize)
{
typeA = Gemm<T>::Types::cudaTypeI;
typeB = Gemm<T>::Types::cudaTypeI;
typeC = Gemm<T>::Types::cudaTypeO;
typeS = Gemm<T>::Types::cudaTypeS;
typeComp = Gemm<T>::Types::cudaTypeCom; // compute
// OPERATION
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescCreate(&operationDesc, typeComp, typeS));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescSetAttribute(
operationDesc, nvinfer1::pluginInternal::CUBLASLT_MATMUL_DESC_TRANSA, &g.opA, sizeof(g.opA)));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulDescSetAttribute(
operationDesc, nvinfer1::pluginInternal::CUBLASLT_MATMUL_DESC_TRANSB, &g.opB, sizeof(g.opB)));
// MAT DESC
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(&Adesc, typeA, g.rA, g.cA, g.ldA));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(&Bdesc, typeB, g.rB, g.cB, g.ldB));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutCreate(&Cdesc, typeC, g.rC, g.cC, g.ldC));
}
void setN(uint64_t n)
{
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutSetAttribute(
Bdesc, nvinfer1::pluginInternal::CUBLASLT_MATRIX_LAYOUT_COLS, &n, sizeof(n)));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatrixLayoutSetAttribute(
Cdesc, nvinfer1::pluginInternal::CUBLASLT_MATRIX_LAYOUT_COLS, &n, sizeof(n)));
}
};
template <typename T>
nvinfer1::pluginInternal::cublasStatus_t cublasLtMatmul(LtContext& ctx, Gemm<T>& g,
nvinfer1::pluginInternal::cublasLtMatmulAlgo_t algo, void* workspace, size_t workspaceSize, cudaStream_t stream)
{
nvinfer1::pluginInternal::CublasLtWrapper& cublasLtWrapper = nvinfer1::pluginInternal::getCublasLtWrapper();
// clang-format off
return cublasLtWrapper.cublasLtMatmul(
ctx.cublas,
ctx.operationDesc,
&g.alpha,
g.A,
ctx.Adesc,
g.B,
ctx.Bdesc,
&g.beta,
g.C,
ctx.Cdesc,
g.C,
ctx.Cdesc,
&algo,
workspace,
workspaceSize,
stream
);
// clang-format on
}
// CAUTION : must match cublasLtMatmulTile_t
char const* const matmulTileName[] = {
"UNDEF",
"8x8",
"8x16",
"16x8",
"8x32",
"16x16",
"32x8",
"8x64",
"16x32",
"32x16",
"64x8",
"32x32",
"32x64",
"64x32",
"32x128",
"64x64",
"128x32",
"64x128",
"128x64",
"64x256",
"128x128",
"256x64",
"64x512",
"128x256",
"256x128",
"512x64",
};
struct AlgoProps
{
int32_t algoId;
int32_t tile;
int32_t swizzle;
int32_t customOption;
int32_t numSplitsK;
int32_t reductionScheme;
uint64_t numericImpl;
void populate(nvinfer1::pluginInternal::cublasLtMatmulAlgo_t const& algo)
{
nvinfer1::pluginInternal::cublasLtMatmulAlgo_t const* matmulAlgo = &algo;
nvinfer1::pluginInternal::CublasLtWrapper& cublasLtWrapper = nvinfer1::pluginInternal::getCublasLtWrapper();
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigGetAttribute(
matmulAlgo, nvinfer1::pluginInternal::CUBLASLT_ALGO_CONFIG_ID, &algoId, sizeof(algoId), nullptr));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigGetAttribute(
matmulAlgo, nvinfer1::pluginInternal::CUBLASLT_ALGO_CONFIG_TILE_ID, &tile, sizeof(tile), nullptr));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigGetAttribute(matmulAlgo,
nvinfer1::pluginInternal::CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &numSplitsK, sizeof(numSplitsK), nullptr));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigGetAttribute(matmulAlgo,
nvinfer1::pluginInternal::CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &reductionScheme, sizeof(reductionScheme),
nullptr));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigGetAttribute(matmulAlgo,
nvinfer1::pluginInternal::CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, &swizzle, sizeof(swizzle), nullptr));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoConfigGetAttribute(matmulAlgo,
nvinfer1::pluginInternal::CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, &customOption, sizeof(customOption),
nullptr));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtMatmulAlgoCapGetAttribute(matmulAlgo,
nvinfer1::pluginInternal::CUBLASLT_ALGO_CAP_NUMERICAL_IMPL_FLAGS, &numericImpl, sizeof(numericImpl),
nullptr));
}
};
template <typename T>
nvinfer1::pluginInternal::cublasLtMatmulAlgo_t gemmSearch(int32_t const m, int32_t const n, int32_t const k,
size_t const workspaceSize, size_t& actualWorkspace, cudaStream_t& stream)
{
Gemm<T> g(m, n, k, false, false);
std::vector<customMatmulPerf_t> perfResults(kNB_ALGO_COMBINATIONS);
bool const useAsync = supportsMemPools();
PLUGIN_CUASSERT(useAsync ? cudaMallocAsync(reinterpret_cast<void**>(&g.A), g.bytesA, stream)
: cudaMalloc(reinterpret_cast<void**>(&g.A), g.bytesA));
PLUGIN_CUASSERT(useAsync ? cudaMallocAsync(reinterpret_cast<void**>(&g.B), g.bytesB, stream)
: cudaMalloc(reinterpret_cast<void**>(&g.B), g.bytesB));
PLUGIN_CUASSERT(useAsync ? cudaMallocAsync(reinterpret_cast<void**>(&g.C), g.bytesC, stream)
: cudaMalloc(reinterpret_cast<void**>(&g.C), g.bytesC));
void* workspace;
PLUGIN_CUASSERT(
useAsync ? cudaMallocAsync(&workspace, workspaceSize, stream) : cudaMalloc(&workspace, workspaceSize));
nvinfer1::pluginInternal::cublasLtHandle_t lt;
nvinfer1::pluginInternal::CublasLtWrapper& cublasLtWrapper = nvinfer1::pluginInternal::getCublasLtWrapper();
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtCreate(&lt));
LtGemmSearch(lt, g, workspace, workspaceSize, perfResults, stream);
PLUGIN_CUASSERT(cudaStreamSynchronize(stream));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtDestroy(lt));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(workspace, stream) : cudaFree(workspace));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(g.A, stream) : cudaFree(g.A));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(g.B, stream) : cudaFree(g.B));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(g.C, stream) : cudaFree(g.C));
actualWorkspace = perfResults[0].workspaceSize;
return perfResults[0].algo;
}
template <typename T>
nvinfer1::pluginInternal::cublasLtMatmulAlgo_t gemmSearch(
Gemm<T>& g, size_t const workspaceSize, size_t& actualWorkspace, cudaStream_t& stream)
{
std::vector<customMatmulPerf_t> perfResults(kNB_ALGO_COMBINATIONS);
bool const useAsync = supportsMemPools();
PLUGIN_CUASSERT(useAsync ? cudaMallocAsync(reinterpret_cast<void**>(&g.A), g.bytesA, stream)
: cudaMalloc(reinterpret_cast<void**>(&g.A), g.bytesA));
PLUGIN_CUASSERT(useAsync ? cudaMallocAsync(reinterpret_cast<void**>(&g.B), g.bytesB, stream)
: cudaMalloc(reinterpret_cast<void**>(&g.B), g.bytesB));
PLUGIN_CUASSERT(useAsync ? cudaMallocAsync(reinterpret_cast<void**>(&g.C), g.bytesC, stream)
: cudaMalloc(reinterpret_cast<void**>(&g.C), g.bytesC));
void* workspace;
PLUGIN_CUASSERT(
useAsync ? cudaMallocAsync(&workspace, workspaceSize, stream) : cudaMalloc(&workspace, workspaceSize));
nvinfer1::pluginInternal::cublasLtHandle_t lt;
nvinfer1::pluginInternal::CublasLtWrapper& cublasLtWrapper = nvinfer1::pluginInternal::getCublasLtWrapper();
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtCreate(&lt));
LtGemmSearch(lt, g, workspace, workspaceSize, perfResults, stream);
PLUGIN_CUASSERT(cudaStreamSynchronize(stream));
PLUGIN_CUBLASASSERT(cublasLtWrapper.cublasLtDestroy(lt));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(workspace, stream) : cudaFree(workspace));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(g.A, stream) : cudaFree(g.A));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(g.B, stream) : cudaFree(g.B));
PLUGIN_CUASSERT(useAsync ? cudaFreeAsync(g.C, stream) : cudaFree(g.C));
actualWorkspace = perfResults[0].workspaceSize;
return perfResults[0].algo;
}
// One of the preferred ways of making TensorRT to be able to see
// our custom layer requires extending IPluginV2 and IPluginCreator classes.
// For requirements for overriden functions, check TensorRT API docs.
class TRT_DEPRECATED_BECAUSE("Superseded by IMatrixMultiplyLayer.") FCPluginDynamic
: public nvinfer1::IPluginV2DynamicExt
{
public:
FCPluginDynamic(
std::string const name, nvinfer1::DataType const type, int32_t const outDim, nvinfer1::Weights const& W);
FCPluginDynamic(std::string const name, void const* data, size_t length);
// It doesn't make sense to make FCPluginDynamic without arguments, so we
// delete default constructor.
FCPluginDynamic() = delete;
// IPluginV2DynamicExt Methods
[[nodiscard]] nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs, int32_t nbInputs,
nvinfer1::IExprBuilder& exprBuilder) noexcept override;
bool supportsFormatCombination(
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override;
// IPluginV2 Methods
char const* getPluginType() const noexcept override;
char const* getPluginVersion() const noexcept override;
int32_t getNbOutputs() const noexcept override;
int32_t initialize() noexcept override;
void terminate() noexcept override;
size_t getSerializationSize() const noexcept override;
void serialize(void* buffer) const noexcept override;
void destroy() noexcept override;
void setPluginNamespace(char const* pluginNamespace) noexcept override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext,
nvinfer1::IGpuAllocator* gpuAllocator) noexcept override;
void detachFromContext() noexcept override;
char const* getPluginNamespace() const noexcept override;
private:
std::string const mLayerName;
std::string mNamespace;
nvinfer1::DataType mType;
size_t mOutDim; // leading dim
size_t mNumParams;
int32_t mNmax;
int32_t mK;
nvinfer1::pluginInternal::cublasLtMatmulAlgo_t mAlgo;
bert::WeightsWithOwnership mW;
bert::cuda_unique_ptr<void> mWdev;
LtContext mLtContext;
cudaStream_t mSharedStream{nullptr};
};
class TRT_DEPRECATED_BECAUSE("Superseded by IMatrixMultiplyLayer.") FCPluginDynamicCreator
: public nvinfer1::IPluginCreator
{
public:
FCPluginDynamicCreator();
char const* getPluginName() const noexcept override;
char const* getPluginVersion() const noexcept override;
nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
nvinfer1::IPluginV2* deserializePlugin(
char const* name, void const* serialData, size_t serialLength) noexcept override;
void setPluginNamespace(char const* pluginNamespace) noexcept override;
char const* getPluginNamespace() const noexcept override;
private:
nvinfer1::PluginFieldCollection mFC;
std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
} // namespace bert
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_FC_PLUGIN_H
#endif // #if CUDA_VERSION >= 10010