From f8ca1e75858a79f04d177dc2884f6cfa305bdccc Mon Sep 17 00:00:00 2001
From: nihuini <nihuini@tencent.com>
Date: Wed, 5 Jan 2022 15:33:03 +0800
Subject: [PATCH] fix pnnx crash on unsupported expression

---
 tools/pnnx/src/pass_level2/torch_mean.cpp     | 22 +++++++++++++++++++
 tools/pnnx/src/pass_level2/torch_sum.cpp      | 22 +++++++++++++++++++
 .../pnnx/src/pass_ncnn/expand_expression.cpp  |  7 ++++++
 3 files changed, 51 insertions(+)

diff --git a/tools/pnnx/src/pass_level2/torch_mean.cpp b/tools/pnnx/src/pass_level2/torch_mean.cpp
index 9176394c3..a43682a8f 100644
--- a/tools/pnnx/src/pass_level2/torch_mean.cpp
+++ b/tools/pnnx/src/pass_level2/torch_mean.cpp
@@ -40,4 +40,26 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_mean, 20)
 
+class torch_mean_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+prim::Constant          op_0        0 1 dtype value=*
+aten::mean              op_1        2 1 input dtype out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.mean";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_mean_1, 20)
+
 } // namespace pnnx
diff --git a/tools/pnnx/src/pass_level2/torch_sum.cpp b/tools/pnnx/src/pass_level2/torch_sum.cpp
index c8418bb42..730ffcce3 100644
--- a/tools/pnnx/src/pass_level2/torch_sum.cpp
+++ b/tools/pnnx/src/pass_level2/torch_sum.cpp
@@ -40,4 +40,26 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_sum, 20)
 
+class torch_sum_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+prim::Constant          op_0        0 1 dtype value=*
+aten::sum               op_1        2 1 input dtype out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.sum";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_sum_1, 20)
+
 } // namespace pnnx
diff --git a/tools/pnnx/src/pass_ncnn/expand_expression.cpp b/tools/pnnx/src/pass_ncnn/expand_expression.cpp
index f19ab5b72..3d417c313 100644
--- a/tools/pnnx/src/pass_ncnn/expand_expression.cpp
+++ b/tools/pnnx/src/pass_ncnn/expand_expression.cpp
@@ -265,6 +265,13 @@ void expand_expression(Graph& graph)
             Operand* old_output_operand = op->outputs[0];
             Operand* new_output_operand = graph.get_operand(op->name + "_" + outname);
 
+            if (!new_output_operand)
+            {
+                // not supported expr
+                nonsupported_expr_ops.insert(op);
+                break;
+            }
+
             for (auto r : op->inputs)
             {
                 r->remove_consumer(op);