limit pnnx onnxruntime threads by OMP_THREAD_LIMIT env (#6049)

.github/workflows/pnnx.yml

@@ -26,6 +26,7 @@ env:
   PROTOBUF_VERSION: 21.12
   ONNXRUNTIME_VERSION: 1.21.1
   CACHE_DATE: 20250512
+  SEGMENT_DOWNLOAD_TIMEOUT_MINS: 15
 
 jobs:
   quick-test:
@@ -236,21 +237,21 @@ jobs:
 
     - name: cache-libtorch
       id: cache-libtorch
-      uses: actions/cache@v4
+      uses: actions/cache/restore@v4
       with:
         path: libtorch-${{ env.LIBTORCH_VERSION }}-install
         key: libtorch-${{ env.LIBTORCH_VERSION }}-linux-install-${{ env.CACHE_DATE }}
 
     - name: cache-torchvision
       id: cache-torchvision
-      uses: actions/cache@v4
+      uses: actions/cache/restore@v4
       with:
         path: torchvision-${{ env.TORCHVISION_VERSION }}-install
         key: torchvision-${{ env.TORCHVISION_VERSION }}-linux-install-${{ env.CACHE_DATE }}
 
     - name: cache-onnxruntime
       id: cache-onnxruntime
-      uses: actions/cache@v4
+      uses: actions/cache/restore@v4
       with:
         path: onnxruntime-${{ env.ONNXRUNTIME_VERSION }}-install
         key: onnxruntime-${{ env.ONNXRUNTIME_VERSION }}-linux-install-${{ env.CACHE_DATE }}

tools/pnnx/src/pass_onnx/fold_constants.cpp

@@ -14,6 +14,7 @@
 
 #include "fold_constants.h"
 
+#include <stdlib.h>
 #include <sstream>
 #include <string>
 #include <unordered_map>
@@ -318,17 +319,20 @@ void fold_constants(onnx::ModelProto& model,
             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));
-        // }
+        const char* omp_thread_limit = std::getenv("OMP_THREAD_LIMIT");
+        const int num_threads = omp_thread_limit ? std::stoi(omp_thread_limit) : 0;
+
+        ort_status = ort_api->SetIntraOpNumThreads(ort_session_opt, num_threads);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort SetIntraOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        ort_status = ort_api->SetInterOpNumThreads(ort_session_opt, num_threads);
+        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);
@@ -901,17 +905,20 @@ void fold_constants_dynamic_shape(onnx::ModelProto& model,
             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));
-        // }
+        const char* omp_thread_limit = std::getenv("OMP_THREAD_LIMIT");
+        const int num_threads = omp_thread_limit ? std::stoi(omp_thread_limit) : 0;
+
+        ort_status = ort_api->SetIntraOpNumThreads(ort_session_opt, num_threads);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort SetIntraOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        ort_status = ort_api->SetInterOpNumThreads(ort_session_opt, num_threads);
+        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);

tools/pnnx/src/pass_onnx/shape_inference.cpp

@@ -14,6 +14,7 @@
 
 #include "shape_inference.h"
 
+#include <stdlib.h>
 #include <sstream>
 #include <string>
 #include <vector>
@@ -220,17 +221,20 @@ void shape_inference(onnx::ModelProto& model,
             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));
-        // }
+        const char* omp_thread_limit = std::getenv("OMP_THREAD_LIMIT");
+        const int num_threads = omp_thread_limit ? std::stoi(omp_thread_limit) : 0;
+
+        ort_status = ort_api->SetIntraOpNumThreads(ort_session_opt, num_threads);
+        if (ort_status)
+        {
+            fprintf(stderr, "ort SetIntraOpNumThreads failed %s\n", ort_api->GetErrorMessage(ort_status));
+        }
+
+        ort_status = ort_api->SetInterOpNumThreads(ort_session_opt, num_threads);
+        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);