[OpenMP] Adding GOMP compatible cancellation

Remove fatal error messages from the cancellation API for GOMP
Add __kmp_barrier_gomp_cancel() to implement cancellation of parallel regions.
This new function uses the linear barrier algorithm with a cancellable
nonsleepable wait loop.

Differential Revision: https://reviews.llvm.org/D57969

llvm-svn: 354367

5 files changed, 246 insertions, 113 deletions

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
index 5ad992aa19b..4dd6cf08559 100644
--- a/openmp/runtime/src/kmp.h
+++ b/openmp/runtime/src/kmp.h
@@ -3560,6 +3560,7 @@ extern int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
                          size_t reduce_size, void *reduce_data,
                          void (*reduce)(void *, void *));
 extern void __kmp_end_split_barrier(enum barrier_type bt, int gtid);
+extern int __kmp_barrier_gomp_cancel(int gtid);
 
 /*!
  * Tell the fork call which compiler generated the fork call, and therefore how
diff --git a/openmp/runtime/src/kmp_barrier.cpp b/openmp/runtime/src/kmp_barrier.cpp
index 332028f7a4e..ecb115ad523 100644
--- a/openmp/runtime/src/kmp_barrier.cpp
+++ b/openmp/runtime/src/kmp_barrier.cpp
@@ -44,7 +44,8 @@ void __kmp_print_structure(void); // Forward declaration
 // ---------------------------- Barrier Algorithms ----------------------------
 
 // Linear Barrier
-static void __kmp_linear_barrier_gather(
+template <bool cancellable = false>
+static bool __kmp_linear_barrier_gather_template(
     enum barrier_type bt, kmp_info_t *this_thr, int gtid, int tid,
     void (*reduce)(void *, void *) USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
   KMP_TIME_DEVELOPER_PARTITIONED_BLOCK(KMP_linear_gather);
@@ -104,7 +105,14 @@ static void __kmp_linear_barrier_gather(
       // Wait for worker thread to arrive
       kmp_flag_64 flag(&other_threads[i]->th.th_bar[bt].bb.b_arrived,
                        new_state);
-      flag.wait(this_thr, FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+      if (cancellable) {
+        bool cancelled = flag.wait_cancellable_nosleep(
+            this_thr, FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+        if (cancelled)
+          return true;
+      } else {
+        flag.wait(this_thr, FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+      }
       ANNOTATE_BARRIER_END(other_threads[i]);
 #if USE_ITT_BUILD && USE_ITT_NOTIFY
       // Barrier imbalance - write min of the thread time and the other thread
@@ -137,9 +145,11 @@ static void __kmp_linear_barrier_gather(
       20,
       ("__kmp_linear_barrier_gather: T#%d(%d:%d) exit for barrier type %d\n",
        gtid, team->t.t_id, tid, bt));
+  return false;
 }
 
-static void __kmp_linear_barrier_release(
+template <bool cancellable = false>
+static bool __kmp_linear_barrier_release_template(
     enum barrier_type bt, kmp_info_t *this_thr, int gtid, int tid,
     int propagate_icvs USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
   KMP_TIME_DEVELOPER_PARTITIONED_BLOCK(KMP_linear_release);
@@ -201,7 +211,15 @@ static void __kmp_linear_barrier_release(
     KA_TRACE(20, ("__kmp_linear_barrier_release: T#%d wait go(%p) == %u\n",
                   gtid, &thr_bar->b_go, KMP_BARRIER_STATE_BUMP));
     kmp_flag_64 flag(&thr_bar->b_go, KMP_BARRIER_STATE_BUMP);
-    flag.wait(this_thr, TRUE USE_ITT_BUILD_ARG(itt_sync_obj));
+    if (cancellable) {
+      bool cancelled = flag.wait_cancellable_nosleep(
+          this_thr, TRUE USE_ITT_BUILD_ARG(itt_sync_obj));
+      if (cancelled) {
+        return true;
+      }
+    } else {
+      flag.wait(this_thr, TRUE USE_ITT_BUILD_ARG(itt_sync_obj));
+    }
     ANNOTATE_BARRIER_END(this_thr);
 #if USE_ITT_BUILD && USE_ITT_NOTIFY
     if ((__itt_sync_create_ptr && itt_sync_obj == NULL) || KMP_ITT_DEBUG) {
@@ -212,7 +230,7 @@ static void __kmp_linear_barrier_release(
       __kmp_itt_task_starting(itt_sync_obj);
 
       if (bt == bs_forkjoin_barrier && TCR_4(__kmp_global.g.g_done))
-        return;
+        return false;
 
       itt_sync_obj = __kmp_itt_barrier_object(gtid, bs_forkjoin_barrier);
       if (itt_sync_obj != NULL)
@@ -222,7 +240,7 @@ static void __kmp_linear_barrier_release(
 #endif /* USE_ITT_BUILD && USE_ITT_NOTIFY */
         // Early exit for reaping threads releasing forkjoin barrier
         if (bt == bs_forkjoin_barrier && TCR_4(__kmp_global.g.g_done))
-      return;
+      return false;
 // The worker thread may now assume that the team is valid.
 #ifdef KMP_DEBUG
     tid = __kmp_tid_from_gtid(gtid);
@@ -239,6 +257,35 @@ static void __kmp_linear_barrier_release(
       20,
       ("__kmp_linear_barrier_release: T#%d(%d:%d) exit for barrier type %d\n",
        gtid, team->t.t_id, tid, bt));
+  return false;
+}
+
+static void __kmp_linear_barrier_gather(
+    enum barrier_type bt, kmp_info_t *this_thr, int gtid, int tid,
+    void (*reduce)(void *, void *) USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
+  __kmp_linear_barrier_gather_template<false>(
+      bt, this_thr, gtid, tid, reduce USE_ITT_BUILD_ARG(itt_sync_obj));
+}
+
+static bool __kmp_linear_barrier_gather_cancellable(
+    enum barrier_type bt, kmp_info_t *this_thr, int gtid, int tid,
+    void (*reduce)(void *, void *) USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
+  return __kmp_linear_barrier_gather_template<true>(
+      bt, this_thr, gtid, tid, reduce USE_ITT_BUILD_ARG(itt_sync_obj));
+}
+
+static void __kmp_linear_barrier_release(
+    enum barrier_type bt, kmp_info_t *this_thr, int gtid, int tid,
+    int propagate_icvs USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
+  __kmp_linear_barrier_release_template<false>(
+      bt, this_thr, gtid, tid, propagate_icvs USE_ITT_BUILD_ARG(itt_sync_obj));
+}
+
+static bool __kmp_linear_barrier_release_cancellable(
+    enum barrier_type bt, kmp_info_t *this_thr, int gtid, int tid,
+    int propagate_icvs USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
+  return __kmp_linear_barrier_release_template<true>(
+      bt, this_thr, gtid, tid, propagate_icvs USE_ITT_BUILD_ARG(itt_sync_obj));
 }
 
 // Tree barrier
@@ -1208,20 +1255,44 @@ static void __kmp_hierarchical_barrier_release(
 
 // End of Barrier Algorithms
 
+// type traits for cancellable value
+// if cancellable is true, then is_cancellable is a normal boolean variable
+// if cancellable is false, then is_cancellable is a compile time constant
+template <bool cancellable> struct is_cancellable {};
+template <> struct is_cancellable<true> {
+  bool value;
+  is_cancellable() : value(false) {}
+  is_cancellable(bool b) : value(b) {}
+  is_cancellable &operator=(bool b) {
+    value = b;
+    return *this;
+  }
+  operator bool() const { return value; }
+};
+template <> struct is_cancellable<false> {
+  is_cancellable &operator=(bool b) { return *this; }
+  constexpr operator bool() const { return false; }
+};
+
 // Internal function to do a barrier.
 /* If is_split is true, do a split barrier, otherwise, do a plain barrier
    If reduce is non-NULL, do a split reduction barrier, otherwise, do a split
    barrier
-   Returns 0 if master thread, 1 if worker thread.  */
-int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
-                  size_t reduce_size, void *reduce_data,
-                  void (*reduce)(void *, void *)) {
+   When cancellable = false,
+     Returns 0 if master thread, 1 if worker thread.
+   When cancellable = true
+     Returns 0 if not cancelled, 1 if cancelled.  */
+template <bool cancellable = false>
+static int __kmp_barrier_template(enum barrier_type bt, int gtid, int is_split,
+                                  size_t reduce_size, void *reduce_data,
+                                  void (*reduce)(void *, void *)) {
   KMP_TIME_PARTITIONED_BLOCK(OMP_plain_barrier);
   KMP_SET_THREAD_STATE_BLOCK(PLAIN_BARRIER);
   int tid = __kmp_tid_from_gtid(gtid);
   kmp_info_t *this_thr = __kmp_threads[gtid];
   kmp_team_t *team = this_thr->th.th_team;
   int status = 0;
+  is_cancellable<cancellable> cancelled;
 #if OMPT_SUPPORT && OMPT_OPTIONAL
   ompt_data_t *my_task_data;
   ompt_data_t *my_parallel_data;
@@ -1305,41 +1376,45 @@ int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
     }
 
     if (KMP_MASTER_TID(tid) && __kmp_tasking_mode != tskm_immediate_exec)
-      __kmp_task_team_setup(
-          this_thr, team,
-          0); // use 0 to only setup the current team if nthreads > 1
-
-    switch (__kmp_barrier_gather_pattern[bt]) {
-    case bp_hyper_bar: {
-      KMP_ASSERT(__kmp_barrier_gather_branch_bits[bt]); // don't set branch bits
-      // to 0; use linear
-      __kmp_hyper_barrier_gather(bt, this_thr, gtid, tid,
-                                 reduce USE_ITT_BUILD_ARG(itt_sync_obj));
-      break;
-    }
-    case bp_hierarchical_bar: {
-      __kmp_hierarchical_barrier_gather(bt, this_thr, gtid, tid,
-                                        reduce USE_ITT_BUILD_ARG(itt_sync_obj));
-      break;
-    }
-    case bp_tree_bar: {
-      KMP_ASSERT(__kmp_barrier_gather_branch_bits[bt]); // don't set branch bits
-      // to 0; use linear
-      __kmp_tree_barrier_gather(bt, this_thr, gtid, tid,
-                                reduce USE_ITT_BUILD_ARG(itt_sync_obj));
-      break;
-    }
-    default: {
-      __kmp_linear_barrier_gather(bt, this_thr, gtid, tid,
+      // use 0 to only setup the current team if nthreads > 1
+      __kmp_task_team_setup(this_thr, team, 0);
+
+    if (cancellable) {
+      cancelled = __kmp_linear_barrier_gather_cancellable(
+          bt, this_thr, gtid, tid, reduce USE_ITT_BUILD_ARG(itt_sync_obj));
+    } else {
+      switch (__kmp_barrier_gather_pattern[bt]) {
+      case bp_hyper_bar: {
+        // don't set branch bits to 0; use linear
+        KMP_ASSERT(__kmp_barrier_gather_branch_bits[bt]);
+        __kmp_hyper_barrier_gather(bt, this_thr, gtid, tid,
+                                   reduce USE_ITT_BUILD_ARG(itt_sync_obj));
+        break;
+      }
+      case bp_hierarchical_bar: {
+        __kmp_hierarchical_barrier_gather(
+            bt, this_thr, gtid, tid, reduce USE_ITT_BUILD_ARG(itt_sync_obj));
+        break;
+      }
+      case bp_tree_bar: {
+        // don't set branch bits to 0; use linear
+        KMP_ASSERT(__kmp_barrier_gather_branch_bits[bt]);
+        __kmp_tree_barrier_gather(bt, this_thr, gtid, tid,
                                   reduce USE_ITT_BUILD_ARG(itt_sync_obj));
-    }
+        break;
+      }
+      default: {
+        __kmp_linear_barrier_gather(bt, this_thr, gtid, tid,
+                                    reduce USE_ITT_BUILD_ARG(itt_sync_obj));
+      }
+      }
     }
 
     KMP_MB();
 
     if (KMP_MASTER_TID(tid)) {
       status = 0;
-      if (__kmp_tasking_mode != tskm_immediate_exec) {
+      if (__kmp_tasking_mode != tskm_immediate_exec && !cancelled) {
         __kmp_task_team_wait(this_thr, team USE_ITT_BUILD_ARG(itt_sync_obj));
       }
 #if USE_DEBUGGER
@@ -1349,10 +1424,13 @@ int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
 #endif
 
 #if OMP_40_ENABLED
-      kmp_int32 cancel_request = KMP_ATOMIC_LD_RLX(&team->t.t_cancel_request);
-      // Reset cancellation flag for worksharing constructs
-      if (cancel_request == cancel_loop || cancel_request == cancel_sections) {
-        KMP_ATOMIC_ST_RLX(&team->t.t_cancel_request, cancel_noreq);
+      if (__kmp_omp_cancellation) {
+        kmp_int32 cancel_request = KMP_ATOMIC_LD_RLX(&team->t.t_cancel_request);
+        // Reset cancellation flag for worksharing constructs
+        if (cancel_request == cancel_loop ||
+            cancel_request == cancel_sections) {
+          KMP_ATOMIC_ST_RLX(&team->t.t_cancel_request, cancel_noreq);
+        }
       }
 #endif
 #if USE_ITT_BUILD
@@ -1416,31 +1494,36 @@ int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
         __kmp_itt_barrier_middle(gtid, itt_sync_obj);
 #endif /* USE_ITT_BUILD */
     }
-    if (status == 1 || !is_split) {
-      switch (__kmp_barrier_release_pattern[bt]) {
-      case bp_hyper_bar: {
-        KMP_ASSERT(__kmp_barrier_release_branch_bits[bt]);
-        __kmp_hyper_barrier_release(bt, this_thr, gtid, tid,
-                                    FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
-        break;
-      }
-      case bp_hierarchical_bar: {
-        __kmp_hierarchical_barrier_release(
+    if ((status == 1 || !is_split) && !cancelled) {
+      if (cancellable) {
+        cancelled = __kmp_linear_barrier_release_cancellable(
             bt, this_thr, gtid, tid, FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
-        break;
-      }
-      case bp_tree_bar: {
-        KMP_ASSERT(__kmp_barrier_release_branch_bits[bt]);
-        __kmp_tree_barrier_release(bt, this_thr, gtid, tid,
-                                   FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
-        break;
-      }
-      default: {
-        __kmp_linear_barrier_release(bt, this_thr, gtid, tid,
+      } else {
+        switch (__kmp_barrier_release_pattern[bt]) {
+        case bp_hyper_bar: {
+          KMP_ASSERT(__kmp_barrier_release_branch_bits[bt]);
+          __kmp_hyper_barrier_release(bt, this_thr, gtid, tid,
+                                      FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+          break;
+        }
+        case bp_hierarchical_bar: {
+          __kmp_hierarchical_barrier_release(
+              bt, this_thr, gtid, tid, FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+          break;
+        }
+        case bp_tree_bar: {
+          KMP_ASSERT(__kmp_barrier_release_branch_bits[bt]);
+          __kmp_tree_barrier_release(bt, this_thr, gtid, tid,
                                      FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+          break;
+        }
+        default: {
+          __kmp_linear_barrier_release(bt, this_thr, gtid, tid,
+                                       FALSE USE_ITT_BUILD_ARG(itt_sync_obj));
+        }
+        }
       }
-      }
-      if (__kmp_tasking_mode != tskm_immediate_exec) {
+      if (__kmp_tasking_mode != tskm_immediate_exec && !cancelled) {
         __kmp_task_team_sync(this_thr, team);
       }
     }
@@ -1506,9 +1589,43 @@ int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
 #endif
   ANNOTATE_BARRIER_END(&team->t.t_bar);
 
+  if (cancellable)
+    return (int)cancelled;
   return status;
 }
 
+// Returns 0 if master thread, 1 if worker thread.
+int __kmp_barrier(enum barrier_type bt, int gtid, int is_split,
+                  size_t reduce_size, void *reduce_data,
+                  void (*reduce)(void *, void *)) {
+  return __kmp_barrier_template<>(bt, gtid, is_split, reduce_size, reduce_data,
+                                  reduce);
+}
+
+#if defined(KMP_GOMP_COMPAT)
+// Returns 1 if cancelled, 0 otherwise
+int __kmp_barrier_gomp_cancel(int gtid) {
+  if (__kmp_omp_cancellation) {
+    int cancelled = __kmp_barrier_template<true>(bs_plain_barrier, gtid, FALSE,
+                                                 0, NULL, NULL);
+    if (cancelled) {
+      int tid = __kmp_tid_from_gtid(gtid);
+      kmp_info_t *this_thr = __kmp_threads[gtid];
+      if (KMP_MASTER_TID(tid)) {
+        // Master does not need to revert anything
+      } else {
+        // Workers need to revert their private b_arrived flag
+        this_thr->th.th_bar[bs_plain_barrier].bb.b_arrived -=
+            KMP_BARRIER_STATE_BUMP;
+      }
+    }
+    return cancelled;
+  }
+  __kmp_barrier(bs_plain_barrier, gtid, FALSE, 0, NULL, NULL);
+  return FALSE;
+}
+#endif
+
 void __kmp_end_split_barrier(enum barrier_type bt, int gtid) {
   KMP_TIME_DEVELOPER_PARTITIONED_BLOCK(KMP_end_split_barrier);
   KMP_SET_THREAD_STATE_BLOCK(PLAIN_BARRIER);
diff --git a/openmp/runtime/src/kmp_ftn_os.h b/openmp/runtime/src/kmp_ftn_os.h
index 6a6ccb31522..b5b32f7d9be 100644
--- a/openmp/runtime/src/kmp_ftn_os.h
+++ b/openmp/runtime/src/kmp_ftn_os.h
@@ -626,8 +626,6 @@
 #define KMP_API_NAME_GOMP_TASKYIELD GOMP_taskyield
 
 // All GOMP_4.0 symbols
-// TODO: As of 2013-10-14, none of the GOMP_4.0 functions are implemented in
-// libomp
 #define KMP_API_NAME_GOMP_BARRIER_CANCEL GOMP_barrier_cancel
 #define KMP_API_NAME_GOMP_CANCEL GOMP_cancel
 #define KMP_API_NAME_GOMP_CANCELLATION_POINT GOMP_cancellation_point
diff --git a/openmp/runtime/src/kmp_gsupport.cpp b/openmp/runtime/src/kmp_gsupport.cpp
index 104e0558813..5d9c44d2ded 100644
--- a/openmp/runtime/src/kmp_gsupport.cpp
+++ b/openmp/runtime/src/kmp_gsupport.cpp
@@ -1544,11 +1544,7 @@ void KMP_EXPAND_NAME(KMP_API_NAME_GOMP_TASKGROUP_END)(void) {
   return;
 }
 
-#ifndef KMP_DEBUG
-static
-#endif /* KMP_DEBUG */
-    kmp_int32
-    __kmp_gomp_to_omp_cancellation_kind(int gomp_kind) {
+static kmp_int32 __kmp_gomp_to_omp_cancellation_kind(int gomp_kind) {
   kmp_int32 cncl_kind = 0;
   switch (gomp_kind) {
   case 1:
@@ -1567,71 +1563,49 @@ static
   return cncl_kind;
 }
 
+// Return true if cancellation should take place, false otherwise
 bool KMP_EXPAND_NAME(KMP_API_NAME_GOMP_CANCELLATION_POINT)(int which) {
-  if (__kmp_omp_cancellation) {
-    KMP_FATAL(NoGompCancellation);
-  }
   int gtid = __kmp_get_gtid();
   MKLOC(loc, "GOMP_cancellation_point");
-  KA_TRACE(20, ("GOMP_cancellation_point: T#%d\n", gtid));
-
+  KA_TRACE(20, ("GOMP_cancellation_point: T#%d which:%d\n", gtid, which));
   kmp_int32 cncl_kind = __kmp_gomp_to_omp_cancellation_kind(which);
-
   return __kmpc_cancellationpoint(&loc, gtid, cncl_kind);
 }
 
-bool KMP_EXPAND_NAME(KMP_API_NAME_GOMP_BARRIER_CANCEL)(void) {
-  if (__kmp_omp_cancellation) {
-    KMP_FATAL(NoGompCancellation);
-  }
-  KMP_FATAL(NoGompCancellation);
-  int gtid = __kmp_get_gtid();
-  MKLOC(loc, "GOMP_barrier_cancel");
-  KA_TRACE(20, ("GOMP_barrier_cancel: T#%d\n", gtid));
-
-  return __kmpc_cancel_barrier(&loc, gtid);
-}
-
+// Return true if cancellation should take place, false otherwise
 bool KMP_EXPAND_NAME(KMP_API_NAME_GOMP_CANCEL)(int which, bool do_cancel) {
-  if (__kmp_omp_cancellation) {
-    KMP_FATAL(NoGompCancellation);
-  } else {
-    return FALSE;
-  }
-
   int gtid = __kmp_get_gtid();
   MKLOC(loc, "GOMP_cancel");
-  KA_TRACE(20, ("GOMP_cancel: T#%d\n", gtid));
-
+  KA_TRACE(20, ("GOMP_cancel: T#%d which:%d do_cancel:%d\n", gtid, which,
+                (int)do_cancel));
   kmp_int32 cncl_kind = __kmp_gomp_to_omp_cancellation_kind(which);
 
   if (do_cancel == FALSE) {
-    return KMP_EXPAND_NAME(KMP_API_NAME_GOMP_CANCELLATION_POINT)(which);
+    return __kmpc_cancellationpoint(&loc, gtid, cncl_kind);
   } else {
     return __kmpc_cancel(&loc, gtid, cncl_kind);
   }
 }
 
+// Return true if cancellation should take place, false otherwise
+bool KMP_EXPAND_NAME(KMP_API_NAME_GOMP_BARRIER_CANCEL)(void) {
+  int gtid = __kmp_get_gtid();
+  KA_TRACE(20, ("GOMP_barrier_cancel: T#%d\n", gtid));
+  return __kmp_barrier_gomp_cancel(gtid);
+}
+
+// Return true if cancellation should take place, false otherwise
 bool KMP_EXPAND_NAME(KMP_API_NAME_GOMP_SECTIONS_END_CANCEL)(void) {
-  if (__kmp_omp_cancellation) {
-    KMP_FATAL(NoGompCancellation);
-  }
   int gtid = __kmp_get_gtid();
-  MKLOC(loc, "GOMP_sections_end_cancel");
   KA_TRACE(20, ("GOMP_sections_end_cancel: T#%d\n", gtid));
-
-  return __kmpc_cancel_barrier(&loc, gtid);
+  return __kmp_barrier_gomp_cancel(gtid);
 }
 
+// Return true if cancellation should take place, false otherwise
 bool KMP_EXPAND_NAME(KMP_API_NAME_GOMP_LOOP_END_CANCEL)(void) {
-  if (__kmp_omp_cancellation) {
-    KMP_FATAL(NoGompCancellation);
-  }
   int gtid = __kmp_get_gtid();
-  MKLOC(loc, "GOMP_loop_end_cancel");
   KA_TRACE(20, ("GOMP_loop_end_cancel: T#%d\n", gtid));
-
-  return __kmpc_cancel_barrier(&loc, gtid);
+  return __kmp_barrier_gomp_cancel(gtid);
 }
 
 // All target functions are empty as of 2014-05-29
diff --git a/openmp/runtime/src/kmp_wait_release.h b/openmp/runtime/src/kmp_wait_release.h
index bea59b3e437..21f36101f3a 100644
--- a/openmp/runtime/src/kmp_wait_release.h
+++ b/openmp/runtime/src/kmp_wait_release.h
@@ -155,8 +155,9 @@ static void __ompt_implicit_task_end(kmp_info_t *this_thr,
    to wake it back up to prevent deadlocks!
 
    NOTE: We may not belong to a team at this point.  */
-template <class C, int final_spin>
-static inline void
+template <class C, int final_spin, bool cancellable = false,
+          bool sleepable = true>
+static inline bool
 __kmp_wait_template(kmp_info_t *this_thr,
                     C *flag USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
 #if USE_ITT_BUILD && USE_ITT_NOTIFY
@@ -176,9 +177,14 @@ __kmp_wait_template(kmp_info_t *this_thr,
   KMP_FSYNC_SPIN_INIT(spin, NULL);
   if (flag->done_check()) {
     KMP_FSYNC_SPIN_ACQUIRED(CCAST(void *, spin));
-    return;
+    return false;
   }
   th_gtid = this_thr->th.th_info.ds.ds_gtid;
+  if (cancellable) {
+    kmp_team_t *team = this_thr->th.th_team;
+    if (team && team->t.t_cancel_request == cancel_parallel)
+      return true;
+  }
 #if KMP_OS_UNIX
   if (final_spin)
     KMP_ATOMIC_ST_REL(&this_thr->th.th_blocking, true);
@@ -400,6 +406,12 @@ final_spin=FALSE)
       KMP_PUSH_PARTITIONED_TIMER(OMP_idle);
     }
 #endif
+    // Check if the barrier surrounding this wait loop has been cancelled
+    if (cancellable) {
+      kmp_team_t *team = this_thr->th.th_team;
+      if (team && team->t.t_cancel_request == cancel_parallel)
+        break;
+    }
 
     // Don't suspend if KMP_BLOCKTIME is set to "infinite"
     if (__kmp_dflt_blocktime == KMP_MAX_BLOCKTIME
@@ -421,6 +433,10 @@ final_spin=FALSE)
     if (KMP_BLOCKING(hibernate_goal, poll_count++))
       continue;
 #endif
+    // Don't suspend if wait loop designated non-sleepable
+    // in template parameters
+    if (!sleepable)
+      continue;
 
 #if OMP_50_ENABLED
     if (__kmp_dflt_blocktime == KMP_MAX_BLOCKTIME &&
@@ -479,6 +495,21 @@ final_spin=FALSE)
     KMP_ATOMIC_ST_REL(&this_thr->th.th_blocking, false);
 #endif
   KMP_FSYNC_SPIN_ACQUIRED(CCAST(void *, spin));
+  if (cancellable) {
+    kmp_team_t *team = this_thr->th.th_team;
+    if (team && team->t.t_cancel_request == cancel_parallel) {
+      if (tasks_completed) {
+        // undo the previous decrement of unfinished_threads so that the
+        // thread can decrement at the join barrier with no problem
+        kmp_task_team_t *task_team = this_thr->th.th_task_team;
+        std::atomic<kmp_int32> *unfinished_threads =
+            &(task_team->tt.tt_unfinished_threads);
+        KMP_ATOMIC_INC(unfinished_threads);
+      }
+      return true;
+    }
+  }
+  return false;
 }
 
 /* Release any threads specified as waiting on the flag by releasing the flag
@@ -796,6 +827,18 @@ public:
       __kmp_wait_template<kmp_flag_64, FALSE>(
           this_thr, this USE_ITT_BUILD_ARG(itt_sync_obj));
   }
+  bool wait_cancellable_nosleep(kmp_info_t *this_thr,
+                                int final_spin
+                                    USE_ITT_BUILD_ARG(void *itt_sync_obj)) {
+    bool retval = false;
+    if (final_spin)
+      retval = __kmp_wait_template<kmp_flag_64, TRUE, true, false>(
+          this_thr, this USE_ITT_BUILD_ARG(itt_sync_obj));
+    else
+      retval = __kmp_wait_template<kmp_flag_64, FALSE, true, false>(
+          this_thr, this USE_ITT_BUILD_ARG(itt_sync_obj));
+    return retval;
+  }
   void release() { __kmp_release_template(this); }
   flag_type get_ptr_type() { return flag64; }
 };