update pnnx ci torch 2.8.0 (#6238)

.github/workflows/pnnx.yml

@@ -21,11 +21,11 @@ permissions:
   contents: read
 
 env:
-  LIBTORCH_VERSION: 2.7.1
-  TORCHVISION_VERSION: 0.22.1
+  LIBTORCH_VERSION: 2.8.0
+  TORCHVISION_VERSION: 0.23.0
   PROTOBUF_VERSION: 21.12
   ONNXRUNTIME_VERSION: 1.22.1
-  CACHE_DATE: 20250723
+  CACHE_DATE: 20250807
   SEGMENT_DOWNLOAD_TIMEOUT_MINS: 15
 
 jobs:
@@ -129,7 +129,6 @@ jobs:
         pip3 install -r requirements.txt --break-system-packages
         patch -p1 -i $GITHUB_WORKSPACE/pnnx-patches/pytorch-v${{ env.LIBTORCH_VERSION }}-fix-mobile-build.patch
         patch -p1 -i $GITHUB_WORKSPACE/pnnx-patches/pytorch-v${{ env.LIBTORCH_VERSION }}-no-link-system-lib.patch
-        patch -p1 -i $GITHUB_WORKSPACE/pnnx-patches/pytorch-v${{ env.LIBTORCH_VERSION }}-fix-pocketfft-build.patch
         mkdir -p build && cd build
         cmake -DCMAKE_INSTALL_PREFIX=$GITHUB_WORKSPACE/libtorch-${{ env.LIBTORCH_VERSION }}-install \
             -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -247,6 +246,7 @@ jobs:
           - { python: '3.12', numpy: '2.2.5',  opencv: '4.10.*', torch: '2.5.0',  torchvision: '0.20.0', torchaudio: '2.5.0+cpu'  }
           - { python: '3.12', numpy: '2.2.5',  opencv: '4.11.*', torch: '2.6.0',  torchvision: '0.21.0', torchaudio: '2.6.0+cpu'  }
           - { python: '3.12', numpy: '2.2.5',  opencv: '4.11.*', torch: '2.7.0',  torchvision: '0.22.0', torchaudio: '2.7.0+cpu'  }
+          - { python: '3.13', numpy: '2.2.5',  opencv: '4.12.*', torch: '2.8.0',  torchvision: '0.23.0', torchaudio: '2.8.0+cpu'  }
 
     name: test-${{ matrix.torch }}-py${{ matrix.python }}
 

tools/pnnx/src/pass_level2/F_layer_norm.cpp

@@ -157,7 +157,11 @@ pnnx.Output             output      1 0 out
             return false;
 
         // dim must be the last N dimensions
-        std::vector<int> dim = captured_params.at("dim").ai;
+        std::vector<int> dim;
+        if (captured_params.at("dim").type == 2)
+            dim.push_back(captured_params.at("dim").i);
+        else // if (captured_params.at("dim").type == 5)
+            dim = captured_params.at("dim").ai;
 
         const int input_rank = (int)inputshape.size();
         const int dim_count = (int)dim.size();
@@ -183,7 +187,11 @@ pnnx.Output             output      1 0 out
     void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
     {
         const std::vector<int>& inputshape = op->inputs[0]->shape;
-        const std::vector<int>& dim = captured_params.at("dim").ai;
+        std::vector<int> dim;
+        if (captured_params.at("dim").type == 2)
+            dim.push_back(captured_params.at("dim").i);
+        else // if (captured_params.at("dim").type == 5)
+            dim = captured_params.at("dim").ai;
         const int input_rank = (int)inputshape.size();
         const int dim_count = (int)dim.size();
 
@@ -240,7 +248,11 @@ pnnx.Output             output      1 0 out
             return false;
 
         // dim must be the last N dimensions
-        std::vector<int> dim = captured_params.at("dim").ai;
+        std::vector<int> dim;
+        if (captured_params.at("dim").type == 2)
+            dim.push_back(captured_params.at("dim").i);
+        else // if (captured_params.at("dim").type == 5)
+            dim = captured_params.at("dim").ai;
 
         const int input_rank = (int)inputshape.size();
         const int dim_count = (int)dim.size();
@@ -278,7 +290,11 @@ pnnx.Output             output      1 0 out
     void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
     {
         const std::vector<int>& inputshape = op->inputs[0]->shape;
-        const std::vector<int>& dim = captured_params.at("dim").ai;
+        std::vector<int> dim;
+        if (captured_params.at("dim").type == 2)
+            dim.push_back(captured_params.at("dim").i);
+        else // if (captured_params.at("dim").type == 5)
+            dim = captured_params.at("dim").ai;
         const int input_rank = (int)inputshape.size();
         const int dim_count = (int)dim.size();
 

tools/pnnx/src/pass_level2/torch_max.cpp

@@ -147,13 +147,16 @@ pnnx.Output             output      2 0 out indices
         if (captured_params.find("op_1.keepdims") == captured_params.end())
             return false;
 
-        if (captured_params.at("op_0.axes").type != 5 || captured_params.at("op_0.axes").ai.size() != 1)
+        if (captured_params.at("op_0.axes").type != 2 && (captured_params.at("op_0.axes").type != 5 || captured_params.at("op_0.axes").ai.size() != 1))
             return false;
 
         if (captured_params.at("op_1.axis").type != 2)
             return false;
 
-        if (captured_params.at("op_0.axes").ai[0] != captured_params.at("op_1.axis").i)
+        if (captured_params.at("op_0.axes").type == 2 && captured_params.at("op_0.axes").i != captured_params.at("op_1.axis").i)
+            return false;
+
+        if (captured_params.at("op_0.axes").type == 5 && captured_params.at("op_0.axes").ai[0] != captured_params.at("op_1.axis").i)
             return false;
 
         if (captured_params.at("op_0.keepdims").i != captured_params.at("op_1.keepdims").i)
@@ -164,7 +167,10 @@ pnnx.Output             output      2 0 out indices
 
     void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
     {
-        op->params["dim"] = captured_params.at("op_0.axes").ai[0];
+        if (captured_params.at("op_0.axes").type == 2)
+            op->params["dim"] = captured_params.at("op_0.axes").i;
+        else
+            op->params["dim"] = captured_params.at("op_0.axes").ai[0];
         op->params["keepdim"] = captured_params.at("op_0.keepdims").i ? true : false;
     }
 };

tools/pnnx/src/pass_level2/torch_min.cpp

@@ -147,13 +147,16 @@ pnnx.Output             output      2 0 out indices
         if (captured_params.find("op_1.keepdims") == captured_params.end())
             return false;
 
-        if (captured_params.at("op_0.axes").type != 5 || captured_params.at("op_0.axes").ai.size() != 1)
+        if (captured_params.at("op_0.axes").type != 2 && (captured_params.at("op_0.axes").type != 5 || captured_params.at("op_0.axes").ai.size() != 1))
             return false;
 
         if (captured_params.at("op_1.axis").type != 2)
             return false;
 
-        if (captured_params.at("op_0.axes").ai[0] != captured_params.at("op_1.axis").i)
+        if (captured_params.at("op_0.axes").type == 2 && captured_params.at("op_0.axes").i != captured_params.at("op_1.axis").i)
+            return false;
+
+        if (captured_params.at("op_0.axes").type == 5 && captured_params.at("op_0.axes").ai[0] != captured_params.at("op_1.axis").i)
             return false;
 
         if (captured_params.at("op_0.keepdims").i != captured_params.at("op_1.keepdims").i)
@@ -164,7 +167,10 @@ pnnx.Output             output      2 0 out indices
 
     void write(Operator* op, const std::map<std::string, Parameter>& captured_params) const
     {
-        op->params["dim"] = captured_params.at("op_0.axes").ai[0];
+        if (captured_params.at("op_0.axes").type == 2)
+            op->params["dim"] = captured_params.at("op_0.axes").i;
+        else
+            op->params["dim"] = captured_params.at("op_0.axes").ai[0];
         op->params["keepdim"] = captured_params.at("op_0.keepdims").i ? true : false;
     }
 };

tools/pnnx/src/pass_level5/fuse_slice_to_tensor_split.cpp

@@ -98,11 +98,17 @@ void fuse_slice_to_tensor_split(Graph& graph)
                     full_dimsize_slice = true;
                     break;
                 }
-                if (!op_in->shape.empty() && end2 == op_in->shape[dim])
+                if (!op_in->shape.empty())
                 {
-                    slice_n_ops.push_back(op2);
-                    full_dimsize_slice = true;
-                    break;
+                    if (dim < 0)
+                        dim += op_in->shape.size();
+
+                    if (end2 == op_in->shape[dim])
+                    {
+                        slice_n_ops.push_back(op2);
+                        full_dimsize_slice = true;
+                        break;
+                    }
                 }
 
                 tensor_split_indices.push_back(end2);

tools/pnnx/src/pass_onnx/fuse_constant_as_attribute.cpp

@@ -28,6 +28,8 @@ static constant_as_attribute caas[] = {
     {"If", 0, "cond"},
     {"Pad", 1, "pads"},
     {"Pad", 2, "value"},
+    {"ReduceL1", 1, "axes"},
+    {"ReduceL2", 1, "axes"},
     {"ReduceLogSumExp", 1, "axes"},
     {"ReduceMax", 1, "axes"},
     {"ReduceMean", 1, "axes"},

tools/pnnx/tests/onnx/test_F_relu.py

@@ -52,7 +52,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x, y, z, w).save("test_F_relu_dynamo.onnx")
+    torch.onnx.export(net, (x, y, z, w), "test_F_relu_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_F_relu_dynamo.onnx inputshape=[16],[2,16],[3,12,16],[5,7,9,11]")

tools/pnnx/tests/onnx/test_convnext_tiny.py

@@ -32,11 +32,11 @@ def test():
     if not torch.allclose(a, b, 1e-4, 1e-4):
         return False
 
-    if version.parse(torch.__version__) < version.parse('2.8'):
+    if version.parse(torch.__version__) < version.parse('2.9'):
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_convnext_tiny_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_convnext_tiny_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_convnext_tiny_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_mobilenet_v2.py

@@ -32,7 +32,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_mobilenet_v2_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_mobilenet_v2_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_mobilenet_v2_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_mobilenet_v3_small.py

@@ -35,7 +35,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_mobilenet_v3_small_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_mobilenet_v3_small_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_mobilenet_v3_small_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_nn_ReLU.py

@@ -54,7 +54,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x, y, z, w).save("test_nn_ReLU_dynamo.onnx")
+    torch.onnx.export(net, (x, y, z, w), "test_nn_ReLU_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_nn_ReLU_dynamo.onnx inputshape=[1,12],[1,12,64],[1,12,24,64],[1,12,24,32,64]")

tools/pnnx/tests/onnx/test_resnet18.py

@@ -32,7 +32,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_resnet18_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_resnet18_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_resnet18_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_shufflenet_v2_x1_0.py

@@ -32,7 +32,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_shufflenet_v2_x1_0_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_shufflenet_v2_x1_0_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_shufflenet_v2_x1_0_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_squeezenet1_1.py

@@ -32,7 +32,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_squeezenet1_1_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_squeezenet1_1_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_squeezenet1_1_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_swin_t.py

@@ -38,11 +38,11 @@ def test():
     if not torch.allclose(a, b, 1e-4, 1e-4):
         return False
 
-    if version.parse(torch.__version__) < version.parse('2.8'):
+    if version.parse(torch.__version__) < version.parse('2.9'):
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_swin_t_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_swin_t_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_swin_t_dynamo.onnx inputshape=[1,3,224,224]")

tools/pnnx/tests/onnx/test_vit_b_32.py

@@ -39,7 +39,7 @@ def test():
         return True
 
     # export dynamo onnx
-    torch.onnx.dynamo_export(net, x).save("test_vit_b_32_dynamo.onnx")
+    torch.onnx.export(net, (x,), "test_vit_b_32_dynamo.onnx", dynamo=True, external_data=False)
 
     # onnx to pnnx
     os.system("../../src/pnnx test_vit_b_32_dynamo.onnx inputshape=[1,3,224,224]")