mobilenet works with trt7(IPluginV2IOExt) ,need to test it with trt8

CMakeLists.txt

@@ -51,13 +51,13 @@ endif()
 
 if(WIN32)
     if(CMAKE_BUILD_TYPE MATCHES Release)
-        set(CMAKE_CXX_FLAGS "/O2 /FS /EHsc")
+        set(CMAKE_CXX_FLAGS "/O2 /FS /EHsc /Md")
-        set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32 -arch=sm_61)
+        set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32)
     endif()
 
     if(CMAKE_BUILD_TYPE MATCHES Debug)
-        set(CMAKE_CXX_FLAGS "/Od /FS /EHsc")
+        set(CMAKE_CXX_FLAGS "/Od /FS /EHsc /MDd")
-        set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32 -arch=sm_61 -G -g)
+        set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} --maxrregcount=32 -G -g)
     endif()
 set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON)
 endif(WIN32)
@@ -92,7 +92,7 @@ include_directories(${CUDNN_INCLUDE_DIR})
 file(GLOB tkdnn_CUSRC "src/kernels/*.cu" "src/sorting.cu" "src/pluginsRT/*.cpp")
 cuda_include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${CUDNN_INCLUDE_DIRS})
 cuda_add_library(kernels SHARED ${tkdnn_CUSRC})
-target_link_libraries(kernels ${CUDA_CUBLAS_LIBRAY} ${CUDA_LIBRARIES} ${CUDNN_LIBRARIES})
+target_link_libraries(kernels ${CUDA_CUBLAS_LIBRARIES} ${CUDA_LIBRARIES} ${CUDNN_LIBRARIES})
 
 
 
@@ -121,7 +121,7 @@ set(tkdnn_LIBS kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} ${CUDNN_LIBRAR
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS} ${OPENCV_INCLUDE_DIRS} ${NVINFER_INCLUDES})
 add_library(tkDNN SHARED ${tkdnn_SRC})
-target_link_libraries(tkDNN ${tkdnn_LIBS} ${CUDA_CUBLAS_LIBRAY})
+target_link_libraries(tkDNN ${tkdnn_LIBS} ${CUDA_CUBLAS_LIBRARIES})
 
 #static 
 #add_library(tkDNN_static STATIC ${tkdnn_SRC})

include/tkDNN/Layer.h

@@ -423,6 +423,8 @@ public:
     virtual layerType_t getLayerType() { return LAYER_FLATTEN; };
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
+
+    int c, h, w, rows, cols;
 };
 
 /**

include/tkDNN/pluginsRT/FlattenConcatRT.h

@@ -3,10 +3,10 @@
 #include <vector>
 #include <utils.h>
 namespace nvinfer1 {
-    class FlattenConcatRT : public IPluginV2 {
+    class FlattenConcatRT : public IPluginV2IOExt {
 
     public:
-        FlattenConcatRT() ;
+        FlattenConcatRT(int c,int h,int w,int rows,int cols) ;
 
         FlattenConcatRT(const void *data, size_t length) ;
 
@@ -16,9 +16,6 @@ namespace nvinfer1 {
 
         Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
 
-        void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
-                            PluginFormat format, int maxBatchSize) NOEXCEPT override ;
-
         int initialize() NOEXCEPT override ;
 
         void terminate() NOEXCEPT override ;
@@ -37,8 +34,6 @@ namespace nvinfer1 {
 
         void destroy() NOEXCEPT override ;
 
-        bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
-
         const char *getPluginType() const NOEXCEPT override ;
 
         const char *getPluginVersion() const NOEXCEPT override;
@@ -47,7 +42,21 @@ namespace nvinfer1 {
 
         void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
 
-        IPluginV2 *clone() const NOEXCEPT override ;
+        IPluginV2IOExt *clone() const NOEXCEPT override ;
+
+        DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
+
+        void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) 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;
+
+        bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT override;
+
+        void detachFromContext() NOEXCEPT override;
 
         int c, h, w;
         int rows, cols;
@@ -65,9 +74,9 @@ namespace nvinfer1 {
 
         const char *getPluginNamespace() const NOEXCEPT override ;
 
-        IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
+        IPluginV2IOExt *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
 
-        IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
+        IPluginV2IOExt *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
 
         const char *getPluginName() const NOEXCEPT override ;
 

src/Flatten.cpp

@@ -15,6 +15,11 @@ Flatten::Flatten(Network *net) : Layer(net) {
     output_dim.w = 1;
     output_dim.l = 1;
 
+    this->h = 1;
+    this->w = 1;
+    this->rows = input_dim.w;
+    this->cols = input_dim.h * input_dim.c;
+    this->c = input_dim.w * input_dim.h * input_dim.c;
 }
 
 Flatten::~Flatten() {

src/NetworkRT.cpp

@@ -18,9 +18,9 @@ using namespace nvinfer1;
 // Logger for info/warning/errors
 class Logger : public ILogger {
     void log(Severity severity, const char* msg) NOEXCEPT override {
-#ifdef DEBUG
+//#ifdef DEBUG
         std::cout <<"TENSORRT LOG: "<< msg << std::endl;
-#endif
+//#endif
     }
 } loggerRT;
 
@@ -473,7 +473,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) {
 
 ILayer* NetworkRT::convert_layer(ITensor *input, Flatten *l) {
 
-    IPluginV2 *plugin = new FlattenConcatRT();
+    IPluginV2IOExt *plugin = new FlattenConcatRT(l->c,l->h,l->w,l->rows,l->cols);
     IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
     checkNULL(lRT);
     return lRT;

src/pluginsRT/FlattenConcatRT.cpp

@@ -4,12 +4,17 @@ using namespace  nvinfer1;
 std::vector<PluginField> FlattenConcatRTPluginCreator::mPluginAttributes;
 PluginFieldCollection FlattenConcatRTPluginCreator::mFC{};
 
-FlattenConcatRT::FlattenConcatRT() {
+FlattenConcatRT::FlattenConcatRT(int c, int h, int w, int rows, int cols) {
     stat = cublasCreate(&handle);
     if (stat != CUBLAS_STATUS_SUCCESS) {
         printf ("CUBLAS initialization failed\n");
         return;
     }
+    this->c = c;
+    this->h = h;
+    this->w = w;
+    this->rows = rows;
+    this->cols = cols;
 }
 
 FlattenConcatRT::FlattenConcatRT(const void *data, size_t length) {
@@ -32,16 +37,6 @@ Dims FlattenConcatRT::getOutputDimensions(int index, const Dims *inputs, int nbI
     return Dims3{ inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2], 1, 1};
 }
 
-void FlattenConcatRT::configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,
-                                          DataType type, PluginFormat format, int maxBatchSize) NOEXCEPT {
-    assert(nbOutputs == 1 && nbInputs ==1);
-    rows = inputDims[0].d[0];
-    cols = inputDims[0].d[1] * inputDims[0].d[2];
-    c = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2];
-    h = 1;
-    w = 1;
-}
-
 int FlattenConcatRT::initialize() NOEXCEPT {
     return 0;
 }
@@ -107,9 +102,7 @@ void FlattenConcatRT::destroy() NOEXCEPT {
     delete this;
 }
 
-bool FlattenConcatRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT {
+
-    return true;
-}
 
 const char *FlattenConcatRT::getPluginType() const NOEXCEPT {
     return "FlattenConcatRT_tkDNN";
@@ -127,12 +120,44 @@ void FlattenConcatRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
     mPluginNamespace = pluginNamespace;
 }
 
-IPluginV2 *FlattenConcatRT::clone() const NOEXCEPT {
+IPluginV2IOExt *FlattenConcatRT::clone() const NOEXCEPT {
-    auto *p = new FlattenConcatRT();
+    auto* p = new FlattenConcatRT(c, h, w, rows, cols);
     p->setPluginNamespace(mPluginNamespace.c_str());
     return p;
 }
 
+DataType FlattenConcatRT::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT
+{
+    return DataType::kFLOAT;
+}
+
+void FlattenConcatRT::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) NOEXCEPT
+{
+}
+
+void FlattenConcatRT::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT
+{
+}
+
+bool FlattenConcatRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT
+{
+    return false;
+}
+
+bool FlattenConcatRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT
+{
+    return false;
+}
+
+bool FlattenConcatRT::supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const NOEXCEPT
+{
+    return true;
+}
+
+void FlattenConcatRT::detachFromContext() NOEXCEPT
+{
+}
+
 FlattenConcatRTPluginCreator::FlattenConcatRTPluginCreator() {
     mPluginAttributes.clear();
     mFC.nbFields = mPluginAttributes.size();
@@ -147,15 +172,21 @@ const char *FlattenConcatRTPluginCreator::getPluginNamespace() const NOEXCEPT {
     return mPluginNamespace.c_str();
 }
 
-IPluginV2 *FlattenConcatRTPluginCreator::deserializePlugin(const char *name, const void *serialData,
+IPluginV2IOExt *FlattenConcatRTPluginCreator::deserializePlugin(const char *name, const void *serialData,
                                                            size_t serialLength) NOEXCEPT {
     auto *pluginObj = new FlattenConcatRT(serialData,serialLength);
     pluginObj->setPluginNamespace(mPluginNamespace.c_str());
     return pluginObj;
 }
 
-IPluginV2 *FlattenConcatRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
+IPluginV2IOExt *FlattenConcatRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
-    auto *pluginObj = new FlattenConcatRT();
+    const PluginField* fields = fc->fields;
+    int c = *(static_cast<const int*>(fields[0].data));
+    int h = *(static_cast<const int*>(fields[1].data));
+    int w = *(static_cast<const int*>(fields[2].data));
+    int rows = *(static_cast<const int*>(fields[3].data));
+    int cols = *(static_cast<const int*>(fields[4].data));
+    auto* pluginObj = new FlattenConcatRT(c, h, w, rows, cols);
     pluginObj->setPluginNamespace(mPluginNamespace.c_str());
     return pluginObj;
 }

src/pluginsRT/ReshapeRT.cpp

@@ -54,10 +54,11 @@ int ReshapeRT::enqueue(int batchSize, const void *const *inputs, void *const *ou
 }
 #elif NV_TENSORRT_MAJOR == 7
 int32_t ReshapeRT::enqueue(int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) {
+    std::cout << new_dim.c << ":" << new_dim.h << std::endl;
     dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
     dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
-
+    std::cout << "C : " << c << "H : " << h << "w :" << w << std::endl;
-    checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
+    checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*new_dim.c*new_dim.h*new_dim.w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
     return 0;
 }
 #endif