NvInferPythonPlugin.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2024-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.
 */
 
#ifndef TRT_PYTHON_IMPL_PLUGIN_H
#define TRT_PYTHON_IMPL_PLUGIN_H
 
#include "NvInfer.h"
 
 
namespace nvinfer1
{
 
enum class PluginArgType : int32_t
{
    kINT = 0,
};
 
enum class PluginArgDataType : int32_t
{
    kINT8 = 0,
    kINT16 = 1,
    kINT32 = 2,
};
class ISymExpr
{
public:
    virtual PluginArgType getType() const noexcept = 0;
    virtual PluginArgDataType getDataType() const noexcept = 0;
    virtual void* getExpr() noexcept = 0;
};
 
class ISymExprsImpl
{
public:
    virtual ISymExpr* getSymExpr(int32_t index) const noexcept = 0;
    virtual bool setSymExpr(int32_t index, ISymExpr* symExpr) noexcept = 0;
    virtual int32_t getNbSymExprs() const noexcept = 0;
    virtual bool setNbSymExprs(int32_t count) noexcept = 0;
 
    virtual ~ISymExprsImpl() noexcept = 0;
};
 
inline ISymExprsImpl::~ISymExprsImpl() noexcept = default;
 
class ISymExprs
{
public:
    ISymExpr* getSymExpr(int32_t index) const noexcept
    {
        return mImpl->getSymExpr(index);
    }
 
    bool setSymExpr(int32_t index, ISymExpr* symExpr) noexcept
    {
        return mImpl->setSymExpr(index, symExpr);
    }
 
    int32_t getNbSymExprs() const noexcept
    {
        return mImpl->getNbSymExprs();
    }
 
    bool setNbSymExprs(int32_t count) noexcept
    {
        return mImpl->setNbSymExprs(count);
    }
 
protected:
    ISymExprsImpl* mImpl{nullptr};
    virtual ~ISymExprs() noexcept = 0;
};
 
inline ISymExprs::~ISymExprs() noexcept = default;
 
enum class QuickPluginCreationRequest : int32_t
{
    kUNKNOWN = 0,
    kPREFER_JIT = 1,
    kPREFER_AOT = 2,
    kSTRICT_JIT = 3,
    kSTRICT_AOT = 4,
};
 
class IKernelLaunchParamsImpl
{
public:
    virtual ISymExpr* getGridX() noexcept = 0;
    virtual bool setGridX(ISymExpr* gridX) noexcept = 0;
 
    virtual ISymExpr* getGridY() noexcept = 0;
    virtual bool setGridY(ISymExpr* gridY) noexcept = 0;
 
    virtual ISymExpr* getGridZ() noexcept = 0;
    virtual bool setGridZ(ISymExpr* gridZ) noexcept = 0;
 
    virtual ISymExpr* getBlockX() noexcept = 0;
    virtual bool setBlockX(ISymExpr* blockX) noexcept = 0;
 
    virtual ISymExpr* getBlockY() noexcept = 0;
    virtual bool setBlockY(ISymExpr* blockY) noexcept = 0;
 
    virtual ISymExpr* getBlockZ() noexcept = 0;
    virtual bool setBlockZ(ISymExpr* blockZ) noexcept = 0;
 
    virtual ISymExpr* getSharedMem() noexcept = 0;
    virtual bool setSharedMem(ISymExpr* sharedMem) noexcept = 0;
 
    virtual ~IKernelLaunchParamsImpl() noexcept = 0;
};
 
inline IKernelLaunchParamsImpl::~IKernelLaunchParamsImpl() noexcept = default;
 
class IKernelLaunchParams
{
public:
    ISymExpr* getGridX() noexcept
    {
        return mImpl->getGridX();
    }
 
    bool setGridX(ISymExpr* gridX) noexcept
    {
        return mImpl->setGridX(gridX);
    }
 
    ISymExpr* getGridY() noexcept
    {
        return mImpl->getGridY();
    }
 
    bool setGridY(ISymExpr* gridY) noexcept
    {
        return mImpl->setGridY(gridY);
    }
 
    ISymExpr* getGridZ() noexcept
    {
        return mImpl->getGridZ();
    }
 
    bool setGridZ(ISymExpr* gridZ) noexcept
    {
        return mImpl->setGridZ(gridZ);
    }
 
    ISymExpr* getBlockX() noexcept
    {
        return mImpl->getBlockX();
    }
 
    bool setBlockX(ISymExpr* blockX) noexcept
    {
        return mImpl->setBlockX(blockX);
    }
 
    ISymExpr* getBlockY() noexcept
    {
        return mImpl->getBlockY();
    }
 
    bool setBlockY(ISymExpr* blockY) noexcept
    {
        return mImpl->setBlockY(blockY);
    }
 
    ISymExpr* getBlockZ() noexcept
    {
        return mImpl->getBlockZ();
    }
 
    bool setBlockZ(ISymExpr* blockZ) noexcept
    {
        return mImpl->setBlockZ(blockZ);
    }
 
    ISymExpr* getSharedMem() noexcept
    {
        return mImpl->getSharedMem();
    }
 
    bool setSharedMem(ISymExpr* sharedMem) noexcept
    {
        return mImpl->setSharedMem(sharedMem);
    }
 
protected:
    IKernelLaunchParamsImpl* mImpl{nullptr};
    virtual ~IKernelLaunchParams() noexcept = 0;
};
 
inline IKernelLaunchParams::~IKernelLaunchParams() noexcept = default;
 
 
namespace v_1_0
{
 
class IPluginV3QuickCore : public IPluginCapability
{
public:
    InterfaceInfo getInterfaceInfo() const noexcept override
    {
        return InterfaceInfo{"PLUGIN_V3QUICK_CORE", 1, 0};
    }
 
    virtual AsciiChar const* getPluginName() const noexcept = 0;
 
    virtual AsciiChar const* getPluginVersion() const noexcept = 0;
 
    virtual AsciiChar const* getPluginNamespace() const noexcept = 0;
};
 
class IPluginV3QuickBuild : public IPluginCapability
{
public:
    InterfaceInfo getInterfaceInfo() const noexcept override
    {
        return InterfaceInfo{"PLUGIN_V3QUICK_BUILD", 1, 0};
    }
 
    virtual int32_t getOutputDataTypes(DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes,
        int32_t const* inputRanks, int32_t nbInputs) const noexcept = 0;
 
    virtual int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
        int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept = 0;
 
    virtual int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs,
        DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept = 0;
 
    virtual int32_t getNbSupportedFormatCombinations(
        DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept = 0;
 
    virtual int32_t getSupportedFormatCombinations(DynamicPluginTensorDesc const* inOut, int32_t nbInputs,
        int32_t nbOutputs, PluginTensorDesc* supportedCombinations, int32_t nbFormatCombinations) noexcept = 0;
 
    virtual int32_t getNbOutputs() const noexcept = 0;
 
    virtual int32_t getAliasedInput(int32_t outputIndex) noexcept
    {
        return -1;
    }
 
    virtual int32_t getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept
    {
        return 0;
    }
 
    virtual int32_t getNbTactics() noexcept
    {
        return 0;
    }
 
    virtual char const* getTimingCacheID() noexcept
    {
        return nullptr;
    }
 
    virtual char const* getMetadataString() noexcept
    {
        return nullptr;
    }
};
 
class IPluginV3QuickAOTBuild : public IPluginV3QuickBuild
{
public:
    InterfaceInfo getInterfaceInfo() const noexcept override
    {
        return InterfaceInfo{"PLUGIN_V3QUICKAOT_BUILD", 1, 0};
    }
 
    virtual int32_t getLaunchParams(DimsExprs const* inputs, DynamicPluginTensorDesc const* inOut, int32_t nbInputs,
        int32_t nbOutputs, IKernelLaunchParams* launchParams, ISymExprs* extraArgs,
        IExprBuilder& exprBuilder) noexcept = 0;
 
    virtual int32_t getKernel(PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out,
        int32_t nbOutputs, const char** kernelName, char** compiledKernel, int32_t* compiledKernelSize) noexcept = 0;
 
    virtual int32_t setTactic(int32_t tactic) noexcept
    {
        return 0;
    }
};
 
class IPluginV3QuickRuntime : public IPluginCapability
{
public:
    InterfaceInfo getInterfaceInfo() const noexcept override
    {
        return InterfaceInfo{"PLUGIN_V3QUICK_RUNTIME", 1, 0};
    }
 
    virtual int32_t setTactic(int32_t tactic) noexcept
    {
        return 0;
    }
 
    virtual int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
        void const* const* inputs, void* const* outputs, Dims const* inputStrides, Dims const* outputStrides,
        int32_t nbInputs, int32_t nbOutputs, cudaStream_t stream) noexcept = 0;
 
    virtual PluginFieldCollection const* getFieldsToSerialize() noexcept = 0;
};
 
class IPluginCreatorV3Quick : public IPluginCreatorInterface
{
public:
    InterfaceInfo getInterfaceInfo() const noexcept override
    {
        return InterfaceInfo{"PLUGIN CREATOR_V3QUICK", 1, 0};
    }
 
    virtual IPluginV3* createPlugin(AsciiChar const* name, AsciiChar const* nspace, PluginFieldCollection const* fc,
        TensorRTPhase phase, QuickPluginCreationRequest quickPluginCreationRequest) noexcept = 0;
 
    virtual PluginFieldCollection const* getFieldNames() noexcept = 0;
 
    virtual AsciiChar const* getPluginName() const noexcept = 0;
 
    virtual AsciiChar const* getPluginVersion() const noexcept = 0;
 
    virtual AsciiChar const* getPluginNamespace() const noexcept = 0;
 
    IPluginCreatorV3Quick() = default;
    virtual ~IPluginCreatorV3Quick() = default;
 
protected:
    IPluginCreatorV3Quick(IPluginCreatorV3Quick const&) = default;
    IPluginCreatorV3Quick(IPluginCreatorV3Quick&&) = default;
    IPluginCreatorV3Quick& operator=(IPluginCreatorV3Quick const&) & = default;
    IPluginCreatorV3Quick& operator=(IPluginCreatorV3Quick&&) & = default;
};
 
} // namespace v_1_0
 
using IPluginV3QuickCore = v_1_0::IPluginV3QuickCore;
 
using IPluginV3QuickBuild = v_1_0::IPluginV3QuickBuild;
 
using IPluginV3QuickAOTBuild = v_1_0::IPluginV3QuickAOTBuild;
 
using IPluginV3QuickRuntime = v_1_0::IPluginV3QuickRuntime;
 
using IPluginCreatorV3Quick = v_1_0::IPluginCreatorV3Quick;
 
} // namespace nvinfer1
 
#endif // TRT_PYTHON_IMPL_PLUGIN_H