Merge pull request #1387 from TianShaoqing/develop

[Backend] update poros

Author: TianShaoqing

.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,42 +57,42 @@ 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")
 
 
 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

@@ -150,15 +150,15 @@ int Compiler::compile(const torch::jit::Module& origin_module,
         opt_graph = graph_and_ivalues.first;
     }
 
-    //cpu的话 过了预处理就返回
+    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(opt_graph, *optimized_module, true);
+        merge_graph_to_module(prewarm_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);
-
     //step2: try to find segments in unfold module
     //划分完子图的模型
     int ret = segment_graph(prewarm_graph);

poros/src/poros/compile/compile.cpp renamed to poros/poros/compile/compile.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/src/poros/compile/compile.h renamed to poros/poros/compile/compile.h

@@ -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/src/poros/compile/graph_prewarm.cpp renamed to poros/poros/compile/graph_prewarm.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/src/poros/compile/graph_prewarm.h renamed to poros/poros/compile/graph_prewarm.h

@@ -135,6 +135,26 @@ class MaxMinConverter : public GpuConverter {
     }
 };
 
+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

poros/src/poros/compile/graph_segment.cpp renamed to poros/poros/compile/graph_segment.cpp

@@ -360,6 +360,9 @@ class ConvertersMap {
     void init_unsupport_op_set() {
         try {
             std::vector<std::string> unsupport_op_vec = PorosGlobalContext::instance().get_poros_options().unsupport_op_list;
+            // 每次设置unsupport_op_list时刷新schema和nodekind set，避免用户下次编译想重新设置还保留之前不支持的op
+            _unsupport_schema_set.clear();
+            _unsupport_nodekind_set.clear();
             for (size_t i = 0; i < unsupport_op_vec.size(); i++) {
                 std::string line = unsupport_op_vec[i];
                 if (line.size() == 0) {

poros/src/poros/compile/graph_segment.h renamed to poros/poros/compile/graph_segment.h

@@ -73,7 +73,7 @@ def compile(module, prewarm_inputs, poros_options):
     Args:
         module (torch.nn.Module / torch.jit.ScriptModule): Source module
         input (list of tensor input): prewarmed data.
-        poros_options(PorosOptions / Dict of settings): compile settings for poros
+        poros_options(PorosOptions): compile settings for poros
     Returns:
         PorosModule: Compiled Module of poros, 
                     when run it will partially execute via inlined engine (which is TensorRT)

poros/src/poros/compile/ivalues_analysis.cpp renamed to poros/poros/compile/ivalues_analysis.cpp

@@ -61,6 +61,8 @@ def convert_prewarm_inputs(prewarm_inputs):
     else:
         raise TypeError("prewarm_inputs for poros should be torch.Tensor or wraped as tuple or inputs-lists, fix it")
     return wraped_prewarm_inputs     
+    # info = poros._C.PreWarmDatas()
+    # info.set_data(prewarm_inputs)
 
 def convert_poros_option(poros_option):
     # type: Dict[str, Any] -> poros._C.PorosOptions
@@ -70,39 +72,6 @@ def convert_poros_option(poros_option):
     option = poros._C.PorosOptions()
     if poros_option is None:
         #default situation. if user do not set the poros_option
-        return option
-    elif isinstance(poros_option, dict):
-        if "debug" in poros_option:
-            assert isinstance(poros_option["debug"], bool)
-            option.debug = poros_option["debug"]
-
-        if "use_fp16" in poros_option:
-            assert isinstance(poros_option["use_fp16"], bool)
-            option.use_fp16 = poros_option["use_fp16"]
-        
-        if "max_workspace_size" in poros_option:
-            assert type(poros_option["max_workspace_size"]) is int
-            option.max_workspace_size = poros_option["max_workspace_size"]
-
-        if "device" in poros_option:
-            option.device = _parse_device(poros_option["device"])
-        
-        if "is_dynamic" in poros_option:
-            assert isinstance(poros_option["is_dynamic"], bool)
-            option.is_dynamic = poros_option["is_dynamic"]
-
-        if "long_to_int" in poros_option:
-            assert isinstance(poros_option["long_to_int"], bool)
-            option.long_to_int = poros_option["long_to_int"]
-        
-        if "device_id" in poros_option:
-            assert type(poros_option["device_id"]) is int
-            option.device_id = poros_option["device_id"]
-
-        if "preprocess_mode" in poros_option:
-            assert type(poros_option["preprocess_mode"]) is int
-            option.preprocess_mode= poros_option["preprocess_mode"]
-
         return option
     elif isinstance(poros_option, PorosOptions):
         return poros_option.to_internal()

poros/src/poros/compile/ivalues_analysis.h renamed to poros/poros/compile/ivalues_analysis.h

@@ -10,11 +10,11 @@ set(GRAPHTEST "graph_test"  )
 file(
         GLOB UT_FILES
         "./op_fuser/*.cpp"
-        "../src/poros/lowering/fuse_*.cpp"
+        "../poros/lowering/fuse_*.cpp"
 )
 list(APPEND UT_FILES
-        "../src/poros/lowering/op_fuse_pass.cpp"
-        "../src/poros/util/graph_test_helper.cpp")
+        "../poros/lowering/op_fuse_pass.cpp"
+        "../poros/util/graph_test_helper.cpp")
 
 add_executable(${GRAPHTEST} ${UT_FILES})
 target_link_libraries(${GRAPHTEST} gtest_main)
@@ -28,8 +28,8 @@ set(UNITTEST "unit_test"  )
 
 file(
         GLOB UT_FILES
-        "../src/poros/*/*.cpp"
-        "../src/poros/converter/*/*.cpp"
+        "../poros/*/*.cpp"
+        "../poros/converter/*/*.cpp"
         "./converter/*.cpp"
 )