split load torchscript sources from pnnx, drop cxxabi hack, link torchscript with whole-archive (#5348)

tools/pnnx/CMakeLists.txt

@@ -25,34 +25,21 @@ include(PNNXPyTorch)
 # c++14 is required for using torch headers
 set(CMAKE_CXX_STANDARD 14)
 
-#set(CMAKE_BUILD_TYPE debug)
+# set(CMAKE_BUILD_TYPE debug)
 #set(CMAKE_BUILD_TYPE relwithdebinfo)
-#set(CMAKE_BUILD_TYPE release)
+# set(CMAKE_BUILD_TYPE release)
 
 option(PNNX_COVERAGE "build for coverage" OFF)
 
 #set(Torch_INSTALL_DIR "/home/nihui/.local/lib/python3.9/site-packages/torch" CACHE STRING "")
 #set(Torch_INSTALL_DIR "/home/nihui/osd/pnnx/pytorch-v1.10.0/build/install" CACHE STRING "")
-#set(Torch_INSTALL_DIR "/home/nihui/osd/pnnx/libtorch" CACHE STRING "")
+# set(Torch_INSTALL_DIR "/home/nihui/osd/pnnx/install" CACHE STRING "")
 set(TorchVision_INSTALL_DIR "/home/nihui/osd/vision/build/install" CACHE STRING "")
 
 #set(Torch_DIR "${Torch_INSTALL_DIR}/share/cmake/Torch")
 set(TorchVision_DIR "${TorchVision_INSTALL_DIR}/share/cmake/TorchVision")
 
-find_package(protobuf CONFIG)
-
-if(protobuf_FOUND)
-    set(PROTOBUF_FOUND ${protobuf_FOUND})
-    set(PROTOBUF_VERSION ${protobuf_VERSION})
-else()
-    # fallback to system
-    find_package(Protobuf)
-    set(PROTOBUF_FOUND ${Protobuf_FOUND})
-    set(PROTOBUF_VERSION ${Protobuf_VERSION})
-    if(TARGET protobuf::protoc)
-        set_target_properties(protobuf::protoc PROPERTIES IMPORTED_LOCATION_RELEASE "${PROTOBUF_PROTOC_EXECUTABLE}")
-    endif()
-endif()
+# test if libtorch and protobuf has the same cxxabi version
 
 find_package(Python3 COMPONENTS Interpreter Development)
 
@@ -75,15 +62,62 @@ if(Torch_VERSION VERSION_GREATER_EQUAL "2.1")
     set(CMAKE_CXX_STANDARD 17)
 endif()
 
-if(TorchVision_FOUND)
+# find torchvision library
+find_library(TORCHVISION_LIBRARY torchvision PATHS "${TorchVision_INSTALL_DIR}/lib" "${TorchVision_INSTALL_DIR}/lib64")
+if(TORCHVISION_LIBRARY)
+    message(STATUS "Found TorchVision: ${TORCHVISION_LIBRARY}")
+    if(APPLE)
+        list(APPEND TORCHVISION_LIBRARY "-Wl,-force_load,${TORCHVISION_LIBRARY}")
+    elseif(MSVC)
+        list(APPEND TORCHVISION_LIBRARY "-WHOLEARCHIVE:${TORCHVISION_LIBRARY}")
+    else()
+        list(APPEND TORCHVISION_LIBRARY "-Wl,--whole-archive ${TORCHVISION_LIBRARY} -Wl,--no-whole-archive")
+    endif()
+    set(TorchVision_FOUND TRUE)
     message(STATUS "Building with TorchVision")
-    add_definitions(-DPNNX_TORCHVISION)
 else()
+    message(WARNING "static library ${TORCHVISION_LIBRARY} not found.")
+    set(TorchVision_FOUND FALSE)
     message(WARNING "Building without TorchVision")
 endif()
 
 include_directories(${TORCH_INCLUDE_DIRS})
 
+if("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
+    include(CheckCXXSourceCompiles)
+    set(CMAKE_REQUIRED_FLAGS "${TORCH_CXX_FLAGS}")
+    check_cxx_source_compiles("#include <cxxabi.h>\n#if _GLIBCXX_USE_CXX11_ABI\nint main() { return 0; }\n#endif" PNNX_TORCH_USE_CXX11_ABI)
+    unset(CMAKE_REQUIRED_FLAGS)
+    check_cxx_source_compiles("#include <cxxabi.h>\n#if _GLIBCXX_USE_CXX11_ABI\nint main() { return 0; }\n#endif" PNNX_COMPILER_USE_CXX11_ABI)
+endif()
+
+if((PNNX_TORCH_USE_CXX11_ABI AND PNNX_COMPILER_USE_CXX11_ABI) OR (NOT PNNX_TORCH_USE_CXX11_ABI AND NOT PNNX_COMPILER_USE_CXX11_ABI))
+    find_package(protobuf CONFIG)
+
+    if(protobuf_FOUND)
+        set(PROTOBUF_FOUND ${protobuf_FOUND})
+        set(PROTOBUF_VERSION ${protobuf_VERSION})
+    else()
+        # fallback to system
+        find_package(Protobuf)
+        set(PROTOBUF_FOUND ${Protobuf_FOUND})
+        set(PROTOBUF_VERSION ${Protobuf_VERSION})
+        if(TARGET protobuf::protoc)
+            set_target_properties(protobuf::protoc PROPERTIES IMPORTED_LOCATION_RELEASE "${PROTOBUF_PROTOC_EXECUTABLE}")
+        endif()
+    endif()
+endif()
+
+# set(onnxruntime_INSTALL_DIR "/home/nihui/osd/ncnn-nihui/tools/pnnx/build/src/test/onnxruntime-1.16.3/build/install" CACHE STRING "")
+# set(onnxruntime_DIR "${onnxruntime_INSTALL_DIR}/lib64/cmake/onnxruntime")
+#
+# find_package(onnxruntime)
+#
+# message(STATUS "onnxruntime_VERSION = ${onnxruntime_VERSION}")
+# message(STATUS "onnxruntime_VERSION_MAJOR = ${onnxruntime_VERSION_MAJOR}")
+# message(STATUS "onnxruntime_VERSION_MINOR = ${onnxruntime_VERSION_MINOR}")
+# message(STATUS "onnxruntime_VERSION_PATCH = ${onnxruntime_VERSION_PATCH}")
+
 add_subdirectory(src)
 
 enable_testing()

tools/pnnx/src/CMakeLists.txt

@@ -572,52 +572,88 @@ if(PROTOBUF_FOUND)
 
     if(Protobuf_FOUND OR protobuf_MODULE_COMPATIBLE)
         protobuf_generate_cpp(ONNX_PROTO_SRCS ONNX_PROTO_HDRS onnx.proto)
+        add_library(onnxproto STATIC ${ONNX_PROTO_SRCS} ${ONNX_PROTO_HDRS})
+        target_include_directories(onnxproto PUBLIC ${PROTOBUF_INCLUDE_DIR} ${CMAKE_CURRENT_BINARY_DIR})
+        target_link_libraries(onnxproto PUBLIC ${PROTOBUF_LIBRARIES})
+    else()
+        add_library(onnxproto STATIC onnx.proto)
+        target_include_directories(onnxproto PUBLIC ${CMAKE_CURRENT_BINARY_DIR})
+        protobuf_generate(TARGET onnxproto)
+        target_link_libraries(onnxproto PUBLIC protobuf::libprotobuf)
+    endif()
+endif()
 
-        add_library(pnnx2onnx STATIC
-            save_onnx.cpp
-            save_onnx_cxxabi_bridge.cpp
-            ${ONNX_PROTO_SRCS} ${ONNX_PROTO_HDRS}
-        )
+set(torch2pnnx_SRCS
+
+    pass_level0.cpp
+    pass_level1.cpp
 
-        target_include_directories(pnnx2onnx PRIVATE ${PROTOBUF_INCLUDE_DIR} ${CMAKE_CURRENT_BINARY_DIR})
-        target_link_libraries(pnnx2onnx PRIVATE ${PROTOBUF_LIBRARIES})
+    ${pnnx_pass_level0_SRCS}
+    ${pnnx_pass_level1_SRCS}
 
-        # libtorch is usually compiled with old cxx11 abi
-        set_source_files_properties(save_onnx_cxxabi_bridge.cpp PROPERTIES COMPILE_FLAGS "${TORCH_CXX_FLAGS}")
+    load_torchscript.cpp
+)
 
-        message(STATUS "Building with onnx-zero")
-    else()
+add_library(torch2pnnx OBJECT ${torch2pnnx_SRCS})
+target_compile_definitions(torch2pnnx PRIVATE BUILD_TORCH2PNNX)
+target_compile_options(torch2pnnx PUBLIC "${TORCH_CXX_FLAGS}")
 
-        add_library(pnnx2onnx STATIC
-            save_onnx.cpp
-            save_onnx_cxxabi_bridge.cpp
-            onnx.proto
-        )
-        set_source_files_properties(save_onnx_cxxabi_bridge.cpp PROPERTIES COMPILE_FLAGS "${TORCH_CXX_FLAGS}")
-        target_include_directories(pnnx2onnx PRIVATE ${CMAKE_CURRENT_BINARY_DIR})
-        protobuf_generate(TARGET pnnx2onnx)
-        target_link_libraries(pnnx2onnx PRIVATE protobuf::libprotobuf)
-        message(STATUS "Building with onnx-zero")
-    endif()
+if(TorchVision_FOUND)
+    set_property(SOURCE load_torchscript.cpp APPEND PROPERTY COMPILE_DEFINITIONS PNNX_TORCHVISION)
+endif()
+
+if(PROTOBUF_FOUND)
+    add_library(pnnx2onnx STATIC
+        save_onnx.cpp
+    )
+    target_link_libraries(pnnx2onnx PRIVATE onnxproto)
+
+    message(STATUS "Building with onnx-zero")
 else()
     message(STATUS "Building without onnx-zero")
 endif()
 
+# if(onnxruntime_FOUND)
+#
+#     set(pnnx_pass_onnx_SRCS
+#         pass_onnx/canonicalize.cpp
+#         pass_onnx/dead_code_elimination.cpp
+#         pass_onnx/fold_constants.cpp
+#         pass_onnx/inline_containers.cpp
+#         pass_onnx/model_stat.cpp
+#         pass_onnx/shape_inference.cpp
+#     )
+#
+#     set(onnx2pnnx_SRCS
+#         ${pnnx_pass_onnx_SRCS}
+#         load_onnx.cpp
+#     )
+#
+#     add_library(onnx2pnnx STATIC ${onnx2pnnx_SRCS})
+#     target_link_libraries(onnx2pnnx PRIVATE onnxproto onnxruntime::onnxruntime)
+#
+#     target_compile_definitions(onnx2pnnx PRIVATE BUILD_ONNX2PNNX)
+#
+#     message(STATUS "Building with dynamo-onnx")
+# else()
+#     message(STATUS "Building without dynamo-onnx")
+# endif()
+
+if(NOT MSVC)
+    add_definitions(-Wall -Wextra)
+endif()
+
 set(pnnx_SRCS
     main.cpp
     ir.cpp
     storezip.cpp
     utils.cpp
 
-    pass_level0.cpp
-    pass_level1.cpp
     pass_level2.cpp
     pass_level3.cpp
     pass_level4.cpp
     pass_level5.cpp
 
-    ${pnnx_pass_level0_SRCS}
-    ${pnnx_pass_level1_SRCS}
     ${pnnx_pass_level2_SRCS}
     ${pnnx_pass_level3_SRCS}
     ${pnnx_pass_level4_SRCS}
@@ -628,36 +664,38 @@ set(pnnx_SRCS
     ${pnnx_pass_ncnn_SRCS}
 )
 
-if(NOT MSVC)
-    add_definitions(-Wall -Wextra)
-endif()
-
 add_executable(pnnx ${pnnx_SRCS})
 
-target_compile_definitions(pnnx PRIVATE BUILD_PNNX)
+set_property(SOURCE main.cpp APPEND PROPERTY COMPILE_DEFINITIONS BUILD_TORCH2PNNX)
+target_link_libraries(pnnx PRIVATE torch2pnnx)
 
-if(PNNX_COVERAGE)
-    target_compile_options(pnnx PUBLIC -coverage -fprofile-arcs -ftest-coverage)
-    target_link_libraries(pnnx PUBLIC -coverage -lgcov)
+if(TorchVision_FOUND)
+    target_link_libraries(pnnx PRIVATE ${TORCHVISION_LIBRARY})
 endif()
 
 if(WIN32)
-    target_compile_definitions(pnnx PUBLIC NOMINMAX)
+    target_link_libraries(pnnx PRIVATE ${TORCH_LIBRARIES})
+else()
+    target_link_libraries(pnnx PRIVATE ${TORCH_LIBRARIES} pthread dl)
 endif()
 
 if(PROTOBUF_FOUND)
-    target_compile_definitions(pnnx PRIVATE BUILD_PNNX2ONNX)
+    set_property(SOURCE main.cpp APPEND PROPERTY COMPILE_DEFINITIONS BUILD_PNNX2ONNX)
     target_link_libraries(pnnx PRIVATE pnnx2onnx)
 endif()
 
-if(TorchVision_FOUND)
-    target_link_libraries(pnnx PRIVATE TorchVision::TorchVision)
+# if(onnxruntime_FOUND)
+#     set_property(SOURCE main.cpp APPEND PROPERTY COMPILE_DEFINITIONS BUILD_ONNX2PNNX)
+#     target_link_libraries(pnnx PRIVATE onnx2pnnx)
+# endif()
+
+if(PNNX_COVERAGE)
+    target_compile_options(pnnx PUBLIC -coverage -fprofile-arcs -ftest-coverage)
+    target_link_libraries(pnnx PUBLIC -coverage -lgcov)
 endif()
 
 if(WIN32)
-    target_link_libraries(pnnx PRIVATE ${TORCH_LIBRARIES})
-else()
-    target_link_libraries(pnnx PRIVATE ${TORCH_LIBRARIES} pthread dl)
+    target_compile_definitions(pnnx PUBLIC NOMINMAX)
 endif()
 
 # set_target_properties(pnnx PROPERTIES COMPILE_FLAGS -fsanitize=address)

tools/pnnx/src/ir.cpp

@@ -23,11 +23,6 @@
 #include <string>
 #include <stack>
 
-#if BUILD_PNNX
-#include <torch/script.h>
-#include <torch/csrc/api/include/torch/version.h>
-#endif
-
 #include "storezip.h"
 #include "utils.h"
 
@@ -135,276 +130,6 @@ static int string_to_type(const char* s)
     return 0; // null
 }
 
-#if BUILD_PNNX
-int get_at_tensor_type(const at::ScalarType& st)
-{
-    if (st == c10::ScalarType::Float) return 1;
-    if (st == c10::ScalarType::Double) return 2;
-    if (st == c10::ScalarType::Half) return 3;
-    if (st == c10::ScalarType::Int) return 4;
-    if (st == c10::ScalarType::QInt32) return 4;
-    if (st == c10::ScalarType::Long) return 5;
-    if (st == c10::ScalarType::Short) return 6;
-    if (st == c10::ScalarType::Char) return 7;
-    if (st == c10::ScalarType::QInt8) return 7;
-    if (st == c10::ScalarType::Byte) return 8;
-    if (st == c10::ScalarType::QUInt8) return 8;
-    if (st == c10::ScalarType::Bool) return 9;
-    if (st == c10::ScalarType::ComplexFloat) return 10;
-    if (st == c10::ScalarType::ComplexDouble) return 11;
-    if (st == c10::ScalarType::ComplexHalf) return 12;
-    return 0; // unknown type
-}
-
-Parameter::Parameter(const torch::jit::Node* value_node)
-{
-    type = 0;
-
-    if (value_node->kind() == c10::prim::Constant)
-    {
-        if (value_node->output()->type()->kind() == c10::TypeKind::NoneType)
-        {
-            type = 0;
-            return;
-        }
-
-        if (!value_node->hasAttribute(torch::jit::attr::value))
-        {
-            fprintf(stderr, "no attribute value\n");
-            value_node->dump();
-            return;
-        }
-
-        switch (value_node->output()->type()->kind())
-        {
-        case c10::TypeKind::NoneType:
-        {
-            type = 0;
-            break;
-        }
-        case c10::TypeKind::BoolType:
-        {
-            type = 1;
-            b = value_node->i(torch::jit::attr::value);
-            break;
-        }
-        case c10::TypeKind::IntType:
-        {
-            type = 2;
-            int64_t i64 = value_node->i(torch::jit::attr::value);
-            if (i64 == std::numeric_limits<int64_t>::max()) i64 = INT_MAX;
-            if (i64 == std::numeric_limits<int64_t>::min()) i64 = INT_MIN;
-            i = (int)i64;
-            break;
-        }
-        case c10::TypeKind::FloatType:
-        {
-            type = 3;
-            f = (float)value_node->f(torch::jit::attr::value);
-            break;
-        }
-        case c10::TypeKind::StringType:
-        {
-            type = 4;
-            s = value_node->s(torch::jit::attr::value);
-            break;
-        }
-        case c10::TypeKind::DeviceObjType:
-        {
-            type = 4;
-            s = value_node->s(torch::jit::attr::value);
-            break;
-        }
-#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 9)
-        case c10::TypeKind::ComplexType:
-        {
-            type = 10;
-            c = std::complex<float>(value_node->c(torch::jit::attr::value));
-            break;
-        }
-#endif
-        case c10::TypeKind::TensorType:
-        {
-            at::Tensor t = value_node->t(torch::jit::attr::value);
-
-            if (t.dim() == 0 && t.numel() == 1)
-            {
-                if (t.scalar_type() == c10::ScalarType::Long)
-                {
-                    type = 2;
-                    int64_t i64 = t.item<int64_t>();
-                    if (i64 == std::numeric_limits<int64_t>::max()) i64 = INT_MAX;
-                    if (i64 == std::numeric_limits<int64_t>::min()) i64 = INT_MIN;
-                    i = (int)i64;
-                }
-                else if (t.scalar_type() == c10::ScalarType::Int)
-                {
-                    type = 2;
-                    i = t.item<int>();
-                }
-                else if (t.scalar_type() == c10::ScalarType::Double)
-                {
-                    type = 3;
-                    f = (float)t.item<double>();
-                }
-                else if (t.scalar_type() == c10::ScalarType::Float)
-                {
-                    type = 3;
-                    f = t.item<float>();
-                }
-                else if (t.scalar_type() == c10::ScalarType::ComplexDouble)
-                {
-                    type = 10;
-                    c = std::complex<float>(t.item<c10::complex<double> >());
-                }
-                else if (t.scalar_type() == c10::ScalarType::ComplexFloat)
-                {
-                    type = 10;
-                    c = std::complex<float>(t.item<c10::complex<float> >());
-                }
-                else
-                {
-                    fprintf(stderr, "unknown Parameter value kind %s of TensorType, t.dim = 0\n", value_node->kind().toDisplayString());
-                }
-            }
-            else
-            {
-                // constant tensor will become pnnx attribute node later
-                type = 8;
-            }
-
-            break;
-        }
-        case c10::TypeKind::ListType:
-        {
-            switch (value_node->output()->type()->containedTypes()[0]->kind())
-            {
-            case c10::TypeKind::IntType:
-            {
-                type = 5;
-                std::vector<int64_t> i64s = value_node->ival(torch::jit::attr::value).toIntVector();
-                for (auto i64 : i64s)
-                {
-                    if (i64 == std::numeric_limits<int64_t>::max()) i64 = INT_MAX;
-                    if (i64 == std::numeric_limits<int64_t>::min()) i64 = INT_MIN;
-                    ai.push_back(i64);
-                }
-                break;
-            }
-            case c10::TypeKind::FloatType:
-            {
-                type = 6;
-                std::vector<double> fs = value_node->ival(torch::jit::attr::value).toDoubleVector();
-                for (auto f : fs)
-                {
-                    af.push_back((float)f);
-                }
-                break;
-            }
-            default:
-            {
-                fprintf(stderr, "unknown Parameter value list element kind %s\n", c10::typeKindToString(value_node->output()->type()->containedTypes()[0]->kind()));
-                break;
-            }
-            }
-            break;
-        }
-        default:
-        {
-            fprintf(stderr, "unknown Parameter value kind %s\n", c10::typeKindToString(value_node->output()->type()->kind()));
-            break;
-        }
-        }
-    }
-    else if (value_node->kind() == c10::prim::ListConstruct)
-    {
-        switch (value_node->output()->type()->cast<c10::ListType>()->getElementType()->kind())
-        {
-        case c10::TypeKind::IntType:
-        {
-            type = 5;
-            for (const auto& x : value_node->inputs())
-            {
-                if (!x->node()->hasAttribute(torch::jit::attr::value))
-                {
-                    fprintf(stderr, "no attribute value in int list\n");
-                    ai.push_back(0);
-                    continue;
-                }
-
-                ai.push_back((int)x->node()->i(torch::jit::attr::value));
-            }
-            break;
-        }
-        case c10::TypeKind::FloatType:
-        {
-            type = 6;
-            for (const auto& x : value_node->inputs())
-            {
-                if (!x->node()->hasAttribute(torch::jit::attr::value))
-                {
-                    fprintf(stderr, "no attribute value in float list\n");
-                    af.push_back(0.f);
-                    continue;
-                }
-
-                af.push_back((float)x->node()->f(torch::jit::attr::value));
-            }
-            break;
-        }
-        case c10::TypeKind::StringType:
-        {
-            type = 7;
-            for (const auto& x : value_node->inputs())
-            {
-                if (!x->node()->hasAttribute(torch::jit::attr::value))
-                {
-                    fprintf(stderr, "no attribute value in string list\n");
-                    as.push_back("");
-                    continue;
-                }
-
-                as.push_back(x->node()->s(torch::jit::attr::value));
-            }
-            break;
-        }
-#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 9)
-        case c10::TypeKind::ComplexType:
-        {
-            type = 11;
-            for (const auto& x : value_node->inputs())
-            {
-                if (!x->node()->hasAttribute(torch::jit::attr::value))
-                {
-                    fprintf(stderr, "no attribute value in complex list\n");
-                    ac.push_back(std::complex<float>(0.f, 0.f));
-                    continue;
-                }
-
-                ac.push_back(std::complex<float>(x->node()->c(torch::jit::attr::value)));
-            }
-            break;
-        }
-#endif
-        default:
-        {
-            fprintf(stderr, "unknown Parameter value list element kind %s\n", c10::typeKindToString(value_node->output()->type()->cast<c10::ListType>()->getElementType()->kind()));
-            break;
-        }
-        }
-    }
-    else
-    {
-        fprintf(stderr, "unknown Parameter value_node kind %s\n", value_node->kind().toDisplayString());
-    }
-}
-
-Parameter::Parameter(const torch::jit::Value* value)
-    : Parameter(value->node())
-{
-}
-#endif // BUILD_PNNX
-
 bool operator==(const Parameter& lhs, const Parameter& rhs)
 {
     if (lhs.type != rhs.type)
@@ -443,59 +168,6 @@ bool operator==(const Parameter& lhs, const Parameter& rhs)
     return false;
 }
 
-#if BUILD_PNNX
-Attribute::Attribute(const at::Tensor& t)
-{
-    type = get_at_tensor_type(t.scalar_type());
-
-    const int ndim = (int)t.dim();
-
-    if (ndim == 0)
-    {
-        shape = {1};
-
-        data.resize(type_to_elemsize(type));
-
-        if (t.scalar_type() == c10::ScalarType::Long)
-        {
-            int64_t i = t.item<int64_t>();
-            memcpy((void*)data.data(), (const void*)&i, data.size());
-        }
-        else if (t.scalar_type() == c10::ScalarType::Int)
-        {
-            int i = t.item<int>();
-            memcpy((void*)data.data(), (const void*)&i, data.size());
-        }
-        else if (t.scalar_type() == c10::ScalarType::Double)
-        {
-            double f = t.item<double>();
-            memcpy((void*)data.data(), (const void*)&f, data.size());
-        }
-        else if (t.scalar_type() == c10::ScalarType::Float)
-        {
-            float f = t.item<float>();
-            memcpy((void*)data.data(), (const void*)&f, data.size());
-        }
-        else
-        {
-            fprintf(stderr, "unknown Attribute tensor scalar type %d\n", type);
-        }
-
-        return;
-    }
-
-    shape.resize(ndim);
-    for (int i = 0; i < ndim; i++)
-        shape[i] = t.size(i);
-
-    if (shape.size() > 0)
-    {
-        data.resize(elemcount() * type_to_elemsize(type));
-        memcpy((void*)data.data(), (const void*)t.cpu().contiguous().data_ptr(), data.size());
-    }
-}
-#endif // BUILD_PNNX
-
 Attribute::Attribute(const std::initializer_list<int>& _shape, const std::vector<float>& t)
 {
     type = 1;
@@ -3129,37 +2801,6 @@ Operator* Graph::new_operator_after(const std::string& type, const std::string&
     return op;
 }
 
-#if BUILD_PNNX
-Operand* Graph::new_operand(const torch::jit::Value* v)
-{
-    Operand* r = new Operand;
-    r->name = v->debugName();
-
-    r->type = -1;
-
-    auto pt = v->type()->cast<c10::TensorType>();
-    if (pt)
-    {
-        if (pt->scalarType().has_value() && pt->dim().has_value())
-        {
-            r->type = get_at_tensor_type(pt->scalarType().value());
-            const int ndim = (int)pt->dim().value();
-            r->shape.resize(ndim);
-            for (int i = 0; i < ndim; i++)
-            {
-                if (pt->sizes()[i].has_value())
-                    r->shape[i] = (int)pt->sizes()[i].value();
-                else
-                    r->shape[i] = -1;
-            }
-        }
-    }
-
-    operands.push_back(r);
-    return r;
-}
-#endif // BUILD_PNNX
-
 Operand* Graph::new_operand(const std::string& name)
 {
     Operand* r = new Operand;

tools/pnnx/src/ir.h

@@ -24,7 +24,7 @@
 #include <string>
 #include <vector>
 
-#if BUILD_PNNX
+#if BUILD_TORCH2PNNX
 namespace torch {
 namespace jit {
 struct Value;
@@ -34,7 +34,14 @@ struct Node;
 namespace at {
 class Tensor;
 }
-#endif // BUILD_PNNX
+#endif // BUILD_TORCH2PNNX
+
+#if BUILD_ONNX2PNNX
+namespace onnx {
+class TensorProto;
+class ValueInfoProto;
+} // namespace onnx
+#endif // BUILD_ONNX2PNNX
 
 namespace pnnx {
 
@@ -102,6 +109,17 @@ public:
         : type(5), ai(_ai)
     {
     }
+    Parameter(const std::vector<int64_t>& _ai)
+        : type(5)
+    {
+        for (const auto& x : _ai)
+        {
+            int64_t _l = x;
+            if (_l == std::numeric_limits<int64_t>::max()) _l = INT_MAX;
+            if (_l == std::numeric_limits<int64_t>::min()) _l = INT_MIN;
+            ai.push_back((int)_l);
+        }
+    }
     Parameter(const std::initializer_list<float>& _af)
         : type(6), af(_af)
     {
@@ -165,10 +183,10 @@ public:
             ac.push_back(std::complex<float>(x));
     }
 
-#if BUILD_PNNX
+#if BUILD_TORCH2PNNX
     Parameter(const torch::jit::Node* value_node);
     Parameter(const torch::jit::Value* value);
-#endif // BUILD_PNNX
+#endif // BUILD_TORCH2PNNX
 
     static Parameter parse_from_string(const std::string& value);
     static std::string encode_to_string(const Parameter& param);
@@ -200,9 +218,12 @@ public:
     {
     }
 
-#if BUILD_PNNX
+#if BUILD_TORCH2PNNX
     Attribute(const at::Tensor& t);
-#endif // BUILD_PNNX
+#endif
+#if BUILD_ONNX2PNNX
+    Attribute(const onnx::TensorProto& t);
+#endif
 
     Attribute(const std::initializer_list<int>& shape, const std::vector<float>& t);
 
@@ -299,9 +320,12 @@ public:
 
     Operator* new_operator_after(const std::string& type, const std::string& name, const Operator* cur);
 
-#if BUILD_PNNX
+#if BUILD_TORCH2PNNX
     Operand* new_operand(const torch::jit::Value* v);
 #endif
+#if BUILD_ONNX2PNNX
+    Operand* new_operand(const onnx::ValueInfoProto& value);
+#endif
 
     Operand* new_operand(const std::string& name);
 

tools/pnnx/src/load_torchscript.cpp

@@ -0,0 +1,552 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "load_torchscript.h"
+
+#if _WIN32
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
+#include <torch/script.h>
+#include <torch/csrc/api/include/torch/version.h>
+#ifdef PNNX_TORCHVISION
+namespace vision {
+int64_t cuda_version();
+} // namespace vision
+#endif
+
+#include "pass_level0.h"
+#include "pass_level1.h"
+
+namespace pnnx {
+
+static int get_at_tensor_type(const at::ScalarType& st)
+{
+    if (st == c10::ScalarType::Float) return 1;
+    if (st == c10::ScalarType::Double) return 2;
+    if (st == c10::ScalarType::Half) return 3;
+    if (st == c10::ScalarType::Int) return 4;
+    if (st == c10::ScalarType::QInt32) return 4;
+    if (st == c10::ScalarType::Long) return 5;
+    if (st == c10::ScalarType::Short) return 6;
+    if (st == c10::ScalarType::Char) return 7;
+    if (st == c10::ScalarType::QInt8) return 7;
+    if (st == c10::ScalarType::Byte) return 8;
+    if (st == c10::ScalarType::QUInt8) return 8;
+    if (st == c10::ScalarType::Bool) return 9;
+    if (st == c10::ScalarType::ComplexFloat) return 10;
+    if (st == c10::ScalarType::ComplexDouble) return 11;
+    if (st == c10::ScalarType::ComplexHalf) return 12;
+    return 0; // unknown type
+}
+
+static size_t type_to_elemsize(int type)
+{
+    if (type == 1) return 4;
+    if (type == 2) return 8;
+    if (type == 3) return 2;
+    if (type == 4) return 4;
+    if (type == 5) return 8;
+    if (type == 6) return 2;
+    if (type == 7) return 1;
+    if (type == 8) return 1;
+    if (type == 9) return 1;
+    if (type == 10) return 8;
+    if (type == 11) return 16;
+    if (type == 12) return 4;
+    return 0; // null
+}
+
+Parameter::Parameter(const torch::jit::Node* value_node)
+{
+    type = 0;
+
+    if (value_node->kind() == c10::prim::Constant)
+    {
+        if (value_node->output()->type()->kind() == c10::TypeKind::NoneType)
+        {
+            type = 0;
+            return;
+        }
+
+        if (!value_node->hasAttribute(torch::jit::attr::value))
+        {
+            fprintf(stderr, "no attribute value\n");
+            value_node->dump();
+            return;
+        }
+
+        switch (value_node->output()->type()->kind())
+        {
+        case c10::TypeKind::NoneType:
+        {
+            type = 0;
+            break;
+        }
+        case c10::TypeKind::BoolType:
+        {
+            type = 1;
+            b = value_node->i(torch::jit::attr::value);
+            break;
+        }
+        case c10::TypeKind::IntType:
+        {
+            type = 2;
+            int64_t i64 = value_node->i(torch::jit::attr::value);
+            if (i64 == std::numeric_limits<int64_t>::max()) i64 = INT_MAX;
+            if (i64 == std::numeric_limits<int64_t>::min()) i64 = INT_MIN;
+            i = (int)i64;
+            break;
+        }
+        case c10::TypeKind::FloatType:
+        {
+            type = 3;
+            f = (float)value_node->f(torch::jit::attr::value);
+            break;
+        }
+        case c10::TypeKind::StringType:
+        {
+            type = 4;
+            s = value_node->s(torch::jit::attr::value);
+            break;
+        }
+        case c10::TypeKind::DeviceObjType:
+        {
+            type = 4;
+            s = value_node->s(torch::jit::attr::value);
+            break;
+        }
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 9)
+        case c10::TypeKind::ComplexType:
+        {
+            type = 10;
+            c = std::complex<float>(value_node->c(torch::jit::attr::value));
+            break;
+        }
+#endif
+        case c10::TypeKind::TensorType:
+        {
+            at::Tensor t = value_node->t(torch::jit::attr::value);
+
+            if (t.dim() == 0 && t.numel() == 1)
+            {
+                if (t.scalar_type() == c10::ScalarType::Long)
+                {
+                    type = 2;
+                    int64_t i64 = t.item<int64_t>();
+                    if (i64 == std::numeric_limits<int64_t>::max()) i64 = INT_MAX;
+                    if (i64 == std::numeric_limits<int64_t>::min()) i64 = INT_MIN;
+                    i = (int)i64;
+                }
+                else if (t.scalar_type() == c10::ScalarType::Int)
+                {
+                    type = 2;
+                    i = t.item<int>();
+                }
+                else if (t.scalar_type() == c10::ScalarType::Double)
+                {
+                    type = 3;
+                    f = (float)t.item<double>();
+                }
+                else if (t.scalar_type() == c10::ScalarType::Float)
+                {
+                    type = 3;
+                    f = t.item<float>();
+                }
+                else if (t.scalar_type() == c10::ScalarType::ComplexDouble)
+                {
+                    type = 10;
+                    c = std::complex<float>(t.item<c10::complex<double> >());
+                }
+                else if (t.scalar_type() == c10::ScalarType::ComplexFloat)
+                {
+                    type = 10;
+                    c = std::complex<float>(t.item<c10::complex<float> >());
+                }
+                else
+                {
+                    fprintf(stderr, "unknown Parameter value kind %s of TensorType, t.dim = 0\n", value_node->kind().toDisplayString());
+                }
+            }
+            else
+            {
+                // constant tensor will become pnnx attribute node later
+                type = 8;
+            }
+
+            break;
+        }
+        case c10::TypeKind::ListType:
+        {
+            switch (value_node->output()->type()->containedTypes()[0]->kind())
+            {
+            case c10::TypeKind::IntType:
+            {
+                type = 5;
+                std::vector<int64_t> i64s = value_node->ival(torch::jit::attr::value).toIntVector();
+                for (auto i64 : i64s)
+                {
+                    if (i64 == std::numeric_limits<int64_t>::max()) i64 = INT_MAX;
+                    if (i64 == std::numeric_limits<int64_t>::min()) i64 = INT_MIN;
+                    ai.push_back(i64);
+                }
+                break;
+            }
+            case c10::TypeKind::FloatType:
+            {
+                type = 6;
+                std::vector<double> fs = value_node->ival(torch::jit::attr::value).toDoubleVector();
+                for (auto f : fs)
+                {
+                    af.push_back((float)f);
+                }
+                break;
+            }
+            default:
+            {
+                fprintf(stderr, "unknown Parameter value list element kind %s\n", c10::typeKindToString(value_node->output()->type()->containedTypes()[0]->kind()));
+                break;
+            }
+            }
+            break;
+        }
+        default:
+        {
+            fprintf(stderr, "unknown Parameter value kind %s\n", c10::typeKindToString(value_node->output()->type()->kind()));
+            break;
+        }
+        }
+    }
+    else if (value_node->kind() == c10::prim::ListConstruct)
+    {
+        switch (value_node->output()->type()->cast<c10::ListType>()->getElementType()->kind())
+        {
+        case c10::TypeKind::IntType:
+        {
+            type = 5;
+            for (const auto& x : value_node->inputs())
+            {
+                if (!x->node()->hasAttribute(torch::jit::attr::value))
+                {
+                    fprintf(stderr, "no attribute value in int list\n");
+                    ai.push_back(0);
+                    continue;
+                }
+
+                ai.push_back((int)x->node()->i(torch::jit::attr::value));
+            }
+            break;
+        }
+        case c10::TypeKind::FloatType:
+        {
+            type = 6;
+            for (const auto& x : value_node->inputs())
+            {
+                if (!x->node()->hasAttribute(torch::jit::attr::value))
+                {
+                    fprintf(stderr, "no attribute value in float list\n");
+                    af.push_back(0.f);
+                    continue;
+                }
+
+                af.push_back((float)x->node()->f(torch::jit::attr::value));
+            }
+            break;
+        }
+        case c10::TypeKind::StringType:
+        {
+            type = 7;
+            for (const auto& x : value_node->inputs())
+            {
+                if (!x->node()->hasAttribute(torch::jit::attr::value))
+                {
+                    fprintf(stderr, "no attribute value in string list\n");
+                    as.push_back("");
+                    continue;
+                }
+
+                as.push_back(x->node()->s(torch::jit::attr::value));
+            }
+            break;
+        }
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 9)
+        case c10::TypeKind::ComplexType:
+        {
+            type = 11;
+            for (const auto& x : value_node->inputs())
+            {
+                if (!x->node()->hasAttribute(torch::jit::attr::value))
+                {
+                    fprintf(stderr, "no attribute value in complex list\n");
+                    ac.push_back(std::complex<float>(0.f, 0.f));
+                    continue;
+                }
+
+                ac.push_back(std::complex<float>(x->node()->c(torch::jit::attr::value)));
+            }
+            break;
+        }
+#endif
+        default:
+        {
+            fprintf(stderr, "unknown Parameter value list element kind %s\n", c10::typeKindToString(value_node->output()->type()->cast<c10::ListType>()->getElementType()->kind()));
+            break;
+        }
+        }
+    }
+    else
+    {
+        fprintf(stderr, "unknown Parameter value_node kind %s\n", value_node->kind().toDisplayString());
+    }
+}
+
+Parameter::Parameter(const torch::jit::Value* value)
+    : Parameter(value->node())
+{
+}
+
+Attribute::Attribute(const at::Tensor& t)
+{
+    type = get_at_tensor_type(t.scalar_type());
+
+    const int ndim = (int)t.dim();
+
+    if (ndim == 0)
+    {
+        shape = {1};
+
+        data.resize(type_to_elemsize(type));
+
+        if (t.scalar_type() == c10::ScalarType::Long)
+        {
+            int64_t i = t.item<int64_t>();
+            memcpy((void*)data.data(), (const void*)&i, data.size());
+        }
+        else if (t.scalar_type() == c10::ScalarType::Int)
+        {
+            int i = t.item<int>();
+            memcpy((void*)data.data(), (const void*)&i, data.size());
+        }
+        else if (t.scalar_type() == c10::ScalarType::Double)
+        {
+            double f = t.item<double>();
+            memcpy((void*)data.data(), (const void*)&f, data.size());
+        }
+        else if (t.scalar_type() == c10::ScalarType::Float)
+        {
+            float f = t.item<float>();
+            memcpy((void*)data.data(), (const void*)&f, data.size());
+        }
+        else
+        {
+            fprintf(stderr, "unknown Attribute tensor scalar type %d\n", type);
+        }
+
+        return;
+    }
+
+    shape.resize(ndim);
+    for (int i = 0; i < ndim; i++)
+        shape[i] = t.size(i);
+
+    if (shape.size() > 0)
+    {
+        data.resize(elemcount() * type_to_elemsize(type));
+        memcpy((void*)data.data(), (const void*)t.cpu().contiguous().data_ptr(), data.size());
+    }
+}
+
+Operand* Graph::new_operand(const torch::jit::Value* v)
+{
+    // Operand* r = new Operand;
+    // r->name = v->debugName();
+
+    Operand* r = new_operand(v->debugName());
+
+    r->type = -1;
+
+    auto pt = v->type()->cast<c10::TensorType>();
+    if (pt)
+    {
+        if (pt->scalarType().has_value() && pt->dim().has_value())
+        {
+            r->type = get_at_tensor_type(pt->scalarType().value());
+            const int ndim = (int)pt->dim().value();
+            r->shape.resize(ndim);
+            for (int i = 0; i < ndim; i++)
+            {
+                if (pt->sizes()[i].has_value())
+                    r->shape[i] = (int)pt->sizes()[i].value();
+                else
+                    r->shape[i] = -1;
+            }
+        }
+    }
+
+    // operands.push_back(r);
+    return r;
+}
+
+static c10::ScalarType input_type_to_c10_ScalarType(const std::string& t)
+{
+    if (t == "c64") return torch::kComplexFloat;
+    if (t == "c32") return torch::kComplexHalf;
+    if (t == "c128") return torch::kComplexDouble;
+    if (t == "f32") return torch::kFloat32;
+    if (t == "f16") return torch::kFloat16;
+    if (t == "f64") return torch::kFloat64;
+    if (t == "i32") return torch::kInt32;
+    if (t == "i16") return torch::kInt16;
+    if (t == "i64") return torch::kInt64;
+    if (t == "i8") return torch::kInt8;
+    if (t == "u8") return torch::kUInt8;
+
+    fprintf(stderr, "unsupported type %s fallback to f32\n", t.c_str());
+    return torch::kFloat32;
+}
+
+const torch::jit::Node* find_node_by_kind(const std::shared_ptr<torch::jit::Graph>& graph, const std::string& kind)
+{
+    for (const auto& n : graph->nodes())
+    {
+        if (n->kind().toDisplayString() == kind)
+            return n;
+    }
+
+    return 0;
+}
+
+int load_torchscript(const std::string& ptpath, Graph& pnnx_graph,
+                     const std::string& device,
+                     const std::vector<std::vector<int64_t> >& input_shapes,
+                     const std::vector<std::string>& input_types,
+                     const std::vector<std::vector<int64_t> >& input_shapes2,
+                     const std::vector<std::string>& input_types2,
+                     const std::vector<std::string>& customop_modules,
+                     const std::vector<std::string>& module_operators,
+                     const std::string& foldable_constants_zippath,
+                     std::set<std::string>& foldable_constants)
+{
+#ifdef PNNX_TORCHVISION
+    // call some vision api to register vision ops  :P
+    (void)vision::cuda_version();
+#endif
+
+    for (auto m : customop_modules)
+    {
+        fprintf(stderr, "load custom module %s\n", m.c_str());
+#if _WIN32
+        HMODULE handle = LoadLibraryExA(m.c_str(), NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
+        if (!handle)
+        {
+            fprintf(stderr, "LoadLibraryExA %s failed %d\n", m.c_str(), GetLastError());
+        }
+#else
+        void* handle = dlopen(m.c_str(), RTLD_LAZY);
+        if (!handle)
+        {
+            fprintf(stderr, "dlopen %s failed %s\n", m.c_str(), dlerror());
+        }
+#endif
+    }
+
+    std::vector<at::Tensor> input_tensors;
+    for (size_t i = 0; i < input_shapes.size(); i++)
+    {
+        const std::vector<int64_t>& shape = input_shapes[i];
+        const std::string& type = input_types[i];
+
+        at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
+        if (device == "gpu")
+            t = t.cuda();
+
+        input_tensors.push_back(t);
+    }
+
+    std::vector<at::Tensor> input_tensors2;
+    for (size_t i = 0; i < input_shapes2.size(); i++)
+    {
+        const std::vector<int64_t>& shape = input_shapes2[i];
+        const std::string& type = input_types2[i];
+
+        at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
+        if (device == "gpu")
+            t = t.cuda();
+
+        input_tensors2.push_back(t);
+    }
+
+    torch::jit::Module mod;
+
+    try
+    {
+        mod = torch::jit::load(ptpath, (device == "gpu") ? c10::kCUDA : c10::kCPU);
+    }
+    catch (const c10::Error& e)
+    {
+        fprintf(stderr, "Load torchscript failed: %s\n", e.what());
+
+        fprintf(stderr, "Please export model to torchscript as follows\n");
+        fprintf(stderr, "------------------------------------------\n");
+        fprintf(stderr, "import torch\n");
+        fprintf(stderr, "import torchvision.models as models\n\n");
+        fprintf(stderr, "net = models.resnet18(pretrained=True)\n");
+        fprintf(stderr, "net = net.eval()\n\n");
+        fprintf(stderr, "x = torch.rand(1, 3, 224, 224)\n");
+        fprintf(stderr, "mod = torch.jit.trace(net, x)\n");
+        fprintf(stderr, "mod.save(\"resnet18.pt\")\n");
+        fprintf(stderr, "------------------------------------------\n");
+
+        return -1;
+    }
+
+    mod.eval();
+
+    //     mod.dump(true, false, false);
+    //     mod.dump(true, true, true);
+
+    auto method = mod.find_method("forward");
+    if (!method)
+    {
+        auto methods = mod.get_methods();
+        if (methods.empty())
+        {
+            fprintf(stderr, "No method in torchscript\n");
+            return -1;
+        }
+
+        method = methods[0];
+        fprintf(stderr, "Use method %s as the entrypoint instead of forward\n", method->name().c_str());
+    }
+
+    auto g = method->graph();
+
+    // g->dump();
+
+    fprintf(stderr, "############# pass_level0\n");
+
+    pnnx::pass_level0(mod, g, input_tensors, input_tensors2, module_operators, ptpath, device, foldable_constants, foldable_constants_zippath);
+
+    // g->dump();
+
+    fprintf(stderr, "############# pass_level1\n");
+
+    pnnx::pass_level1(mod, g, module_operators, pnnx_graph);
+
+    return 0;
+}
+
+} // namespace pnnx

tools/pnnx/src/load_torchscript.h

@@ -0,0 +1,35 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#ifndef PNNX_LOAD_TORCHSCRIPT_H
+#define PNNX_LOAD_TORCHSCRIPT_H
+
+#include "ir.h"
+
+namespace pnnx {
+
+int load_torchscript(const std::string& ptpath, Graph& g,
+                     const std::string& device,
+                     const std::vector<std::vector<int64_t> >& input_shapes,
+                     const std::vector<std::string>& input_types,
+                     const std::vector<std::vector<int64_t> >& input_shapes2,
+                     const std::vector<std::string>& input_types2,
+                     const std::vector<std::string>& customop_modules,
+                     const std::vector<std::string>& module_operators,
+                     const std::string& foldable_constants_zippath,
+                     std::set<std::string>& foldable_constants);
+
+} // namespace pnnx
+
+#endif // PNNX_LOAD_TORCHSCRIPT_H

tools/pnnx/src/main.cpp

@@ -13,31 +13,25 @@
 // specific language governing permissions and limitations under the License.
 
 #include <stdio.h>
+#include <string.h>
 
-#if _WIN32
-#include <windows.h>
-#else
-#include <dlfcn.h>
-#endif
-
+#include <algorithm>
 #include <string>
 #include <vector>
 
-#include <torch/script.h>
-
-#ifdef PNNX_TORCHVISION
-// register torchvision ops via including headers
-#include <torchvision/vision.h>
-#endif
-
 #include "ir.h"
-#include "pass_level0.h"
-#include "pass_level1.h"
 #include "pass_level2.h"
 #include "pass_level3.h"
 #include "pass_level4.h"
 #include "pass_level5.h"
 
+#if BUILD_TORCH2PNNX
+#include "load_torchscript.h"
+#endif
+#if BUILD_ONNX2PNNX
+#include "load_onnx.h"
+#endif
+
 #include "pass_ncnn.h"
 #include "save_ncnn.h"
 
@@ -160,24 +154,6 @@ static void print_shape_list(const std::vector<std::vector<int64_t> >& shapes, c
     }
 }
 
-static c10::ScalarType input_type_to_c10_ScalarType(const std::string& t)
-{
-    if (t == "c64") return torch::kComplexFloat;
-    if (t == "c32") return torch::kComplexHalf;
-    if (t == "c128") return torch::kComplexDouble;
-    if (t == "f32") return torch::kFloat32;
-    if (t == "f16") return torch::kFloat16;
-    if (t == "f64") return torch::kFloat64;
-    if (t == "i32") return torch::kInt32;
-    if (t == "i16") return torch::kInt16;
-    if (t == "i64") return torch::kInt64;
-    if (t == "i8") return torch::kInt8;
-    if (t == "u8") return torch::kUInt8;
-
-    fprintf(stderr, "unsupported type %s fallback to f32\n", t.c_str());
-    return torch::kFloat32;
-}
-
 static void show_usage()
 {
     fprintf(stderr, "Usage: pnnx [model.pt] [(key=value)...]\n");
@@ -314,114 +290,17 @@ int main(int argc, char** argv)
         fprintf(stderr, "\n");
     }
 
-#ifdef PNNX_TORCHVISION
-    // call some vision api to register vision ops  :P
-    (void)vision::cuda_version();
-#endif
-
-    for (auto m : customop_modules)
-    {
-        fprintf(stderr, "load custom module %s\n", m.c_str());
-#if _WIN32
-        HMODULE handle = LoadLibraryExA(m.c_str(), NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
-        if (!handle)
-        {
-            fprintf(stderr, "LoadLibraryExA %s failed %d\n", m.c_str(), GetLastError());
-        }
-#else
-        void* handle = dlopen(m.c_str(), RTLD_LAZY);
-        if (!handle)
-        {
-            fprintf(stderr, "dlopen %s failed %s\n", m.c_str(), dlerror());
-        }
-#endif
-    }
-
-    std::vector<at::Tensor> input_tensors;
-    for (size_t i = 0; i < input_shapes.size(); i++)
-    {
-        const std::vector<int64_t>& shape = input_shapes[i];
-        const std::string& type = input_types[i];
-
-        at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
-        if (device == "gpu")
-            t = t.cuda();
-
-        input_tensors.push_back(t);
-    }
-
-    std::vector<at::Tensor> input_tensors2;
-    for (size_t i = 0; i < input_shapes2.size(); i++)
-    {
-        const std::vector<int64_t>& shape = input_shapes2[i];
-        const std::string& type = input_types2[i];
-
-        at::Tensor t = torch::ones(shape, input_type_to_c10_ScalarType(type));
-        if (device == "gpu")
-            t = t.cuda();
-
-        input_tensors2.push_back(t);
-    }
-
-    torch::jit::Module mod;
-
-    try
-    {
-        mod = torch::jit::load(ptpath, (device == "gpu") ? c10::kCUDA : c10::kCPU);
-    }
-    catch (const c10::Error& e)
-    {
-        fprintf(stderr, "Load torchscript failed: %s\n", e.what());
-
-        fprintf(stderr, "Please export model to torchscript as follows\n");
-        fprintf(stderr, "------------------------------------------\n");
-        fprintf(stderr, "import torch\n");
-        fprintf(stderr, "import torchvision.models as models\n\n");
-        fprintf(stderr, "net = models.resnet18(pretrained=True)\n");
-        fprintf(stderr, "net = net.eval()\n\n");
-        fprintf(stderr, "x = torch.rand(1, 3, 224, 224)\n");
-        fprintf(stderr, "mod = torch.jit.trace(net, x)\n");
-        fprintf(stderr, "mod.save(\"resnet18.pt\")\n");
-        fprintf(stderr, "------------------------------------------\n");
-
-        return -1;
-    }
-
-    mod.eval();
-
-    //     mod.dump(true, false, false);
-    //     mod.dump(true, true, true);
-
-    auto method = mod.find_method("forward");
-    if (!method)
-    {
-        auto methods = mod.get_methods();
-        if (methods.empty())
-        {
-            fprintf(stderr, "No method in torchscript\n");
-            return -1;
-        }
-
-        method = methods[0];
-        fprintf(stderr, "Use method %s as the entrypoint instead of forward\n", method->name().c_str());
-    }
-
-    auto g = method->graph();
-
-    //     g->dump();
-
-    fprintf(stderr, "############# pass_level0\n");
-
     std::set<std::string> foldable_constants;
     std::string foldable_constants_zippath = ptbase + ".foldable_constants.zip";
-    pnnx::pass_level0(mod, g, input_tensors, input_tensors2, module_operators, ptpath, device, foldable_constants, foldable_constants_zippath);
-
-    //     g->dump();
-
-    fprintf(stderr, "############# pass_level1\n");
 
     pnnx::Graph pnnx_graph;
-    pnnx::pass_level1(mod, g, module_operators, pnnx_graph);
+    load_torchscript(ptpath, pnnx_graph,
+                     device, input_shapes, input_types,
+                     input_shapes2, input_types2,
+                     customop_modules, module_operators,
+                     foldable_constants_zippath, foldable_constants);
+
+    // load_onnx(ptpath.c_str(), pnnx_graph);
 
     //     g->dump();
 

tools/pnnx/src/onnx.proto

@@ -3,8 +3,8 @@
 //
 
 
-// Copyright (c) ONNX Project Contributors.
-// Licensed under the MIT license.
+// SPDX-License-Identifier: Apache-2.0
+
 
 syntax = "proto2";
 
@@ -20,23 +20,16 @@ package onnx;
 //
 // This document describes the syntax of models and their computation graphs,
 // as well as the standard data types. Together, they are referred to as the ONNX
-// Intermediate Representation, or 'IR' for short. 
+// Intermediate Representation, or 'IR' for short.
 //
 // The normative semantic specification of the ONNX IR is found in docs/IR.md.
 // Definitions of the built-in neural network operators may be found in docs/Operators.md.
 
 // Notes
 //
-// Release
-//
-// We are still in the very early stage of defining ONNX. The current
-// version of ONNX is a starting point. While we are actively working
-// towards a complete spec, we would like to get the community involved
-// by sharing our working version of ONNX.
-//
 // Protobuf compatibility
-// 
-// To simplify framework compatibility, ONNX is defined using the subset of protobuf 
+//
+// To simplify framework compatibility, ONNX is defined using the subset of protobuf
 // that is compatible with both protobuf v2 and v3. This means that we do not use any
 // protobuf features that are only available in one of the two versions.
 //
@@ -60,8 +53,8 @@ enum Version {
   _START_VERSION = 0;
   // The version field is always serialized and we will use it to store the
   // version that the  graph is generated from. This helps us set up version
-  // control. 
-  // For the IR, we are using simple numbers starting with with 0x00000001, 
+  // control.
+  // For the IR, we are using simple numbers starting with 0x00000001,
   // which was the version we published on Oct 10, 2017.
   IR_VERSION_2017_10_10 = 0x0000000000000001;
 
@@ -84,7 +77,36 @@ enum Version {
   // IR VERSION 5 published on March 18, 2019
   // - Add message TensorAnnotation.
   // - Add quantization annotation in GraphProto to map tensor with its scale and zero point quantization parameters.
-  IR_VERSION = 0x0000000000000005;
+  IR_VERSION_2019_3_18 = 0x0000000000000005;
+
+  // IR VERSION 6 published on Sep 19, 2019
+  // - Add support for sparse tensor constants stored in model.
+  //   - Add message SparseTensorProto
+  //   - Add sparse initializers
+  IR_VERSION_2019_9_19 = 0x0000000000000006;
+
+  // IR VERSION 7 published on May 8, 2020
+  // - Add support to allow function body graph to rely on multiple external opreator sets.
+  // - Add a list to promote inference graph's initializers to global and
+  //   mutable variables. Global variables are visible in all graphs of the
+  //   stored models.
+  // - Add message TrainingInfoProto to store initialization
+  //   method and training algorithm. The execution of TrainingInfoProto
+  //   can modify the values of mutable variables.
+  // - Implicitly add inference graph into each TrainingInfoProto's algorithm.
+  IR_VERSION_2020_5_8 = 0x0000000000000007;
+
+  // IR VERSION 8 published on July 30, 2021
+  // Introduce TypeProto.SparseTensor
+  // Introduce TypeProto.Optional
+  // Added a list of FunctionProtos local to the model
+  // Deprecated since_version and operator status from FunctionProto
+  IR_VERSION_2021_7_30 = 0x0000000000000008;
+
+  // IR VERSION 9 published on May 5, 2023
+  // Added AttributeProto to FunctionProto so that default attribute values can be set.
+  // Added FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ.
+  IR_VERSION = 0x0000000000000009;
 }
 
 // Attributes
@@ -94,6 +116,8 @@ enum Version {
 // An AttributeProto MUST contain the name field, and *only one* of the
 // following content fields, effectively enforcing a C/C++ union equivalent.
 message AttributeProto {
+  reserved 12, 16 to 19;
+  reserved "v";
 
   // Note: this enum is structurally identical to the OpSchema::AttrType
   // enum defined in schema.h.  If you rev one, you likely need to rev the other.
@@ -104,17 +128,21 @@ message AttributeProto {
     STRING = 3;
     TENSOR = 4;
     GRAPH = 5;
+    SPARSE_TENSOR = 11;
+    TYPE_PROTO = 13;
 
     FLOATS = 6;
     INTS = 7;
     STRINGS = 8;
     TENSORS = 9;
     GRAPHS = 10;
+    SPARSE_TENSORS = 12;
+    TYPE_PROTOS = 14;
   }
 
   // The name field MUST be present for this version of the IR.
   optional string name = 1;           // namespace Attribute
- 
+
   // if ref_attr_name is not empty, ref_attr_name is the attribute name in parent function.
   // In this case, this AttributeProto does not contain data, and it's a reference of attribute
   // in parent scope.
@@ -126,7 +154,7 @@ message AttributeProto {
 
   // The type field MUST be present for this version of the IR.
   // For 0.0.1 versions of the IR, this field was not defined, and
-  // implementations needed to use has_field hueristics to determine
+  // implementations needed to use has_field heuristics to determine
   // which value field was in use.  For IR_VERSION 0.0.2 or later, this
   // field MUST be set and match the f|i|s|t|... field in use.  This
   // change was made to accommodate proto3 implementations.
@@ -138,14 +166,18 @@ message AttributeProto {
   optional bytes s = 4;               // UTF-8 string
   optional TensorProto t = 5;         // tensor value
   optional GraphProto g = 6;          // graph
+  optional SparseTensorProto sparse_tensor = 22;  // sparse tensor value
   // Do not use field below, it's deprecated.
   // optional ValueProto v = 12;         // value - subsumes everything but graph
+  optional TypeProto tp = 14;          // type proto
 
   repeated float floats = 7;          // list of floats
   repeated int64 ints = 8;            // list of ints
   repeated bytes strings = 9;         // list of UTF-8 strings
   repeated TensorProto tensors = 10;  // list of tensors
   repeated GraphProto graphs = 11;    // list of graph
+  repeated SparseTensorProto sparse_tensors = 23; // list of sparse tensors
+  repeated TypeProto type_protos = 15;// list of type protos
 }
 
 // Defines information on value, including the name, the type, and
@@ -153,7 +185,8 @@ message AttributeProto {
 message ValueInfoProto {
   // This field MUST be present in this version of the IR.
   optional string name = 1;     // namespace Value
-  // This field MUST be present in this version of the IR.
+  // This field MUST be present in this version of the IR for
+  // inputs and outputs of the top-level graph.
   optional TypeProto type = 2;
   // A human-readable documentation for this value. Markdown is allowed.
   optional string doc_string = 3;
@@ -164,7 +197,7 @@ message ValueInfoProto {
 // Computation graphs are made up of a DAG of nodes, which represent what is
 // commonly called a "layer" or "pipeline stage" in machine learning frameworks.
 //
-// For example, it can be a node of type "Conv" that takes in an image, a filter 
+// For example, it can be a node of type "Conv" that takes in an image, a filter
 // tensor and a bias tensor, and produces the convolved output.
 message NodeProto {
   repeated string input = 1;    // namespace Value
@@ -186,12 +219,130 @@ message NodeProto {
   optional string doc_string = 6;
 }
 
+// Training information
+// TrainingInfoProto stores information for training a model.
+// In particular, this defines two functionalities: an initialization-step
+// and a training-algorithm-step. Initialization resets the model
+// back to its original state as if no training has been performed.
+// Training algorithm improves the model based on input data.
+//
+// The semantics of the initialization-step is that the initializers
+// in ModelProto.graph and in TrainingInfoProto.algorithm are first
+// initialized as specified by the initializers in the graph, and then
+// updated by the "initialization_binding" in every instance in
+// ModelProto.training_info.
+//
+// The field "algorithm" defines a computation graph which represents a
+// training algorithm's step. After the execution of a
+// TrainingInfoProto.algorithm, the initializers specified by "update_binding"
+// may be immediately updated. If the targeted training algorithm contains
+// consecutive update steps (such as block coordinate descent methods),
+// the user needs to create a TrainingInfoProto for each step.
+message TrainingInfoProto {
+  // This field describes a graph to compute the initial tensors
+  // upon starting the training process. Initialization graph has no input
+  // and can have multiple outputs. Usually, trainable tensors in neural
+  // networks are randomly initialized. To achieve that, for each tensor,
+  // the user can put a random number operator such as RandomNormal or
+  // RandomUniform in TrainingInfoProto.initialization.node and assign its
+  // random output to the specific tensor using "initialization_binding".
+  // This graph can also set the initializers in "algorithm" in the same
+  // TrainingInfoProto; a use case is resetting the number of training
+  // iteration to zero.
+  //
+  // By default, this field is an empty graph and its evaluation does not
+  // produce any output. Thus, no initializer would be changed by default.
+  optional GraphProto initialization = 1;
+
+  // This field represents a training algorithm step. Given required inputs,
+  // it computes outputs to update initializers in its own or inference graph's
+  // initializer lists. In general, this field contains loss node, gradient node,
+  // optimizer node, increment of iteration count.
+  //
+  // An execution of the training algorithm step is performed by executing the
+  // graph obtained by combining the inference graph (namely "ModelProto.graph")
+  // and the "algorithm" graph. That is, the actual
+  // input/initializer/output/node/value_info/sparse_initializer list of
+  // the training graph is the concatenation of
+  // "ModelProto.graph.input/initializer/output/node/value_info/sparse_initializer"
+  // and "algorithm.input/initializer/output/node/value_info/sparse_initializer"
+  // in that order. This combined graph must satisfy the normal ONNX conditions.
+  // Now, let's provide a visualization of graph combination for clarity.
+  // Let the inference graph (i.e., "ModelProto.graph") be
+  //    tensor_a, tensor_b -> MatMul -> tensor_c -> Sigmoid -> tensor_d
+  // and the "algorithm" graph be
+  //    tensor_d -> Add -> tensor_e
+  // The combination process results
+  //    tensor_a, tensor_b -> MatMul -> tensor_c -> Sigmoid -> tensor_d -> Add -> tensor_e
+  //
+  // Notice that an input of a node in the "algorithm" graph may reference the
+  // output of a node in the inference graph (but not the other way round). Also, inference
+  // node cannot reference inputs of "algorithm". With these restrictions, inference graph
+  // can always be run independently without training information.
+  //
+  // By default, this field is an empty graph and its evaluation does not
+  // produce any output. Evaluating the default training step never
+  // update any initializers.
+  optional GraphProto algorithm = 2;
+
+  // This field specifies the bindings from the outputs of "initialization" to
+  // some initializers in "ModelProto.graph.initializer" and
+  // the "algorithm.initializer" in the same TrainingInfoProto.
+  // See "update_binding" below for details.
+  //
+  // By default, this field is empty and no initializer would be changed
+  // by the execution of "initialization".
+  repeated StringStringEntryProto initialization_binding = 3;
+
+  // Gradient-based training is usually an iterative procedure. In one gradient
+  // descent iteration, we apply
+  //
+  // x = x - r * g
+  //
+  // where "x" is the optimized tensor, "r" stands for learning rate, and "g" is
+  // gradient of "x" with respect to a chosen loss. To avoid adding assignments
+  // into the training graph, we split the update equation into
+  //
+  // y = x - r * g
+  // x = y
+  //
+  // The user needs to save "y = x - r * g" into TrainingInfoProto.algorithm. To
+  // tell that "y" should be assigned to "x", the field "update_binding" may
+  // contain a key-value pair of strings, "x" (key of StringStringEntryProto)
+  // and "y" (value of StringStringEntryProto).
+  // For a neural network with multiple trainable (mutable) tensors, there can
+  // be multiple key-value pairs in "update_binding".
+  //
+  // The initializers appears as keys in "update_binding" are considered
+  // mutable variables. This implies some behaviors
+  // as described below.
+  //
+  //  1. We have only unique keys in all "update_binding"s so that two
+  //     variables may not have the same name. This ensures that one
+  //     variable is assigned up to once.
+  //  2. The keys must appear in names of "ModelProto.graph.initializer" or
+  //     "TrainingInfoProto.algorithm.initializer".
+  //  3. The values must be output names of "algorithm" or "ModelProto.graph.output".
+  //  4. Mutable variables are initialized to the value specified by the
+  //     corresponding initializer, and then potentially updated by
+  //     "initializer_binding"s and "update_binding"s in "TrainingInfoProto"s.
+  //
+  // This field usually contains names of trainable tensors
+  // (in ModelProto.graph), optimizer states such as momentums in advanced
+  // stochastic gradient methods (in TrainingInfoProto.graph),
+  // and number of training iterations (in TrainingInfoProto.graph).
+  //
+  // By default, this field is empty and no initializer would be changed
+  // by the execution of "algorithm".
+  repeated StringStringEntryProto update_binding = 4;
+}
+
 // Models
 //
 // ModelProto is a top-level file/container format for bundling a ML model and
 // associating its computation graph with metadata.
 //
-// The semantics of the model are described by the associated GraphProto.
+// The semantics of the model are described by the associated GraphProto's.
 message ModelProto {
   // The version of the IR this model targets. See Version enum above.
   // This field MUST be present.
@@ -236,13 +387,42 @@ message ModelProto {
 
   // Named metadata values; keys should be distinct.
   repeated StringStringEntryProto metadata_props = 14;
+
+  // Training-specific information. Sequentially executing all stored
+  // `TrainingInfoProto.algorithm`s and assigning their outputs following
+  // the corresponding `TrainingInfoProto.update_binding`s is one training
+  // iteration. Similarly, to initialize the model
+  // (as if training hasn't happened), the user should sequentially execute
+  // all stored `TrainingInfoProto.initialization`s and assigns their outputs
+  // using `TrainingInfoProto.initialization_binding`s.
+  //
+  // If this field is empty, the training behavior of the model is undefined.
+  repeated TrainingInfoProto training_info = 20;
+
+  // A list of function protos local to the model.
+  //
+  // Name of the function "FunctionProto.name" should be unique within the domain "FunctionProto.domain".
+  // In case of any conflicts the behavior (whether the model local functions are given higher priority,
+  // or standard operator sets are given higher priotity or this is treated as error) is defined by
+  // the runtimes.
+  //
+  // The operator sets imported by FunctionProto should be compatible with the ones
+  // imported by ModelProto and other model local FunctionProtos.
+  // Example, if same operator set say 'A' is imported by a FunctionProto and ModelProto
+  // or by 2 FunctionProtos then versions for the operator set may be different but,
+  // the operator schema returned for op_type, domain, version combination
+  // for both the versions should be same for every node in the function body.
+  //
+  // One FunctionProto can reference other FunctionProto in the model, however, recursive reference
+  // is not allowed.
+  repeated FunctionProto functions = 25;
 };
 
 // StringStringEntryProto follows the pattern for cross-proto-version maps.
 // See https://developers.google.com/protocol-buffers/docs/proto3#maps
 message StringStringEntryProto {
   optional string key = 1;
-  optional string value= 2;
+  optional string value = 2;
 };
 
 message TensorAnnotation {
@@ -258,7 +438,7 @@ message TensorAnnotation {
 
 // Graphs
 //
-// A graph defines the computational logic of a model and is comprised of a parameterized 
+// A graph defines the computational logic of a model and is comprised of a parameterized
 // list of nodes that form a directed acyclic graph based on their inputs and outputs.
 // This is the equivalent of the "network" or "graph" in many deep learning
 // frameworks.
@@ -270,10 +450,14 @@ message GraphProto {
   optional string name = 2;   // namespace Graph
 
   // A list of named tensor values, used to specify constant inputs of the graph.
-  // Each TensorProto entry must have a distinct name (within the list) that
-  // MAY also appear in the input list.
+  // Each initializer (both TensorProto as well SparseTensorProto) MUST have a name.
+  // The name MUST be unique across both initializer and sparse_initializer,
+  // but the name MAY also appear in the input list.
   repeated TensorProto initializer = 5;
 
+  // Initializers (see above) stored in sparse format.
+  repeated SparseTensorProto sparse_initializer = 15;
+
   // A human-readable documentation for this graph. Markdown is allowed.
   optional string doc_string = 10;
 
@@ -291,13 +475,8 @@ message GraphProto {
   // which means, tensor 'a_scale' and tensor 'a_zero_point' are scale and zero point of tensor 'a' in the model.
   repeated TensorAnnotation quantization_annotation = 14;
 
-  // DO NOT USE the following fields, they were deprecated from earlier versions.
-  // repeated string input = 3;
-  // repeated string output = 4;
-  // optional int64 ir_version = 6;
-  // optional int64 producer_version = 7;
-  // optional string producer_tag = 8;
-  // optional string domain = 9;
+  reserved 3, 4, 6 to 9;
+  reserved "ir_version", "producer_version", "producer_tag", "domain";
 }
 
 // Tensors
@@ -332,6 +511,17 @@ message TensorProto {
     // This format has 1 sign bit, 8 exponent bits, and 7 mantissa bits.
     BFLOAT16 = 16;
 
+    // Non-IEEE floating-point format based on papers
+    // FP8 Formats for Deep Learning, https://arxiv.org/abs/2209.05433,
+    // 8-bit Numerical Formats For Deep Neural Networks, https://arxiv.org/pdf/2206.02915.pdf.
+    // Operators supported FP8 are Cast, CastLike, QuantizeLinear, DequantizeLinear.
+    // The computation usually happens inside a block quantize / dequantize
+    // fused by the runtime.
+    FLOAT8E4M3FN = 17;    // float 8, mostly used for coefficients, supports nan, not inf
+    FLOAT8E4M3FNUZ = 18;  // float 8, mostly used for coefficients, supports nan, not inf, no negative zero
+    FLOAT8E5M2 = 19;      // follows IEEE 754, supports nan, inf, mostly used for gradients
+    FLOAT8E5M2FNUZ = 20;  // follows IEEE 754, supports nan, inf, mostly used for gradients, no negative zero
+
     // Future extensions go here.
   }
 
@@ -359,17 +549,17 @@ message TensorProto {
   // For float and complex64 values
   // Complex64 tensors are encoded as a single array of floats,
   // with the real components appearing in odd numbered positions,
-  // and the corresponding imaginary component apparing in the
+  // and the corresponding imaginary component appearing in the
   // subsequent even numbered position. (e.g., [1.0 + 2.0i, 3.0 + 4.0i]
   // is encoded as [1.0, 2.0 ,3.0 ,4.0]
   // When this field is present, the data_type field MUST be FLOAT or COMPLEX64.
   repeated float float_data = 4 [packed = true];
 
-  // For int32, uint8, int8, uint16, int16, bool, and float16 values
-  // float16 values must be bit-wise converted to an uint16_t prior
+  // For int32, uint8, int8, uint16, int16, bool, float8, and float16 values
+  // float16 and float8 values must be bit-wise converted to an uint16_t prior
   // to writing to the buffer.
   // When this field is present, the data_type field MUST be
-  // INT32, INT16, INT8, UINT16, UINT8, BOOL, or FLOAT16
+  // INT32, INT16, INT8, UINT16, UINT8, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ
   repeated int32 int32_data = 5 [packed = true];
 
   // For strings.
@@ -431,7 +621,7 @@ message TensorProto {
   // For double
   // Complex128 tensors are encoded as a single array of doubles,
   // with the real components appearing in odd numbered positions,
-  // and the corresponding imaginary component apparing in the
+  // and the corresponding imaginary component appearing in the
   // subsequent even numbered position. (e.g., [1.0 + 2.0i, 3.0 + 4.0i]
   // is encoded as [1.0, 2.0 ,3.0 ,4.0]
   // When this field is present, the data_type field MUST be DOUBLE or COMPLEX128
@@ -443,6 +633,30 @@ message TensorProto {
   repeated uint64 uint64_data = 11 [packed = true];
 }
 
+// A serialized sparse-tensor value
+message SparseTensorProto {
+  // The sequence of non-default values are encoded as a tensor of shape [NNZ].
+  // The default-value is zero for numeric tensors, and empty-string for string tensors.
+  // values must have a non-empty name present which serves as a name for SparseTensorProto
+  // when used in sparse_initializer list.
+  optional TensorProto values = 1;
+
+  // The indices of the non-default values, which may be stored in one of two formats.
+  // (a) Indices can be a tensor of shape [NNZ, rank] with the [i,j]-th value
+  // corresponding to the j-th index of the i-th value (in the values tensor).
+  // (b) Indices can be a tensor of shape [NNZ], in which case the i-th value
+  // must be the linearized-index of the i-th value (in the values tensor).
+  // The linearized-index can be converted into an index tuple (k_1,...,k_rank)
+  // using the shape provided below.
+  // The indices must appear in ascending order without duplication.
+  // In the first format, the ordering is lexicographic-ordering:
+  // e.g., index-value [1,4] must appear before [2,1]
+  optional TensorProto indices = 2;
+
+  // The shape of the underlying dense-tensor: [dim_1, dim_2, ... dim_rank]
+  repeated int64 dims = 3;
+}
+
 // Defines a tensor shape. A dimension can be either an integer value
 // or a symbolic variable. A symbolic variable represents an unknown
 // dimension.
@@ -455,7 +669,7 @@ message TensorShapeProto {
     // Standard denotation can optionally be used to denote tensor
     // dimensions with standard semantic descriptions to ensure
     // that operations are applied to the correct axis of a tensor.
-    // Refer to https://github.com/onnx/onnx/blob/master/docs/DimensionDenotation.md#denotation-definition
+    // Refer to https://github.com/onnx/onnx/blob/main/docs/DimensionDenotation.md#denotation-definition
     // for pre-defined dimension denotations.
     optional string denotation = 3;
   };
@@ -475,16 +689,68 @@ message TypeProto {
     optional TensorShapeProto shape = 2;
   }
 
+  // repeated T
+  message Sequence {
+    // The type and optional shape of each element of the sequence.
+    // This field MUST be present for this version of the IR.
+    optional TypeProto elem_type = 1;
+  };
+
+  // map<K,V>
+  message Map {
+    // This field MUST have a valid TensorProto.DataType value
+    // This field MUST be present for this version of the IR.
+    // This field MUST refer to an integral type ([U]INT{8|16|32|64}) or STRING
+    optional int32 key_type = 1;
+    // This field MUST be present for this version of the IR.
+    optional TypeProto value_type = 2;
+  };
+
+  // wrapper for Tensor, Sequence, or Map
+  message Optional {
+    // The type and optional shape of the element wrapped.
+    // This field MUST be present for this version of the IR.
+    // Possible values correspond to OptionalProto.DataType enum
+    optional TypeProto elem_type = 1;
+  };
+
+
+  message SparseTensor {
+    // This field MUST NOT have the value of UNDEFINED
+    // This field MUST have a valid TensorProto.DataType value
+    // This field MUST be present for this version of the IR.
+    optional int32 elem_type = 1;
+    optional TensorShapeProto shape = 2;
+  }
+
 
   oneof value {
     // The type of a tensor.
     Tensor tensor_type = 1;
 
+    // NOTE:  DNN-only implementations of ONNX MAY elect to not support non-tensor values
+    //        as input and output to graphs and nodes. These types are needed to naturally
+    //        support classical ML operators.  DNN operators SHOULD restrict their input
+    //        and output types to tensors.
+
+    // The type of a sequence.
+    Sequence sequence_type = 4;
+
+    // The type of a map.
+    Map map_type = 5;
+
+    // The type of an optional.
+    Optional optional_type = 9;
+
+
+    // Type of the sparse tensor
+    SparseTensor sparse_tensor_type = 8;
+
   }
 
-  // An optional denotation can be used to denote the whole 
-  // type with a standard semantic description as to what is 
-  // stored inside. Refer to https://github.com/onnx/onnx/blob/master/docs/TypeDenotation.md#type-denotation-definition
+  // An optional denotation can be used to denote the whole
+  // type with a standard semantic description as to what is
+  // stored inside. Refer to https://github.com/onnx/onnx/blob/main/docs/TypeDenotation.md#type-denotation-definition
   // for pre-defined type denotations.
   optional string denotation = 6;
 }
@@ -503,3 +769,68 @@ message OperatorSetIdProto {
   // This field MUST be present in this version of the IR.
   optional int64 version = 2;
 }
+
+// Operator/function status.
+enum OperatorStatus {
+    EXPERIMENTAL = 0;
+    STABLE = 1;
+}
+
+message FunctionProto {
+  // The name of the function, similar usage of op_type in OperatorProto.
+  // Combined with FunctionProto.domain, this forms the unique identity of
+  // the FunctionProto.
+  optional string name = 1;
+
+  // Deprecated since IR Version 8
+  // optional int64 since_version = 2;
+  reserved 2;
+  reserved "since_version";
+
+  // Deprecated since IR Version 8
+  // optional OperatorStatus status = 3;
+  reserved 3;
+  reserved "status";
+
+  // The inputs and outputs of the function.
+  repeated string input = 4;
+  repeated string output = 5;
+
+  // The attribute parameters of the function.
+  // It is for function parameters without default values.
+  repeated string attribute = 6;
+
+  // The attribute protos of the function.
+  // It is for function attributes with default values.
+  // A function attribute shall be represented either as
+  // a string attribute or an AttributeProto, not both.
+  repeated AttributeProto attribute_proto = 11;
+
+  // The nodes in the function.
+  repeated NodeProto node = 7;
+  // A human-readable documentation for this function. Markdown is allowed.
+  optional string doc_string = 8;
+
+  // The OperatorSets this function body (graph) relies on.
+  //
+  // All nodes in the function body (graph) will bind against the operator
+  // with the same-domain/same-op_type operator with the HIGHEST version
+  // in the referenced operator sets. This means at most one version can be relied
+  // for one domain.
+  //
+  // The operator sets imported by FunctionProto should be compatible with the ones
+  // imported by ModelProto. Example, if same operator set say 'A' is imported by FunctionProto
+  // and ModelProto then versions for the operator set may be different but,
+  // the operator schema returned for op_type, domain, version combination
+  // for both the versions should be same.
+
+  repeated OperatorSetIdProto opset_import = 9;
+
+  // The domain which this function belongs to. Combined with FunctionProto.name, this forms the unique identity of
+  // the FunctionProto.
+  optional string domain = 10;
+}
+
+// For using protobuf-lite
+option optimize_for = LITE_RUNTIME;
+

tools/pnnx/src/pass_level2/torch_permute.cpp

@@ -14,7 +14,7 @@
 
 #include "pass_level2.h"
 
-#include <torch/csrc/api/include/torch/torch.h>
+#include <torch/csrc/api/include/torch/version.h>
 
 namespace pnnx {
 

tools/pnnx/src/pass_level2/torch_repeat_interleave.cpp

@@ -14,8 +14,6 @@
 
 #include "pass_level2.h"
 
-#include <torch/csrc/api/include/torch/torch.h>
-
 namespace pnnx {
 
 class torch_repeat_interleave : public GraphRewriterPass

tools/pnnx/src/pass_level5/fuse_layernorm.cpp

@@ -19,7 +19,7 @@
 #include <math.h>
 #include <string.h>
 
-#include <torch/csrc/api/include/torch/torch.h>
+#include <torch/csrc/api/include/torch/version.h>
 
 namespace pnnx {
 

tools/pnnx/src/pass_level5/fuse_scaled_dot_product_attention.cpp

@@ -19,7 +19,7 @@
 #include <math.h>
 #include <string.h>
 
-#include <torch/csrc/api/include/torch/torch.h>
+#include <torch/csrc/api/include/torch/version.h>
 
 namespace pnnx {
 

tools/pnnx/src/save_onnx.cpp

@@ -24,17 +24,6 @@
 
 namespace pnnx {
 
-// from cxxabi bridge
-extern const char* get_operand_name(const Operand* x);
-extern const char* get_operator_type(const Operator* op);
-extern const char* get_operator_name(const Operator* op);
-extern std::vector<const char*> get_operator_params_keys(const Operator* op);
-extern std::vector<const char*> get_operator_attrs_keys(const Operator* op);
-extern const Parameter& get_operator_param(const Operator* op, const char* key);
-extern const Attribute& get_operator_attr(const Operator* op, const char* key);
-extern const char* get_param_s(const Parameter& p);
-extern std::vector<const char*> get_param_as(const Parameter& p);
-
 int save_onnx(const Graph& g, const char* onnxpath, int fp16)
 {
     onnx::ModelProto model;
@@ -45,7 +34,7 @@ int save_onnx(const Graph& g, const char* onnxpath, int fp16)
     {
         onnx::ValueInfoProto* vip = gp->add_value_info();
 
-        vip->set_name(get_operand_name(x));
+        vip->set_name(x->name);
 
         onnx::TypeProto* tp = vip->mutable_type();
 
@@ -108,27 +97,26 @@ int save_onnx(const Graph& g, const char* onnxpath, int fp16)
     {
         onnx::NodeProto* np = gp->add_node();
 
-        np->set_op_type(get_operator_type(op));
-        np->set_name(get_operator_name(op));
+        np->set_op_type(op->type);
+        np->set_name(op->name);
 
         for (const Operand* oprand : op->inputs)
         {
-            np->add_input(get_operand_name(oprand));
+            np->add_input(oprand->name);
         }
 
         for (const Operand* oprand : op->outputs)
         {
-            np->add_output(get_operand_name(oprand));
+            np->add_output(oprand->name);
         }
 
-        std::vector<const char*> params_keys = get_operator_params_keys(op);
-        for (const char* param_name : params_keys)
+        for (const auto& it : op->params)
         {
-            const Parameter& param = get_operator_param(op, param_name);
+            const Parameter& param = it.second;
 
             onnx::AttributeProto* ap = np->add_attribute();
 
-            ap->set_name(param_name);
+            ap->set_name(it.first);
 
             if (param.type == 0)
             {
@@ -156,7 +144,7 @@ int save_onnx(const Graph& g, const char* onnxpath, int fp16)
             if (param.type == 4)
             {
                 ap->set_type(onnx::AttributeProto::STRING);
-                ap->set_s(get_param_s(param));
+                ap->set_s(param.s);
             }
             if (param.type == 5)
             {
@@ -177,24 +165,22 @@ int save_onnx(const Graph& g, const char* onnxpath, int fp16)
             if (param.type == 7)
             {
                 ap->set_type(onnx::AttributeProto::STRINGS);
-                std::vector<const char*> as = get_param_as(param);
-                for (auto s : as)
+                for (auto s : param.as)
                 {
                     ap->add_strings(s);
                 }
             }
         }
 
-        std::vector<const char*> attrs_keys = get_operator_attrs_keys(op);
-        for (const char* attr_name : attrs_keys)
+        for (const auto& it : op->attrs)
         {
             onnx::TensorProto* tp = gp->add_initializer();
 
-            tp->set_name(std::string(get_operator_name(op)) + "." + attr_name);
+            tp->set_name(op->name + "." + it.first);
 
-            np->add_input(std::string(get_operator_name(op)) + "." + attr_name);
+            np->add_input(op->name + "." + it.first);
 
-            const Attribute& attr = get_operator_attr(op, attr_name);
+            const Attribute& attr = it.second;
             for (auto s : attr.shape)
             {
                 tp->add_dims(s);

tools/pnnx/src/save_onnx_cxxabi_bridge.cpp

@@ -1,81 +0,0 @@
-// Tencent is pleased to support the open source community by making ncnn available.
-//
-// Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
-//
-// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
-// in compliance with the License. You may obtain a copy of the License at
-//
-// https://opensource.org/licenses/BSD-3-Clause
-//
-// Unless required by applicable law or agreed to in writing, software distributed
-// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
-// CONDITIONS OF ANY KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations under the License.
-
-#include "ir.h"
-
-namespace pnnx {
-
-const char* get_operand_name(const Operand* x)
-{
-    return x->name.c_str();
-}
-
-const char* get_operator_type(const Operator* op)
-{
-    return op->type.c_str();
-}
-
-const char* get_operator_name(const Operator* op)
-{
-    return op->name.c_str();
-}
-
-std::vector<const char*> get_operator_params_keys(const Operator* op)
-{
-    std::vector<const char*> keys;
-    for (const auto& it : op->params)
-    {
-        const std::string& key = it.first;
-        keys.push_back(key.c_str());
-    }
-    return keys;
-}
-
-std::vector<const char*> get_operator_attrs_keys(const Operator* op)
-{
-    std::vector<const char*> keys;
-    for (const auto& it : op->attrs)
-    {
-        const std::string& key = it.first;
-        keys.push_back(key.c_str());
-    }
-    return keys;
-}
-
-const Parameter& get_operator_param(const Operator* op, const char* key)
-{
-    return op->params.at(key);
-}
-
-const Attribute& get_operator_attr(const Operator* op, const char* key)
-{
-    return op->attrs.at(key);
-}
-
-const char* get_param_s(const Parameter& p)
-{
-    return p.s.c_str();
-}
-
-std::vector<const char*> get_param_as(const Parameter& p)
-{
-    std::vector<const char*> as;
-    for (const auto& s : p.as)
-    {
-        as.push_back(s.c_str());
-    }
-    return as;
-}
-
-} // namespace pnnx

tools/pnnx/src/utils.cpp

@@ -16,17 +16,6 @@
 
 namespace pnnx {
 
-const torch::jit::Node* find_node_by_kind(const std::shared_ptr<torch::jit::Graph>& graph, const std::string& kind)
-{
-    for (const auto& n : graph->nodes())
-    {
-        if (n->kind().toDisplayString() == kind)
-            return n;
-    }
-
-    return 0;
-}
-
 unsigned short float32_to_float16(float value)
 {
     // 1 : 8 : 23

tools/pnnx/src/utils.h

@@ -15,12 +15,21 @@
 #ifndef PNNX_UTILS_H
 #define PNNX_UTILS_H
 
-#include <torch/script.h>
-#include <torch/csrc/jit/api/module.h>
+#if BUILD_TORCH2PNNX
+#include <memory>
+namespace torch {
+namespace jit {
+struct Graph;
+struct Node;
+} // namespace jit
+} // namespace torch
+#endif
 
 namespace pnnx {
 
+#if BUILD_TORCH2PNNX
 const torch::jit::Node* find_node_by_kind(const std::shared_ptr<torch::jit::Graph>& graph, const std::string& kind);
+#endif
 
 unsigned short float32_to_float16(float value);