write shape, fuse sam image encoder attention (#4792)

* write shape, fuse sam image encoder attention

* set more dynamic shape as static

* less warning for constant tensor node

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;

tools/pnnx/src/pass_level2.cpp

@@ -100,6 +100,50 @@ void GraphRewriterPass::write(Operator* op, const std::map<std::string, Paramete
 
         op->params[x.first] = Parameter::parse_from_string(str);
     }
+
+    for (size_t i = 0; i < op->inputs.size(); i++)
+    {
+        Operand* operand = op->inputs[i];
+        std::vector<int>& 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<int>& 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<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const

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]);
     }

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<int> build_shape(const std::string& expr)
+static void build_shape(const std::string& expr, std::vector<int>& shape, std::vector<std::string>& expr_tokens)
 {
     std::string listexpr = expr.substr(1, expr.size() - 2);
 
-    std::vector<int> 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<int> 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<int> 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<std::string>& 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<int> shape = build_shape(expr);
+            std::vector<int> outshape = op->outputs[0]->shape;
+            if (outshape.empty())
+                continue;
+
+            std::vector<int> shape;
+            std::vector<std::string> expr_tokens;
+            build_shape(expr, shape, expr_tokens);
 
             // replace -1 with static dim-size
-            std::vector<int> 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<int> outshape = op->outputs[0]->shape;
+        if (outshape.empty())
+            continue;
+
+        std::vector<int> 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

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";
     }

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

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<std::string, Parameter>& 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
 }
 

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<std::string, Parameter>& captured_params) const
     {
-        // (116,2,1024)
-        // (1,0,2)
-        // (2,-1,116,32,32)
-        const std::vector<int>& shape = captured_params.at("shape").ai;
-        const std::vector<int>& dims = captured_params.at("dims").ai;
-        const std::vector<int>& shape2 = captured_params.at("shape2").ai;
-
-        if (dims != std::vector<int>{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<std::string, Parameter>& captured_params) const
+    void write(const std::map<std::string, Operator*>& ops, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& captured_attrs) const
     {
-        const std::vector<int>& 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<int>{-233, -233};
-            slice->params["1"] = 0;
-
-            break;
-        }
-
-        if (!matched)
-            break;
-    }
 }
 
 } // namespace ncnn