diff --git a/.github/workflows/pnnx.yml b/.github/workflows/pnnx.yml index c8471f0ca..067c4e1b7 100644 --- a/.github/workflows/pnnx.yml +++ b/.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 }} diff --git a/tools/pnnx/src/pass_onnx/fold_constants.cpp b/tools/pnnx/src/pass_onnx/fold_constants.cpp index 462f72d11..135be2b3a 100644 --- a/tools/pnnx/src/pass_onnx/fold_constants.cpp +++ b/tools/pnnx/src/pass_onnx/fold_constants.cpp @@ -14,6 +14,7 @@ #include "fold_constants.h" +#include #include #include #include @@ -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); diff --git a/tools/pnnx/src/pass_onnx/shape_inference.cpp b/tools/pnnx/src/pass_onnx/shape_inference.cpp index cc40a3cc3..e01532ed3 100644 --- a/tools/pnnx/src/pass_onnx/shape_inference.cpp +++ b/tools/pnnx/src/pass_onnx/shape_inference.cpp @@ -14,6 +14,7 @@ #include "shape_inference.h" +#include #include #include #include @@ -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);