diff --git a/.ci/pnnx.yml b/.ci/pnnx.yml index ede82a87a..e0d1987cc 100644 --- a/.ci/pnnx.yml +++ b/.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: diff --git a/tools/pnnx/src/pass_level0/inline_block.cpp b/tools/pnnx/src/pass_level0/inline_block.cpp index 6f98a5898..57dda9ca8 100644 --- a/tools/pnnx/src/pass_level0/inline_block.cpp +++ b/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 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 new_outputs = torch::jit::insertGraph(*to_replace->owningGraph(), *(toGraphFunction(*callee).graph()), to_replace->inputs(), value_map); #else std::vector 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 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 { 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 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); diff --git a/tools/pnnx/src/pass_level1.cpp b/tools/pnnx/src/pass_level1.cpp index 0aaf4d897..7b79713a8 100644 --- a/tools/pnnx/src/pass_level1.cpp +++ b/tools/pnnx/src/pass_level1.cpp @@ -224,7 +224,7 @@ void pass_level1(const torch::jit::Module& mod, const std::shared_ptr= 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_ptrname = 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); diff --git a/tools/pnnx/src/pass_level2/F_normalize.cpp b/tools/pnnx/src/pass_level2/F_normalize.cpp index 9717282dd..b898fbbf6 100644 --- a/tools/pnnx/src/pass_level2/F_normalize.cpp +++ b/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 diff --git a/tools/pnnx/src/pass_level2/torch_norm.cpp b/tools/pnnx/src/pass_level2/torch_norm.cpp index 08ffda043..7464cc167 100644 --- a/tools/pnnx/src/pass_level2/torch_norm.cpp +++ b/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