pnnx onnx with inputshape (#5480)

tools/pnnx/src/CMakeLists.txt

@@ -363,6 +363,7 @@ set(pnnx_pass_level5_SRCS
     pass_level5/fuse_slice_copy.cpp
     pass_level5/fuse_slice_indices.cpp
     pass_level5/fuse_slice_to_tensor_split.cpp
+    pass_level5/fuse_slice_squeeze_to_select.cpp
     pass_level5/fuse_static_batchnorm.cpp
     pass_level5/fuse_static_conv.cpp
     pass_level5/fuse_static_convtranspose.cpp
@@ -371,6 +372,7 @@ set(pnnx_pass_level5_SRCS
     pass_level5/fuse_static_instancenorm.cpp
     pass_level5/fuse_static_layernorm.cpp
     pass_level5/fuse_static_linear.cpp
+    pass_level5/fuse_swish.cpp
     pass_level5/normalize_einsum_equation.cpp
     pass_level5/unroll_rnn_op.cpp
 )

tools/pnnx/src/ir.cpp

@@ -1004,11 +1004,20 @@ static std::string expand_expression(const Operator* op)
         {
             std::string a = exprstack.top();
             exprstack.pop();
-            std::string b = exprstack.top();
-            exprstack.pop();
 
-            std::string r = a + ".size(" + b + ")";
-            exprstack.push(r);
+            if (exprstack.empty())
+            {
+                std::string r = a + ".shape";
+                exprstack.push(r);
+            }
+            else
+            {
+                std::string b = exprstack.top();
+                exprstack.pop();
+
+                std::string r = a + ".size(" + b + ")";
+                exprstack.push(r);
+            }
         }
         else if (t == "int"
                  || t == "abs"

tools/pnnx/src/load_onnx.cpp

@@ -23,6 +23,8 @@
 
 #include <fstream>
 
+#include <onnxruntime_c_api.h>
+
 #include "ir.h"
 
 #include "pass_onnx/canonicalize.h"
@@ -304,7 +306,14 @@ Operand* Graph::new_operand(const onnx::ValueInfoProto& value)
     r->shape.resize(tensor_shape.dim_size());
     for (int z = 0; z < tensor_shape.dim_size(); z++)
     {
-        r->shape[z] = tensor_shape.dim(z).dim_value();
+        if (!tensor_shape.dim(z).has_dim_value())
+        {
+            r->shape[z] = -1;
+        }
+        else
+        {
+            r->shape[z] = tensor_shape.dim(z).dim_value();
+        }
     }
 
     operands.push_back(r);
@@ -366,7 +375,147 @@ static double get_current_time()
     return usec.count() / 1000.0;
 }
 
-int load_onnx(const std::string& onnxpath, Graph& pnnx_graph)
+static const char* get_onnx_tensor_elem_data_type_str(ONNXTensorElementDataType type)
+{
+    switch (type)
+    {
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8:
+        return "i8";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8:
+        return "u8";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16:
+        return "i16";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16:
+        return "u16";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
+        return "i32";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32:
+        return "u32";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64:
+        return "i64";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64:
+        return "u64";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16:
+        return "f16";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
+        return "f32";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE:
+        return "f64";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16:
+        return "bf16";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64:
+        return "c64";
+    case ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128:
+        return "c128";
+    default:
+        break;
+    }
+
+    // unknown
+    fprintf(stderr, "unsupported tensor elem data type %d\n", (int)type);
+    return "";
+}
+
+static bool check_input_shape(onnx::ModelProto& model, const std::vector<std::vector<int64_t> >& input_shapes, const std::vector<std::string>& input_types)
+{
+    const onnx::GraphProto& graph = model.graph();
+
+    if (!input_shapes.empty() && (int)input_shapes.size() != graph.input_size())
+    {
+        fprintf(stderr, "input_shape expect %d tensors but got %d\n", graph.input_size(), (int)input_shapes.size());
+        return false;
+    }
+
+    for (int i = 0; i < graph.input_size(); i++)
+    {
+        const onnx::ValueInfoProto& value = graph.input(i);
+        const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape();
+
+        ONNXTensorElementDataType datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type();
+
+        bool matched = true;
+
+        if (input_shapes.empty())
+        {
+            // dynamic dimension size
+            for (int j = 0; j < tsp.dim_size(); j++)
+            {
+                if (!tsp.dim(j).has_dim_value())
+                {
+                    matched = false;
+                    break;
+                }
+            }
+        }
+        else
+        {
+            if ((int)input_shapes[i].size() != tsp.dim_size())
+            {
+                matched = false;
+            }
+            else
+            {
+                for (int j = 0; j < tsp.dim_size(); j++)
+                {
+                    if (!tsp.dim(j).has_dim_value())
+                        continue;
+
+                    int64_t ds = tsp.dim(j).dim_value();
+                    if (input_shapes[i][j] != ds)
+                        matched = false;
+                }
+            }
+
+            if (input_types[i] != get_onnx_tensor_elem_data_type_str(datatype))
+                matched = false;
+        }
+
+        if (!matched)
+        {
+            fprintf(stderr, "input_shapes[%d] expect [", i);
+            for (int j = 0; j < tsp.dim_size(); j++)
+            {
+                if (tsp.dim(j).has_dim_value())
+                {
+                    int64_t ds = tsp.dim(j).dim_value();
+                    fprintf(stderr, "%ld", ds);
+                }
+                else
+                {
+                    fprintf(stderr, "?");
+                }
+                if (j + 1 != tsp.dim_size())
+                    fprintf(stderr, ",");
+            }
+            fprintf(stderr, "]%s but got ", get_onnx_tensor_elem_data_type_str(datatype));
+            if (input_shapes.empty())
+            {
+                fprintf(stderr, "nothing\n");
+            }
+            else
+            {
+                fprintf(stderr, "[");
+                for (size_t j = 0; j < input_shapes[i].size(); j++)
+                {
+                    fprintf(stderr, "%ld", input_shapes[i][j]);
+                    if (j + 1 != input_shapes[i].size())
+                        fprintf(stderr, ",");
+                }
+                fprintf(stderr, "]%s\n", input_types[i].c_str());
+            }
+
+            return false;
+        }
+    }
+
+    return true;
+}
+
+int load_onnx(const std::string& onnxpath, Graph& pnnx_graph,
+              const std::vector<std::vector<int64_t> >& input_shapes,
+              const std::vector<std::string>& input_types,
+              const std::vector<std::vector<int64_t> >& input_shapes2,
+              const std::vector<std::string>& input_types2)
 {
     onnx::ModelProto model;
 
@@ -377,6 +526,16 @@ int load_onnx(const std::string& onnxpath, Graph& pnnx_graph)
         return -1;
     }
 
+    // input shape sanity check
+    if (!check_input_shape(model, input_shapes, input_types))
+    {
+        return -1;
+    }
+    if (!input_shapes2.empty() && !check_input_shape(model, input_shapes2, input_types2))
+    {
+        return -1;
+    }
+
     fprintf(stderr, "############# pass_level0 onnx \n");
 
     onnx2pnnx::ModelStat oldstat = onnx2pnnx::get_model_stat(model);
@@ -415,7 +574,7 @@ int load_onnx(const std::string& onnxpath, Graph& pnnx_graph)
 
     t0 = get_current_time();
 
-    onnx2pnnx::fold_constants(model);
+    onnx2pnnx::fold_constants(model, input_shapes, input_types, input_shapes2, input_types2);
 
     t1 = get_current_time();
 
@@ -445,7 +604,7 @@ int load_onnx(const std::string& onnxpath, Graph& pnnx_graph)
 
     t0 = get_current_time();
 
-    onnx2pnnx::shape_inference(model);
+    onnx2pnnx::shape_inference(model, input_shapes, input_types, input_shapes2, input_types2);
 
     t1 = get_current_time();
 

tools/pnnx/src/load_onnx.h

@@ -19,7 +19,11 @@
 
 namespace pnnx {
 
-int load_onnx(const std::string& onnxpath, Graph& g);
+int load_onnx(const std::string& onnxpath, Graph& g,
+              const std::vector<std::vector<int64_t> >& input_shapes,
+              const std::vector<std::string>& input_types,
+              const std::vector<std::vector<int64_t> >& input_shapes2,
+              const std::vector<std::string>& input_types2);
 
 } // namespace pnnx
 

tools/pnnx/src/main.cpp

@@ -316,7 +316,9 @@ int main(int argc, char** argv)
 #if BUILD_ONNX2PNNX
     if (!model_file_maybe_torchscript(ptpath))
     {
-        load_onnx(ptpath.c_str(), pnnx_graph);
+        load_onnx(ptpath.c_str(), pnnx_graph,
+                  input_shapes, input_types,
+                  input_shapes2, input_types2);
     }
     else
 #endif

tools/pnnx/src/pass_level2/F_avg_pool2d.cpp

@@ -43,7 +43,7 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d, 10)
 
-class F_avg_pool2d_1 : public GraphRewriterPass
+class F_avg_pool2d_onnx : public GraphRewriterPass
 {
 public:
     const char* match_pattern_graph() const
@@ -83,6 +83,22 @@ pnnx.Output             output      1 0 out
     }
 };
 
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d_1, 10)
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d_onnx, 10)
+
+class F_avg_pool2d_onnx_1 : public F_avg_pool2d_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input_0     0 1 input
+AveragePool             op_0        1 1 input out kernel_shape=%kernel_shape strides=%strides pads=%pads ceil_mode=%ceil_mode count_include_pad=%count_include_pad
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_avg_pool2d_onnx_1, 10)
 
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_batch_norm.cpp

@@ -104,4 +104,29 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_batch_norm_onnx, 10)
 
+class F_batch_norm_onnx_1 : 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 weight
+pnnx.Input              input_2     0 1 bias
+pnnx.Input              input_3     0 1 running_mean
+pnnx.Input              input_4     0 1 running_var
+BatchNormalization      op_0        5 1 input weight bias running_mean running_var out epsilon=%eps momentum=*
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.batch_norm";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_batch_norm_onnx_1, 10)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_conv2d.cpp

@@ -227,6 +227,34 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv2d_onnx, 10)
 
+class F_conv2d_onnx_0 : public F_conv2d_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 weight
+Conv                    op_0        2 1 input weight out kernel_shape=%kernel_shape strides=%strides pads=%pads dilations=%dilations group=%group
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* replace_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 weight
+F.conv2d                conv        2 1 input weight out bias=None
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_conv2d_onnx_0, 10)
+
 class F_conv2d_onnx_1 : public GraphRewriterPass
 {
 public:

tools/pnnx/src/pass_level2/F_gelu.cpp

@@ -214,7 +214,7 @@ public:
 pnnx.Input              input       0 1 input
 prim::Constant          op_0        0 1 sqrt2 value=%sqrt2
 aten::div               op_1        2 1 input sqrt2 16
-Erf                     op_2        1 1 16 17
+aten::erf               op_2        1 1 16 17
 prim::Constant          op_3        0 1 one value=%1
 aten::add               op_4        2 1 17 one 22
 aten::mul               op_5        2 1 input 22 25

tools/pnnx/src/pass_level2/F_hardswish.cpp

@@ -200,6 +200,34 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_hardswish_6, 8)
 
+class F_hardswish_7 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+10 9
+pnnx.Input              input       0 1 input
+prim::Constant          op_0        0 1 25 value=3
+aten::add               op_1        2 1 input 25 a
+prim::Constant          op_2        0 1 min value=0
+prim::Constant          op_3        0 1 max value=6
+aten::clamp             op_4        3 1 a min max b
+aten::mul               op_5        2 1 input b c
+prim::Constant          op_6        0 1 49 value=6
+aten::div               op_7        2 1 c 49 out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.hardswish";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_hardswish_7, 8)
+
 class F_hardswish_onnx : public GraphRewriterPass
 {
 public:

tools/pnnx/src/pass_level2/F_interpolate.cpp

@@ -997,4 +997,118 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_interpolate_7, 10)
 
+class F_interpolate_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_0        0 1 size @data
+Resize                  op_1        2 1 input size out coordinate_transformation_mode=* mode=nearest nearest_mode=floor cubic_coeff_a=*
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.interpolate";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& /*captured_params*/, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        auto size = captured_attrs.at("op_0.data");
+        if (size.type != 5)
+            return false;
+
+        int size_count = size.data.size() / sizeof(int64_t);
+        if (size_count < 3 || size_count > 5)
+            return false;
+
+        const std::vector<int>& input_shape = matched_operators.at("op_1")->inputs[0]->shape;
+        if (input_shape.size() < 3 || input_shape.size() > 5)
+            return false;
+
+        const int64_t* ps = (const int64_t*)size.data.data();
+        if (input_shape[0] != ps[0] || input_shape[1] != ps[1])
+            return false;
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& /*captured_params*/, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        auto size = captured_attrs.at("op_0.data");
+        int size_count = size.data.size() / sizeof(float);
+        const int64_t* ps = (const int64_t*)size.data.data();
+
+        op->params["mode"] = "nearest";
+        op->params["align_corners"] = false;
+        if (size_count == 3)
+            op->params["size"] = {ps[2]};
+        if (size_count == 4)
+            op->params["size"] = {ps[2], ps[3]};
+        if (size_count == 5)
+            op->params["size"] = {ps[2], ps[3], ps[4]};
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_interpolate_onnx, 10)
+
+class F_interpolate_onnx_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+pnnx.Attribute          op_0        0 1 scale_factor @data
+Resize                  op_1        2 1 input scale_factor out coordinate_transformation_mode=* mode=nearest nearest_mode=floor cubic_coeff_a=*
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.interpolate";
+    }
+
+    bool match(const std::map<std::string, Parameter>& /*captured_params*/, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        auto scale_factor = captured_attrs.at("op_0.data");
+        if (scale_factor.type != 1)
+            return false;
+
+        int scale_factor_count = scale_factor.data.size() / sizeof(float);
+        if (scale_factor_count < 3 || scale_factor_count > 5)
+            return false;
+
+        const float* ps = (const float*)scale_factor.data.data();
+        if (ps[0] != 1.f || ps[1] != 1.f)
+            return false;
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& /*captured_params*/, const std::map<std::string, Attribute>& captured_attrs) const
+    {
+        auto scale_factor = captured_attrs.at("op_0.data");
+        int scale_factor_count = scale_factor.data.size() / sizeof(float);
+        const float* ps = (const float*)scale_factor.data.data();
+
+        op->params["mode"] = "nearest";
+        op->params["recompute_scale_factor"] = true;
+        if (scale_factor_count == 3)
+            op->params["scale_factor"] = {ps[2]};
+        if (scale_factor_count == 4)
+            op->params["scale_factor"] = {ps[2], ps[3]};
+        if (scale_factor_count == 5)
+            op->params["scale_factor"] = {ps[2], ps[3], ps[4]};
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_interpolate_onnx_1, 10)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_linear.cpp

@@ -198,7 +198,7 @@ public:
     const char* match_pattern_graph() const
     {
         return R"PNNXIR(7767517
-5 4
+4 3
 pnnx.Input              input_0     0 1 input
 pnnx.Attribute          weight      0 1 weight @data=(%in_features,%out_features)f32
 MatMul                  matmul      2 1 input weight out
@@ -209,7 +209,7 @@ pnnx.Output             output      1 0 out
     const char* replace_pattern_graph() const
     {
         return R"PNNXIR(7767517
-5 4
+4 3
 pnnx.Input              input_0     0 1 input
 pnnx.Attribute          weight      0 1 weight
 F.linear                linear      2 1 input weight out bias=None $weight=weight
@@ -220,4 +220,62 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_linear_onnx_2, 9)
 
+class F_linear_onnx_3 : public F_linear_onnx_1
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+6 5
+pnnx.Input              input_0     0 1 input
+pnnx.Attribute          weight      0 1 weight @data=(%in_features,%out_features)f32
+pnnx.Attribute          bias        0 1 bias @data=(%out_features)f32
+MatMul                  matmul      2 1 input weight mm
+aten::add               add         2 1 mm bias out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* replace_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input_0     0 1 input
+pnnx.Attribute          weight      0 1 weight
+pnnx.Attribute          bias        0 1 bias
+F.linear                linear      3 1 input weight bias out $weight=weight $bias=bias
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    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
+    {
+        F_linear_onnx_1::write(ops, captured_params, captured_attrs);
+
+        Operator* op_bias = ops.at("bias");
+        op_bias->attrs["data"] = captured_attrs.at("bias.data");
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_linear_onnx_3, 8)
+
+class F_linear_onnx_4 : public F_linear_onnx_3
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+6 5
+pnnx.Input              input_0     0 1 input
+pnnx.Attribute          weight      0 1 weight @data=(%in_features,%out_features)f32
+pnnx.Attribute          bias        0 1 bias @data=(%out_features)f32
+MatMul                  matmul      2 1 input weight mm
+aten::add               add         2 1 bias mm out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_linear_onnx_4, 8)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_pad.cpp

@@ -259,4 +259,26 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_pad_8, 10)
 
+class F_pad_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 pad
+Pad                     op_0        2 1 input pad out mode=constant
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.pad";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_pad_onnx, 10)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/F_softmax.cpp

@@ -60,4 +60,25 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_softmax_1, 10)
 
+class F_softmax_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input_0     0 1 input
+Softmax                 op_0        1 1 input out axis=%dim
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.softmax";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(F_softmax_onnx, 10)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/Tensor_reshape.cpp

@@ -44,10 +44,11 @@ public:
     const char* match_pattern_graph() const
     {
         return R"PNNXIR(7767517
-4 3
+5 4
 pnnx.Input              input_0     0 1 input
 pnnx.Input              input_1     0 1 shape
-Reshape                 op_0        2 1 input shape out allowzero=*
+aten::cat               op_0        1 1 shape cat dim=0
+Reshape                 op_1        2 1 input cat out allowzero=*
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
@@ -58,9 +59,27 @@ pnnx.Output             output      1 0 out
     }
 };
 
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx, 20)
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx, 19)
+
+class Tensor_reshape_onnx_1 : public Tensor_reshape_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+5 4
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 shape
+aten::cat               op_0        1 1 shape cat dim=0
+Reshape                 op_1        2 1 input cat out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_1, 19)
 
-class Tensor_reshape_onnx_1 : public GraphRewriterPass
+class Tensor_reshape_onnx_2 : public Tensor_reshape_onnx
 {
 public:
     const char* match_pattern_graph() const
@@ -69,17 +88,29 @@ public:
 4 3
 pnnx.Input              input_0     0 1 input
 pnnx.Input              input_1     0 1 shape
-Reshape                 op_0        2 1 input shape out
+Reshape                 op_1        2 1 input shape out allowzero=*
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
+};
 
-    const char* type_str() const
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_2, 20)
+
+class Tensor_reshape_onnx_3 : public Tensor_reshape_onnx
+{
+public:
+    const char* match_pattern_graph() const
     {
-        return "Tensor.reshape";
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 shape
+Reshape                 op_1        2 1 input shape out
+pnnx.Output             output      1 0 out
+)PNNXIR";
     }
 };
 
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_1, 20)
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_reshape_onnx_3, 20)
 
 } // namespace pnnx

tools/pnnx/src/pass_level2/Tensor_select.cpp

@@ -46,7 +46,7 @@ public:
     {
         return R"PNNXIR(7767517
 3 2
-pnnx.Input              input_0     0 1 input
+pnnx.Input              input       0 1 input
 aten::select            op_0        1 1 input out dim=%dim index=%index
 pnnx.Output             output      1 0 out
 )PNNXIR";
@@ -60,4 +60,26 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_select_1, 20)
 
+class Tensor_select_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+prim::Constant          op_0        0 1 indices value=%index
+Gather                  op_1        2 1 input indices out axis=%dim
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Tensor.select";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_select_onnx, 20)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/Tensor_to.cpp

@@ -119,4 +119,47 @@ REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to, 20)
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to_1, 20)
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to_2, 20)
 
+class Tensor_to_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 input
+Cast                    op_0        1 1 input out to=%to
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Tensor.to";
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        const int to = captured_params.at("to").i;
+
+        op->params["non_blocking"] = false;
+        op->params["copy"] = false;
+        op->params["memory_format"] = "torch.preserve_format";
+
+        if (to == 1) op->params["dtype"] = "torch.float";
+        if (to == 2) op->params["dtype"] = "torch.uint8";
+        if (to == 3) op->params["dtype"] = "torch.int8";
+        if (to == 5) op->params["dtype"] = "torch.short";
+        if (to == 6) op->params["dtype"] = "torch.int";
+        if (to == 7) op->params["dtype"] = "torch.long";
+        if (to == 9) op->params["dtype"] = "torch.bool";
+        if (to == 10) op->params["dtype"] = "torch.half";
+        if (to == 11) op->params["dtype"] = "torch.double";
+        if (to == 14) op->params["dtype"] = "torch.complex64";
+        if (to == 15) op->params["dtype"] = "torch.complex128";
+        if (to == 16) op->params["dtype"] = "torch.bfloat16";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(Tensor_to_onnx, 20)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/torch_clamp.cpp

@@ -41,36 +41,13 @@ REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp, 20)
 
 class torch_clamp_onnx : public GraphRewriterPass
 {
-public:
-    const char* match_pattern_graph() const
-    {
-        return R"PNNXIR(7767517
-5 4
-pnnx.Input              input_0     0 1 input
-pnnx.Input              input_1     0 1 min
-pnnx.Input              input_2     0 1 max
-Clip                    op_0        3 1 input min max out
-pnnx.Output             output      1 0 out
-)PNNXIR";
-    }
-
-    const char* type_str() const
-    {
-        return "torch.clamp";
-    }
-};
-
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp_onnx, 20)
-
-class torch_clamp_onnx_1 : public GraphRewriterPass
-{
 public:
     const char* match_pattern_graph() const
     {
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-Clip                    op_0        1 1 input out min=%min max=%max
+aten::clamp             op_0        1 1 input out min=%min max=%max
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
@@ -81,6 +58,6 @@ pnnx.Output             output      1 0 out
     }
 };
 
-REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp_onnx_1, 20)
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_clamp_onnx, 20)
 
 } // namespace pnnx

tools/pnnx/src/pass_level2/torch_split.cpp

@@ -105,4 +105,21 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_split_onnx, 20)
 
+class torch_split_onnx_1 : public torch_split_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 tensor
+pnnx.Input              input_1     0 1 split_size_or_sections
+aten::split             op_0        2 1 tensor split_size_or_sections out dim=%dim indices=None
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_split_onnx_1, 20)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/torch_squeeze.cpp

@@ -80,4 +80,42 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_squeeze_0, 20)
 
+class torch_squeeze_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 dim
+Squeeze                 op_0        2 1 input dim out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.squeeze";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_squeeze_onnx, 20)
+
+class torch_squeeze_onnx_1 : public torch_squeeze_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 input
+Squeeze                 op_0        1 1 input out axes=%dim
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_squeeze_onnx_1, 20)
+
 } // namespace pnnx

tools/pnnx/src/pass_level2/torch_unsqueeze.cpp

@@ -59,4 +59,42 @@ pnnx.Output             output      1 0 out
 
 REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_unsqueeze_01, 20)
 
+class torch_unsqueeze_onnx : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 dim
+Unsqueeze               op_0        2 1 input dim out
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "torch.unsqueeze";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_unsqueeze_onnx, 20)
+
+class torch_unsqueeze_onnx_1 : public torch_unsqueeze_onnx
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 input
+Unsqueeze               op_0        1 1 input out axes=%dim
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_unsqueeze_onnx_1, 20)
+
 } // namespace pnnx

tools/pnnx/src/pass_level3/fuse_expression.cpp

@@ -54,7 +54,7 @@ static bool operand_maybe_tensor(const Operand* operand)
 
     if (op->type == "aten::size")
     {
-        return false;
+        return op->inputs.size() == 1;
     }
 
     if (op->type == "aten::Int")
@@ -508,8 +508,11 @@ static void fuse_expression(Graph& graph, Operand* operand, std::string& expr, s
     {
         expr += "size(";
         fuse_expression(graph, op->inputs[0], expr, inputs, foldable_constants, zip);
-        expr += ",";
-        fuse_expression(graph, op->inputs[1], expr, inputs, foldable_constants, zip);
+        if (op->inputs.size() == 2)
+        {
+            expr += ",";
+            fuse_expression(graph, op->inputs[1], expr, inputs, foldable_constants, zip);
+        }
         expr += ")";
     }
     else if (op->type == "aten::Int")

tools/pnnx/src/pass_level5.cpp

@@ -49,6 +49,7 @@
 #include "pass_level5/fuse_slice_copy.h"
 #include "pass_level5/fuse_slice_indices.h"
 #include "pass_level5/fuse_slice_to_tensor_split.h"
+#include "pass_level5/fuse_slice_squeeze_to_select.h"
 #include "pass_level5/fuse_static_batchnorm.h"
 #include "pass_level5/fuse_static_conv.h"
 #include "pass_level5/fuse_static_convtranspose.h"
@@ -57,6 +58,7 @@
 #include "pass_level5/fuse_static_instancenorm.h"
 #include "pass_level5/fuse_static_layernorm.h"
 #include "pass_level5/fuse_static_linear.h"
+#include "pass_level5/fuse_swish.h"
 #include "pass_level5/normalize_einsum_equation.h"
 #include "pass_level4/dead_code_elimination.h"
 #include "pass_level4/canonicalize.h"
@@ -77,6 +79,8 @@ void pass_level5(Graph& g, const std::set<std::string>& foldable_constants, cons
 
     eliminate_noop_slice(g);
 
+    fuse_slice_squeeze_to_select(g);
+
     fuse_slice_indices(g);
 
     normalize_einsum_equation(g);
@@ -140,6 +144,8 @@ void pass_level5(Graph& g, const std::set<std::string>& foldable_constants, cons
     fuse_multiheadattention(g);
     fuse_scaled_dot_product_attention(g);
 
+    fuse_swish(g);
+
     fuse_index_expression(g);
 
     dead_code_elimination(g);

tools/pnnx/src/pass_level5/eval_expression.cpp

@@ -119,40 +119,49 @@ static std::string eval_expression(const Operator* op)
         {
             std::string a = exprstack.top();
             exprstack.pop();
-            std::string b = exprstack.top();
-            exprstack.pop();
 
-            if (token_is_argument(a) && token_is_literal(b))
+            if (exprstack.empty())
             {
-                int input_index = std::stoi(a.substr(1));
-                if (op->inputs[input_index]->shape.empty())
-                {
-                    std::string r = std::string("size(") + a + "," + b + ")";
-                    exprstack.push(r);
-                }
-                else
+                std::string r = std::string("size(") + a + ")";
+                exprstack.push(r);
+            }
+            else
+            {
+                std::string b = exprstack.top();
+                exprstack.pop();
+
+                if (token_is_argument(a) && token_is_literal(b))
                 {
-                    int bi = std::stoi(b);
-                    if (bi < 0)
-                        bi = op->inputs[input_index]->shape.size() + bi;
-                    int r = op->inputs[input_index]->shape[bi];
-                    if (r == -1)
+                    int input_index = std::stoi(a.substr(1));
+                    if (op->inputs[input_index]->shape.empty())
                     {
-                        // do not evaluate dynamic size info as -1
-                        // just keep the size expression
                         std::string r = std::string("size(") + a + "," + b + ")";
                         exprstack.push(r);
                     }
                     else
                     {
-                        exprstack.push(std::to_string(r));
+                        int bi = std::stoi(b);
+                        if (bi < 0)
+                            bi = op->inputs[input_index]->shape.size() + bi;
+                        int r = op->inputs[input_index]->shape[bi];
+                        if (r == -1)
+                        {
+                            // do not evaluate dynamic size info as -1
+                            // just keep the size expression
+                            std::string r = std::string("size(") + a + "," + b + ")";
+                            exprstack.push(r);
+                        }
+                        else
+                        {
+                            exprstack.push(std::to_string(r));
+                        }
                     }
                 }
-            }
-            else
-            {
-                std::string r = std::string("size(") + a + "," + b + ")";
-                exprstack.push(r);
+                else
+                {
+                    std::string r = std::string("size(") + a + "," + b + ")";
+                    exprstack.push(r);
+                }
             }
         }
         else if (t == "int"

tools/pnnx/src/pass_level5/fuse_layernorm.cpp

@@ -65,12 +65,112 @@ pnnx.Output             output      1 0 out
     }
 };
 
+class fuse_layernorm_pass_1 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        // clang-format off
+        // *INDENT-OFF*
+
+        return R"PNNXIR(7767517
+9 8
+pnnx.Input              input       0 1 input #input=(1,?,%c)f32
+pnnx.Attribute          op_0        0 1 weight @data #weight=(%c)f32
+pnnx.Attribute          op_1        0 1 bias @data #bias=(%c)f32
+torch.mean              op_2        1 1 input mean dim=(-1) keepdim=True
+pnnx.Expression         op_3        2 1 input mean 173 expr=sub(@0,@1)
+pnnx.Expression         op_4        1 1 173 174 expr=pow(@0,2.000000e+00)
+torch.mean              op_5        1 1 174 var dim=(-1) keepdim=True
+pnnx.Expression         op_6        4 1 173 var weight bias out expr=add(mul(div(@0,sqrt(add(@1,%eps))),@2),@3)
+pnnx.Output             output      1 0 out
+)PNNXIR";
+
+        // *INDENT-ON*
+        // clang-format on
+    }
+
+    const char* replace_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 input
+nn.LayerNorm            ln          1 1 input out elementwise_affine=True eps=%eps normalized_shape=(%c) @weight=%op_0.data @bias=%op_1.data
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+class fuse_layernorm_pass_2 : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        // clang-format off
+        // *INDENT-OFF*
+
+        return R"PNNXIR(7767517
+7 6
+pnnx.Input              input       0 1 input #input=(1,?,%c)f32
+torch.mean              op_0        1 1 input mean dim=(-1) keepdim=True
+pnnx.Expression         op_1        2 1 input mean 25 expr=sub(@0,@1)
+pnnx.Expression         op_2        1 1 25 26 expr=pow(@0,2.000000e+00)
+torch.mean              op_3        1 1 26 var dim=(-1) keepdim=True
+pnnx.Expression         op_4        2 1 25 var out expr=div(@0,sqrt(add(@1,%eps)))
+pnnx.Output             output      1 0 out
+)PNNXIR";
+
+        // *INDENT-ON*
+        // clang-format on
+    }
+
+    const char* replace_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+3 2
+pnnx.Input              input       0 1 input
+nn.LayerNorm            ln          1 1 input out elementwise_affine=False eps=%eps normalized_shape=(%c)
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+class fuse_layernorm_pass_2_1 : public fuse_layernorm_pass_2
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        // clang-format off
+        // *INDENT-OFF*
+
+        return R"PNNXIR(7767517
+7 6
+pnnx.Input              input       0 1 input #input=(1,?,?,%c)f32
+torch.mean              op_0        1 1 input mean dim=(-1) keepdim=True
+pnnx.Expression         op_1        2 1 input mean 32 expr=sub(@0,@1)
+pnnx.Expression         op_2        1 1 32 33 expr=pow(@0,2.000000e+00)
+torch.mean              op_3        1 1 33 var dim=(-1) keepdim=True
+pnnx.Expression         op_4        2 1 32 var out expr=div(@0,sqrt(add(@1,%eps)))
+pnnx.Output             output      1 0 out
+)PNNXIR";
+
+        // *INDENT-ON*
+        // clang-format on
+    }
+};
+
 void fuse_layernorm(Graph& graph)
 {
     fuse_layernorm_pass a;
+    fuse_layernorm_pass_1 b;
+    fuse_layernorm_pass_2 c;
+    fuse_layernorm_pass_2_1 c1;
     int opindex = 0;
 
     pnnx_graph_rewrite(graph, &a, opindex);
+    pnnx_graph_rewrite(graph, &b, opindex);
+    pnnx_graph_rewrite(graph, &c, opindex);
+    pnnx_graph_rewrite(graph, &c1, opindex);
 }
 
 } // namespace pnnx

tools/pnnx/src/pass_level5/fuse_multiheadattention.cpp

@@ -51,7 +51,7 @@ torch.unbind            op_3        1 3 3 4 5 6 dim=0
 pnnx.Expression         op_4        1 1 4 7 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
 Tensor.permute          op_5        1 1 5 8 dims=(0,1,3,2)
 torch.matmul            op_6        2 1 7 8 9
-F.softmax               op_7        1 1 9 10 dim=-1
+F.softmax               softmax     1 1 9 10 dim=%softmax_dim
 torch.matmul            op_8        2 1 10 6 11
 Tensor.permute          op_9        1 1 11 12 dims=(0,2,1,3)
 Tensor.reshape          op_10       1 1 12 13 shape=(%batch,%size,%embed_dim)
@@ -70,13 +70,14 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int qkv_out_features = captured_params.at("qkv_out_features").i;
         const int num_heads = captured_params.at("num_heads").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;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
 
         if (qkv_out_features != embed_dim * 3)
             return false;
@@ -87,6 +88,10 @@ pnnx.Output             output      1 0 out
         if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001))
             return false;
 
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
         return true;
     }
 
@@ -157,7 +162,7 @@ Tensor.permute          op_6        1 1 5 9 dims=(0,2,1,3)
 torch.transpose         op_7        1 1 8 10 dim0=-1 dim1=-2
 torch.matmul            op_8        2 1 9 10 11
 pnnx.Expression         op_9        1 1 11 12 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
-nn.Softmax              op_10       1 1 12 13 dim=-1
+nn.Softmax              softmax     1 1 12 13 dim=%softmax_dim
 Tensor.permute          op_11       1 1 7 14 dims=(0,2,1,3)
 torch.matmul            op_12       2 1 13 14 15
 Tensor.permute          op_13       1 1 15 16 dims=(0,2,1,3)
@@ -168,6 +173,129 @@ pnnx.Output             output      1 0 out
     }
 };
 
+class fuse_multiheadattention_pass_11_0 : public fuse_multiheadattention_pass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+22 21
+pnnx.Input              input       0 1 input
+Tensor.permute          op_1        1 1 input 13 dims=(1,0,2)
+nn.Linear               op_0        1 1 13 14 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight
+Tensor.slice            op_2        1 1 14 15 dim=-1 end=%embed_dim start=0 step=1
+Tensor.slice            op_3        1 1 14 16 dim=-1 end=%embed_dim2 start=%embed_dim step=1
+Tensor.slice            op_4        1 1 14 17 dim=-1 end=%qkv_out_features start=%embed_dim2 step=1
+Tensor.reshape          op_5        1 1 15 18 shape=(%size,%num_heads,%feat_per_head)
+Tensor.permute          op_6        1 1 18 19 dims=(1,0,2)
+Tensor.reshape          op_7        1 1 16 20 shape=(%size,%num_heads,%feat_per_head)
+Tensor.reshape          op_8        1 1 17 21 shape=(%size,%num_heads,%feat_per_head)
+Tensor.permute          op_9        1 1 21 22 dims=(1,0,2)
+pnnx.Expression         op_10       1 1 19 23 expr=div(@0,%sqrt_embed_dim_per_head)
+Tensor.permute          op_11       1 1 20 24 dims=(1,2,0)
+torch.matmul            op_12       2 1 23 24 25
+F.softmax               softmax     1 1 25 26 dim=%softmax_dim
+torch.matmul            op_14       2 1 26 22 27
+Tensor.permute          op_15       1 1 27 28 dims=(1,0,2)
+Tensor.reshape          op_16       1 1 28 29 shape=(%size,%embed_dim)
+nn.Linear               out_proj    1 1 29 30 bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
+Tensor.reshape          op_24       1 1 30 31 shape=(%size,%batch,%embed_dim)
+Tensor.permute          op_25       1 1 31 out dims=(1,0,2)
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
+    {
+        const int embed_dim = captured_params.at("embed_dim").i;
+        const int embed_dim2 = captured_params.at("embed_dim2").i;
+        const int qkv_out_features = captured_params.at("qkv_out_features").i;
+        const int num_heads = captured_params.at("num_heads").i;
+        const int feat_per_head = captured_params.at("feat_per_head").i;
+        const float sqrt_embed_dim_per_head = captured_params.at("sqrt_embed_dim_per_head").f;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
+
+        if (qkv_out_features != embed_dim * 3 || embed_dim2 != embed_dim * 2)
+            return false;
+
+        if (embed_dim != num_heads * feat_per_head)
+            return false;
+
+        if (!NearlyEqual(sqrt_embed_dim_per_head, sqrt(feat_per_head), 0.001))
+            return false;
+
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
+        return true;
+    }
+};
+
+class fuse_multiheadattention_pass_11_1 : public fuse_multiheadattention_pass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+28 27
+pnnx.Input              input       0 1 input
+Tensor.permute          op_1        1 1 input 9 dims=(1,0,2)
+nn.Linear               op_0        1 1 9 10 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight
+Tensor.reshape          op_2        1 1 10 11 shape=(1,%size,%batch,3,%embed_dim)
+Tensor.permute          op_3        1 1 11 12 dims=(3,1,2,0,4)
+torch.squeeze           op_4        1 1 12 13 dim=3
+torch.unbind            op_5        1 3 13 14 15 16 dim=0
+Tensor.reshape          op_6        1 1 14 17 shape=(%size,%num_heads,%feat_per_head)
+Tensor.permute          op_7        1 1 17 18 dims=(1,0,2)
+Tensor.reshape          op_8        1 1 15 19 shape=(%size,%num_heads,%feat_per_head)
+Tensor.permute          op_9        1 1 19 20 dims=(1,0,2)
+Tensor.reshape          op_10       1 1 16 21 shape=(%size,%num_heads,%feat_per_head)
+Tensor.permute          op_11       1 1 21 22 dims=(1,0,2)
+Tensor.reshape          op_12       1 1 18 23 shape=(%batch,%num_heads,%size,%feat_per_head)
+Tensor.reshape          op_13       1 1 20 24 shape=(%batch,%num_heads,%size,%feat_per_head)
+Tensor.reshape          op_14       1 1 22 25 shape=(%batch,%num_heads,%size,%feat_per_head)
+Tensor.permute          op_15       1 1 24 26 dims=(0,1,3,2)
+pnnx.Expression         op_16       1 1 23 27 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
+pnnx.Expression         op_17       1 1 26 28 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
+torch.matmul            op_18       2 1 27 28 29
+F.softmax               softmax     1 1 29 30 dim=%softmax_dim
+torch.matmul            op_20       2 1 30 25 31
+Tensor.permute          op_21       1 1 31 32 dims=(2,0,1,3)
+Tensor.reshape          op_22       1 1 32 33 shape=(%size,%embed_dim)
+nn.Linear               out_proj    1 1 33 34 bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
+Tensor.reshape          op_24       1 1 34 35 shape=(%size,%batch,%embed_dim)
+Tensor.permute          op_25       1 1 35 out dims=(1,0,2)
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
+    {
+        const int embed_dim = captured_params.at("embed_dim").i;
+        const int qkv_out_features = captured_params.at("qkv_out_features").i;
+        const int num_heads = captured_params.at("num_heads").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;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
+
+        if (qkv_out_features != embed_dim * 3)
+            return false;
+
+        if (embed_dim != num_heads * feat_per_head)
+            return false;
+
+        if (!NearlyEqual(inv_sqrt_embed_dim_per_head, sqrt(1.f / sqrt(feat_per_head)), 0.001))
+            return false;
+
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
+        return true;
+    }
+};
+
 class fuse_multiheadattention_pass_sameqkv : public GraphRewriterPass
 {
 public:
@@ -189,7 +317,7 @@ Tensor.permute          op_9        1 1 35 40 dims=(0,2,1,3)
 Tensor.reshape          op_10       1 1 40 41 shape=(%num_heads,%size,%feat_per_head)
 Tensor.permute          op_11       1 1 39 42 dims=(0,2,1)
 torch.matmul            op_12       2 1 41 42 43
-F.softmax               op_13       1 1 43 44 dim=-1
+F.softmax               softmax     1 1 43 44 dim=%softmax_dim
 Tensor.permute          op_14       1 1 37 45 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 45 46 shape=(%num_heads,%size,%feat_per_head)
 torch.matmul            op_16       2 1 44 46 47
@@ -211,12 +339,13 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").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;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
 
         if (embed_dim != num_heads * feat_per_head)
             return false;
@@ -224,6 +353,10 @@ pnnx.Output             output      1 0 out
         if (!NearlyEqual(inv_sqrt_embed_dim_per_head, 1.f / sqrt(feat_per_head), 0.001))
             return false;
 
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
         return true;
     }
 
@@ -326,7 +459,7 @@ Tensor.permute          op_9        1 1 36 41 dims=(0,2,1,3)
 Tensor.reshape          op_10       1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head)
 Tensor.permute          op_11       1 1 40 43 dims=(0,2,1)
 torch.matmul            op_12       2 1 42 43 44
-F.softmax               op_13       1 1 44 45 dim=-1
+F.softmax               softmax     1 1 44 45 dim=%softmax_dim
 Tensor.permute          op_14       1 1 38 46 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.matmul            op_16       2 1 45 47 48
@@ -461,7 +594,7 @@ Tensor.permute          op_9        1 1 36 41 dims=(0,2,1,3)
 Tensor.reshape          op_10       1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head)
 Tensor.permute          op_11       1 1 40 43 dims=(0,2,1)
 torch.matmul            op_12       2 1 42 43 44
-F.softmax               op_13       1 1 44 45 dim=-1
+F.softmax               softmax     1 1 44 45 dim=%softmax_dim
 Tensor.permute          op_14       1 1 38 46 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.matmul            op_16       2 1 45 47 48
@@ -505,7 +638,7 @@ Tensor.permute          op_8        1 1 35 40 dims=(0,2,1,3)
 Tensor.reshape          op_9        1 1 40 41 shape=(%num_heads,%size,%feat_per_head)
 torch.einsum            op_10       2 1 41 39 42 equation=ijl,ikl->ijk
 pnnx.Expression         op_11       1 1 42 43 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
-F.softmax               op_12       1 1 43 44 dim=-1
+F.softmax               softmax     1 1 43 44 dim=%softmax_dim
 Tensor.permute          op_13       1 1 37 45 dims=(0,2,1,3)
 Tensor.reshape          op_14       1 1 45 46 shape=(%num_heads,%size,%feat_per_head)
 torch.einsum            op_15       2 1 44 46 47 equation=ijl,ilk->ijk
@@ -537,7 +670,7 @@ Tensor.permute          op_7        1 1 53 54 dims=(0,1,3,2)
 torch.transpose         op_8        1 1 50 55 dim0=1 dim1=2
 torch.matmul            op_9        2 1 55 54 56
 pnnx.Expression         op_10       1 1 56 57 expr=div(@0,%sqrt_feat_per_head)
-F.softmax               op_11       1 1 57 58 dim=-1
+F.softmax               softmax     1 1 57 58 dim=%softmax_dim
 torch.transpose         op_12       1 1 52 59 dim0=1 dim1=2
 torch.matmul            op_13       2 1 58 59 60
 torch.transpose         op_14       1 1 60 61 dim0=1 dim1=2
@@ -547,12 +680,13 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").i;
         const int feat_per_head = captured_params.at("feat_per_head").i;
         const float sqrt_feat_per_head = captured_params.at("sqrt_feat_per_head").f;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
 
         if (embed_dim != num_heads * feat_per_head)
             return false;
@@ -560,6 +694,10 @@ pnnx.Output             output      1 0 out
         if (!NearlyEqual(sqrt_feat_per_head, sqrt(feat_per_head), 0.001))
             return false;
 
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
         return true;
     }
 };
@@ -585,7 +723,7 @@ Tensor.permute          op_7        1 1 53 54 dims=(0,1,3,2)
 torch.transpose         op_8        1 1 50 55 dim0=1 dim1=2
 torch.matmul            op_9        2 1 55 54 56
 pnnx.Expression         op_10       1 1 56 57 expr=div(@0,%sqrt_feat_per_head)
-F.softmax               op_11       1 1 57 58 dim=-1
+F.softmax               softmax     1 1 57 58 dim=%softmax_dim
 torch.transpose         op_12       1 1 52 59 dim0=1 dim1=2
 torch.matmul            op_13       2 1 58 59 60
 torch.transpose         op_14       1 1 60 61 dim0=1 dim1=2
@@ -595,12 +733,13 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").i;
         const int feat_per_head = captured_params.at("feat_per_head").i;
         const float sqrt_feat_per_head = captured_params.at("sqrt_feat_per_head").f;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
 
         if (embed_dim != num_heads * feat_per_head)
             return false;
@@ -608,6 +747,10 @@ pnnx.Output             output      1 0 out
         if (!NearlyEqual(sqrt_feat_per_head, sqrt(feat_per_head), 0.001))
             return false;
 
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
         return true;
     }
 };
@@ -634,7 +777,7 @@ Tensor.permute          op_8        1 1 36 41 dims=(0,2,1,3)
 Tensor.reshape          op_9        1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head)
 torch.einsum            op_10       2 1 42 40 43 equation=ijl,ikl->ijk
 pnnx.Expression         op_11       1 1 43 44 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
-F.softmax               op_12       1 1 44 45 dim=-1
+F.softmax               softmax     1 1 44 45 dim=%softmax_dim
 Tensor.permute          op_13       1 1 38 46 dims=(0,2,1,3)
 Tensor.reshape          op_14       1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.einsum            op_15       2 1 45 47 48 equation=ijl,ilk->ijk
@@ -668,7 +811,7 @@ Tensor.permute          op_8        1 1 36 41 dims=(0,2,1,3)
 Tensor.reshape          op_9        1 1 41 42 shape=(%num_heads,%qsize,%feat_per_head)
 torch.einsum            op_10       2 1 42 40 43 equation=ijl,ikl->ijk
 pnnx.Expression         op_11       1 1 43 44 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
-F.softmax               op_12       1 1 44 45 dim=-1
+F.softmax               softmax     1 1 44 45 dim=%softmax_dim
 Tensor.permute          op_13       1 1 38 46 dims=(0,2,1,3)
 Tensor.reshape          op_14       1 1 46 47 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.einsum            op_15       2 1 45 47 48 equation=ijl,ilk->ijk
@@ -702,7 +845,7 @@ Tensor.reshape          op_9        1 1 41 42 shape=(%num_heads,%size,%feat_per_
 pnnx.Attribute          op_10       0 1 43 @data
 torch.transpose         op_11       1 1 42 44 dim0=-1 dim1=-2
 torch.baddbmm           op_12       3 1 43 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0
-F.softmax               op_13       1 1 45 46 dim=-1
+F.softmax               softmax     1 1 45 46 dim=%softmax_dim
 Tensor.permute          op_14       1 1 38 47 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 47 48 shape=(%num_heads,%size,%feat_per_head)
 torch.bmm               op_16       2 1 46 48 49
@@ -739,7 +882,7 @@ pnnx.Expression         op_10       2 1 40 42 43 expr=%expr_zero_shape
 torch.empty             op_11       1 1 43 zeros
 torch.transpose         op_12       1 1 42 44 dim0=-1 dim1=-2
 torch.baddbmm           op_13       3 1 zeros 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0
-F.softmax               op_14       1 1 45 46 dim=-1
+F.softmax               softmax     1 1 45 46 dim=%softmax_dim
 Tensor.permute          op_15       1 1 38 47 dims=(0,2,1,3)
 Tensor.reshape          op_16       1 1 47 48 shape=(%num_heads,%size,%feat_per_head)
 torch.bmm               op_17       2 1 46 48 49
@@ -777,7 +920,7 @@ Tensor.reshape          op_9        1 1 41 42 shape=(%num_heads,%kvsize,%feat_pe
 pnnx.Attribute          op_10       0 1 43 @data
 torch.transpose         op_11       1 1 42 44 dim0=-1 dim1=-2
 torch.baddbmm           op_12       3 1 43 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0
-F.softmax               op_13       1 1 45 46 dim=-1
+F.softmax               softmax     1 1 45 46 dim=%softmax_dim
 Tensor.permute          op_14       1 1 38 47 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.bmm               op_16       2 1 46 48 49
@@ -814,7 +957,7 @@ Tensor.reshape          op_9        1 1 41 42 shape=(%num_heads,%kvsize,%feat_pe
 pnnx.Attribute          op_10       0 1 43 @data
 torch.transpose         op_11       1 1 42 44 dim0=-1 dim1=-2
 torch.baddbmm           op_12       3 1 43 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0
-F.softmax               op_13       1 1 45 46 dim=-1
+F.softmax               softmax     1 1 45 46 dim=%softmax_dim
 Tensor.permute          op_14       1 1 38 47 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.bmm               op_16       2 1 46 48 49
@@ -853,7 +996,7 @@ pnnx.Expression         op_10       1 1 40 43 expr=%expr_zero_shape
 torch.empty             op_11       1 1 43 zeros
 torch.transpose         op_12       1 1 42 44 dim0=-1 dim1=-2
 torch.baddbmm           op_13       3 1 zeros 40 44 45 alpha=%inv_sqrt_embed_dim_per_head beta=0
-F.softmax               op_14       1 1 45 46 dim=-1
+F.softmax               softmax     1 1 45 46 dim=%softmax_dim
 Tensor.permute          op_15       1 1 38 47 dims=(0,2,1,3)
 Tensor.reshape          op_16       1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.bmm               op_17       2 1 46 48 49
@@ -891,7 +1034,7 @@ pnnx.Expression         op_10       1 1 40 43 expr=%expr_zero_shape
 torch.empty             op_11       1 1 43 zeros
 torch.transpose         op_12       1 1 42 44 dim0=-1 dim1=-2
 torch.baddbmm           op_13       3 1 zeros 40 44 45 alpha=%alpha beta=0
-F.softmax               op_14       1 1 45 46 dim=-1
+F.softmax               softmax     1 1 45 46 dim=%softmax_dim
 Tensor.permute          op_15       1 1 38 47 dims=(0,2,1,3)
 Tensor.reshape          op_16       1 1 47 48 shape=(%num_heads,%kvsize,%feat_per_head)
 torch.bmm               op_17       2 1 46 48 49
@@ -929,7 +1072,7 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").i;
@@ -995,7 +1138,7 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").i;
@@ -1034,7 +1177,7 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& /*matched_operators*/, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").i;
@@ -1070,7 +1213,7 @@ Tensor.reshape          op_11       1 1 10 12 shape=(%num_heads,%batch_mul_size,
 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
+F.softmax               softmax     1 1 15 16 dim=%softmax_dim
 torch.bmm               op_16       2 1 16 17 18
 Tensor.view             op_17       1 1 18 19 shape=(%batch,%num_heads,%size,%feat_per_head)
 torch.transpose         op_18       1 1 19 20 dim0=1 dim1=2
@@ -1108,7 +1251,7 @@ Tensor.view             op_15       1 1 16 17 shape=(%batch,%num_heads,%size,%si
 pnnx.Attribute          attn_mask   0 1 attn_mask @data=(1,1,%size,%size)f32
 pnnx.Expression         op_16       2 1 17 attn_mask 18 expr=add(@0,@1)
 Tensor.view             op_17       1 1 18 19 shape=(%num_heads,%size,%size)
-F.softmax               op_18       1 1 19 20 dim=-1
+F.softmax               softmax     1 1 19 20 dim=%softmax_dim
 torch.bmm               op_19       2 1 20 21 22
 Tensor.view             op_20       1 1 22 23 shape=(%batch,%num_heads,%size,%feat_per_head)
 torch.transpose         op_21       1 1 23 24 dim0=1 dim1=2
@@ -1160,7 +1303,7 @@ torch.transpose         op_9        1 1 37 39 dim0=-1 dim1=-2
 torch.matmul            op_10       2 1 38 39 40
 pnnx.Attribute          attn_mask   0 1 attn_mask @data=(1,1,1,%size)f32
 pnnx.Expression         op_11       2 1 40 attn_mask 41 expr=add(div(@0,%sqrt_feat_per_head),@1)
-F.softmax               op_12       1 1 41 42 dim=-1
+F.softmax               softmax     1 1 41 42 dim=%softmax_dim
 torch.matmul            op_13       2 1 42 43 44
 Tensor.permute          op_14       1 1 44 45 dims=(0,2,1,3)
 Tensor.reshape          op_15       1 1 45 46 shape=(%batch,%size,%embed_dim)
@@ -1180,12 +1323,13 @@ pnnx.Output             output      1 0 out
 )PNNXIR";
     }
 
-    bool match(const std::map<std::string, Parameter>& captured_params) const
+    bool match(const std::map<std::string, const Operator*>& matched_operators, const std::map<std::string, Parameter>& captured_params, const std::map<std::string, Attribute>& /*captured_attrs*/) const
     {
         const int embed_dim = captured_params.at("embed_dim").i;
         const int num_heads = captured_params.at("num_heads").i;
         const int feat_per_head = captured_params.at("feat_per_head").i;
         const float sqrt_feat_per_head = captured_params.at("sqrt_feat_per_head").f;
+        const int softmax_dim = captured_params.at("softmax_dim").i;
 
         if (embed_dim != num_heads * feat_per_head)
             return false;
@@ -1193,6 +1337,10 @@ pnnx.Output             output      1 0 out
         if (!NearlyEqual(sqrt_feat_per_head, sqrt(feat_per_head), 0.001))
             return false;
 
+        int softmax_input_rank = (int)matched_operators.at("softmax")->inputs[0]->shape.size();
+        if (softmax_dim != -1 && softmax_dim != softmax_input_rank - 1)
+            return false;
+
         return true;
     }
 
@@ -1237,7 +1385,7 @@ torch.transpose         op_5        1 1 12 15 dim0=-2 dim1=-1
 torch.matmul            op_6        2 1 14 15 16
 pnnx.Attribute          attn_mask   0 1 attn_mask @data=(1,%num_heads,%size,%size)f32
 pnnx.Expression         op_7        2 1 16 attn_mask 18 expr=add(@0,@1)
-F.softmax               op_8        1 1 18 19 dim=-1
+F.softmax               softmax     1 1 18 19 dim=%softmax_dim
 torch.matmul            op_9        2 1 19 13 20
 torch.transpose         op_10       1 1 20 21 dim0=1 dim1=2
 Tensor.reshape          op_11       1 1 21 22 shape=(%batch,%size,%embed_dim)
@@ -1281,6 +1429,33 @@ pnnx.Output             output      1 0 out
     }
 };
 
+class fuse_multiheadattention_pass_17_1 : public fuse_multiheadattention_pass_17
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+17 16
+pnnx.Input              input_0     0 1 input
+nn.Linear               op_0        1 1 input 8 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight
+Tensor.reshape          op_1        1 1 8 9 shape=(%batch,%size,3,%num_heads,%feat_per_head)
+Tensor.permute          op_2        1 1 9 10 dims=(2,0,3,1,4)
+torch.unbind            op_3        1 3 10 11 12 13 dim=0
+pnnx.Expression         op_4        1 1 11 14 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
+Tensor.permute          op_5        1 1 12 15 dims=(0,1,3,2)
+torch.matmul            op_6        2 1 14 15 16
+pnnx.Attribute          attn_mask   0 1 attn_mask @data=(1,%num_heads,%size,%size)f32
+pnnx.Expression         op_7        2 1 16 attn_mask 18 expr=add(@0,@1)
+F.softmax               softmax     1 1 18 19 dim=%softmax_dim
+torch.matmul            op_9        2 1 19 13 20
+Tensor.permute          op_10       1 1 20 21 dims=(0,2,1,3)
+Tensor.reshape          op_11       1 1 21 22 shape=(%batch,%size,%embed_dim)
+nn.Linear               out_proj    1 1 22 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
 class fuse_multiheadattention_pass_18 : public fuse_multiheadattention_pass
 {
 public:
@@ -1302,7 +1477,7 @@ Tensor.view             op_8        1 1 35 36 shape=(1,%batch,%num_heads,%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=(%batch,%num_heads,%size,%size)
-F.softmax               op_12       1 1 39 40 dim=-1
+F.softmax               softmax     1 1 39 40 dim=%softmax_dim
 torch.matmul            op_13       2 1 40 30 41
 torch.transpose         op_14       1 1 41 42 dim0=1 dim1=2
 Tensor.reshape          op_15       1 1 42 43 shape=(%batch,%size,%embed_dim)
@@ -1368,11 +1543,44 @@ pnnx.Output             output      1 0 out
     }
 };
 
+class fuse_multiheadattention_pass_18_1 : public fuse_multiheadattention_pass_18
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+20 19
+pnnx.Input              input_0     0 1 input
+nn.Linear               op_0        1 1 input 25 bias=%qkvbias in_features=%embed_dim out_features=%qkv_out_features @bias @weight
+Tensor.reshape          op_1        1 1 25 26 shape=(%batch,%size,3,%num_heads,%feat_per_head)
+Tensor.permute          op_2        1 1 26 27 dims=(2,0,3,1,4)
+torch.unbind            op_3        1 3 27 28 29 30 dim=0
+pnnx.Expression         op_4        1 1 28 31 expr=mul(@0,%inv_sqrt_embed_dim_per_head)
+Tensor.permute          op_5        1 1 29 32 dims=(0,1,3,2)
+torch.matmul            op_6        2 1 31 32 33
+pnnx.Attribute          attn_mask   0 1 attn_mask @data=(1,%num_heads,%size,%size)f32
+pnnx.Expression         op_7        2 1 33 attn_mask 35 expr=add(@0,@1)
+Tensor.reshape          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.reshape          op_11       1 1 38 39 shape=(%batch,%num_heads,%size,%size)
+F.softmax               softmax     1 1 39 40 dim=%softmax_dim
+torch.matmul            op_13       2 1 40 30 41
+Tensor.permute          op_14       1 1 41 42 dims=(0,2,1,3)
+Tensor.reshape          op_15       1 1 42 43 shape=(%batch,%size,%embed_dim)
+nn.Linear               out_proj    1 1 43 out bias=%outbias in_features=%embed_dim out_features=%embed_dim @bias @weight
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
 void fuse_multiheadattention(Graph& graph)
 {
 #if TORCH_VERSION_MAJOR >= 2 || (TORCH_VERSION_MAJOR >= 1 && TORCH_VERSION_MINOR >= 9)
     fuse_multiheadattention_pass a;
     fuse_multiheadattention_pass_11 a1;
+    fuse_multiheadattention_pass_11_0 a2;
+    fuse_multiheadattention_pass_11_1 a3;
     fuse_multiheadattention_pass_sameqkv b;
     fuse_multiheadattention_pass_qkv c;
     fuse_multiheadattention_pass_q_samekv d;
@@ -1395,11 +1603,15 @@ void fuse_multiheadattention(Graph& graph)
     fuse_multiheadattention_pass_16 o;
     fuse_multiheadattention_pass_16_1 o1;
     fuse_multiheadattention_pass_17 p;
+    fuse_multiheadattention_pass_17_1 p1;
     fuse_multiheadattention_pass_18 q;
+    fuse_multiheadattention_pass_18_1 q1;
     int opindex = 0;
 
     pnnx_graph_rewrite(graph, &a, opindex);
     pnnx_graph_rewrite(graph, &a1, opindex);
+    pnnx_graph_rewrite(graph, &a2, opindex);
+    pnnx_graph_rewrite(graph, &a3, opindex);
     pnnx_graph_rewrite(graph, &b, opindex);
     pnnx_graph_rewrite(graph, &c, opindex);
     pnnx_graph_rewrite(graph, &d, opindex);
@@ -1422,7 +1634,9 @@ void fuse_multiheadattention(Graph& graph)
     pnnx_graph_rewrite(graph, &o, opindex);
     pnnx_graph_rewrite(graph, &o1, opindex);
     pnnx_graph_rewrite(graph, &p, opindex);
+    pnnx_graph_rewrite(graph, &p1, opindex);
     pnnx_graph_rewrite(graph, &q, opindex);
+    pnnx_graph_rewrite(graph, &q1, opindex);
 #endif
 }
 

tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.cpp

@@ -0,0 +1,86 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "fuse_slice_squeeze_to_select.h"
+
+#include "pass_level2.h"
+
+namespace pnnx {
+
+class fuse_slice_squeeze_to_select_pass : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+Tensor.slice            op_0        1 1 input a dim=%dim start=%start end=%end step=1
+torch.squeeze           op_1        1 1 a out dim=%squeeze_dim
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "Tensor.select";
+    }
+
+    const char* name_str() const
+    {
+        return "select";
+    }
+
+    bool match(const std::map<std::string, Parameter>& captured_params) const
+    {
+        const int start = captured_params.at("start").i;
+        const int end = captured_params.at("end").i;
+        if (end != start + 1)
+            return false;
+
+        const int dim = captured_params.at("dim").i;
+        int squeeze_dim;
+        if (captured_params.at("squeeze_dim").type == 2)
+        {
+            squeeze_dim = captured_params.at("squeeze_dim").i;
+        }
+        else // if (captured_params.at("squeeze_dim").type == 5)
+        {
+            if (captured_params.at("squeeze_dim").ai.size() != 1)
+                return false;
+
+            squeeze_dim = captured_params.at("squeeze_dim").ai[0];
+        }
+        if (squeeze_dim != dim)
+            return false;
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
+    {
+        op->params["dim"] = captured_params.at("dim");
+        op->params["index"] = captured_params.at("start");
+    }
+};
+
+void fuse_slice_squeeze_to_select(Graph& graph)
+{
+    fuse_slice_squeeze_to_select_pass a;
+    int opindex = 0;
+
+    pnnx_graph_rewrite(graph, &a, opindex);
+}
+
+} // namespace pnnx

tools/pnnx/src/pass_level5/fuse_slice_squeeze_to_select.h

@@ -0,0 +1,21 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "ir.h"
+
+namespace pnnx {
+
+void fuse_slice_squeeze_to_select(Graph& graph);
+
+} // namespace pnnx

tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp

@@ -35,19 +35,19 @@ void fuse_slice_to_tensor_split(Graph& graph)
 
             Operand* op_in = op->inputs[0];
 
-            if (op->params.find("dims") == op->params.end()
-                    || op->params.find("starts") == op->params.end()
-                    || op->params.find("ends") == op->params.end()
-                    || op->params.find("steps") == op->params.end())
+            if ((!op->has_param("dim") && !op->has_param("dims"))
+                    || (!op->has_param("start") && !op->has_param("starts"))
+                    || (!op->has_param("end") && !op->has_param("ends"))
+                    || (!op->has_param("step") && !op->has_param("steps")))
                 continue;
 
-            if (op->params.at("dims").ai.size() != 1)
+            if (!op->has_param("dim") && op->params.at("dims").ai.size() != 1)
                 continue;
 
-            int dim = op->params.at("dims").ai[0];
-            int start = op->params.at("starts").ai[0];
-            int end = op->params.at("ends").ai[0];
-            int step = op->params.at("steps").ai[0];
+            int dim = op->has_param("dim") ? op->params.at("dim").i : op->params.at("dims").ai[0];
+            int start = op->has_param("start") ? op->params.at("start").i : op->params.at("starts").ai[0];
+            int end = op->has_param("end") ? op->params.at("end").i : op->params.at("ends").ai[0];
+            int step = op->has_param("step") ? op->params.at("step").i : op->params.at("steps").ai[0];
             if (start != 0 || step != 1)
                 continue;
 
@@ -74,18 +74,18 @@ void fuse_slice_to_tensor_split(Graph& graph)
                     if (x->inputs[0] != op_in)
                         continue;
 
-                    if (x->params.find("dims") == x->params.end()
-                            || x->params.find("starts") == x->params.end()
-                            || x->params.find("ends") == x->params.end()
-                            || x->params.find("steps") == x->params.end())
+                    if ((!x->has_param("dim") && !x->has_param("dims"))
+                            || (!x->has_param("start") && !x->has_param("starts"))
+                            || (!x->has_param("end") && !x->has_param("ends"))
+                            || (!x->has_param("step") && !x->has_param("steps")))
                         continue;
 
-                    if (x->params.at("dims").ai.size() != 1)
+                    if (!x->has_param("dim") && x->params.at("dims").ai.size() != 1)
                         continue;
 
-                    int dim2 = x->params.at("dims").ai[0];
-                    int start2 = x->params.at("starts").ai[0];
-                    int step2 = x->params.at("steps").ai[0];
+                    int dim2 = x->has_param("dim") ? x->params.at("dim").i : x->params.at("dims").ai[0];
+                    int start2 = x->has_param("start") ? x->params.at("start").i : x->params.at("starts").ai[0];
+                    int step2 = x->has_param("step") ? x->params.at("step").i : x->params.at("steps").ai[0];
                     if (step2 != 1)
                         continue;
 
@@ -102,7 +102,7 @@ void fuse_slice_to_tensor_split(Graph& graph)
                 if (std::find(graph.ops.begin(), graph.ops.end(), op2) < std::find(graph.ops.begin(), graph.ops.end(), cur))
                     cur = op2;
 
-                int end2 = op2->params.at("ends").ai[0];
+                int end2 = op2->has_param("end") ? op2->params.at("end").i : op2->params.at("ends").ai[0];
                 if (end2 == INT_MAX)
                 {
                     slice_n_ops.push_back(op2);

tools/pnnx/src/pass_level5/fuse_static_groupnorm.cpp

@@ -41,7 +41,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.GroupNorm            group_norm  1 1 input out num_channels=%num_channels num_groups=%num_groups eps=%eps affine=True @weight=%op_weight.data @bias=%op_bias.data
+nn.GroupNorm            gn          1 1 input out num_channels=%num_channels num_groups=%num_groups eps=%eps affine=True @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }

tools/pnnx/src/pass_level5/fuse_static_instancenorm.cpp

@@ -41,7 +41,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.InstanceNorm1d       instance_norm 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data
+nn.InstanceNorm1d       in          1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
@@ -73,7 +73,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.InstanceNorm2d       instance_norm 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data
+nn.InstanceNorm2d       in          1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }
@@ -99,7 +99,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.InstanceNorm3d       instance_norm 1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data
+nn.InstanceNorm3d       in          1 1 input out num_features=%num_features eps=%eps affine=True track_running_stats=False @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }

tools/pnnx/src/pass_level5/fuse_static_layernorm.cpp

@@ -41,7 +41,7 @@ pnnx.Output             output      1 0 out
         return R"PNNXIR(7767517
 3 2
 pnnx.Input              input       0 1 input
-nn.LayerNorm            layer_norm  1 1 input out normalized_shape=%normalized_shape eps=%eps elementwise_affine=True @weight=%op_weight.data @bias=%op_bias.data
+nn.LayerNorm            ln          1 1 input out normalized_shape=%normalized_shape eps=%eps elementwise_affine=True @weight=%op_weight.data @bias=%op_bias.data
 pnnx.Output             output      1 0 out
 )PNNXIR";
     }

tools/pnnx/src/pass_level5/fuse_swish.cpp

@@ -0,0 +1,71 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2023 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "fuse_swish.h"
+
+#include "pass_level2.h"
+
+namespace pnnx {
+
+class fuse_swish_pass : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+F.sigmoid               op_0        1 1 input a
+pnnx.Expression         op_1        2 1 input a out expr=mul(@0,@1)
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "F.swish";
+    }
+
+    const char* name_str() const
+    {
+        return "swish";
+    }
+};
+
+class fuse_swish_pass_1 : public fuse_swish_pass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+4 3
+pnnx.Input              input       0 1 input
+nn.Sigmoid              op_0        1 1 input a
+pnnx.Expression         op_1        2 1 input a out expr=mul(@0,@1)
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+};
+
+void fuse_swish(Graph& graph)
+{
+    fuse_swish_pass a;
+    fuse_swish_pass_1 b;
+    int opindex = 0;
+
+    pnnx_graph_rewrite(graph, &a, opindex);
+    pnnx_graph_rewrite(graph, &b, opindex);
+}
+
+} // namespace pnnx

tools/pnnx/src/pass_level5/fuse_swish.h

@@ -0,0 +1,21 @@
+// Tencent is pleased to support the open source community by making ncnn available.
+//
+// Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
+//
+// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// https://opensource.org/licenses/BSD-3-Clause
+//
+// Unless required by applicable law or agreed to in writing, software distributed
+// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+// CONDITIONS OF ANY KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations under the License.
+
+#include "ir.h"
+
+namespace pnnx {
+
+void fuse_swish(Graph& graph);
+
+} // namespace pnnx

tools/pnnx/src/pass_onnx.cpp

@@ -571,6 +571,62 @@ static void fuse_list_unpack(Graph& graph)
     }
 }
 
+static void constant_unpooling(Graph& graph)
+{
+    while (1)
+    {
+        bool matched = false;
+
+        for (size_t i = 0; i < graph.ops.size(); i++)
+        {
+            Operator* op = graph.ops[i];
+
+            if (op->type != "prim::Constant" && op->type != "pnnx.Expression" && op->type != "pnnx.Attribute")
+                continue;
+
+            Operand* op_out = op->outputs[0];
+            if (op_out->consumers.size() == 1)
+                continue;
+
+            matched = true;
+
+            // create shadow node for all consumers
+            for (size_t j = 1; j < op_out->consumers.size(); j++)
+            {
+                Operator* op1 = op_out->consumers[j];
+
+                Operator* op0 = graph.new_operator_before(op->type, op->name + "_pnnxshadow" + std::to_string(j), op1);
+                op0->inputnames = op->inputnames;
+                op0->params = op->params;
+                op0->attrs = op->attrs;
+
+                Operand* op0_out = graph.new_operand(op_out->name + "_pnnxshadow" + std::to_string(j));
+                op0_out->type = op_out->type;
+                op0_out->shape = op_out->shape;
+                op0_out->params = op_out->params;
+
+                op0_out->producer = op0;
+                op0->outputs.push_back(op0_out);
+
+                for (size_t k = 0; k < op1->inputs.size(); k++)
+                {
+                    if (op1->inputs[k] == op_out)
+                        op1->inputs[k] = op0_out;
+                }
+
+                op0_out->consumers.push_back(op1);
+            }
+
+            op_out->consumers.resize(1);
+
+            break;
+        }
+
+        if (!matched)
+            break;
+    }
+}
+
 void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
 {
     onnx2pnnx::OnnxModelProxy modelproxy(model);
@@ -634,6 +690,11 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
                 sim_op_type = "aten::split";
             }
 
+            if (op_type == "Shape")
+            {
+                sim_op_type = "aten::size";
+            }
+
             // unaryop
             if (op_type == "Abs") sim_op_type = "aten::abs";
             if (op_type == "Acos") sim_op_type = "aten::acos";
@@ -645,12 +706,14 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
             if (op_type == "Ceil") sim_op_type = "aten::ceil";
             if (op_type == "Cos") sim_op_type = "aten::cos";
             if (op_type == "Cosh") sim_op_type = "aten::cosh";
+            if (op_type == "Erf") sim_op_type = "aten::erf";
             if (op_type == "Exp") sim_op_type = "aten::exp";
             if (op_type == "Floor") sim_op_type = "aten::floor";
             if (op_type == "Log") sim_op_type = "aten::log";
             if (op_type == "Neg") sim_op_type = "aten::neg";
             if (op_type == "Reciprocal") sim_op_type = "aten::reciprocal";
             if (op_type == "Round") sim_op_type = "aten::round";
+            if (op_type == "Sigmoid") sim_op_type = "aten::sigmoid";
             if (op_type == "Sign") sim_op_type = "aten::sign";
             if (op_type == "Sin") sim_op_type = "aten::sin";
             if (op_type == "Sinh") sim_op_type = "aten::sinh";
@@ -666,6 +729,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
             if (op_type == "Max") sim_op_type = "aten::max";
             if (op_type == "Min") sim_op_type = "aten::min";
             if (op_type == "Pow") sim_op_type = "aten::pow";
+
+            // trinaryop
+            if (op_type == "Clip") sim_op_type = "aten::clamp";
         }
         else if (string_starts_with(op_type, "aten_"))
         {
@@ -716,7 +782,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
             if (input.empty())
                 continue;
 
-            if (modelproxy.has_initializer(input))
+            Operand* op_in = pnnx_graph.get_operand(input);
+
+            if (!op_in && modelproxy.has_initializer(input))
             {
                 // skip function weight
                 if (is_function_op)
@@ -733,6 +801,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
                     if (sim_op_type == "Reshape" && j == 1)
                         is_attr_list = true;
 
+                    if (sim_op_type == "Pad" && j == 1)
+                        is_attr_list = true;
+
                     if (sim_op_type == "ReduceMean" && j == 1)
                         is_attr_list = true;
                 }
@@ -893,9 +964,9 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
 
                     op_const->attrs["data"] = tensor;
                 }
-            }
 
-            Operand* op_in = pnnx_graph.get_operand(input);
+                op_in = pnnx_graph.get_operand(input);
+            }
 
             op_in->consumers.push_back(op);
             op->inputs.push_back(op_in);
@@ -1039,6 +1110,8 @@ void pass_onnx(const onnx::ModelProto& model, Graph& pnnx_graph)
 
     // post process
     fuse_list_unpack(pnnx_graph);
+
+    constant_unpooling(pnnx_graph);
 }
 
 } // namespace pnnx

tools/pnnx/src/pass_onnx/eliminate_noop.cpp

@@ -33,7 +33,7 @@ void eliminate_noop(onnx::ModelProto& model)
         const onnx::NodeProto& node = graph->node(i);
         const std::string& op_type = node.op_type();
 
-        if (op_type == "Identity" || op_type == "aten_copy")
+        if (op_type == "Identity" || op_type == "Dropout" || op_type == "aten_copy")
         {
             const std::string& input_name = node.input(0);
             const std::string& output_name = node.output(0);

tools/pnnx/src/pass_onnx/fold_constants.cpp

@@ -62,8 +62,36 @@ static size_t sizeof_onnx_datatype(ONNXTensorElementDataType type)
     return 0;
 }
 
-void fold_constants(onnx::ModelProto& model)
+static ONNXTensorElementDataType get_onnx_tensor_elem_data_type(const std::string& type)
 {
+    if (type == "i8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8;
+    if (type == "u8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8;
+    if (type == "i16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16;
+    if (type == "u16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16;
+    if (type == "i32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32;
+    if (type == "u32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32;
+    if (type == "i64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64;
+    if (type == "u64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64;
+    if (type == "f16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
+    if (type == "f32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
+    if (type == "f64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE;
+    if (type == "bf16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16;
+    if (type == "c64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64;
+    if (type == "c128") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128;
+
+    // unknown
+    fprintf(stderr, "unsupported tensor elem data type %s\n", type.c_str());
+    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
+}
+
+void fold_constants(onnx::ModelProto& model,
+                    const std::vector<std::vector<int64_t> >& input_shapes,
+                    const std::vector<std::string>& input_types,
+                    const std::vector<std::vector<int64_t> >& input_shapes2,
+                    const std::vector<std::string>& input_types2)
+{
+    bool ignore_aten_size = input_shapes2.empty();
+
     // collect initializers
     std::unordered_set<std::string> initializers;
     {
@@ -114,13 +142,13 @@ void fold_constants(onnx::ModelProto& model)
                     || op_type == "aten_ones_like"
                     || op_type == "aten_zeros_like")
             {
-                is_outputs_foldable = true;
+                is_outputs_foldable = ignore_aten_size;
             }
 
             // TODO whitelist for static shape
             if (op_type == "Shape")
             {
-                is_outputs_foldable = true;
+                is_outputs_foldable = ignore_aten_size;
             }
 
             // TODO whitelist for static type
@@ -263,14 +291,22 @@ void fold_constants(onnx::ModelProto& model)
             const onnx::ValueInfoProto& value = graph->input(i);
 
             std::vector<int64_t> shape;
-            const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape();
-            for (int k = 0; k < tsp.dim_size(); k++)
+            ONNXTensorElementDataType datatype;
+            if (!input_shapes.empty())
             {
-                // TODO has_dim_value ?
-                shape.push_back(tsp.dim(k).dim_value());
+                shape = input_shapes[i];
+                datatype = get_onnx_tensor_elem_data_type(input_types[i]);
             }
+            else
+            {
+                const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape();
+                for (int k = 0; k < tsp.dim_size(); k++)
+                {
+                    shape.push_back(tsp.dim(k).dim_value());
+                }
 
-            ONNXTensorElementDataType datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type();
+                datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type();
+            }
 
             OrtValue* ort_val = 0;
             ort_status = ort_api->CreateTensorAsOrtValue(ort_allocator, (const int64_t*)shape.data(), shape.size(), datatype, &ort_val);

tools/pnnx/src/pass_onnx/fold_constants.h

@@ -18,7 +18,11 @@ namespace pnnx {
 
 namespace onnx2pnnx {
 
-void fold_constants(onnx::ModelProto& model);
+void fold_constants(onnx::ModelProto& model,
+                    const std::vector<std::vector<int64_t> >& input_shapes,
+                    const std::vector<std::string>& input_types,
+                    const std::vector<std::vector<int64_t> >& input_shapes2,
+                    const std::vector<std::string>& input_types2);
 
 } // namespace onnx2pnnx
 

tools/pnnx/src/pass_onnx/shape_inference.cpp

@@ -24,12 +24,38 @@ namespace pnnx {
 
 namespace onnx2pnnx {
 
+static ONNXTensorElementDataType get_onnx_tensor_elem_data_type(const std::string& type)
+{
+    if (type == "i8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8;
+    if (type == "u8") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8;
+    if (type == "i16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16;
+    if (type == "u16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16;
+    if (type == "i32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32;
+    if (type == "u32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32;
+    if (type == "i64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64;
+    if (type == "u64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64;
+    if (type == "f16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16;
+    if (type == "f32") return ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
+    if (type == "f64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE;
+    if (type == "bf16") return ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16;
+    if (type == "c64") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX64;
+    if (type == "c128") return ONNX_TENSOR_ELEMENT_DATA_TYPE_COMPLEX128;
+
+    // unknown
+    fprintf(stderr, "unsupported tensor elem data type %s\n", type.c_str());
+    return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
+}
+
 static bool string_starts_with(const std::string& s, const std::string& s2)
 {
     return strncmp(s.c_str(), s2.c_str(), s2.size()) == 0;
 }
 
-void shape_inference(onnx::ModelProto& model)
+void shape_inference(onnx::ModelProto& model,
+                     const std::vector<std::vector<int64_t> >& input_shapes,
+                     const std::vector<std::string>& input_types,
+                     const std::vector<std::vector<int64_t> >& input_shapes2,
+                     const std::vector<std::string>& input_types2)
 {
     onnx::GraphProto* graph = model.mutable_graph();
 
@@ -166,14 +192,22 @@ void shape_inference(onnx::ModelProto& model)
             const onnx::ValueInfoProto& value = graph->input(i);
 
             std::vector<int64_t> shape;
-            const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape();
-            for (int k = 0; k < tsp.dim_size(); k++)
+            ONNXTensorElementDataType datatype;
+            if (!input_shapes.empty())
             {
-                // TODO has_dim_value ?
-                shape.push_back(tsp.dim(k).dim_value());
+                shape = input_shapes[i];
+                datatype = get_onnx_tensor_elem_data_type(input_types[i]);
             }
+            else
+            {
+                const onnx::TensorShapeProto& tsp = value.type().tensor_type().shape();
+                for (int k = 0; k < tsp.dim_size(); k++)
+                {
+                    shape.push_back(tsp.dim(k).dim_value());
+                }
 
-            ONNXTensorElementDataType datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type();
+                datatype = (ONNXTensorElementDataType)value.type().tensor_type().elem_type();
+            }
 
             OrtValue* ort_val = 0;
             ort_status = ort_api->CreateTensorAsOrtValue(ort_allocator, (const int64_t*)shape.data(), shape.size(), datatype, &ort_val);
@@ -333,6 +367,232 @@ void shape_inference(onnx::ModelProto& model)
         ort_api->ReleaseSessionOptions(ort_session_opt);
         ort_api->ReleaseEnv(ort_env);
     }
+
+    // onnxrt inference for input_shapes2
+    if (!input_shapes2.empty())
+    {
+        const OrtApi* ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION);
+
+        OrtStatus* ort_status = 0;
+
+        OrtEnv* ort_env = 0;
+        ort_status = ort_api->CreateEnv(ORT_LOGGING_LEVEL_WARNING, "pnnx", &ort_env);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort CreateEnv failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        OrtSessionOptions* ort_session_opt = 0;
+        ort_status = ort_api->CreateSessionOptions(&ort_session_opt);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort CreateSessionOptions failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        ort_status = ort_api->SetSessionGraphOptimizationLevel(ort_session_opt, ORT_DISABLE_ALL);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort SetSessionGraphOptimizationLevel failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        // ort_status = ort_api->SetIntraOpNumThreads(ort_session_opt, 4);
+        // if (ort_status)
+        // {
+        //     fprintf(stderr, "ort SetIntraOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status));
+        // }
+        //
+        // ort_status = ort_api->SetInterOpNumThreads(ort_session_opt, 4);
+        // if (ort_status)
+        // {
+        //     fprintf(stderr, "ort SetInterOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status));
+        // }
+
+        OrtSession* ort_session = 0;
+        ort_status = ort_api->CreateSessionFromArray(ort_env, (const void*)tmp_onnx_data.data(), tmp_onnx_data.size(), ort_session_opt, &ort_session);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort CreateSession failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        OrtRunOptions* ort_run_opt = 0;
+        ort_status = ort_api->CreateRunOptions(&ort_run_opt);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort CreateRunOptions failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        OrtAllocator* ort_allocator = 0;
+        ort_status = ort_api->GetAllocatorWithDefaultOptions(&ort_allocator);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort GetAllocatorWithDefaultOptions failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        std::vector<const char*> input_names;
+        std::vector<OrtValue*> inputs;
+        for (int i = 0; i < graph->input_size(); i++)
+        {
+            const onnx::ValueInfoProto& value = graph->input(i);
+
+            std::vector<int64_t> shape = input_shapes2[i];
+            ONNXTensorElementDataType datatype = get_onnx_tensor_elem_data_type(input_types2[i]);
+
+            OrtValue* ort_val = 0;
+            ort_status = ort_api->CreateTensorAsOrtValue(ort_allocator, (const int64_t*)shape.data(), shape.size(), datatype, &ort_val);
+            if (ort_status)
+            {
+                fprintf(stderr, "ort CreateTensorAsOrtValue failed %s\n", ort_api->GetErrorMessage(ort_status));
+            }
+
+            input_names.push_back(value.name().c_str());
+            inputs.push_back(ort_val);
+        }
+
+        std::vector<const char*> output_names;
+        std::vector<OrtValue*> outputs;
+        for (size_t i = 0; i < intermediates.size(); i++)
+        {
+            output_names.push_back(intermediates[i].c_str());
+            outputs.push_back(0);
+        }
+
+        ort_status = ort_api->Run(ort_session, ort_run_opt,
+                                  input_names.data(), inputs.data(), input_names.size(),
+                                  output_names.data(), output_names.size(), outputs.data());
+        if (ort_status)
+        {
+            fprintf(stderr, "ort Run failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        // TODO get output data
+
+        // graph->clear_output();
+
+        for (size_t i = 0; i < output_names.size(); i++)
+        {
+            OrtTypeInfo* type_info = 0;
+            ort_status = ort_api->GetTypeInfo(outputs[i], &type_info);
+            if (ort_status)
+            {
+                fprintf(stderr, "ort GetTypeInfo failed %s\n", ort_api->GetErrorMessage(ort_status));
+            }
+
+            ONNXType type = ONNX_TYPE_UNKNOWN;
+            if (type_info)
+            {
+                ort_status = ort_api->GetOnnxTypeFromTypeInfo(type_info, &type);
+                if (ort_status)
+                {
+                    fprintf(stderr, "ort GetOnnxTypeFromTypeInfo failed %s\n", ort_api->GetErrorMessage(ort_status));
+                }
+            }
+
+            if (type == ONNX_TYPE_TENSOR)
+            {
+                OrtTensorTypeAndShapeInfo* info = 0;
+                ort_status = ort_api->GetTensorTypeAndShape(outputs[i], &info);
+                if (ort_status)
+                {
+                    fprintf(stderr, "ort GetTensorTypeAndShape failed %s\n", ort_api->GetErrorMessage(ort_status));
+                }
+
+                ONNXTensorElementDataType datatype = ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED;
+                ort_status = ort_api->GetTensorElementType(info, &datatype);
+                if (ort_status)
+                {
+                    fprintf(stderr, "ort GetTensorElementType failed %s\n", ort_api->GetErrorMessage(ort_status));
+                }
+
+                size_t out_dims = 0;
+                ort_status = ort_api->GetDimensionsCount(info, &out_dims);
+                if (ort_status)
+                {
+                    fprintf(stderr, "ort GetDimensionsCount failed %s\n", ort_api->GetErrorMessage(ort_status));
+                }
+
+                // fprintf(stderr, "   out_dims = %lu\n", out_dims);
+
+                std::vector<int64_t> out_shape;
+                out_shape.resize(out_dims);
+                ort_status = ort_api->GetDimensions(info, out_shape.data(), out_dims);
+                if (ort_status)
+                {
+                    fprintf(stderr, "ort GetDimensions failed %s\n", ort_api->GetErrorMessage(ort_status));
+                }
+
+                // fprintf(stderr, "%16s = ", output_names[i]);
+                // for (size_t j = 0; j < out_dims; j++)
+                // {
+                //     fprintf(stderr, "%lu ", out_shape[j]);
+                // }
+                // fprintf(stderr, "\n");
+
+                // assign value info
+                {
+                    onnx::ValueInfoProto* value = 0;
+
+                    // maybe output
+                    for (size_t j = 0; j < orig_outputs.size(); j++)
+                    {
+                        if (orig_outputs[j] == output_names[i])
+                        {
+                            value = graph->mutable_output(j);
+                            break;
+                        }
+                    }
+                    if (!value)
+                    {
+                        for (int j = 0; j < graph->value_info_size(); j++)
+                        {
+                            if (graph->mutable_value_info(j)->name() == output_names[i])
+                            {
+                                value = graph->mutable_value_info(j);
+                                break;
+                            }
+                        }
+                    }
+
+                    // fprintf(stderr, "assign value info2 %s\n", value->name().c_str());
+
+                    value->mutable_type()->mutable_tensor_type()->set_elem_type((int32_t)datatype);
+
+                    onnx::TensorShapeProto* tsp = value->mutable_type()->mutable_tensor_type()->mutable_shape();
+
+                    // tsp->clear_dim();
+                    for (size_t j = 0; j < out_dims; j++)
+                    {
+                        if (tsp->dim(j).dim_value() == out_shape[j])
+                            continue;
+
+                        // dynamic dim size
+                        tsp->mutable_dim(j)->clear_dim_value();
+                    }
+                }
+
+                ort_api->ReleaseTensorTypeAndShapeInfo(info);
+            }
+
+            if (type_info)
+            {
+                ort_api->ReleaseTypeInfo(type_info);
+            }
+        }
+
+        for (size_t i = 0; i < input_names.size(); i++)
+        {
+            ort_api->ReleaseValue(inputs[i]);
+        }
+
+        for (size_t i = 0; i < output_names.size(); i++)
+        {
+            ort_api->ReleaseValue(outputs[i]);
+        }
+
+        ort_api->ReleaseRunOptions(ort_run_opt);
+        ort_api->ReleaseSession(ort_session);
+        ort_api->ReleaseSessionOptions(ort_session_opt);
+        ort_api->ReleaseEnv(ort_env);
+    }
 }
 
 } // namespace onnx2pnnx

tools/pnnx/src/pass_onnx/shape_inference.h

@@ -18,7 +18,11 @@ namespace pnnx {
 
 namespace onnx2pnnx {
 
-void shape_inference(onnx::ModelProto& model);
+void shape_inference(onnx::ModelProto& model,
+                     const std::vector<std::vector<int64_t> >& input_shapes,
+                     const std::vector<std::string>& input_types,
+                     const std::vector<std::vector<int64_t> >& input_shapes2,
+                     const std::vector<std::string>& input_types2);
 
 } // namespace onnx2pnnx
 

tools/pnnx/tests/onnx/CMakeLists.txt

@@ -13,3 +13,5 @@ pnnx_onnx_add_test(mobilenet_v3_small)
 pnnx_onnx_add_test(resnet18)
 pnnx_onnx_add_test(shufflenet_v2_x1_0)
 pnnx_onnx_add_test(squeezenet1_1)
+pnnx_onnx_add_test(swin_t)
+pnnx_onnx_add_test(vit_b_32)

tools/pnnx/tests/onnx/test_swin_t.py

@@ -0,0 +1,65 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torchvision
+import torchvision.models as models
+from packaging import version
+
+def test():
+    if version.parse(torchvision.__version__) < version.parse('0.13'):
+        return True
+
+    net = models.swin_t()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 3, 224, 224)
+
+    a = net(x)
+
+    # export onnx
+    torch.onnx.export(net, (x,), "test_swin_t.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_swin_t.onnx inputshape=[1,3,224,224]")
+
+    # pnnx inference
+    import test_swin_t_pnnx
+    b = test_swin_t_pnnx.test_inference()
+
+    if not torch.allclose(a, b, 1e-4, 1e-4):
+        return False
+
+    if version.parse(torch.__version__) < version.parse('2.4'):
+        return True
+
+    # export dynamo onnx
+    torch.onnx.dynamo_export(net, x).save("test_swin_t_dynamo.onnx")
+
+    # onnx to pnnx
+    os.system("../../src/pnnx test_swin_t_dynamo.onnx inputshape=[1,3,224,224]")
+
+    # pnnx inference
+    import test_swin_t_dynamo_pnnx
+    b = test_swin_t_dynamo_pnnx.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)

tools/pnnx/tests/onnx/test_vit_b_32.py

@@ -0,0 +1,68 @@
+# Tencent is pleased to support the open source community by making ncnn available.
+#
+# Copyright (C) 2024 THL A29 Limited, a Tencent company. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+import torch
+import torchvision
+import torchvision.models as models
+from packaging import version
+
+def test():
+    if version.parse(torchvision.__version__) < version.parse('0.12'):
+        return True
+
+    if version.parse(torch.__version__) >= version.parse('2.0') and version.parse(torch.__version__) < version.parse('2.1'):
+        return True
+
+    net = models.vit_b_32()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(1, 3, 224, 224)
+
+    a = net(x)
+
+    # export onnx
+    torch.onnx.export(net, (x,), "test_vit_b_32.onnx")
+
+    # onnx to pnnx
+    import os
+    os.system("../../src/pnnx test_vit_b_32.onnx inputshape=[1,3,224,224]")
+
+    # pnnx inference
+    import test_vit_b_32_pnnx
+    b = test_vit_b_32_pnnx.test_inference()
+
+    if not torch.allclose(a, b, 1e-4, 1e-4):
+        return False
+
+    if version.parse(torch.__version__) < version.parse('2.4'):
+        return True
+
+    # export dynamo onnx
+    torch.onnx.dynamo_export(net, x).save("test_vit_b_32_dynamo.onnx")
+
+    # onnx to pnnx
+    os.system("../../src/pnnx test_vit_b_32_dynamo.onnx inputshape=[1,3,224,224]")
+
+    # pnnx inference
+    import test_vit_b_32_dynamo_pnnx
+    b = test_vit_b_32_dynamo_pnnx.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)