include/tkDNN/Layer.h

@@ -535,12 +535,13 @@ typedef enum {
 
 class Padding : public Layer {
 public:
-    Padding(Network *net,int32_t pad_h,int32_t pad_w,tkdnnPaddingMode_t padding_mode);
+    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;
 
 };
 

include/tkDNN/kernels.h

@@ -51,4 +51,8 @@ void scalAdd(dnnType* dstData, int size, float alpha, float beta, int inc, cudaS
 
 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/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,99 @@
+#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;
+    };
+};
+#endif
+

src/NetworkRT.cpp

@@ -456,12 +456,42 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input,Padding *l){
+
+    float rt_ver =  float(NV_TENSORRT_MAJOR) +
+                    float(NV_TENSORRT_MINOR)/10 +
+                    float(NV_TENSORRT_PATCH)/100;
+
+#if ((NV_TENSORRT_MAJOR == 8 && NV_TENSORRT_MINOR >= 2) || NV_TENSORRT_MAJOR > 8)
     auto *lRT = networkRT->addSlice(*input,Dims3{0,0,0},Dims3{l->output_dim.c,l->output_dim.h,l->output_dim.w},Dims3{0,0,0});
     if(l->padding_mode == PADDING_MODE_REFLECTION){
         lRT->setMode(SliceMode::kREFLECT);
+    }else if(l->padding_mode == PADDING_MODE_CONSTANT || l->padding_mode == PADDING_MODE_ZERO){
+        lRT->setMode(SliceMode::kFILL);
+        lRT->setInput(4, reinterpret_cast<ITensor &>(l->constant));
     }
     checkNULL(lRT);
     return lRT;
+#else
+    //todo  add PADDING_MODE_CONSTANT AND PADDING_MODE_ZERO for tensorrt versions < 8.2
+    if(l->padding_mode == PADDING_MODE_REFLECTION){
+    auto creator = getPluginRegistry()->getPluginCreator("ReflectionPaddingRT_tkDNN","1");
+    std::vector<PluginField> mPluginAttributes;
+    PluginFieldCollection mFC{};
+    mPluginAttributes.emplace_back(PluginField("padH",&l->paddingH,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("padW",&l->paddingW,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("inputH",&l->input_dim.h,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("inputW",&l->input_dim.w,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("outputH",&l->output_dim.h,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("outputW",&l->output_dim.w,PluginFieldType::kINT32,1));
+    mPluginAttributes.emplace_back(PluginField("n",&l->input_dim.n,PluginFieldType::kINT32,1));
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+    auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC);
+    }
+    auto *lRT = networkRT->addPluginV2(&input, 1, *plugin);
+    checkNULL(lRT);
+    return lRT;
+#endif
 }
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {

src/Padding.cpp

@@ -7,7 +7,7 @@
 #include "kernels.h"
 
 namespace tk{ namespace dnn {
-    Padding::Padding(Network *net, int32_t pad_h, int32_t pad_w, tkdnnPaddingMode_t padding_mode) : Layer(net) {
+    Padding::Padding(Network *net, int32_t pad_h, int32_t pad_w, tkdnnPaddingMode_t padding_mode,float constant) : Layer(net) {
         this->paddingH = pad_h;
         this->paddingW = pad_w;
         this->padding_mode = padding_mode;
@@ -15,6 +15,11 @@ namespace tk{ namespace dnn {
         output_dim.n = input_dim.n;
         output_dim.h = input_dim.h + 2 * (this->paddingH);
         output_dim.w = input_dim.w + 2 * (this->paddingW);
+        if(padding_mode == tkdnnPaddingMode_t::PADDING_MODE_CONSTANT){
+            this->constant = constant;
+        }else{
+            this->constant = 0;
+        }
         checkCuda(cudaMalloc(&dstData,output_dim.tot()*sizeof(dnnType)));
     }
 
@@ -28,6 +33,11 @@ namespace tk{ namespace dnn {
             reflection_pad2d_out_forward(paddingH, paddingW, srcData, dstData, input_dim.h, input_dim.w, input_dim.c,
                                          input_dim.n);
         }
+        else if(padding_mode == tkdnnPaddingMode_t::PADDING_MODE_CONSTANT){
+            constant_pad2d_forward(srcData,dstData,input_dim.h,input_dim.w,output_dim.h,output_dim.w,input_dim.c,
+                                   input_dim.n,paddingH,paddingW,constant);
+        }
+
         dim = output_dim;
         return dstData;
     }

src/kernels/padding.cu

@@ -2,7 +2,9 @@
 #include <thrust/pair.h>
 #include <stdio.h>
 
-
+/*
+ * Reflection padding is from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/ReflectionPad.cu
+ */
 __device__
 inline thrust::pair<int32_t,int32_t> get_index_mapping2d(
         int32_t input_dim_x,int32_t input_dim_y,int32_t output_dim_x,
@@ -66,3 +68,43 @@ void reflection_pad2d_out_forward(int32_t pad_h,int32_t pad_w,float *srcData,flo
     }
 
 }
+
+/*
+ * constant padding is inspired from https://github.com/apache/incubator-mxnet/blob/master/src/operator/pad.cu
+ */
+
+__global__
+void constant_pad2d_kernel(dnnType *srcData,dnnType *dstData,const int32_t padT,const int32_t padL,float constant,int32_t n,int32_t c,int32_t i_h,int32_t i_w,int32_t o_h,int32_t o_w){
+    int outputPointId = threadIdx.x + blockIdx.x * blockDim.x;
+    if(outputPointId >= o_h*o_w){
+        return ;
+    }
+
+    int Ny = i_h;
+    int Nx = i_w;
+
+    int plane =   blockIdx.y;
+    int batch = blockIdx.z;
+    int outputPointX = outputPointId % o_w;
+    int outputPointY = outputPointId / o_w;
+    int checkT       = max(0, outputPointY - padT + 1);
+    int checkB       = max(0, padT + Ny - outputPointY);
+    int checkL       = max(0, outputPointX - padL + 1);
+    int checkR       = max(0, padL + Nx - outputPointX);
+    int inputPointX  = min(max(outputPointX - padL, 0), Nx - 1);
+    int inputPointY  = min(max(outputPointY - padT, 0), Ny - 1);
+    int need_pad     = !(checkT * checkB * checkL * checkR);
+    float value_to_copy = srcData[batch*c*i_h*i_w + plane*i_h*i_w + inputPointY*i_w + inputPointX];
+    dstData[batch*c*o_w*o_h + plane*o_h*o_w + outputPointY*o_w + outputPointX] = value_to_copy * (!need_pad) + need_pad*constant;
+
+}
+
+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){
+    int32_t output_plane_size = output_h*output_w;
+    dim3 block_size(output_plane_size>256 ?256:output_plane_size);
+    dim3 grid_size(ceilDiv(output_plane_size,static_cast<int32_t>(256)),c,n);
+    constant_pad2d_kernel<<<grid_size,block_size,0,cudaStream>>>(srcData,dstData,padT,padL,constant,n,c,input_h,input_w,output_h,output_w);
+
+}
+

src/pluginsRT/ReflectionPadding.cpp

@@ -0,0 +1,202 @@
+#include <tkDNN/pluginsRT/ReflectionPadding.h>
+using namespace nvinfer1;
+
+std::vector<PluginField> ReflectionPaddingRTPluginCreator::mPluginAttributes;
+PluginFieldCollection ReflectionPaddingRTPluginCreator::mFC{};
+
+static const char* REFLECTIONPADDINGRT_PLUGIN_VERSION{"1"};
+static const char* REFLECTIONPADDINGRT_PLUGIN_NAME{"ReflectionPaddingRT_tkDNN"};
+
+ReflectionPaddingRT::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) {
+    this->padH = padH;
+    this->padW = padW;
+    this->input_h  = input_h;
+    this->input_w = input_w;
+    this->output_h = output_h;
+    this->output_w = output_w;
+    this->n = n;
+    this->c = c;
+}
+
+ReflectionPaddingRT::ReflectionPaddingRT(const void *data, size_t length) {
+    const char* buf = reinterpret_cast<const char*>(data),*bufcheck=buf;
+    padH = readBUF<int32_t>(buf);
+    padW = readBUF<int32_t>(buf);
+    input_h = readBUF<int32_t>(buf);
+    input_w = readBUF<int32_t>(buf);
+    output_h = readBUF<int32_t>(buf);
+    output_w = readBUF<int32_t>(buf);
+    n = readBUF<int32_t>(buf);
+    c = readBUF<int32_t>(buf);
+    assert(buf = bufcheck + length);
+}
+
+ReflectionPaddingRT::~ReflectionPaddingRT() {}
+
+int ReflectionPaddingRT::getNbOutputs() const NOEXCEPT {
+    return 1;
+}
+
+Dims ReflectionPaddingRT::getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT {
+    return Dims3{c,output_h,output_w};
+}
+
+int ReflectionPaddingRT::initialize() NOEXCEPT {
+    return 0;
+}
+
+void ReflectionPaddingRT::terminate() NOEXCEPT {
+
+}
+
+size_t ReflectionPaddingRT::getWorkspaceSize(int maxBatchSize) const NOEXCEPT {
+    return 0;
+}
+
+#if NV_TENSORRT_MAJOR > 7
+int ReflectionPaddingRT::enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
+                                 cudaStream_t stream) NOEXCEPT {
+    dnnType* srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+    dnnType* dstData = reinterpret_cast<dnnType*>(outputs[0]);
+    reflection_pad2d_out_forward(padH,padW,srcData,dstData,input_h,input_w,c,n,stream);
+    return 0;
+}
+
+#elif NV_TENSORRT_MAJOR <= 7
+int32_t ReflectionPaddingRT::enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream){
+    dnnType* srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+    dnnType* dstData = reinterpret_cast<dnnType*>(outputs[0]);
+    reflection_pad2d_out_forward(padH,padW,srcData,dstData,input_h,input_w,c,n,stream);
+    return 0;
+}
+#endif
+
+
+size_t ReflectionPaddingRT::getSerializationSize() const NOEXCEPT {
+    return 8*sizeof(int32_t);
+}
+
+void ReflectionPaddingRT::serialize(void *buffer) const NOEXCEPT {
+    char *buf = reinterpret_cast<char*>(buffer),*a=buf;
+    writeBUF(buf,padH);
+    writeBUF(buf,padW);
+    writeBUF(buf,input_h);
+    writeBUF(buf,input_w);
+    writeBUF(buf,output_h);
+    writeBUF(buf,output_w);
+    writeBUF(buf,n);
+    writeBUF(buf,c);
+}
+
+void ReflectionPaddingRT::destroy() NOEXCEPT {
+    delete this;
+}
+
+const char *ReflectionPaddingRT::getPluginType() const NOEXCEPT {
+    return REFLECTIONPADDINGRT_PLUGIN_NAME;
+}
+
+const char *ReflectionPaddingRT::getPluginVersion() const NOEXCEPT {
+    return REFLECTIONPADDINGRT_PLUGIN_VERSION;
+}
+
+const char *ReflectionPaddingRT::getPluginNamespace() const NOEXCEPT {
+    return mPluginNamespace.c_str();
+}
+
+void ReflectionPaddingRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
+    mPluginNamespace = pluginNamespace;
+}
+
+IPluginV2Ext *ReflectionPaddingRT::clone() const NOEXCEPT {
+    auto *p = new ReflectionPaddingRT(padH,padW,input_h,input_w,output_h,output_w,c,n);
+    p->setPluginNamespace(mPluginNamespace.c_str());
+    return p;
+}
+
+DataType
+ReflectionPaddingRT::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const NOEXCEPT {
+    return DataType::kFLOAT;
+}
+
+void ReflectionPaddingRT::attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
+                                          IGpuAllocator *gpuAllocator) NOEXCEPT {
+}
+
+bool ReflectionPaddingRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted,
+                                                       int nbInputs) const NOEXCEPT {
+    return false;
+}
+
+bool ReflectionPaddingRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT {
+    return false;
+}
+
+void
+ReflectionPaddingRT::configurePlugin(const Dims *inputDims, int32_t nbInputs, const Dims *outputDims, int32_t nbOutputs,
+                                     const DataType *inputTypes, const DataType *outputTypes,
+                                     const bool *inputIsBroadcast, const bool *outputIsBroadcast,
+                                     PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT {
+
+}
+
+void ReflectionPaddingRT::detachFromContext() NOEXCEPT {
+
+}
+
+bool ReflectionPaddingRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT {
+    return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
+}
+
+
+ReflectionPaddingRTPluginCreator::ReflectionPaddingRTPluginCreator() {
+    mPluginAttributes.clear();
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+void ReflectionPaddingRTPluginCreator::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
+    mPluginNamespace = pluginNamespace;
+}
+
+const char *ReflectionPaddingRTPluginCreator::getPluginNamespace() const NOEXCEPT {
+    return mPluginNamespace.c_str();
+}
+
+IPluginV2Ext *ReflectionPaddingRTPluginCreator::deserializePlugin(const char *name, const void *serialData,
+                                                                  size_t serialLength) NOEXCEPT {
+    auto *pluginObj = new ReflectionPaddingRT(serialData,serialLength);
+    pluginObj->setPluginNamespace(mPluginNamespace.c_str());
+    return pluginObj;
+}
+
+IPluginV2Ext *
+ReflectionPaddingRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
+    const PluginField *fields = fc->fields;
+    int padH = *(static_cast<const int32_t*>(fields[0].data));
+    int padW = *(static_cast<const int32_t*>(fields[1].data));
+    int inputH = *(static_cast<const int32_t*>(fields[2].data));
+    int inputW = *(static_cast<const int32_t*>(fields[3].data));
+    int outputH = *(static_cast<const int32_t*>(fields[4].data));
+    int outputW = *(static_cast<const int32_t*>(fields[5].data));
+    int n = *(static_cast<const int32_t*>(fields[6].data));
+    int c = *(static_cast<const int32_t*>(fields[7].data));
+    auto *pluginObj = new ReflectionPaddingRT(padH,padW,inputH,inputW,outputH,outputW,c,n);
+    pluginObj->setPluginNamespace(mPluginNamespace.c_str());
+    return pluginObj;
+}
+
+const char *ReflectionPaddingRTPluginCreator::getPluginName() const NOEXCEPT {
+    return REFLECTIONPADDINGRT_PLUGIN_NAME;
+}
+
+const char *ReflectionPaddingRTPluginCreator::getPluginVersion() const NOEXCEPT {
+    return REFLECTIONPADDINGRT_PLUGIN_VERSION;
+}
+
+const PluginFieldCollection *ReflectionPaddingRTPluginCreator::getFieldNames() NOEXCEPT {
+    return &mFC;
+}
+
+