mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
311 lines
8.3 KiB
Smarty
311 lines
8.3 KiB
Smarty
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2022 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 "plugin.h"
|
|
#include <iostream>
|
|
|
|
using namespace nvinfer1;
|
|
using nvinfer1::plugin::[[ plugin_name ]]Creator;
|
|
using nvinfer1::plugin::[[ plugin_name ]];
|
|
|
|
PluginFieldCollection [[ plugin_name ]]Creator::mFC{};
|
|
std::vector<PluginField> [[ plugin_name ]]Creator::mPluginAttributes;
|
|
static bool triton_kernels_loaded = false;
|
|
|
|
// constructor
|
|
[[ plugin_name ]]::[[ plugin_name ]]( [[ construct_arg_list ]] )
|
|
{
|
|
{% for arg in params -%}
|
|
this->[[arg.name]] = [[arg.name]];
|
|
{% endfor %}
|
|
}
|
|
|
|
|
|
// Parameterized constructor
|
|
[[ plugin_name ]]::[[ plugin_name ]](const void* data, size_t length)
|
|
{
|
|
const char *d = static_cast<const char*>(data), *a = d;
|
|
|
|
{% for arg in params -%}
|
|
read(d, [[arg.name]]);
|
|
{% endfor %}
|
|
TLLM_CHECK(d == a + getSerializationSize());
|
|
}
|
|
|
|
|
|
nvinfer1::IPluginV2DynamicExt* [[ plugin_name ]]::clone() const noexcept
|
|
{
|
|
auto* plugin = new [[plugin_name]]([[', '.join(param_names)]]);
|
|
plugin->setPluginNamespace(mNamespace.c_str());
|
|
return plugin;
|
|
}
|
|
|
|
|
|
nvinfer1::DimsExprs [[ plugin_name ]]::getOutputDimensions(
|
|
int outputIndex, const nvinfer1::DimsExprs* inputDims, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
|
|
{
|
|
[[ getOutputDimensions_body ]]
|
|
}
|
|
|
|
bool [[ plugin_name ]]::supportsFormatCombination(
|
|
int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept
|
|
{
|
|
PLUGIN_ASSERT(nbInputs + nbOutputs == [[io_count]]);
|
|
PLUGIN_ASSERT(0 <= pos && pos < nbInputs + nbOutputs);
|
|
PLUGIN_ASSERT(nbInputs == [[input_count]]);
|
|
PLUGIN_ASSERT(nbOutputs == [[output_count]]);
|
|
|
|
{% for arg in inputs %}
|
|
if (pos == [[loop.index0]]) {
|
|
{%- if arg.is_tensor -%}
|
|
return inOut[pos].type == DataType::[[ arg.dtype.dtype.to('trt') ]] && inOut[pos].format == TensorFormat::kLINEAR;
|
|
{%- else -%}
|
|
return inOut[pos].type == DataType::[[ arg.dtype.dtype.to('trt') ]];
|
|
{%- endif -%}
|
|
}
|
|
{% endfor %}
|
|
|
|
{% for arg in outputs %}
|
|
if (pos == nbInputs + [[loop.index0]])
|
|
return inOut[pos].type == DataType::[[ arg.dtype.dtype.to('trt') ]] && inOut[pos].format == TensorFormat::kLINEAR;
|
|
{% endfor %}
|
|
}
|
|
|
|
|
|
void [[ plugin_name ]]::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs,
|
|
const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept
|
|
{
|
|
[[ configurePlugin_body ]]
|
|
}
|
|
|
|
|
|
size_t [[ plugin_name ]]::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs,
|
|
const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept
|
|
{
|
|
[[ getWorkspaceSize_body ]]
|
|
}
|
|
|
|
int [[ plugin_name ]]::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
|
|
const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
|
{
|
|
// input arguments
|
|
{% for arg in inputs %}
|
|
const auto* [[arg.name]]_buf = reinterpret_cast<const [[arg.dtype.dtype.to("c")]] *>(inputs[ [[arg.offset]] ]);
|
|
{% if arg.dtype.is_tensor -%}
|
|
auto [[arg.name]] = reinterpret_cast<CUdeviceptr>([[arg.name]]_buf);
|
|
{% else %}
|
|
const auto [[arg.name]] = * [[arg.name]]_buf;
|
|
{% endif -%}
|
|
{% endfor %}
|
|
|
|
// outputs
|
|
{% for arg in outputs %}
|
|
auto* [[arg.name]]_buf = reinterpret_cast<const [[arg.dtype.dtype.to("c")]] *>(outputs[ [[arg.offset]] ]);
|
|
auto [[arg.name]] = reinterpret_cast<CUdeviceptr>([[arg.name]]_buf);
|
|
{% endfor %}
|
|
|
|
// dim size arguments
|
|
{%- for arg in dim_size_args -%}
|
|
{# code field is dedicated for DimSizeArg #}
|
|
[[arg.dtype.dtype.to("c")]] [[arg.name]] = [[arg.code]];
|
|
{%- endfor %}
|
|
|
|
// TODO: Check result code
|
|
[[kernel_name]]([[enqueue_body_arg_list]]);
|
|
|
|
return 0;
|
|
}
|
|
|
|
|
|
nvinfer1::DataType [[ plugin_name ]]::getOutputDataType(
|
|
int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept
|
|
{
|
|
{% for arg in outputs %}
|
|
if (index == [[loop.index0]]) {
|
|
return DataType::[[arg.dtype.dtype.to('trt')]];
|
|
}
|
|
{% endfor %}
|
|
}
|
|
|
|
|
|
const char* [[ plugin_name ]]::getPluginType() const noexcept
|
|
{
|
|
return "[[ plugin_name ]]";
|
|
}
|
|
|
|
const char* [[ plugin_name ]]::getPluginVersion() const noexcept
|
|
{
|
|
return "[[ kernel_version ]]";
|
|
}
|
|
|
|
int [[ plugin_name ]]::getNbOutputs() const noexcept
|
|
{
|
|
return [[outputs|length]];
|
|
}
|
|
|
|
int [[ plugin_name ]]::initialize() noexcept
|
|
{
|
|
if (triton_kernels_loaded) {
|
|
return 0;
|
|
}
|
|
load_[[kernel_name]]();
|
|
triton_kernels_loaded = true;
|
|
return 0;
|
|
}
|
|
|
|
void [[ plugin_name ]]::terminate() noexcept {
|
|
if (!triton_kernels_loaded) {
|
|
return;
|
|
}
|
|
unload_[[kernel_name]]();
|
|
triton_kernels_loaded = false;
|
|
}
|
|
|
|
size_t [[ plugin_name ]]::getSerializationSize() const noexcept
|
|
{
|
|
size_t ret = 0;
|
|
|
|
{% for arg in params -%}
|
|
ret += sizeof([[arg.dtype.dtype.to('c')]]);
|
|
{% endfor %}
|
|
|
|
return ret;
|
|
|
|
}
|
|
|
|
void [[ plugin_name ]]::serialize(void* buffer) const noexcept
|
|
{
|
|
|
|
char *d = static_cast<char*>(buffer), *a = d;
|
|
|
|
{% for arg in params -%}
|
|
write(d, [[arg.name]]);
|
|
{% endfor %}
|
|
TLLM_CHECK(d == a + getSerializationSize());
|
|
|
|
}
|
|
|
|
void [[ plugin_name ]]::destroy() noexcept
|
|
{
|
|
// This gets called when the network containing plugin is destroyed
|
|
delete this;
|
|
}
|
|
|
|
void [[ plugin_name ]]::setPluginNamespace(const char* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
const char* [[ plugin_name ]]::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
|
|
|
|
[[ plugin_name ]]Creator::[[ plugin_name ]]Creator()
|
|
{
|
|
// Fill PluginFieldCollection with PluginField arguments metadata
|
|
mPluginAttributes.clear();
|
|
|
|
{% for arg in params %}
|
|
mPluginAttributes.emplace_back(PluginField("[[arg.name]]", nullptr, PluginFieldType::[[arg.dtype.dtype.to('trt_plugin')]]));
|
|
{% endfor %}
|
|
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
|
|
const char* [[ plugin_name ]]Creator::getPluginName() const noexcept
|
|
{
|
|
return "[[ plugin_name ]]";
|
|
}
|
|
|
|
const char* [[ plugin_name ]]Creator::getPluginVersion() const noexcept
|
|
{
|
|
return "[[ plugin_version ]]";
|
|
}
|
|
|
|
const PluginFieldCollection* [[ plugin_name ]]Creator::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
}
|
|
|
|
IPluginV2* [[ plugin_name ]]Creator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
|
|
{
|
|
const PluginField* fields = fc->fields;
|
|
|
|
// declare parameters
|
|
{% for arg in params %}
|
|
[[arg.dtype.dtype.to('c')]] [[arg.name]];
|
|
{% endfor %}
|
|
|
|
for (int i = 0; i < fc->nbFields; ++i) {
|
|
const char* attrName = fields[i].name;
|
|
{% for arg in params %}
|
|
if (!strcmp(attrName, "[[arg.name]]"))
|
|
{
|
|
PLUGIN_ASSERT(fields[i].type == PluginFieldType::[[arg.dtype.dtype.to('trt_plugin')]]);
|
|
[[arg.name]] = static_cast<[[arg.dtype.dtype.to('c')]]>(*(static_cast<const [[arg.dtype.dtype.to('c')]]*>(fields[i].data)));
|
|
}
|
|
{% endfor %}
|
|
}
|
|
|
|
try
|
|
{
|
|
auto* obj = new [[plugin_name]]([[ ', '.join(param_names) ]]);
|
|
obj->setPluginNamespace(mNamespace.c_str());
|
|
return obj;
|
|
}
|
|
catch (const std::exception& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
|
|
}
|
|
|
|
|
|
IPluginV2* [[ plugin_name ]]Creator::deserializePlugin(
|
|
const char* name, const void* serialData, size_t serialLength) noexcept
|
|
{
|
|
// This object will be deleted when the network is destroyed, which will
|
|
// call [[ plugin_name ]]::destroy()
|
|
try
|
|
{
|
|
auto* obj = new [[ plugin_name ]](serialData, serialLength);
|
|
obj->setPluginNamespace("tensorrt_llm");
|
|
return obj;
|
|
}
|
|
catch (const std::exception& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void [[plugin_name]]Creator::setPluginNamespace(const char* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
const char* [[plugin_name]]Creator::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|