NvOnnxParser.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#ifndef NV_ONNX_PARSER_H
#define NV_ONNX_PARSER_H
 
#include "NvInfer.h"
#include <stddef.h>
 
 
#define NV_ONNX_PARSER_MAJOR 0
#define NV_ONNX_PARSER_MINOR 2
#define NV_ONNX_PARSER_PATCH 0
 
static constexpr int32_t NV_ONNX_PARSER_VERSION
    = ((NV_ONNX_PARSER_MAJOR * 10000) + (NV_ONNX_PARSER_MINOR * 100) + NV_ONNX_PARSER_PATCH);
 
namespace nvonnxparser
{
 
template <typename T>
constexpr int32_t EnumMax() noexcept = delete;
 
enum class ErrorCode : int
{
    kSUCCESS = 0,
    kINTERNAL_ERROR = 1,
    kMEM_ALLOC_FAILED = 2,
    kMODEL_DESERIALIZE_FAILED = 3,
    kINVALID_VALUE = 4,
    kINVALID_GRAPH = 5,
    kINVALID_NODE = 6,
    kUNSUPPORTED_GRAPH = 7,
    kUNSUPPORTED_NODE = 8,
    kUNSUPPORTED_NODE_ATTR = 9,
    kUNSUPPORTED_NODE_INPUT = 10,
    kUNSUPPORTED_NODE_DATATYPE = 11,
    kUNSUPPORTED_NODE_DYNAMIC = 12,
    kUNSUPPORTED_NODE_SHAPE = 13,
    kREFIT_FAILED = 14
};
 
template <>
constexpr int32_t EnumMax<ErrorCode>() noexcept
{
    return 14;
}
 
using OnnxParserFlags
    = uint32_t;
 
enum class OnnxParserFlag : int32_t
{
    kNATIVE_INSTANCENORM = 0,
    kENABLE_UINT8_AND_ASYMMETRIC_QUANTIZATION_DLA = 1,
    // This flag is set to be OFF by default.
    kREPORT_CAPABILITY_DLA = 2,
    kENABLE_PLUGIN_OVERRIDE = 3,
    kADJUST_FOR_DLA = 4
};
 
template <>
constexpr int32_t EnumMax<OnnxParserFlag>() noexcept
{
    return 4;
}
 
class IParserError
{
public:
    virtual ErrorCode code() const = 0;
    virtual char const* desc() const = 0;
    virtual char const* file() const = 0;
    virtual int line() const = 0;
    virtual char const* func() const = 0;
    virtual int node() const = 0;
    virtual char const* nodeName() const = 0;
    virtual char const* nodeOperator() const = 0;
    virtual char const* const* localFunctionStack() const = 0;
    //              a stack size of 0.
    virtual int32_t localFunctionStackSize() const = 0;
 
protected:
    virtual ~IParserError() {}
};
 
class IRefitterObserver;
 
class IParser
{
public:
    virtual bool parse(
        void const* serialized_onnx_model, size_t serialized_onnx_model_size, const char* model_path = nullptr) noexcept
        = 0;
 
    virtual bool parseFromFile(char const* onnxModelFile, int verbosity) noexcept = 0;
 
    virtual bool supportsOperator(const char* op_name) const noexcept = 0;
 
    virtual int getNbErrors() const noexcept = 0;
 
    virtual IParserError const* getError(int index) const noexcept = 0;
 
    virtual void clearErrors() noexcept = 0;
 
    virtual ~IParser() noexcept = default;
 
    virtual char const* const* getUsedVCPluginLibraries(int64_t& nbPluginLibs) const noexcept = 0;
 
    virtual void setFlags(OnnxParserFlags onnxParserFlags) noexcept = 0;
 
    virtual OnnxParserFlags getFlags() const noexcept = 0;
 
    virtual void clearFlag(OnnxParserFlag onnxParserFlag) noexcept = 0;
 
    virtual void setFlag(OnnxParserFlag onnxParserFlag) noexcept = 0;
 
    virtual bool getFlag(OnnxParserFlag onnxParserFlag) const noexcept = 0;
 
    virtual nvinfer1::ITensor const* getLayerOutputTensor(char const* name, int64_t i) noexcept = 0;
 
    virtual bool supportsModelV2(
        void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
 
    virtual int64_t getNbSubgraphs() noexcept = 0;
 
    virtual bool isSubgraphSupported(int64_t const index) noexcept = 0;
 
    virtual int64_t* getSubgraphNodes(int64_t const index, int64_t& subgraphLength) noexcept = 0;
 
    virtual bool loadModelProto(
        void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
 
    virtual bool loadInitializer(char const* name, void const* data, size_t size) noexcept = 0;
 
    virtual bool parseModelProto() noexcept = 0;
 
    virtual bool setBuilderConfig(const nvinfer1::IBuilderConfig* const builderConfig) noexcept = 0;
 
    virtual void setRefitObserver(IRefitterObserver* observer) noexcept = 0;
};
 
enum class RefitTransformKind : int32_t
{
    kIDENTITY = 0,
 
    kDOUBLE_TO_FLOAT = 1,
 
    kBATCH_NORM_FOLD_SCALE = 2,
 
    kBATCH_NORM_FOLD_BIAS = 3,
 
    kCONSTANT_NODE = 4,
 
    kCONSTANT_OF_SHAPE = 5,
};
 
template <>
constexpr int32_t EnumMax<RefitTransformKind>() noexcept
{
    return 5;
}
 
struct RefitRecord
{
    char const* trtName;
 
    RefitTransformKind kind;
 
    int32_t onnxDtype;
 
    nvinfer1::DataType trtDtype;
 
    int64_t count;
 
    int32_t nbSources;
 
    char const* const* sourceOnnxNames;
 
    float epsilon;
 
    void const* fixedData;
 
    size_t fixedDataSize;
};
 
class IRefitterObserver
{
public:
    virtual ~IRefitterObserver() noexcept = default;
 
    virtual void onRefittableWeight(RefitRecord const& record) noexcept = 0;
};
 
class IParserRefitter
{
public:
    virtual bool refitFromBytes(
        void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept
        = 0;
 
    virtual bool refitFromFile(char const* onnxModelFile) noexcept = 0;
 
    virtual int32_t getNbErrors() const noexcept = 0;
 
    virtual IParserError const* getError(int32_t index) const noexcept = 0;
 
    virtual void clearErrors() = 0;
 
    virtual ~IParserRefitter() noexcept = default;
 
    virtual bool loadModelProto(
        void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
 
    virtual bool loadInitializer(char const* name, void const* data, size_t size) noexcept = 0;
 
    virtual bool refitModelProto() noexcept = 0;
 
    virtual void setRefitObserver(IRefitterObserver* observer) noexcept = 0;
};
 
} // namespace nvonnxparser
 
extern "C" TENSORRTAPI void* createNvOnnxParser_INTERNAL(void* network, void* logger, int version) noexcept;
extern "C" TENSORRTAPI void* createNvOnnxParserRefitter_INTERNAL(
    void* refitter, void* logger, int32_t version) noexcept;
extern "C" TENSORRTAPI int getNvOnnxParserVersion() noexcept;
 
namespace nvonnxparser
{
 
namespace
{
 
inline IParser* createParser(nvinfer1::INetworkDefinition& network, nvinfer1::ILogger& logger) noexcept
{
    return static_cast<IParser*>(createNvOnnxParser_INTERNAL(&network, &logger, NV_ONNX_PARSER_VERSION));
}
 
inline IParserRefitter* createParserRefitter(nvinfer1::IRefitter& refitter, nvinfer1::ILogger& logger) noexcept
{
    return static_cast<IParserRefitter*>(
        createNvOnnxParserRefitter_INTERNAL(&refitter, &logger, NV_ONNX_PARSER_VERSION));
}
 
} // namespace
 
} // namespace nvonnxparser
 
#endif // NV_ONNX_PARSER_H