Merge pull request #1387 from TianShaoqing/develop

[Backend] update poros

.gitmodules

@@ -0,0 +1,6 @@
+[submodule "poros/third_party/googletest"]
+	path = poros/third_party/googletest
+	url = https://github.com/google/googletest.git
+[submodule "poros/third_party/gflags"]
+	path = poros/third_party/gflags
+	url = https://github.com/gflags/gflags.git

poros/CMakeLists.txt

@@ -1,8 +1,6 @@
 cmake_minimum_required(VERSION 3.21)
 project(poros)
 set(CMAKE_CXX_STANDARD 14)
-add_definitions(-D_GLIBCXX_USE_CXX11_ABI=1)
-
 
 option(BUILD_STATIC "build lib${PROJECT_NAME}.a static lib" OFF)
 option(BUILD_KERNEL "build lib${PROJECT_NAME}-kernel.so shared lib" OFF)
@@ -11,7 +9,17 @@ option(BUILD_TOOL "build ${PROJECT_NAME}-tool, an executable binary output" OFF)
 option(TEST "build for test. copy '.so' to site-packages automatically after compile" OFF)
 option(DEBUG "build for debug. add '-g' flag to gcc for detailed debug information" ON)
 option(UT "build for unit test" OFF)
+option(ABI "build ${PROJECT_NAME} with application binary interface (ABI) is on" OFF)
 
+# abi configuration
+if (NOT ABI)
+    add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)
+else ()
+    if (BUILD_TOOL OR UT)
+        message(FATAL_ERROR "${PROJECT_NAME}-tool or unit test are not supported when abi is on.")
+    endif()
+    add_definitions(-D_GLIBCXX_USE_CXX11_ABI=1)
+endif ()
 
 # minimum requirements
 set(PYTHON_MINIMUM_VERSION 3.6)
@@ -49,33 +57,33 @@ find_package(CUDNN ${CUDNN_MINIMUM_VERSION} REQUIRED)
 
 ## release headers
 # engine
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/engine/iengine.h" "${PROJECT_SOURCE_DIR}/src/poros/engine/engine_context.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/engine/iengine.h" "${PROJECT_SOURCE_DIR}/poros/engine/engine_context.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/engine")
 # compile
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/compile/poros_module.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/compile/poros_module.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/compile")
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/compile/compile.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/compile/compile.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/compile")
 # converter
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/converter/iconverter.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/converter/iconverter.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/converter")
 # iplugin
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/iplugin/*.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/iplugin/*.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/iplugin")
 ## context
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/context/*.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/context/*.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/context")
 ## context
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/context/*.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/context/*.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/context")
 ## lowering
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/lowering/*.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/lowering/*.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/lowering")
 ## util
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/util/*.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/util/*.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/util")
 ## log
-file(GLOB headers "${PROJECT_SOURCE_DIR}/src/poros/log/*.h")
+file(GLOB headers "${PROJECT_SOURCE_DIR}/poros/log/*.h")
 file(COPY ${headers} DESTINATION "${PROJECT_SOURCE_DIR}/build/include/poros/log")
 
 
@@ -83,8 +91,8 @@ include_directories(${TORCH_INCLUDE_DIRS})
 include_directories(${TensorRT_INCLUDE_DIRS})
 include_directories(${CUDA_INCLUDE_DIRS})
 include_directories(${CUDNN_INCLUDE_PATH})
-include_directories(src)
-include_directories(src/poros/compile)
+include_directories(${PROJECT_SOURCE_DIR})
+include_directories(poros/compile)
 
 
 add_compile_options(-D__const__= -D_GNU_SOURCE)
@@ -109,9 +117,9 @@ add_compile_options(
 
 file(
         GLOB POROS_CPP_FILES
-        "./src/poros/*/*.cpp"
-        "./src/poros/converter/*/*.cpp"
-        "./src/poros/converter/gpu/plugins/*.cpp"
+        "./poros/*/*.cpp"
+        "./poros/converter/*/*.cpp"
+        "./poros/converter/gpu/plugins/*.cpp"
 )
 
 
@@ -183,13 +191,13 @@ endif ()
 # kernel
 file(
         GLOB POROS_KERNEL_CPP_FILES
-        ./src/poros/compile/*.cpp
-        ./src/poros/context/*.cpp
-        ./src/poros/iplugin/*.cpp
-        ./src/poros/log/*.cpp
-        ./src/poros/lowering/*.cpp
-        ./src/poros/util/*.cpp
-        ./src/poros/engine/engine.cpp
+        ./poros/compile/*.cpp
+        ./poros/context/*.cpp
+        ./poros/iplugin/*.cpp
+        ./poros/log/*.cpp
+        ./poros/lowering/*.cpp
+        ./poros/util/*.cpp
+        ./poros/engine/engine.cpp
 )
 
 # kernel SHARED

poros/README.md

@@ -43,7 +43,7 @@ get Poros source code:
 ```shell
 git clone https://github.com/PaddlePaddle/FastDeploy.git
 cd poros
-git submodule update --init --recursive --jobs 0 -f
+git submodule update --init --recursive
 ```
 
 We strongly recommend you to prepare the building environment with anaconda3:

poros/poros/compile/compile.cpp

@@ -150,15 +150,15 @@ int Compiler::compile(const torch::jit::Module& origin_module,
         opt_graph = graph_and_ivalues.first;
     }
 
-    //cpu的话 过了预处理就返回
-    if (_options.device == Device::CPU) {
-        merge_graph_to_module(opt_graph, *optimized_module, true);
-        return 0;
-    }
-
     std::shared_ptr<torch::jit::Graph> prewarm_graph = graph_prewarm(opt_graph, prewarm_datas);
     GRAPH_DUMP("prewarmed_module graph:", prewarm_graph);
 
+    //cpu的话，预热后就返回
+    if (_options.device == Device::CPU) {
+        merge_graph_to_module(prewarm_graph, *optimized_module, true);
+        return 0;
+    }
+
     //step2: try to find segments in unfold module
     //划分完子图的模型
     int ret = segment_graph(prewarm_graph);

poros/poros/compile/graph_segment.cpp

@@ -102,6 +102,24 @@ struct PorosGraphSegment {
             }
         }
 
+        // aten::__getitem__ idx参数不支持非constant类型
+        if (node->kind() == torch::jit::aten::__getitem__) {
+            if (node->inputs().size() == 2 && 
+                node->input(1)->node()->kind() != torch::jit::prim::Constant) {
+                LOG(WARNING) << "The index input of aten::__getitem__ is not supported as non-constant type.";
+                return false;
+            }
+        }
+
+        // aten::_set_item idx参数不支持非constant类型
+        if (node->kind() == torch::jit::aten::_set_item) {
+            if (node->inputs().size() == 3 && 
+                node->input(1)->node()->kind() != torch::jit::prim::Constant) {
+                LOG(WARNING) << "The index input of aten::_set_item is not supported as non-constant type.";
+                return false;
+            }
+        }
+
         if (node->kind() == kind_ || engine_->is_node_supported(node)) {
             return true;
         }

poros/poros/converter/gpu/aten_eval.cpp

@@ -66,6 +66,8 @@ bool AppendConverter::converter(TensorrtEngine* engine, const torch::jit::Node *
 bool GetitemConverter::converter(TensorrtEngine* engine, const torch::jit::Node *node) {
     at::ArrayRef<const torch::jit::Value*> inputs = node->inputs();
     POROS_CHECK_TRUE((inputs.size() == 2), "invaid inputs size for GetitemConverter");
+    POROS_CHECK_TRUE((inputs[1]->node()->kind() == torch::jit::prim::Constant),
+        "inputs[1] for GetitemConverter is not come from prim::Constant as expected");
 
     if (node->outputs()[0]->type()->str() == "Tensor") {
         //extract list
@@ -126,6 +128,8 @@ bool GetitemConverter::converter(TensorrtEngine* engine, const torch::jit::Node
 bool SetitemConverter::converter(TensorrtEngine* engine, const torch::jit::Node *node) {
     at::ArrayRef<const torch::jit::Value*> inputs = node->inputs();
     POROS_CHECK_TRUE((inputs.size() == 3), "invaid inputs size for SetitemConverter");
+    POROS_CHECK_TRUE((inputs[1]->node()->kind() == torch::jit::prim::Constant),
+        "inputs[1] for SetitemConverter is not come from prim::Constant as expected");
 
     size_t idx = engine->context().get_constant(inputs[1]).toInt();
 

poros/poros/converter/gpu/reduce.cpp

@@ -271,10 +271,75 @@ bool MaxMinConverter::converter(TensorrtEngine* engine, const torch::jit::Node *
     return true;
 }
 
+/*
+"aten::argmax(Tensor self, int? dim=None, bool keepdim=False) -> (Tensor)"*/
+bool ArgmaxArgminConverter::converter(TensorrtEngine* engine, const torch::jit::Node *node) {
+    at::ArrayRef<const torch::jit::Value*> inputs = node->inputs();
+    POROS_CHECK_TRUE((inputs[0]->type()->isSubtypeOf(c10::TensorType::get())), 
+        "input[0] for ArgmaxArgminConverter is not Tensor as expected");
+
+    // TODO: to imp dim=None
+    POROS_CHECK_TRUE((inputs[1]->type()->isSubtypeOf(c10::IntType::get())), 
+        "input[1] for ArgmaxArgminConverter is not int as expected");
+
+    //extract self
+    auto in_tensor = engine->context().get_tensor(inputs[0]);
+    POROS_CHECK_TRUE((in_tensor != nullptr), "Unable to init input tensor for node: " << *node);
+    auto in_dims = nvdim_to_sizes(in_tensor->getDimensions());
+
+    bool is_dynamic = check_nvtensor_is_dynamic(in_tensor);
+
+    POROS_CHECK_TRUE((in_dims.size() > 1), 
+        "Converter aten::argmax error: At least 2 dimensions are required for input[0].");
+    nvinfer1::ITensor* output_indices = nullptr;
+
+    int64_t dim = 0;
+    dim = engine->context().get_constant(inputs[1]).toInt();
+    dim = dim < 0 ? in_dims.size() + dim : dim;
+    bool keep_dim = engine->context().get_constant(inputs[2]).toBool();
+    uint32_t shiftDim = 1 << dim;
+
+    // nvinfer1::TopKOperation noly support kFLOAT, so this is transfer kINT32 to kFLOAT
+    if (in_tensor->getType() == nvinfer1::DataType::kINT32) {
+        auto id_layer = engine->network()->addIdentity(*in_tensor);
+        id_layer->setOutputType(0, nvinfer1::DataType::kFLOAT);
+        id_layer->setName((layer_info(node) + "_IIdentityLayer_int32_to_float").c_str());
+        in_tensor = id_layer->getOutput(0);
+    }
+
+    nvinfer1::TopKOperation topk_option = (node->kind() == torch::jit::aten::argmax) ?
+                                            nvinfer1::TopKOperation::kMAX : 
+                                            nvinfer1::TopKOperation::kMIN;
+    nvinfer1::ITopKLayer* topk_layer =  engine->network()->addTopK(*in_tensor, topk_option, 1, shiftDim);
+    POROS_CHECK(topk_layer, "Unable to create TopK layer from node: " << *node);
+    topk_layer->setName((layer_info(node) + "_ITopKLayer").c_str());
+    output_indices = topk_layer->getOutput(1);
+
+    // squeeze output dim
+    if (in_tensor->getDimensions().nbDims > 1 && !keep_dim) {
+        auto shuffle_layer = engine->network()->addShuffle(*output_indices);
+        if (is_dynamic) {
+            nvinfer1::ITensor* self_shape_tensor = engine->network()->addShape(*in_tensor)->getOutput(0);
+            nvinfer1::ITensor* squeeze_output_shape = squeeze_nv_shapetensor(engine, self_shape_tensor, dim);
+            shuffle_layer->setInput(1, *squeeze_output_shape);
+        } else {
+            in_dims.erase(in_dims.begin() + dim);
+            nvinfer1::Dims squeeze_output_dims = sizes_to_nvdim(in_dims);
+            shuffle_layer->setReshapeDimensions(squeeze_output_dims);
+        }
+        output_indices = shuffle_layer->getOutput(0);
+    }
+    engine->context().set_tensor(node->outputs()[0], output_indices);
+    LOG(INFO) << "Output tensor shape: " << output_indices->getDimensions();
+    return true;
+}
+
+
 POROS_REGISTER_CONVERTER(TensorrtEngine, MeanConverter);
 POROS_REGISTER_CONVERTER(TensorrtEngine, SumConverter);
 POROS_REGISTER_CONVERTER(TensorrtEngine, ProdConverter);
 POROS_REGISTER_CONVERTER(TensorrtEngine, MaxMinConverter);
+POROS_REGISTER_CONVERTER(TensorrtEngine, ArgmaxArgminConverter);
 
 }  // namespace poros 
 }  // namespace mirana

poros/poros/converter/gpu/reduce.h

@@ -135,6 +135,26 @@ public:
     }
 };
 
+class ArgmaxArgminConverter : public GpuConverter {
+public:
+    ArgmaxArgminConverter() {}
+    virtual ~ArgmaxArgminConverter() {}
+
+    bool converter(TensorrtEngine* engine, const torch::jit::Node *node);
+
+    const std::vector<std::string> schema_string() {
+        return {"aten::argmax(Tensor self, int dim, bool keepdim=False) -> (Tensor)",
+                "aten::argmax(Tensor self, int? dim=None, bool keepdim=False) -> (Tensor)",
+                "aten::argmin(Tensor self, int dim, bool keepdim=False) -> (Tensor)",
+                };
+    }
+
+    const std::vector<torch::jit::NodeKind> node_kind() {
+        return {torch::jit::aten::argmax,
+                torch::jit::aten::argmin};
+    }
+};
+
 }  // namespace poros 
 }  // namespace mirana
 }  // namespace baidu