pnnx torch 2.0 (#4579)

* fix build with torch-2.0

* torch 2.0 new patterns

* add torch 2.0 ci

.ci/pnnx.yml

@@ -44,6 +44,10 @@ jobs:
             torchvision-version: 0.14.0
             torchvision-cache-key: '0_14_0'
 
+          - torch-version: 2.0.0
+            torchvision-version: 0.15.1
+            torchvision-cache-key: '0_15_1'
+
     runs-on:
       pool-name: docker
       container:

tools/pnnx/src/pass_level0/inline_block.cpp

@@ -27,7 +27,7 @@ static void inlineCallTo(torch::jit::Node* to_replace, torch::jit::Function* cal
     torch::jit::WithInsertPoint guard(to_replace);
 
     std::unordered_map<torch::jit::Value*, torch::jit::Value*> value_map;
-#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11)
     std::vector<torch::jit::Value*> new_outputs = torch::jit::insertGraph(*to_replace->owningGraph(), *(toGraphFunction(*callee).graph()), to_replace->inputs(), value_map);
 #else
     std::vector<torch::jit::Value*> new_outputs = torch::jit::insertGraph(*to_replace->owningGraph(), *(callee->graph()), to_replace->inputs(), value_map);
@@ -56,7 +56,7 @@ static void inlineCalls(torch::jit::Block* block, const std::vector<std::string>
             if (!fun_type->function()->isGraphFunction())
                 continue;
 
-#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11)
             inlineCalls(toGraphFunction(*(fun_type->function())).graph()->block(), module_operators, inlined_modules, inside_module_op);
 #else
             inlineCalls(fun_type->function()->graph()->block(), module_operators, inlined_modules, inside_module_op);
@@ -87,7 +87,7 @@ static void inlineCalls(torch::jit::Block* block, const std::vector<std::string>
             {
                 if (std::find(module_operators.begin(), module_operators.end(), class_type_str_no_torch_prefix) != module_operators.end())
                 {
-#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11)
                     inlineCalls(toGraphFunction(function).graph()->block(), module_operators, inlined_modules, true);
 #else
                     inlineCalls(function.graph()->block(), module_operators, inlined_modules, true);
@@ -110,7 +110,7 @@ static void inlineCalls(torch::jit::Block* block, const std::vector<std::string>
                     continue;
             }
 
-#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11)
             inlineCalls(toGraphFunction(function).graph()->block(), module_operators, inlined_modules, inside_module_op);
 #else
             inlineCalls(function.graph()->block(), module_operators, inlined_modules, inside_module_op);

tools/pnnx/src/pass_level1.cpp

@@ -224,7 +224,7 @@ void pass_level1(const torch::jit::Module& mod, const std::shared_ptr<torch::jit
                 {
                     int pnnx_moduleop_unknown_index = 0;
 
-#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11)
                     torch::jit::Block* moduleop_block = toGraphFunction(function).graph()->block();
 #else
                     torch::jit::Block* moduleop_block = function.graph()->block();
@@ -339,7 +339,7 @@ void pass_level1(const torch::jit::Module& mod, const std::shared_ptr<torch::jit
 
                     op->name = wrapped_name;
 
-#if TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11
+#if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 11)
                     ow->write(op, toGraphFunction(function).graph(), sub_mod);
 #else
                     ow->write(op, function.graph(), sub_mod);

tools/pnnx/src/pass_level2/F_normalize.cpp

@@ -45,4 +45,34 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_normalize, 10)
 
+class F_normalize_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+12 11
+pnnx.Input              input       0 1 input
+prim::Constant          op_0        0 1 dtype value=*
+prim::Constant          op_1        0 1 keepdim value=True
+prim::Constant          op_2        0 1 p value=%p
+prim::Constant          op_3        0 1 dim value=%dim
+prim::Constant          op_4        0 1 eps value=%eps
+prim::ListConstruct     op_5        1 1 dim dims
+aten::linalg_vector_norm op_6       5 1 input p dims keepdim dtype 10
+aten::clamp_min         op_7        2 1 10 eps 13
+aten::expand_as         op_8        2 1 13 input denorm
+aten::div               op_9        2 1 input denorm out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.normalize";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_normalize_1, 10)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/torch_norm.cpp

@@ -69,4 +69,29 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_norm_1, 20)
 
+class torch_norm_2 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+7 6
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 dim
+prim::Constant          op_0        0 1 p value=%p
+prim::Constant          op_1        0 1 keepdim value=%keepdim
+prim::Constant          op_2        0 1 dtype value=*
+aten::linalg_vector_norm op_3       5 1 input p dim keepdim dtype out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.norm";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_norm_2, 20)
+
 } // namespace pnnx