NvInfer.h

Go to the documentation of this file.

/*
 * Copyright 1993-2018 NVIDIA Corporation.  All rights reserved.
 *
 * NOTICE TO LICENSEE:
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 *
 * These Licensed Deliverables contained herein is PROPRIETARY and
 * CONFIDENTIAL to NVIDIA and is being provided under the terms and
 * conditions of a form of NVIDIA software license agreement by and
 * between NVIDIA and Licensee ("License Agreement") or electronically
 * accepted by Licensee.  Notwithstanding any terms or conditions to
 * the contrary in the License Agreement, reproduction or disclosure
 * of the Licensed Deliverables to any third party without the express
 * written consent of NVIDIA is prohibited.
 *
 * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
 * LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
 * SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE.  IT IS
 * PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
 * NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
 * DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
 * NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
 * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
 * LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
 * SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
 * DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
 * WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
 * ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
 * OF THESE LICENSED DELIVERABLES.
 *
 * U.S. Government End Users.  These Licensed Deliverables are a
 * "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
 * 1995), consisting of "commercial computer software" and "commercial
 * computer software documentation" as such terms are used in 48
 * C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
 * only as a commercial end item.  Consistent with 48 C.F.R.12.212 and
 * 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
 * U.S. Government End Users acquire the Licensed Deliverables with
 * only those rights set forth herein.
 *
 * Any use of the Licensed Deliverables in individual and commercial
 * software must include, in the user documentation and internal
 * comments to the code, the above Disclaimer and U.S. Government End
 * Users Notice.
 */

#ifndef NV_INFER_H
#define NV_INFER_H

#include <cstddef>
#include <cstdint>

#define NV_TENSORRT_MAJOR 4 
#define NV_TENSORRT_MINOR 0 
#define NV_TENSORRT_PATCH 1 
#define NV_TENSORRT_BUILD 3 

#define NV_TENSORRT_SONAME_MAJOR 4 
#define NV_TENSORRT_SONAME_MINOR 1 
#define NV_TENSORRT_SONAME_PATCH 2 

#if __cplusplus > 201103L
#define _TENSORRT_FINAL final
#else
#define _TENSORRT_FINAL
#endif

#ifdef TENSORRT_BUILD_LIB
#define TENSORRTAPI __attribute__ ((visibility ("default")))
#else
#define TENSORRTAPI
#endif



// forward declare some CUDA types to avoid an include dependency

struct cublasContext;
struct cudnnContext;

typedef struct CUstream_st* cudaStream_t; 
typedef struct CUevent_st* cudaEvent_t;   

static const int NV_TENSORRT_VERSION = (NV_TENSORRT_MAJOR * 1000) + (NV_TENSORRT_MINOR * 100) + NV_TENSORRT_PATCH; // major, minor, patch

namespace nvinfer1
{

template <typename T>
inline int EnumMax(); 

enum class DataType : int
{
    kFLOAT = 0, 
    kHALF = 1,  
    kINT8 = 2,  
    kINT32 = 3  
};

template <>
inline int EnumMax<DataType>()
{
    return 4;
} 

enum class DimensionType : int
{
    kSPATIAL = 0, 
    kCHANNEL = 1, 
    kINDEX = 2,   
    kSEQUENCE = 3 
};

template <>
inline int EnumMax<DimensionType>()
{
    return 4;
} 

class Dims
{
public:
    static const int MAX_DIMS = 8; 
    int nbDims;                    
    int d[MAX_DIMS];               
    DimensionType type[MAX_DIMS];  
};

class Dims2 : public Dims
{
public:
    Dims2()
    {
        nbDims = 2;
        d[0] = d[1] = 0;
    }

    Dims2(int d0, int d1)
    {
        nbDims = 2;
        d[0] = d0;
        d[1] = d1;
    }
};

class DimsHW : public Dims2
{
public:
    DimsHW()
        : Dims2()
    {
        type[0] = type[1] = DimensionType::kSPATIAL;
    }

    DimsHW(int height, int width)
        : Dims2(height, width)
    {
        type[0] = type[1] = DimensionType::kSPATIAL;
    }

    int& h() { return d[0]; }

    int h() const { return d[0]; }

    int& w() { return d[1]; }

    int w() const { return d[1]; }
};

class Dims3 : public Dims
{
public:
    Dims3()
    {
        nbDims = 3;
        d[0] = d[1] = d[2] = 0;
    }

    Dims3(int d0, int d1, int d2)
    {
        nbDims = 3;
        d[0] = d0;
        d[1] = d1;
        d[2] = d2;
    }
};

class DimsCHW : public Dims3
{
public:
    DimsCHW()
        : Dims3()
    {
        type[0] = DimensionType::kCHANNEL;
        type[1] = type[2] = DimensionType::kSPATIAL;
    }

    DimsCHW(int channels, int height, int width)
        : Dims3(channels, height, width)
    {
        type[0] = DimensionType::kCHANNEL;
        type[1] = type[2] = DimensionType::kSPATIAL;
    }

    int& c() { return d[0]; }

    int c() const { return d[0]; }

    int& h() { return d[1]; }

    int h() const { return d[1]; }

    int& w() { return d[2]; }

    int w() const { return d[2]; }
};

class Dims4 : public Dims
{
public:
    Dims4()
    {
        nbDims = 4;
        d[0] = d[1] = d[2] = d[3] = 0;
    }

    Dims4(int d0, int d1, int d2, int d3)
    {
        nbDims = 4;
        d[0] = d0;
        d[1] = d1;
        d[2] = d2;
        d[3] = d3;
    }
};

class DimsNCHW : public Dims4
{
public:
    DimsNCHW()
        : Dims4()
    {
        type[0] = DimensionType::kINDEX;
        type[1] = DimensionType::kCHANNEL;
        type[2] = type[3] = DimensionType::kSPATIAL;
    }

    DimsNCHW(int batchSize, int channels, int height, int width)
        : Dims4(batchSize, channels, height, width)
    {
        type[0] = DimensionType::kINDEX;
        type[1] = DimensionType::kCHANNEL;
        type[2] = type[3] = DimensionType::kSPATIAL;
    }

    int& n() { return d[0]; }

    int n() const { return d[0]; }

    int& c() { return d[1]; }

    int c() const { return d[1]; }

    int& h() { return d[2]; }

    int h() const { return d[2]; }

    int& w() { return d[3]; }

    int w() const { return d[3]; }
};

class Weights
{
public:
    DataType type;      
    const void* values; 
    int64_t count;      
};

class IHostMemory
{
public:
    virtual void* data() const = 0;       
    virtual std::size_t size() const = 0; 
    virtual DataType type() const = 0;    
    virtual void destroy() = 0;           
protected:
    virtual ~IHostMemory() {}
};

enum class LayerType : int
{
    kCONVOLUTION = 0,     
    kFULLY_CONNECTED = 1, 
    kACTIVATION = 2,      
    kPOOLING = 3,         
    kLRN = 4,             
    kSCALE = 5,           
    kSOFTMAX = 6,         
    kDECONVOLUTION = 7,   
    kCONCATENATION = 8,   
    kELEMENTWISE = 9,     
    kPLUGIN = 10,         
    kRNN = 11,            
    kUNARY = 12,          
    kPADDING = 13,        
    kSHUFFLE = 14,        
    kREDUCE = 15,         
    kTOPK = 16,           
    kGATHER = 17,         
    kMATRIX_MULTIPLY = 18,
    kRAGGED_SOFTMAX = 19, 
    kCONSTANT = 20,       
    kRNN_V2 = 21          
};

template <>
inline int EnumMax<LayerType>()
{
    return 22;
} 

enum class TensorLocation : int
{
    kDEVICE = 0, 
    kHOST = 1    
};

template <>
inline int EnumMax<TensorLocation>()
{
    return 2;
} 

class ITensor
{
public:
    virtual void setName(const char* name) = 0;

    virtual const char* getName() const = 0;

    virtual void setDimensions(Dims dimensions) = 0; // only valid for input tensors

    virtual Dims getDimensions() const = 0;

    virtual void setType(DataType type) = 0;

    virtual DataType getType() const = 0;

    virtual bool isNetworkInput() const = 0;

    virtual bool isNetworkOutput() const = 0;

protected:
    virtual ~ITensor() {}

public:
    virtual void setBroadcastAcrossBatch(bool broadcastAcrossBatch) = 0;

    virtual bool getBroadcastAcrossBatch() const = 0;

    virtual TensorLocation getLocation() const = 0;

    virtual void setLocation(TensorLocation location) = 0;
};

class ILayer
{
public:
    virtual LayerType getType() const = 0;

    virtual void setName(const char* name) = 0;

    virtual const char* getName() const = 0;

    virtual int getNbInputs() const = 0;

    virtual ITensor* getInput(int index) const = 0;

    virtual int getNbOutputs() const = 0;

    virtual ITensor* getOutput(int index) const = 0;

protected:
    virtual ~ILayer() {}
};

class IConvolutionLayer : public ILayer
{
public:
    virtual void setKernelSize(DimsHW kernelSize) = 0;

    virtual DimsHW getKernelSize() const = 0;

    virtual void setNbOutputMaps(int nbOutputMaps) = 0;

    virtual int getNbOutputMaps() const = 0;

    virtual void setStride(DimsHW stride) = 0;

    virtual DimsHW getStride() const = 0;

    virtual void setPadding(DimsHW padding) = 0;

    virtual DimsHW getPadding() const = 0; // padding defaults to 0

    virtual void setNbGroups(int nbGroups) = 0;

    virtual int getNbGroups() const = 0;

    virtual void setKernelWeights(Weights weights) = 0;

    virtual Weights getKernelWeights() const = 0;

    virtual void setBiasWeights(Weights weights) = 0;

    virtual Weights getBiasWeights() const = 0;

    virtual void setDilation(DimsHW dims) = 0;

    virtual DimsHW getDilation() const = 0;

protected:
    virtual ~IConvolutionLayer() {}
};

class IFullyConnectedLayer : public ILayer
{
public:
    virtual void setNbOutputChannels(int nbOutputs) = 0;

    virtual int getNbOutputChannels() const = 0;

    virtual void setKernelWeights(Weights weights) = 0;

    virtual Weights getKernelWeights() const = 0;

    virtual void setBiasWeights(Weights weights) = 0;

    virtual Weights getBiasWeights() const = 0;

protected:
    virtual ~IFullyConnectedLayer() {}
};

enum class ActivationType : int
{
    kRELU = 0,    
    kSIGMOID = 1, 
    kTANH = 2     
};

template <>
inline int EnumMax<ActivationType>()
{
    return 3;
} 

class IActivationLayer : public ILayer
{
public:
    virtual void setActivationType(ActivationType type) = 0;

    virtual ActivationType getActivationType() const = 0;

protected:
    virtual ~IActivationLayer() {}
};

enum class PoolingType : int
{
    kMAX = 0,              // Maximum over elements
    kAVERAGE = 1,          // Average over elements. If the tensor is padded, the count includes the padding
    kMAX_AVERAGE_BLEND = 2 // Blending between the max pooling and average pooling: (1-blendFactor)*maxPool + blendFactor*avgPool
};

template <>
inline int EnumMax<PoolingType>()
{
    return 3;
} 

class IPoolingLayer : public ILayer
{
public:
    virtual void setPoolingType(PoolingType type) = 0;

    virtual PoolingType getPoolingType() const = 0;

    virtual void setWindowSize(DimsHW windowSize) = 0;

    virtual DimsHW getWindowSize() const = 0;

    virtual void setStride(DimsHW stride) = 0;

    virtual DimsHW getStride() const = 0;

    virtual void setPadding(DimsHW padding) = 0;

    virtual DimsHW getPadding() const = 0;

    virtual void setBlendFactor(float blendFactor) = 0;

    virtual float getBlendFactor() const = 0;

    virtual void setAverageCountExcludesPadding(bool exclusive) = 0;

    virtual bool getAverageCountExcludesPadding() const = 0;

protected:
    virtual ~IPoolingLayer() {}
};

class ILRNLayer : public ILayer
{
public:
    virtual void setWindowSize(int windowSize) = 0;

    virtual int getWindowSize() const = 0;

    virtual void setAlpha(float alpha) = 0;

    virtual float getAlpha() const = 0;

    virtual void setBeta(float beta) = 0;

    virtual float getBeta() const = 0;

    virtual void setK(float k) = 0;

    virtual float getK() const = 0;

protected:
    virtual ~ILRNLayer() {}
};

enum class ScaleMode : int
{
    kUNIFORM = 0,    
    kCHANNEL = 1,    
    kELEMENTWISE = 2 
};

template <>
inline int EnumMax<ScaleMode>()
{
    return 3;
} 

class IScaleLayer : public ILayer
{
public:
    virtual void setMode(ScaleMode mode) = 0;

    virtual ScaleMode getMode() const = 0;

    virtual void setShift(Weights shift) = 0;

    virtual Weights getShift() const = 0;

    virtual void setScale(Weights scale) = 0;

    virtual Weights getScale() const = 0;

    virtual void setPower(Weights power) = 0;

    virtual Weights getPower() const = 0;

protected:
    virtual ~IScaleLayer() {}
};

class ISoftMaxLayer : public ILayer
{
protected:
    virtual ~ISoftMaxLayer() {}
public:
    virtual void setAxes(uint32_t axes) = 0;

    virtual uint32_t getAxes() const = 0;
};

class IConcatenationLayer : public ILayer
{
protected:
    virtual ~IConcatenationLayer() {}

public:
    virtual void setAxis(int axis) = 0;

    virtual int getAxis() const = 0;
};

class IDeconvolutionLayer : public ILayer
{
public:
    virtual void setKernelSize(DimsHW kernelSize) = 0;

    virtual DimsHW getKernelSize() const = 0;

    virtual void setNbOutputMaps(int nbOutputMaps) = 0;

    virtual int getNbOutputMaps() const = 0;

    virtual void setStride(DimsHW stride) = 0;

    virtual DimsHW getStride() const = 0;

    virtual void setPadding(DimsHW padding) = 0;

    virtual DimsHW getPadding() const = 0; // padding defaults to 0

    virtual void setNbGroups(int nbGroups) = 0;

    virtual int getNbGroups() const = 0;

    virtual void setKernelWeights(Weights weights) = 0;

    virtual Weights getKernelWeights() const = 0;

    virtual void setBiasWeights(Weights weights) = 0;

    virtual Weights getBiasWeights() const = 0;

protected:
    virtual ~IDeconvolutionLayer() {}
};

enum class ElementWiseOperation : int
{
    kSUM = 0,  
    kPROD = 1, 
    kMAX = 2,  
    kMIN = 3,  
    kSUB = 4,  
    kDIV = 5,  
    kPOW = 6   
};

template <>
inline int EnumMax<ElementWiseOperation>()
{
    return 7;
} 

class IElementWiseLayer : public ILayer
{
public:
    virtual void setOperation(ElementWiseOperation type) = 0;

    virtual ElementWiseOperation getOperation() const = 0;

protected:
    virtual ~IElementWiseLayer() {}
};

class IGatherLayer : public ILayer
{
public:
    virtual void setGatherAxis(int axis) = 0;

    virtual int getGatherAxis() const = 0;

protected:
    virtual ~IGatherLayer() {}
};

enum class RNNOperation : int
{
    kRELU = 0, 
    kTANH = 1, 
    kLSTM = 2, 
    kGRU = 3   
};

template <>
inline int EnumMax<RNNOperation>()
{
    return 4;
} 

enum class RNNDirection : int
{
    kUNIDIRECTION = 0, 
    kBIDIRECTION = 1   
};

template <>
inline int EnumMax<RNNDirection>()
{
    return 2;
} 

enum class RNNInputMode : int
{
    kLINEAR = 0, 
    kSKIP = 1    
};

template <>
inline int EnumMax<RNNInputMode>()
{
    return 2;
} 

class IRNNLayer : public ILayer
{
public:
    virtual unsigned getLayerCount() const = 0;

    virtual std::size_t getHiddenSize() const = 0;

    virtual int getSeqLength() const = 0;

    virtual void setOperation(RNNOperation op) = 0;

    virtual RNNOperation getOperation() const = 0;

    virtual void setInputMode(RNNInputMode op) = 0;

    virtual RNNInputMode getInputMode() const = 0;

    virtual void setDirection(RNNDirection op) = 0;

    virtual RNNDirection getDirection() const = 0;

    virtual void setWeights(Weights weights) = 0;

    virtual Weights getWeights() const = 0;

    virtual void setBias(Weights bias) = 0;

    virtual Weights getBias() const = 0;

    virtual int getDataLength() const = 0;

    virtual void setHiddenState(ITensor& hidden) = 0;

    virtual ITensor* getHiddenState() const = 0;

    virtual void setCellState(ITensor& cell) = 0;

    virtual ITensor* getCellState() const = 0;

protected:
    virtual ~IRNNLayer() {}
};

enum class RNNGateType : int
{
    kINPUT  = 0, 
    kOUTPUT = 1, 
    kFORGET = 2, 
    kUPDATE = 3, 
    kRESET  = 4, 
    kCELL   = 5, 
    kHIDDEN = 6  
};

template <>
inline int EnumMax<RNNGateType>()
{
    return 7;
} 

class IRNNv2Layer : public ILayer
{
public:
    virtual int32_t getLayerCount() const = 0; //< Get the layer count of the RNN
    virtual int32_t getHiddenSize() const = 0; //< Get the hidden size of the RNN
    virtual int32_t getMaxSeqLength() const = 0; //< Get the maximum sequence length of the RNN
    virtual int32_t getDataLength() const = 0; //< Get the maximum data length of the RNN

    virtual void setSequenceLengths(ITensor& seqLengths) = 0;

    virtual ITensor *getSequenceLengths() const = 0;

    virtual void setOperation(RNNOperation op) = 0;

    virtual RNNOperation getOperation() const = 0;

    virtual void setInputMode(RNNInputMode op) = 0;

    virtual RNNInputMode getInputMode() const = 0;

    virtual void setDirection(RNNDirection op) = 0;

    virtual RNNDirection getDirection() const = 0;

    virtual void setWeightsForGate(int layerIndex, RNNGateType gate, bool isW, Weights weights) = 0;

    virtual Weights getWeightsForGate(int layerIndex, RNNGateType gate, bool isW) const = 0;

    virtual void setBiasForGate(int layerIndex, RNNGateType gate, bool isW, Weights bias) = 0;

    virtual Weights getBiasForGate(int layerIndex, RNNGateType gate, bool isW) const = 0;

    virtual void setHiddenState(ITensor& hidden) = 0;

    virtual ITensor* getHiddenState() const = 0;

    virtual void setCellState(ITensor& cell) = 0;

    virtual ITensor* getCellState() const = 0;
protected:
    virtual ~IRNNv2Layer() {}
};

class IOutputDimensionsFormula
{
public:
    virtual DimsHW compute(DimsHW inputDims, DimsHW kernelSize, DimsHW stride, DimsHW padding, DimsHW dilation, const char* layerName) = 0;

protected:
    virtual ~IOutputDimensionsFormula() {}
};

enum class PluginFormat: uint8_t
{
    kNCHW        = 0, 
    kNC2HW2      = 1, 
    kNHWC8       = 2  
};

template <>
inline int EnumMax<PluginFormat>()
{
    return 3;
} 

class IPlugin
{
public:
    virtual int getNbOutputs() const = 0;

    virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) = 0;

    virtual void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) = 0;

    virtual int initialize() = 0;

    virtual void terminate() = 0;

    virtual size_t getWorkspaceSize(int maxBatchSize) const = 0;

    virtual int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) = 0;

    virtual size_t getSerializationSize() = 0;

    virtual void serialize(void* buffer) = 0;
protected:
    virtual ~IPlugin() {}
};

class IPluginExt : public IPlugin
{
public:
    virtual int getTensorRTVersion() const
    {
        return NV_TENSORRT_VERSION;
    }

    virtual bool supportsFormat(DataType type, PluginFormat format) const = 0;

    virtual void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) = 0;

protected:
    void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) _TENSORRT_FINAL {}
    virtual ~IPluginExt() {}
};

class IPluginLayer : public ILayer
{
public:
    virtual IPlugin& getPlugin() = 0;

protected:
    virtual ~IPluginLayer() {}
};

enum class UnaryOperation : int
{
    kEXP = 0,   
    kLOG = 1,   
    kSQRT = 2,  
    kRECIP = 3, 
    kABS = 4,   
    kNEG = 5,   
};

template <>
inline int EnumMax<UnaryOperation>()
{
    return 6;
} 

class IUnaryLayer : public ILayer
{
public:
    virtual void setOperation(UnaryOperation op) = 0;

    virtual UnaryOperation getOperation() const = 0;

protected:
    virtual ~IUnaryLayer() {}
};

enum class ReduceOperation : int
{
    kSUM = 0,
    kPROD = 1,
    kMAX = 2,
    kMIN = 3,
    kAVG = 4
};

template <>
inline int EnumMax<ReduceOperation>()
{
    return 5;
} 

class IReduceLayer : public ILayer
{
public:
    virtual void setOperation(ReduceOperation op) = 0;

    virtual ReduceOperation getOperation() const = 0;

    virtual void setReduceAxes(uint32_t reduceAxes) = 0;

    virtual uint32_t getReduceAxes() const = 0;

    virtual void setKeepDimensions(bool keepDimensions) = 0;

    virtual bool getKeepDimensions() const = 0;

protected:
    virtual ~IReduceLayer() {}
};

class IPaddingLayer : public ILayer
{
public:
    virtual void setPrePadding(DimsHW padding) = 0;

    virtual DimsHW getPrePadding() const = 0;

    virtual void setPostPadding(DimsHW padding) = 0;

    virtual DimsHW getPostPadding() const = 0;

protected:
    virtual ~IPaddingLayer() {}
};

struct Permutation
{
    int order[Dims::MAX_DIMS];
};

class IShuffleLayer : public ILayer
{
public:
    virtual void setFirstTranspose(Permutation permutation) = 0;

    virtual Permutation getFirstTranspose() const = 0;

    virtual void setReshapeDimensions(Dims dimensions) = 0;

    virtual Dims getReshapeDimensions() const = 0;

    virtual void setSecondTranspose(Permutation permutation) = 0;

    virtual Permutation getSecondTranspose() const = 0;

protected:
    virtual ~IShuffleLayer() {}
};

enum class TopKOperation : int
{
    kMAX = 0,  
    kMIN = 1,  
};

template <>
inline int EnumMax<TopKOperation>()
{
    return 2;
} 

class ITopKLayer : public ILayer
{
public:
    virtual void setOperation(TopKOperation op) = 0;

    virtual TopKOperation getOperation() const = 0;

    virtual void setK(int k) = 0;

    virtual int getK() const = 0;

    virtual void setReduceAxes(uint32_t reduceAxes) = 0;

    virtual uint32_t getReduceAxes() const = 0;

protected:
    virtual ~ITopKLayer() {}
};

class IMatrixMultiplyLayer : public ILayer
{
public:
    virtual void setTranspose(int index, bool val) = 0;

    virtual bool getTranspose(int index) const = 0;

protected:
    virtual ~IMatrixMultiplyLayer() {}
};

class IRaggedSoftMaxLayer : public ILayer
{
protected:
    virtual ~IRaggedSoftMaxLayer() {}
};

class IConstantLayer : public ILayer
{
public:
    virtual void setWeights(Weights weights) = 0;

    virtual Weights getWeights() const = 0;

    virtual void setDimensions(Dims dimensions) = 0;

    virtual Dims getDimensions() const = 0;

protected:
    virtual ~IConstantLayer() {}
};

class INetworkDefinition
{
public:
    virtual ITensor* addInput(const char* name, DataType type, Dims dimensions) = 0;

    virtual void markOutput(ITensor& tensor) = 0;

    virtual IConvolutionLayer* addConvolution(ITensor& input, int nbOutputMaps, DimsHW kernelSize, Weights kernelWeights, Weights biasWeights) = 0;

    virtual IFullyConnectedLayer* addFullyConnected(ITensor& input, int nbOutputs, Weights kernelWeights, Weights biasWeights) = 0;

    virtual IActivationLayer* addActivation(ITensor& input, ActivationType type) = 0;

    virtual IPoolingLayer* addPooling(ITensor& input, PoolingType type, DimsHW windowSize) = 0;

    virtual ILRNLayer* addLRN(ITensor& input, int window, float alpha, float beta, float k) = 0;

    virtual IScaleLayer* addScale(ITensor& input, ScaleMode mode, Weights shift, Weights scale, Weights power) = 0;

    virtual ISoftMaxLayer* addSoftMax(ITensor& input) = 0;

    virtual IConcatenationLayer* addConcatenation(ITensor* const* inputs, int nbInputs) = 0;

    virtual IDeconvolutionLayer* addDeconvolution(ITensor& input, int nbOutputMaps, DimsHW kernelSize, Weights kernelWeights, Weights biasWeights) = 0;

    virtual IElementWiseLayer* addElementWise(ITensor& input1, ITensor& input2, ElementWiseOperation op) = 0;

    virtual IRNNLayer* addRNN(ITensor& inputs, int layerCount, std::size_t hiddenSize, int maxSeqLen, RNNOperation op, RNNInputMode mode, RNNDirection dir, Weights weights, Weights bias) = 0;

    virtual IPluginLayer* addPlugin(ITensor* const* inputs, int nbInputs, IPlugin& plugin) = 0;

    virtual IUnaryLayer* addUnary(ITensor& input, UnaryOperation operation) = 0;

    virtual IPaddingLayer* addPadding(ITensor& input, DimsHW prePadding, DimsHW postPadding) = 0;

    virtual IShuffleLayer* addShuffle(ITensor& input) = 0;

    virtual void setPoolingOutputDimensionsFormula(IOutputDimensionsFormula* formula) = 0;

    virtual IOutputDimensionsFormula& getPoolingOutputDimensionsFormula() const = 0;

    virtual void setConvolutionOutputDimensionsFormula(IOutputDimensionsFormula* formula) = 0;

    virtual IOutputDimensionsFormula& getConvolutionOutputDimensionsFormula() const = 0;

    virtual void setDeconvolutionOutputDimensionsFormula(IOutputDimensionsFormula* formula) = 0;

    virtual IOutputDimensionsFormula& getDeconvolutionOutputDimensionsFormula() const = 0;

    virtual int getNbLayers() const = 0;

    virtual ILayer* getLayer(int index) const = 0;

    virtual int getNbInputs() const = 0;

    virtual ITensor* getInput(int index) const = 0; // adding inputs invalidates indexing here

    virtual int getNbOutputs() const = 0;

    virtual ITensor* getOutput(int index) const = 0; // adding outputs invalidates indexing here

    virtual void destroy() = 0;

protected:
    virtual ~INetworkDefinition() {}

public:
    virtual IReduceLayer* addReduce(ITensor& input, ReduceOperation operation, uint32_t reduceAxes, bool keepDimensions) = 0;

    virtual ITopKLayer* addTopK(ITensor& input, TopKOperation op, int k, uint32_t reduceAxes) = 0;

    virtual IGatherLayer* addGather(ITensor& data, ITensor& indices, int axis) = 0;

    virtual IRaggedSoftMaxLayer* addRaggedSoftMax(ITensor &input, ITensor &bounds) = 0;

    virtual IMatrixMultiplyLayer* addMatrixMultiply(ITensor& input0, bool transpose0, ITensor& input1, bool transpose1) = 0;

    virtual IConstantLayer* addConstant(Dims dimensions, Weights weights) = 0;

    virtual IRNNv2Layer *addRNNv2(ITensor &input, int32_t layerCount, int32_t hiddenSize, int32_t maxSeqLen, RNNOperation op) = 0;

    virtual IPluginLayer* addPluginExt(ITensor* const* inputs, int nbInputs, IPluginExt& plugin) = 0;
};

class IProfiler
{
public:
    virtual void reportLayerTime(const char* layerName, float ms) = 0;

protected:
    virtual ~IProfiler() {}
};

class ICudaEngine;

class IExecutionContext
{
public:
    virtual bool execute(int batchSize, void** bindings) = 0;

    virtual bool enqueue(int batchSize, void** bindings, cudaStream_t stream, cudaEvent_t* inputConsumed) = 0;

    virtual void setDebugSync(bool sync) = 0;

    virtual bool getDebugSync() const = 0;

    virtual void setProfiler(IProfiler*) = 0;

    virtual IProfiler* getProfiler() const = 0;

    virtual const ICudaEngine& getEngine() const = 0;

    virtual void destroy() = 0;

protected:
    virtual ~IExecutionContext() {}

public:
    virtual void setName(const char* name) = 0;

    virtual const char* getName() const = 0;

    virtual void setDeviceMemory(void* memory) = 0;
};

class ICudaEngine
{
public:
    virtual int getNbBindings() const = 0;

    virtual int getBindingIndex(const char* name) const = 0;

    virtual const char* getBindingName(int bindingIndex) const = 0;

    virtual bool bindingIsInput(int bindingIndex) const = 0;

    virtual Dims getBindingDimensions(int bindingIndex) const = 0;

    virtual DataType getBindingDataType(int bindingIndex) const = 0;

    virtual int getMaxBatchSize() const = 0;

    virtual int getNbLayers() const = 0;

    virtual std::size_t getWorkspaceSize() const = 0;

    virtual IHostMemory* serialize() const = 0;

    virtual IExecutionContext* createExecutionContext() = 0;

    virtual void destroy() = 0;

    virtual TensorLocation getLocation(int bindingIndex) const = 0;

protected:
    virtual ~ICudaEngine() {}

public:
    virtual IExecutionContext* createExecutionContextWithoutDeviceMemory() = 0;

    virtual size_t getDeviceMemorySize() const = 0;
};

enum class CalibrationAlgoType : int
{
    kLEGACY_CALIBRATION = 0,
    kENTROPY_CALIBRATION = 1
};

template <>
inline int EnumMax<CalibrationAlgoType>()
{
    return 2;
} 

class IInt8Calibrator
{
public:
    virtual int getBatchSize() const = 0;

    virtual bool getBatch(void* bindings[], const char* names[], int nbBindings) = 0; // get a pointer to the input batch

    virtual const void* readCalibrationCache(std::size_t& length) = 0;

    virtual void writeCalibrationCache(const void* ptr, std::size_t length) = 0;

    virtual CalibrationAlgoType getAlgorithm() = 0;

protected:
    virtual ~IInt8Calibrator() {}
};

class IInt8EntropyCalibrator : public IInt8Calibrator
{
    virtual CalibrationAlgoType getAlgorithm() { return CalibrationAlgoType::kENTROPY_CALIBRATION; }
protected:
    virtual ~IInt8EntropyCalibrator() {}
};

class IInt8LegacyCalibrator : public IInt8Calibrator
{
public:
    virtual CalibrationAlgoType getAlgorithm() { return CalibrationAlgoType::kLEGACY_CALIBRATION; }

    virtual double getQuantile() const = 0;

    virtual double getRegressionCutoff() const = 0;

    virtual const void* readHistogramCache(std::size_t& length) = 0;

    virtual void writeHistogramCache(const void* ptr, std::size_t length) = 0;

protected:
    virtual ~IInt8LegacyCalibrator() {}
};

class IGpuAllocator
{
public:
    virtual void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) = 0;

    virtual void free(void* memory) = 0;
};

class IBuilder
{
public:
    virtual nvinfer1::INetworkDefinition* createNetwork() = 0;

    virtual void setMaxBatchSize(int batchSize) = 0;

    virtual int getMaxBatchSize() const = 0;

    virtual void setMaxWorkspaceSize(std::size_t workspaceSize) = 0;

    virtual std::size_t getMaxWorkspaceSize() const = 0;

    virtual void setHalf2Mode(bool mode) = 0;

    virtual bool getHalf2Mode() const = 0;

    virtual void setDebugSync(bool sync) = 0;

    virtual bool getDebugSync() const = 0;

    virtual void setMinFindIterations(int minFind) = 0;

    virtual int getMinFindIterations() const = 0;

    virtual void setAverageFindIterations(int avgFind) = 0;

    virtual int getAverageFindIterations() const = 0;

    virtual nvinfer1::ICudaEngine* buildCudaEngine(nvinfer1::INetworkDefinition& network) = 0;

    virtual bool platformHasFastFp16() const = 0;

    virtual bool platformHasFastInt8() const = 0;

    virtual void destroy() = 0;

    virtual void setInt8Mode(bool mode) = 0;

    virtual bool getInt8Mode() const = 0;

    virtual void setInt8Calibrator(IInt8Calibrator* calibrator) = 0;

protected:
    virtual ~IBuilder() {}

public:
    virtual void setGpuAllocator(IGpuAllocator* allocator) = 0;

    virtual void setFp16Mode(bool mode) = 0;

    virtual bool getFp16Mode() const = 0;
};

class IPluginFactory
{
public:
    virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) = 0;
};

class IRuntime
{
public:
    virtual nvinfer1::ICudaEngine* deserializeCudaEngine(const void* blob, std::size_t size, IPluginFactory* pluginFactory) = 0;

    virtual void destroy() = 0;

protected:
    virtual ~IRuntime() {}

public:
    virtual void setGpuAllocator(IGpuAllocator* allocator) = 0;
};

class ILogger
{
public:
    enum class Severity
    {
        kINTERNAL_ERROR = 0, 
        kERROR = 1,          
        kWARNING = 2,        
        kINFO = 3            
    };

    virtual void log(Severity severity, const char* msg) = 0;

protected:
    virtual ~ILogger() {}
};

template <>
inline int EnumMax<ILogger::Severity>()
{
    return 4;
} 

} // namespace nvinfer1

extern "C" TENSORRTAPI void* createInferBuilder_INTERNAL(void* logger, int version); 
extern "C" TENSORRTAPI void* createInferRuntime_INTERNAL(void* logger, int version); 

extern "C" TENSORRTAPI nvinfer1::ILogger* getLogger();

extern "C" TENSORRTAPI int getInferLibVersion();

namespace nvinfer1
{
namespace // unnamed namespace in case the compiler doesn't inline these
{
inline IBuilder* createInferBuilder(ILogger& logger)
{
    return static_cast<IBuilder*>(createInferBuilder_INTERNAL(&logger, NV_TENSORRT_VERSION));
}

inline IRuntime* createInferRuntime(ILogger& logger)
{
    return static_cast<IRuntime*>(createInferRuntime_INTERNAL(&logger, NV_TENSORRT_VERSION));
}
}
}

#endif