[OpenMP] Corrects the setting of the Static_Steal schedule with `omp_set_schedule` (#227236)

Static-steal is considered a valid OMP scheduler by `kmp_runtime.cpp`'s
`__kmp_set_schedule` function. But when trying pass it to the setter, we
end up with a failure because it was not added to the `__kmp_sch_map` in
`kmp_global.cpp`.

This PR adds a test that each schedule in `kmp.h` can be set and gotten
by `omp_set_schedule` and `omp_get_schedule`, which fails with static
steal and a fix for `__kmp_sch_map`. To deal with static-steal set this
way with `ordered` loops, I added similar handling as with setting
static steal as an environment variable: gets converted to dynamic,
nonmonotonic when initialising the loop, so when combined with `ordered`
simply dynamic scheduling is used. A test that static-steal with
`ordered` works was added, which used to hang before adding this
handling.

This PR was made with AI assitance. All code was reviewed and where
necessary modified by myself.

GitOrigin-RevId: 1f48c2f8f4021f016a62fe0cebf877c4c8f5dffd
diff --git a/runtime/src/kmp_dispatch.cpp b/runtime/src/kmp_dispatch.cpp
index 3b4a1f3..57e4bcb 100644
--- a/runtime/src/kmp_dispatch.cpp
+++ b/runtime/src/kmp_dispatch.cpp
@@ -265,6 +265,13 @@
       // Use the scheduling specified by OMP_SCHEDULE (or __kmp_sch_default if
       // not specified)
       schedule = team->t.t_sched.r_sched_type;
+#if KMP_STATIC_STEAL_ENABLED
+      // Treat static_steal set by omp_set_schedule() like
+      // OMP_SCHEDULE=static_steal, i.e. as nonmonotonic dynamic
+      if (SCHEDULE_WITHOUT_MODIFIERS(schedule) == kmp_sch_static_steal)
+        schedule = (enum sched_type)(kmp_sch_dynamic_chunked |
+                                     kmp_sch_modifier_nonmonotonic);
+#endif
       monotonicity = __kmp_get_monotonicity(loc, schedule, use_hier);
       schedule = SCHEDULE_WITHOUT_MODIFIERS(schedule);
       if (pr->flags.ordered) // correct monotonicity for ordered loop if needed
diff --git a/runtime/src/kmp_global.cpp b/runtime/src/kmp_global.cpp
index 15b9bab..78ae429 100644
--- a/runtime/src/kmp_global.cpp
+++ b/runtime/src/kmp_global.cpp
@@ -235,9 +235,10 @@
     kmp_sch_dynamic_chunked, // ==> kmp_sched_dynamic           = 2
     kmp_sch_guided_chunked, // ==> kmp_sched_guided            = 3
     kmp_sch_auto, // ==> kmp_sched_auto              = 4
-    kmp_sch_trapezoidal // ==> kmp_sched_trapezoidal       = 101
-    // will likely not be used, introduced here just to debug the code
-    // of public intel extension schedules
+    kmp_sch_trapezoidal, // ==> kmp_sched_trapezoidal       = 101
+#if KMP_STATIC_STEAL_ENABLED
+    kmp_sch_static_steal // ==> kmp_sched_static_steal      = 102
+#endif
 };
 
 #if KMP_MIC_SUPPORTED
diff --git a/runtime/test/worksharing/for/kmp_sched_set_get.c b/runtime/test/worksharing/for/kmp_sched_set_get.c
new file mode 100644
index 0000000..f1f1ae3
--- /dev/null
+++ b/runtime/test/worksharing/for/kmp_sched_set_get.c
@@ -0,0 +1,81 @@
+// RUN: %libomp-compile-and-run
+
+// The test checks that every schedule kind accepted by omp_set_schedule(),
+// i.e. that is not considered out of range in __kmp_set_schedule in
+// kmp_runtime.cpp, can be set, and we can afterwards get the schedule back with
+// omp_get_schedule
+
+#include <omp.h>
+#include <stdio.h>
+
+// ---------------------------------------------------------------------------
+// As with kmp.h, static steal is by default enabled, but can be manually
+// disabled. If libomp is built with -DKMP_STATIC_STEAL_ENABLED=0, pass the
+// same flag to the tests (for example through OPENMP_TEST_FLAGS).
+#ifndef KMP_STATIC_STEAL_ENABLED
+#define KMP_STATIC_STEAL_ENABLED 1
+#endif
+
+// These definitions need to match kmp_sched_t in kmp.h.
+#ifndef KMP_SCHED_TYPE_DEFINED
+#define KMP_SCHED_TYPE_DEFINED
+typedef enum kmp_sched {
+  kmp_sched_lower = 0,
+  kmp_sched_static = 1,
+  kmp_sched_dynamic = 2,
+  kmp_sched_guided = 3,
+  kmp_sched_auto = 4,
+  kmp_sched_upper_std = 5,
+  kmp_sched_lower_ext = 100,
+  kmp_sched_trapezoidal = 101,
+#if KMP_STATIC_STEAL_ENABLED
+  kmp_sched_static_steal = 102,
+#endif
+  kmp_sched_upper,
+  kmp_sched_default = kmp_sched_static,
+  kmp_sched_monotonic = 0x80000000
+} kmp_sched_t;
+#endif
+
+// ---------------------------------------------------------------------------
+
+int main() {
+  const int chunk = 5;
+  int err = 0;
+  unsigned k;
+
+  // Visit every kind that __kmp_set_schedule() accepts: the standard kinds
+  // between lower and upper_std, and the extension kinds between lower_ext
+  // and upper.
+  for (k = kmp_sched_lower + 1; k < kmp_sched_upper; ++k) {
+    omp_sched_t kind_get;
+    int chunk_get;
+
+    if (k >= kmp_sched_upper_std && k <= kmp_sched_lower_ext) {
+      continue;
+    }
+
+#ifdef DEBUG
+    printf("checking kind %u\n", k);
+#endif
+
+    omp_set_schedule((omp_sched_t)k, chunk);
+    omp_get_schedule(&kind_get, &chunk_get);
+
+    // Check kind and chunk match.
+    // The chunk size is ignored for auto, so allow a mismatched chunk size.
+    if (kind_get != (omp_sched_t)k ||
+        (k != kmp_sched_auto && chunk_get != chunk)) {
+      printf("Error: kind %u: schedule: (%d, %d) is not equal to (%u, %d)\n", k,
+             (int)kind_get, chunk_get, k, chunk);
+      ++err;
+    }
+  }
+
+  if (err > 0) {
+    printf("Failed\n");
+    return 1;
+  }
+  printf("Passed\n");
+  return 0;
+}
diff --git a/runtime/test/worksharing/for/kmp_sched_static_steal_ordered.c b/runtime/test/worksharing/for/kmp_sched_static_steal_ordered.c
new file mode 100644
index 0000000..d92d110
--- /dev/null
+++ b/runtime/test/worksharing/for/kmp_sched_static_steal_ordered.c
@@ -0,0 +1,63 @@
+// RUN: %libomp-compile-and-run
+
+// Tests that an ordered loop with schedule(runtime) when the run-time
+// schedule is set to static_steal with omp_set_schedule(), does not
+// use the unordered static-steal scheduling nor hang.
+// Such loops must fall back to dynamic, as they do with the environment
+// variable OMP_SCHEDULE=static_steal.
+
+#include <omp.h>
+#include "omp_testsuite.h"
+
+// ---------------------------------------------------------------------------
+// Definition copied from OpenMP RTL (kmp_sched_t in kmp.h).
+enum { kmp_sched_static_steal = 102 };
+// End of definition copied from OpenMP RTL.
+// ---------------------------------------------------------------------------
+
+static int last_i = 0;
+
+/* Utility function to check that i is increasing monotonically
+   with each call */
+static int check_i_islarger(int i) {
+  int islarger;
+  islarger = (i > last_i);
+  last_i = i;
+  return (islarger);
+}
+
+int test_static_steal_ordered() {
+  int sum;
+  int is_larger = 1;
+  int known_sum;
+  int i;
+
+  last_i = 0;
+  sum = 0;
+  omp_set_schedule((omp_sched_t)kmp_sched_static_steal, 4);
+
+  // num_threads(4): static_steal needs more than one thread.
+#pragma omp parallel for schedule(runtime) ordered num_threads(4)
+  for (i = 1; i <= LOOPCOUNT; i++) {
+#pragma omp ordered
+    {
+      is_larger = check_i_islarger(i) && is_larger;
+      sum = sum + i;
+    }
+  }
+
+  known_sum = (LOOPCOUNT * (LOOPCOUNT + 1)) / 2;
+  return (known_sum == sum) && is_larger;
+}
+
+int main() {
+  int i;
+  int num_failed = 0;
+
+  for (i = 0; i < REPETITIONS; i++) {
+    if (!test_static_steal_ordered()) {
+      num_failed++;
+    }
+  }
+  return num_failed;
+}
\ No newline at end of file