New features and fixes

nvdsinfer_custom_impl_Yolo/yoloPlugins.h

@@ -38,57 +38,68 @@
   }                                                                                                                        \
 }
 
+#if NV_TENSORRT_MAJOR >= 8
+  #define INT int32_t
+#else
+  #define INT int
+#endif
+
 namespace {
   const char* YOLOLAYER_PLUGIN_VERSION {"1"};
   const char* YOLOLAYER_PLUGIN_NAME {"YoloLayer_TRT"};
 } // namespace
 
-class YoloLayer : public nvinfer1::IPluginV2 {
+class YoloLayer : public nvinfer1::IPluginV2DynamicExt {
   public:
     YoloLayer(const void* data, size_t length);
 
     YoloLayer(const uint& netWidth, const uint& netHeight, const uint& numClasses, const uint& newCoords,
         const std::vector<TensorInfo>& yoloTensors, const uint64_t& outputSize);
 
-    const char* getPluginType() const noexcept override { return YOLOLAYER_PLUGIN_NAME; }
-
-    const char* getPluginVersion() const noexcept override { return YOLOLAYER_PLUGIN_VERSION; }
-
-    int getNbOutputs() const noexcept override { return 1; }
-
-    nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) noexcept override;
-
-    bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format) const noexcept override;
-
-    void configureWithFormat(const nvinfer1::Dims* inputDims, int nbInputs, const nvinfer1::Dims* outputDims, int nbOutputs,
-        nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) noexcept override;
+    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
 
     int initialize() noexcept override { return 0; }
 
     void terminate() noexcept override {}
 
-    size_t getWorkspaceSize(int maxBatchSize) const noexcept override {
-      return maxBatchSize * sizeof(int);
-    }
-
-#ifdef LEGACY
-    int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override;
-#else
-    int32_t enqueue(int batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream)
-        noexcept override;
-#endif
+    void destroy() noexcept override { delete this; }
 
     size_t getSerializationSize() const noexcept override;
 
     void serialize(void* buffer) const noexcept override;
 
-    void destroy() noexcept override { delete this; }
+    int getNbOutputs() const noexcept override { return 3; }
 
-    nvinfer1::IPluginV2* clone() const noexcept override;
+    nvinfer1::DimsExprs getOutputDimensions(INT index, const nvinfer1::DimsExprs* inputs, INT nbInputDims,
+        nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+
+    size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, INT nbInputs,
+        const nvinfer1::PluginTensorDesc* outputs, INT nbOutputs) const noexcept override { return 0; }
+
+    bool supportsFormatCombination(INT pos, const nvinfer1::PluginTensorDesc* inOut, INT nbInputs, INT nbOutputs) noexcept
+        override;
+
+    const char* getPluginType() const noexcept override { return YOLOLAYER_PLUGIN_NAME; }
+
+    const char* getPluginVersion() const noexcept override { return YOLOLAYER_PLUGIN_VERSION; }
 
     void setPluginNamespace(const char* pluginNamespace) noexcept override { m_Namespace = pluginNamespace; }
 
-    virtual const char* getPluginNamespace() const noexcept override { return m_Namespace.c_str(); }
+    const char* getPluginNamespace() const noexcept override { return m_Namespace.c_str(); }
+
+    nvinfer1::DataType getOutputDataType(INT index, const nvinfer1::DataType* inputTypes, INT nbInputs) const noexcept
+        override;
+    
+    void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, nvinfer1::IGpuAllocator* gpuAllocator)
+        noexcept override {}
+
+    void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, INT nbInput,
+        const nvinfer1::DynamicPluginTensorDesc* out, INT nbOutput) noexcept override;
+
+    void detachFromContext() noexcept override {}
+    
+    INT enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc*  outputDesc,
+        void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
 
   private:
     std::string m_Namespace {""};
@@ -115,12 +126,14 @@ class YoloLayerPluginCreator : public nvinfer1::IPluginCreator {
       return nullptr;
     }
 
-    nvinfer1::IPluginV2* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept override {
+    nvinfer1::IPluginV2DynamicExt* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept
+        override {
       std::cerr<< "YoloLayerPluginCreator::getFieldNames is not implemented";
       return nullptr;
     }
 
-    nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override {
+    nvinfer1::IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
+        override {
       std::cout << "Deserialize yoloLayer plugin: " << name << std::endl;
       return new YoloLayer(serialData, serialLength);
     }