diff --git a/tools/pnnx/src/ir.cpp b/tools/pnnx/src/ir.cpp index 8eae58264..e6d902c05 100644 --- a/tools/pnnx/src/ir.cpp +++ b/tools/pnnx/src/ir.cpp @@ -269,10 +269,8 @@ Parameter::Parameter(const torch::jit::Node* value_node) } else { - const int ndim = (int)t.dim(); - + // constant tensor will become pnnx attribute node later type = 8; - fprintf(stderr, "unknown Parameter value kind %s of TensorType, t.dim = %d\n", value_node->kind().toDisplayString(), ndim); } break; diff --git a/tools/pnnx/src/pass_level2.cpp b/tools/pnnx/src/pass_level2.cpp index 199861774..41f366d5e 100644 --- a/tools/pnnx/src/pass_level2.cpp +++ b/tools/pnnx/src/pass_level2.cpp @@ -100,6 +100,50 @@ void GraphRewriterPass::write(Operator* op, const std::mapparams[x.first] = Parameter::parse_from_string(str); } + + for (size_t i = 0; i < op->inputs.size(); i++) + { + Operand* operand = op->inputs[i]; + std::vector& shape = operand->shape; + for (size_t j = 0; j < shape.size(); j++) + { + int ai = shape[j]; + if (ai == -233) + { + std::string key = operand->params.at(std::string("__shape_") + std::to_string(j)).s; + + if (captured_params.find(key) == captured_params.end()) + { + fprintf(stderr, "replace pattern param %%%s missing captured\n", key.c_str()); + return; + } + + shape[j] = captured_params.at(key).i; + } + } + } + + for (size_t i = 0; i < op->outputs.size(); i++) + { + Operand* operand = op->outputs[i]; + std::vector& shape = operand->shape; + for (size_t j = 0; j < shape.size(); j++) + { + int ai = shape[j]; + if (ai == -233) + { + std::string key = operand->params.at(std::string("__shape_") + std::to_string(j)).s; + + if (captured_params.find(key) == captured_params.end()) + { + fprintf(stderr, "replace pattern param %%%s missing captured\n", key.c_str()); + return; + } + + shape[j] = captured_params.at(key).i; + } + } + } } void GraphRewriterPass::write(Operator* op, const std::map& captured_params, const std::map& captured_attrs) const diff --git a/tools/pnnx/src/pass_level3/fuse_expression.cpp b/tools/pnnx/src/pass_level3/fuse_expression.cpp index b23b0ae43..a86c1afea 100644 --- a/tools/pnnx/src/pass_level3/fuse_expression.cpp +++ b/tools/pnnx/src/pass_level3/fuse_expression.cpp @@ -100,7 +100,8 @@ static bool operand_maybe_tensor(const Operand* operand) || op->type == "aten::div" || op->type == "aten::floor_divide" || op->type == "aten::mul" - || op->type == "aten::pow") + || op->type == "aten::pow" + || op->type == "aten::remainder") { return operand_maybe_tensor(op->inputs[0]) || operand_maybe_tensor(op->inputs[1]); } diff --git a/tools/pnnx/src/pass_level5/eliminate_reshape_shape_expression.cpp b/tools/pnnx/src/pass_level5/eliminate_reshape_shape_expression.cpp index 744d29016..e51fa240f 100644 --- a/tools/pnnx/src/pass_level5/eliminate_reshape_shape_expression.cpp +++ b/tools/pnnx/src/pass_level5/eliminate_reshape_shape_expression.cpp @@ -31,13 +31,12 @@ static bool token_is_interger_literal(const std::string& t) return iss.eof() && !iss.fail(); } -static std::vector build_shape(const std::string& expr) +static void build_shape(const std::string& expr, std::vector& shape, std::vector& expr_tokens) { std::string listexpr = expr.substr(1, expr.size() - 2); - std::vector shape; - std::string t; + std::string et; int level = 0; for (size_t i = 0; i < listexpr.size(); i++) { @@ -47,21 +46,26 @@ static std::vector build_shape(const std::string& expr) { level += 1; t = "-1"; + et += ch; } else if (ch == ')' || ch == ']') { level -= 1; t = "-1"; + et += ch; } else if (level == 0 && ch == ',') { int dimsize = token_is_interger_literal(t) ? std::stoi(t) : -1; shape.push_back(dimsize); + expr_tokens.push_back(et); t.clear(); + et.clear(); } else { t += ch; + et += ch; } } @@ -71,7 +75,26 @@ static std::vector build_shape(const std::string& expr) shape.push_back(dimsize); } - return shape; + if (level == 0 && !et.empty()) + { + expr_tokens.push_back(et); + } +} + +static std::string build_expr(const std::vector& expr_tokens) +{ + std::string expr; + + expr += '['; + for (int i = 0; i < (int)expr_tokens.size(); i++) + { + expr += expr_tokens[i]; + if (i != (int)expr_tokens.size() - 1) + expr += ','; + } + expr += ']'; + + return expr; } void eliminate_reshape_shape_expression(Graph& graph) @@ -98,18 +121,21 @@ void eliminate_reshape_shape_expression(Graph& graph) if (expr.empty() || expr[0] != '[') continue; - std::vector shape = build_shape(expr); + std::vector outshape = op->outputs[0]->shape; + if (outshape.empty()) + continue; + + std::vector shape; + std::vector expr_tokens; + build_shape(expr, shape, expr_tokens); // replace -1 with static dim-size - std::vector outshape = op->outputs[0]->shape; - if (!outshape.empty()) + for (size_t j = 0; j < outshape.size(); j++) { - for (size_t j = 0; j < outshape.size(); j++) + if (outshape[j] != -1) { - if (outshape[j] != -1) - { - shape[j] = outshape[j]; - } + shape[j] = outshape[j]; + expr_tokens[j] = std::to_string(outshape[j]); } } @@ -124,7 +150,10 @@ void eliminate_reshape_shape_expression(Graph& graph) } if (dynamic_dim_count > 1) + { + op_expr->params["expr"] = build_expr(expr_tokens); continue; + } matched = true; @@ -156,6 +185,34 @@ void eliminate_reshape_shape_expression(Graph& graph) if (!matched) break; } + + for (size_t i = 0; i < graph.ops.size(); i++) + { + Operator* op = graph.ops[i]; + + if (op->type != "Tensor.view" && op->type != "Tensor.reshape") + continue; + + if (op->inputs.size() != 1) + continue; + + std::vector outshape = op->outputs[0]->shape; + if (outshape.empty()) + continue; + + std::vector shape = op->params.at("shape").ai; + + // replace -1 with static dim-size + for (size_t j = 0; j < outshape.size(); j++) + { + if (outshape[j] != -1) + { + shape[j] = outshape[j]; + } + } + + op->params["shape"] = shape; + } } } // namespace pnnx diff --git a/tools/pnnx/src/pass_level5/fuse_channel_shuffle.cpp b/tools/pnnx/src/pass_level5/fuse_channel_shuffle.cpp index 6b79bc059..3a38f5941 100644 --- a/tools/pnnx/src/pass_level5/fuse_channel_shuffle.cpp +++ b/tools/pnnx/src/pass_level5/fuse_channel_shuffle.cpp @@ -56,7 +56,7 @@ public: pnnx.Input input 0 1 input Tensor.view op_0 1 1 input 13 shape=(%batch,%groups,%channels_per_group,%h,%w) torch.transpose op_1 1 1 13 14 dim0=1 dim1=2 -Tensor.reshape op_2 1 1 14 out shape=(%batch,-1,%h,%w) +Tensor.reshape op_2 1 1 14 out shape=(%batch,%channels,%h,%w) pnnx.Output output 1 0 out )PNNXIR"; } diff --git a/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp b/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp index 4dd004264..336ad9dfb 100644 --- a/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp +++ b/tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp @@ -1060,14 +1060,14 @@ nn.Linear op_1 1 1 input 4 bias=%kbias in_features=%embed_d nn.Linear op_2 1 1 input 6 bias=%vbias in_features=%embed_dim out_features=%embed_dim @bias @weight pnnx.Expression op_3 1 1 2 3 expr=mul(@0,%inv_sqrt_embed_dim_per_head) Tensor.view op_4 1 1 3 8 shape=(%batch,%size,%num_heads,%feat_per_head) -Tensor.view op_5 1 1 4 5 shape=(%batch,-1,%num_heads,%feat_per_head) -Tensor.view op_6 1 1 6 7 shape=(%batch,-1,%num_heads,%feat_per_head) +Tensor.view op_5 1 1 4 5 shape=(%batch,%size,%num_heads,%feat_per_head) +Tensor.view op_6 1 1 6 7 shape=(%batch,%size,%num_heads,%feat_per_head) torch.transpose op_7 1 1 8 9 dim0=1 dim1=2 torch.transpose op_8 1 1 5 10 dim0=1 dim1=2 torch.transpose op_9 1 1 7 11 dim0=1 dim1=2 -Tensor.reshape op_10 1 1 9 14 shape=(%num_heads,-1,%feat_per_head) -Tensor.reshape op_11 1 1 10 12 shape=(%num_heads,-1,%feat_per_head) -Tensor.reshape op_12 1 1 11 17 shape=(%num_heads,-1,%feat_per_head) +Tensor.reshape op_10 1 1 9 14 shape=(%num_heads,%batch_mul_size,%feat_per_head) +Tensor.reshape op_11 1 1 10 12 shape=(%num_heads,%batch_mul_size,%feat_per_head) +Tensor.reshape op_12 1 1 11 17 shape=(%num_heads,%batch_mul_size,%feat_per_head) torch.transpose op_13 1 1 12 13 dim0=1 dim1=2 torch.bmm op_14 2 1 14 13 15 F.softmax op_15 1 1 15 16 dim=-1 @@ -1094,14 +1094,14 @@ nn.Linear op_1 1 1 input 5 bias=%kbias in_features=%embed_d nn.Linear op_2 1 1 input 7 bias=%vbias in_features=%embed_dim out_features=%embed_dim @bias @weight pnnx.Expression op_3 1 1 3 4 expr=mul(@0,%inv_sqrt_embed_dim_per_head) Tensor.view op_4 1 1 4 9 shape=(%batch,%size,%num_heads,%feat_per_head) -Tensor.view op_5 1 1 5 6 shape=(%batch,-1,%num_heads,%feat_per_head) -Tensor.view op_6 1 1 7 8 shape=(%batch,-1,%num_heads,%feat_per_head) +Tensor.view op_5 1 1 5 6 shape=(%batch,%size,%num_heads,%feat_per_head) +Tensor.view op_6 1 1 7 8 shape=(%batch,%size,%num_heads,%feat_per_head) torch.transpose op_7 1 1 9 10 dim0=1 dim1=2 torch.transpose op_8 1 1 6 11 dim0=1 dim1=2 torch.transpose op_9 1 1 8 12 dim0=1 dim1=2 -Tensor.reshape op_10 1 1 10 15 shape=(%num_heads,-1,%feat_per_head) -Tensor.reshape op_11 1 1 11 13 shape=(%num_heads,-1,%feat_per_head) -Tensor.reshape op_12 1 1 12 21 shape=(%num_heads,-1,%feat_per_head) +Tensor.reshape op_10 1 1 10 15 shape=(%num_heads,%batch_mul_size,%feat_per_head) +Tensor.reshape op_11 1 1 11 13 shape=(%num_heads,%batch_mul_size,%feat_per_head) +Tensor.reshape op_12 1 1 12 21 shape=(%num_heads,%batch_mul_size,%feat_per_head) torch.transpose op_13 1 1 13 14 dim0=1 dim1=2 torch.bmm op_14 2 1 15 14 16 Tensor.view op_15 1 1 16 17 shape=(%batch,%num_heads,%size,%size) @@ -1301,7 +1301,7 @@ pnnx.Expression op_7 2 1 33 attn_mask 35 expr=add(@0,@1) Tensor.view op_8 1 1 35 36 shape=(1,%batch,%num_heads,%size,%size) pnnx.Attribute op_9 0 1 37 @data=(1,%batch,1,%size,%size)f32 pnnx.Expression op_10 2 1 36 37 38 expr=add(@0,@1) -Tensor.view op_11 1 1 38 39 shape=(-1,%num_heads,%size,%size) +Tensor.view op_11 1 1 38 39 shape=(%batch,%num_heads,%size,%size) F.softmax op_12 1 1 39 40 dim=-1 torch.matmul op_13 2 1 40 30 41 torch.transpose op_14 1 1 41 42 dim0=1 dim1=2 diff --git a/tools/pnnx/src/pass_level5/fuse_scaled_dot_product_attention.cpp b/tools/pnnx/src/pass_level5/fuse_scaled_dot_product_attention.cpp index ab5eb886e..e6f1489a9 100644 --- a/tools/pnnx/src/pass_level5/fuse_scaled_dot_product_attention.cpp +++ b/tools/pnnx/src/pass_level5/fuse_scaled_dot_product_attention.cpp @@ -79,13 +79,73 @@ pnnx.Output output 1 0 out } }; +class fuse_scaled_dot_product_attention_pass_1 : public GraphRewriterPass +{ +public: + const char* match_pattern_graph() const + { + return R"PNNXIR(7767517 +14 13 +pnnx.Input input_0 0 1 query #query=(%batch,%qsize,%feat_per_head)f32 +pnnx.Input input_1 0 1 key #key=(%batch,%kvsize,%feat_per_head)f32 +pnnx.Input input_2 0 1 value #value=(%batch,%kvsize,%feat_per_head)f32 +pnnx.Input input_Rh 0 1 Rh #Rh=(%batch,%h,%w,%h,1)f32 +pnnx.Input input_Rw 0 1 Rw #Rw=(%batch,%h,%w,1,%w)f32 +pnnx.Expression op_0 1 1 query 17 expr=mul(@0,%inv_sqrt_embed_dim_per_head) +torch.transpose op_1 1 1 key 22 dim0=-2 dim1=-1 +torch.matmul op_2 2 1 17 22 23 +Tensor.view op_3 1 1 23 24 shape=(%batch,%h,%w,%h,%w) +pnnx.Expression op_4 3 1 24 Rh Rw 28 expr=add(add(@0,@1),@2) +Tensor.view op_5 1 1 28 29 shape=(%batch,%qsize,%qsize) +F.softmax op_6 1 1 29 30 dim=-1 +torch.matmul op_7 2 1 30 value out +pnnx.Output output 1 0 out +)PNNXIR"; + } + + const char* replace_pattern_graph() const + { + return R"PNNXIR(7767517 +9 8 +pnnx.Input input_0 0 1 query +pnnx.Input input_1 0 1 key +pnnx.Input input_2 0 1 value +pnnx.Input input_Rh 0 1 Rh +pnnx.Input input_Rw 0 1 Rw +pnnx.Expression RhRw 2 1 Rh Rw RhRw expr=add(@0,@1) #RhRw=(%batch,%h,%w,%h,%w)f32 +Tensor.reshape attn_mask 1 1 RhRw attn_mask shape=(%batch,%qsize,%qsize) #attn_mask=(%batch,%qsize,%qsize)f32 +F.scaled_dot_product_attention op_0 4 1 query key value attn_mask out dropout_p=0.0 is_causal=False $attn_mask=attn_mask +pnnx.Output output 1 0 out +)PNNXIR"; + } + + bool match(const std::map& captured_params) const + { + const int qsize = captured_params.at("qsize").i; + const int h = captured_params.at("h").i; + const int w = captured_params.at("w").i; + const int feat_per_head = captured_params.at("feat_per_head").i; + const float inv_sqrt_embed_dim_per_head = captured_params.at("inv_sqrt_embed_dim_per_head").f; + + if (qsize != h * w) + return false; + + if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001)) + return false; + + return true; + } +}; + void fuse_scaled_dot_product_attention(Graph& graph) { #if TORCH_VERSION_MAJOR >= 2 fuse_scaled_dot_product_attention_pass a; + fuse_scaled_dot_product_attention_pass_1 b; int opindex = 0; pnnx_graph_rewrite(graph, &a, opindex); + pnnx_graph_rewrite(graph, &b, opindex); #endif } diff --git a/tools/pnnx/src/pass_ncnn/fuse_convert_shufflechannel_slice.cpp b/tools/pnnx/src/pass_ncnn/fuse_convert_shufflechannel_slice.cpp index 9e7265c15..d71313bd3 100644 --- a/tools/pnnx/src/pass_ncnn/fuse_convert_shufflechannel_slice.cpp +++ b/tools/pnnx/src/pass_ncnn/fuse_convert_shufflechannel_slice.cpp @@ -35,51 +35,58 @@ public: { return R"PNNXIR(7767517 6 6 -pnnx.Input input 0 1 input -Tensor.reshape op_0 1 1 input a shape=%shape -torch.permute op_1 1 1 a b dims=%dims -Tensor.reshape op_2 1 1 b c shape=%shape2 +pnnx.Input input 0 1 input #input=(%batch,%c,%h,%w)f32 +Tensor.reshape op_0 1 1 input a shape=(%batch_mul_ch_per_group,%groups,%h_mul_w) +torch.permute op_1 1 1 a b dims=(1,0,2) +Tensor.reshape op_2 1 1 b c shape=(%groups,%batch,%ch_per_group,%h,%w) torch.unbind op_3 1 2 c out0 out1 dim=0 pnnx.Output output 2 0 out0 out1 )PNNXIR"; } - const char* type_str() const - { - return "ncnn._shufflechannel_slice"; - } - - const char* name_str() const + const char* replace_pattern_graph() const { - return "shufflechannel_slice"; + return R"PNNXIR(7767517 +4 4 +pnnx.Input input 0 1 input +ShuffleChannel shufflechannel 1 1 input a 0=%groups 1=1 #a=(%batch,%c,%h,%w)f32 +Slice slice 1 2 a out0 out1 0=(-233,-233) 1=0 +pnnx.Output output 2 0 out0 out1 +)PNNXIR"; } bool match(const std::map& captured_params) const { - // (116,2,1024) - // (1,0,2) - // (2,-1,116,32,32) - const std::vector& shape = captured_params.at("shape").ai; - const std::vector& dims = captured_params.at("dims").ai; - const std::vector& shape2 = captured_params.at("shape2").ai; - - if (dims != std::vector{1, 0, 2}) + const int groups = captured_params.at("groups").i; + const int batch = captured_params.at("batch").i; + const int batch_mul_ch_per_group = captured_params.at("batch_mul_ch_per_group").i; + const int ch_per_group = captured_params.at("ch_per_group").i; + const int h_mul_w = captured_params.at("h_mul_w").i; + const int c = captured_params.at("c").i; + const int h = captured_params.at("h").i; + const int w = captured_params.at("w").i; + + if (groups != 2 || groups * ch_per_group != c) return false; - if (shape[0] != shape2[2] || shape[1] != shape2[0] || shape[2] != shape2[3] * shape2[4] || shape[1] != 2 || shape2[1] != -1) + if (batch_mul_ch_per_group != batch * ch_per_group) + return false; + + if (h_mul_w != h * w) return false; return true; } - void write(Operator* op, const std::map& captured_params) const + void write(const std::map& ops, const std::map& captured_params, const std::map& captured_attrs) const { - const std::vector& shape = captured_params.at("shape").ai; + GraphRewriterPass::write(ops, captured_params, captured_attrs); - int groups = shape[1]; + const int batch_index = ops.at("shufflechannel")->inputs[0]->params["__batch_index"].i; - op->params["0"] = groups; - op->params["1"] = 1; + ops.at("slice")->inputs[0]->params["__batch_index"] = batch_index; + ops.at("slice")->outputs[0]->params["__batch_index"] = batch_index; + ops.at("slice")->outputs[1]->params["__batch_index"] = batch_index; } }; @@ -90,10 +97,10 @@ public: { return R"PNNXIR(7767517 6 6 -pnnx.Input input 0 1 input -Tensor.reshape op_0 1 1 input a shape=%shape -Tensor.permute op_1 1 1 a b dims=%dims -Tensor.reshape op_2 1 1 b c shape=%shape2 +pnnx.Input input 0 1 input #input=(%batch,%c,%h,%w)f32 +Tensor.reshape op_0 1 1 input a shape=(%batch_mul_ch_per_group,%groups,%h_mul_w) +Tensor.permute op_1 1 1 a b dims=(1,0,2) +Tensor.reshape op_2 1 1 b c shape=(%groups,%batch,%ch_per_group,%h,%w) torch.unbind op_3 1 2 c out0 out1 dim=0 pnnx.Output output 2 0 out0 out1 )PNNXIR"; @@ -108,57 +115,6 @@ void fuse_convert_shufflechannel_slice(Graph& graph) pnnx_graph_rewrite(graph, &a, opindex); pnnx_graph_rewrite(graph, &b, opindex); - - int op_index = 0; - - while (1) - { - bool matched = false; - - for (Operator* op : graph.ops) - { - if (op->type != "ncnn._shufflechannel_slice") - continue; - - matched = true; - - const int batch_index = op->inputs[0]->params["__batch_index"].i; - - op->type = "ShuffleChannel"; - op->name = std::string("shufflechannel_") + std::to_string(op_index++); - - Operand* out0 = op->outputs[0]; - Operand* out1 = op->outputs[1]; - - Operator* slice = graph.new_operator_after("Slice", op->name + "_slice", op); - - Operand* slice_in = graph.new_operand(op->name + "_slice_in"); - - slice_in->params["__batch_index"] = batch_index; - out0->params["__batch_index"] = batch_index; - out1->params["__batch_index"] = batch_index; - - slice->inputs.push_back(slice_in); - slice->outputs.push_back(out0); - slice->outputs.push_back(out1); - - op->outputs.clear(); - op->outputs.push_back(slice_in); - - out0->producer = slice; - out1->producer = slice; - slice_in->producer = op; - slice_in->consumers.push_back(slice); - - slice->params["0"] = std::vector{-233, -233}; - slice->params["1"] = 0; - - break; - } - - if (!matched) - break; - } } } // namespace ncnn