try / catch in ThreadedApply

masahi · Apr 8, 2022 · 476129b · 476129b
1 parent d8226ff
commit 476129b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 20 deletions.
diff --git a/src/meta_schedule/utils.h b/src/meta_schedule/utils.h
@@ -35,6 +35,7 @@
 #include <tvm/meta_schedule/tune_context.h>
 #include <tvm/node/node.h>
 #include <tvm/node/serialization.h>
+#include <tvm/runtime/container/optional.h>
 #include <tvm/support/parallel_for.h>
 #include <tvm/tir/schedule/schedule.h>
 
@@ -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
@@ -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