Merge pull request #3 from perseusdg/tensorrt8

push tensorrt8 commits to the rds branch

include/tkDNN/Layer.h

@@ -31,7 +31,8 @@ enum layerType_t {
     LAYER_SHORTCUT,
     LAYER_UPSAMPLE,
     LAYER_REGION,
-    LAYER_YOLO
+    LAYER_YOLO,
+    LAYER_PADDING
 };
 
 #define TKDNN_BN_MIN_EPSILON 1e-5
@@ -87,6 +88,7 @@ public:
             case LAYER_UPSAMPLE:            return "Upsample";
             case LAYER_REGION:              return "Region";
             case LAYER_YOLO:                return "Yolo";
+            case LAYER_PADDING:             return "Padding";
             default:                        return "unknown";
         }
     }
@@ -520,9 +522,33 @@ protected:
     bool poolOn3d;
 };
 
+/**
+ * Padding Layers
+ * tkDNN supports reflection,constant and zero padding
+ */
+
+typedef enum {
+    PADDING_MODE_CONSTANT = 0,
+    PADDING_MODE_ZERO = 1,
+    PADDING_MODE_REFLECTION = 2
+} tkdnnPaddingMode_t;
+
+class Padding : public Layer {
+public:
+    Padding(Network *net,int32_t pad_h,int32_t pad_w,tkdnnPaddingMode_t padding_mode,float constant = 0.0);
+    virtual ~Padding();
+    virtual layerType_t getLayerType(){return LAYER_PADDING ;};
+    virtual dnnType* infer(dataDim_t& dim,dnnType* srcData);
+    int32_t paddingH,paddingW;
+    tkdnnPaddingMode_t padding_mode;
+    float constant;
+
+};
+
 /**
     Softmax layer
 */
+
 class Softmax : public Layer {
 
 public:

include/tkDNN/NetworkRT.h

@@ -23,6 +23,8 @@
 #include <pluginsRT/ShortcutRT.h>
 #include <pluginsRT/UpsampleRT.h>
 #include <pluginsRT/YoloRT.h>
+#include <pluginsRT/ConstantPaddingRT.h>
+#include <pluginsRT/ReflectionPadding.h>
 
 
 
@@ -93,10 +95,16 @@ public:
     nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Region *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l);
     nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Yolo *l);
-    nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Upsample *l);
+    nvinfer1::IResizeLayer* convert_layer(nvinfer1::ITensor *input, Upsample *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeformConv2d *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input,Padding *l);
 
+#if NV_TENSORRT_MAJOR > 5 && NV_TENSORRT_MAJOR < 8
     bool serialize(const char *filename);
+#else
+    bool serialize(const char *filename,nvinfer1::IHostMemory *ptr);
+#endif
+
     bool deserialize(const char *filename);
     void destroy();
 

include/tkDNN/kernels.h

@@ -48,4 +48,11 @@ void dcnV2CudaForward(cublasStatus_t stat, cublasHandle_t handle,
                          const int dst_dim, cudaStream_t stream = cudaStream_t(0));
 
 void scalAdd(dnnType* dstData, int size, float alpha, float beta, int inc, cudaStream_t stream = cudaStream_t(0));
+
+void reflection_pad2d_out_forward(int32_t pad_h,int32_t pad_w,float *srcData,float *dstData,int32_t input_h,int32_t input_w,int32_t plane_dim,int32_t n_batch,cudaStream_t cudaStream = cudaStream_t(0));
+
+void constant_pad2d_forward(dnnType *srcData,dnnType *dstData,int32_t input_h,int32_t input_w,int32_t output_h,
+                            int32_t output_w,int32_t c,int32_t n,int32_t padT,int32_t padL,dnnType constant,cudaStream_t cudaStream = cudaStream_t(0));
+
+
 #endif //KERNELS_H

include/tkDNN/pluginsRT/ConstantPaddingRT.h

@@ -0,0 +1,109 @@
+//
+// Created by perseusdg on 1/7/22.
+//
+
+#ifndef _CONSTANTPADDINGRT_PLUGIN_H
+#define _CONSTANTPADDINGRT_PLUGIN_H
+
+#include<cassert>
+#include <NvInfer.h>
+#include <vector>
+#include <utils.h>
+#include <kernels.h>
+
+namespace nvinfer1{
+    class ConstantPaddingRT : public IPluginV2Ext {
+    public:
+        ConstantPaddingRT(int32_t padH,int32_t padW,int32_t n,int32_t c,int32_t i_h,int32_t i_w,int32_t o_h,int32_t o_w,float constant);
+
+        ConstantPaddingRT(const void *data,size_t length);
+
+        ~ConstantPaddingRT();
+
+        int getNbOutputs() const NOEXCEPT override;
+
+        Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
+
+        int initialize() NOEXCEPT override ;
+
+        void terminate() NOEXCEPT override ;
+
+        size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
+
+
+#if NV_TENSORRT_MAJOR > 7
+        int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
+#elif NV_TENSORRT_MAJOR <= 7
+        int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
+#endif
+
+        size_t getSerializationSize() const NOEXCEPT override ;
+
+        void serialize(void *buffer) const NOEXCEPT override ;
+
+        void destroy() NOEXCEPT override ;
+
+        const char *getPluginType() const NOEXCEPT override ;
+
+        const char *getPluginVersion() const NOEXCEPT override;
+
+        const char *getPluginNamespace() const NOEXCEPT override ;
+
+        void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
+
+        IPluginV2Ext *clone() const NOEXCEPT override ;
+
+        DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
+
+        void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
+
+        bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
+
+        bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
+
+        void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
+                              int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
+                              bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
+                              int32_t maxBatchSize) NOEXCEPT override;
+
+        void detachFromContext() NOEXCEPT override;
+
+        bool supportsFormat (DataType type, PluginFormat format) const NOEXCEPT override;
+
+        int32_t i_h,i_w,o_h,o_w,n,c,padH,padW;
+        float constant;
+    private:
+        std::string mPluginNamespace;
+
+    };
+
+    class ConstantPaddingRTPluginCreator : public IPluginCreator {
+    public:
+        ConstantPaddingRTPluginCreator();
+
+        void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override;
+
+        const char *getPluginNamespace() const NOEXCEPT override;
+
+        IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
+
+        IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
+
+        const char *getPluginName() const NOEXCEPT override ;
+
+        const char *getPluginVersion() const NOEXCEPT override;
+
+        const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
+
+    private:
+        static PluginFieldCollection mFC;
+        static std::vector<PluginField> mPluginAttributes;
+        std::string mPluginNamespace;
+
+    };
+
+    REGISTER_TENSORRT_PLUGIN(ConstantPaddingRTPluginCreator);
+};
+
+
+#endif //TKDNN_CONSTANTPADDINGRT_H

include/tkDNN/pluginsRT/FlattenConcatRT.h

@@ -1,3 +1,6 @@
+#ifndef _FLATTENCONCATRT_PLUGIN_H
+#define _FLATTENCONCATRT_PLUGIN_H
+
 #include<cassert>
 #include <NvInfer.h>
 #include <vector>
@@ -93,4 +96,5 @@ namespace nvinfer1 {
     };
 
     REGISTER_TENSORRT_PLUGIN(FlattenConcatRTPluginCreator);
-};
+};
+#endif

include/tkDNN/pluginsRT/ReflectionPadding.h

@@ -0,0 +1,101 @@
+#ifndef _REFLECTIONPADDINGRT_PLUGIN_H
+#define _REFLECTIONPADDINGRT_PLUGIN_H
+
+#include<cassert>
+#include <NvInfer.h>
+#include <vector>
+#include <utils.h>
+#include <kernels.h>
+
+namespace nvinfer1{
+    class ReflectionPaddingRT : public IPluginV2Ext {
+    public:
+        ReflectionPaddingRT(int32_t padH,int32_t padW,int32_t input_h,int32_t input_w,int32_t output_h,int32_t output_w,int32_t c,int32_t n);
+
+        ReflectionPaddingRT(const void *data,size_t length);
+
+        ~ReflectionPaddingRT();
+
+        int getNbOutputs() const NOEXCEPT override;
+
+        Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
+
+        int initialize() NOEXCEPT override ;
+
+        void terminate() NOEXCEPT override ;
+
+        size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
+
+#if NV_TENSORRT_MAJOR > 7
+        int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
+#elif NV_TENSORRT_MAJOR <= 7
+        int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
+#endif
+
+        size_t getSerializationSize() const NOEXCEPT override ;
+
+        void serialize(void *buffer) const NOEXCEPT override ;
+
+        void destroy() NOEXCEPT override ;
+
+        const char *getPluginType() const NOEXCEPT override ;
+
+        const char *getPluginVersion() const NOEXCEPT override;
+
+        const char *getPluginNamespace() const NOEXCEPT override ;
+
+        void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
+
+        IPluginV2Ext *clone() const NOEXCEPT override ;
+
+        DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
+
+        void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
+
+        bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
+
+        bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
+
+        void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
+                              int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
+                              bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
+                              int32_t maxBatchSize) NOEXCEPT override;
+
+        void detachFromContext() NOEXCEPT override;
+
+        bool supportsFormat (DataType type, PluginFormat format) const NOEXCEPT override;
+
+        int32_t padH,padW,input_h,input_w,output_h,output_w,n,c;
+    private:
+        std::string mPluginNamespace;
+
+    };
+
+    class ReflectionPaddingRTPluginCreator : public IPluginCreator {
+    public:
+        ReflectionPaddingRTPluginCreator();
+
+        void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
+
+        const char *getPluginNamespace() const NOEXCEPT override ;
+
+        IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
+
+        IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
+
+        const char *getPluginName() const NOEXCEPT override ;
+
+        const char *getPluginVersion() const NOEXCEPT override;
+
+        const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
+
+    private:
+        static PluginFieldCollection mFC;
+        static std::vector<PluginField> mPluginAttributes;
+        std::string mPluginNamespace;
+    };
+
+    REGISTER_TENSORRT_PLUGIN(ReflectionPaddingRTPluginCreator);
+};
+#endif
+

include/tkDNN/utils.h

@@ -6,6 +6,8 @@
 #include <fstream>
 #include <iomanip>
 #include <stdlib.h>
+#include <yaml-cpp/yaml.h>
+
 
 #include "cuda.h"
 #include "cuda_runtime_api.h"
@@ -16,7 +18,6 @@
 
 #ifdef __linux__
 #include <unistd.h>
-
 #endif
 
 #include <ios>
@@ -161,5 +162,19 @@ static inline bool isCudaPointer(void *data) {
   return cudaPointerGetAttributes(&attr, data) == 0;
 }
 
+inline YAML::Node YAMLloadConf(const std::string& conf_file) {
+    std::cerr<<"Loading YAML: "<<conf_file<<"\n";
+    return YAML::LoadFile(conf_file);
+}
+
+template<typename T>
+inline T YAMLgetConf(YAML::Node conf, std::string key, T defaultVal) {
+    T val = defaultVal;
+    if(conf && conf[key]) {
+        val = conf[key].as<T>();
+    }
+    return val;
+}
+
 
 #endif //UTILS_H