diff --git a/src/meta_schedule/utils.h b/src/meta_schedule/utils.h index 2ee18a8668be7..3dfde3be721fc 100644 --- a/src/meta_schedule/utils.h +++ b/src/meta_schedule/utils.h @@ -35,6 +35,7 @@ #include #include #include +#include #include #include @@ -307,12 +308,19 @@ struct ThreadedTraceApply { /*rand_state=*/ForkSeed(rand_state), /*debug_mode=*/0, /*error_render_level=*/tir::ScheduleErrorRenderLevel::kNone); + trace->ApplyToSchedule(sch, /*remove_postproc=*/true); sch->EnterPostproc(); + for (int i = 0; i < n_; ++i) { Item& item = items_[i]; - if (!item.postproc->Apply(sch)) { - ++item.fail_counter; + try { + if (!item.postproc->Apply(sch)) { + ++item.fail_counter; + return NullOpt; + } + } catch (const std::exception& e) { + LOG(WARNING) << "ThreadedTraceApply::Apply failed with error " << e.what(); return NullOpt; } } diff --git a/src/support/parallel_for.cc b/src/support/parallel_for.cc index f025934383b57..e90967562d163 100644 --- a/src/support/parallel_for.cc +++ b/src/support/parallel_for.cc @@ -118,36 +118,26 @@ void parallel_for_dynamic(int begin, int end, int num_threads, futures.emplace_back(task.get_future()); threads.emplace_back(std::move(task), thread_id); } - // Step 2.2. Launch worker 0 inplace - - auto try_await_futures = [&futures]() { - try { - for (auto&& future : futures) { - future.get(); - } - } catch (const std::exception& e) { - LOG(WARNING) << "RuntimeError: parallel_for_dynamic error on future get with " << e.what(); - } - }; - try { worker(0); } catch (const std::exception& e) { for (auto&& thread : threads) { thread.join(); } - LOG(WARNING) << "RuntimeError: parallel_for_dynamic error on thread join with " << e.what(); - try_await_futures(); - return; + LOG(FATAL) << "RuntimeError: parallel_for_dynamic error with " << e.what(); } - // Step 3. Join threads and check exceptions for (auto&& thread : threads) { thread.join(); } - - try_await_futures(); + try { + for (auto&& future : futures) { + future.get(); + } + } catch (const std::exception& e) { + LOG(FATAL) << "RuntimeError: parallel_for_dynamic error with " << e.what(); + } } } // namespace support