Added constant and zero padding with ISliceLayer,need to add them for tensorrt versions less than 8.2 using IPluginV2Ext instead of ISliceLayer since they dont seem to support reflect and zero

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;
+}
+
+