[OpenMP] libomp: implement OpenMP 5.1 nteams-var and teams-thread-limit-var ICVs

[AndreyChurbanov]

Summary

The change includes OMP_NUM_TEAMS, OMP_TEAMS_THREAD_LIMIT env variables,
omp_set_num_teams, omp_get_max_teams, omp_set_teams_thread_limit,
omp_get_teams_thread_limit routines.

[jdoerfert]

We should replace the OMP_NUM_TEAMS handling in the plugins with this, right? Separate patch or can we do it right away?

[hbae]

The scope of nteams-var and teams-thread-limit-var is device, so I believe proper handling in plugin should be using extended variable name such as OMP_NUM_TEAMS_DEV[_<device>] as described in OpenMP 5.1 page 75.
I think it is best to leave this up to contributors to each plugins.

1 Each ICV that does not have global scope (see Table 2.3) has a set of device-specific environment
2 variables that extend the variables defined in Table 2.1 with the following syntax:
3 <ENVIRONMENT VARIABLE>_DEV[_<device>]
4 where <ENVIRONMENT VARIABLE> is one of the variables from Table 2.1 and <device> is the
5 device number as specified in the device clause (see Section 2.14).

[hbae]

Looks good to me.

openmp/runtime/src/dllexports

kmp_set_disp_num_buffers                    890
     ompc_set_affinity_format                752
     ompc_get_affinity_format                753
     ompc_display_affinity                   754
     ompc_capture_affinity                   755
     omp_pause_resource                      756
     omp_pause_resource_all                  757
     omp_get_supported_active_levels         758
     omp_fulfill_event                       759
+    omp_set_num_teams                       802
+    omp_get_max_teams                       803
+    omp_set_teams_thread_limit              804
+    omp_get_teams_thread_limit              805
     omp_display_env                         733
     omp_calloc                              776
     omp_realloc                             777
 
     omp_null_allocator                     DATA
     omp_default_mem_alloc                  DATA
     omp_large_cap_mem_alloc                DATA
     omp_const_mem_alloc                    DATA

openmp/runtime/src/include/omp.h.var

typedef enum omp_pause_resource_t {
       omp_pause_soft = 1,
       omp_pause_hard = 2
     } omp_pause_resource_t;
     extern int __KAI_KMPC_CONVENTION omp_pause_resource(omp_pause_resource_t, int);
     extern int __KAI_KMPC_CONVENTION omp_pause_resource_all(omp_pause_resource_t);
 
     extern int __KAI_KMPC_CONVENTION omp_get_supported_active_levels(void);
 
+    /* OpenMP 5.1 */
+    extern void __KAI_KMPC_CONVENTION omp_set_num_teams(int num_teams);
+    extern int __KAI_KMPC_CONVENTION omp_get_max_teams(void);
+    extern void __KAI_KMPC_CONVENTION omp_set_teams_thread_limit(int limit);
+    extern int __KAI_KMPC_CONVENTION omp_get_teams_thread_limit(void);
+
     /* OpenMP 5.1 Display Environment */
     extern void omp_display_env(int verbose);
 
 #   undef __KAI_KMPC_CONVENTION
 #   undef __KMP_IMP
 
     /* Warning:
        The following typedefs are not standard, deprecated and will be removed in a future release.

openmp/runtime/src/include/omp_lib.h.var

!DIR$ ENDIF
 
         function omp_capture_affinity(buffer, format)
           import
           character (len=*) :: format
           character (len=*) :: buffer
           integer (kind=kmp_size_t_kind) :: omp_capture_affinity
         end function omp_capture_affinity
 
+        subroutine omp_set_num_teams(num_teams) bind(c)
+          import
+          integer (kind=omp_integer_kind), value :: num_teams
+        end subroutine omp_set_num_teams
+
+        function omp_get_max_teams() bind(c)
+          import
+          integer (kind=omp_integer_kind) omp_get_max_teams
+        end function omp_get_max_teams
+
+        subroutine omp_set_teams_thread_limit(thread_limit) bind(c)
+          import
+          use omp_lib_kinds
+          integer (kind=omp_integer_kind), value :: thread_limit
+        end subroutine omp_set_teams_thread_limit
+
+        function omp_get_teams_thread_limit() bind(c)
+          import
+          use omp_lib_kinds
+          integer (kind=omp_integer_kind) omp_get_teams_thread_limit
+        end function omp_get_teams_thread_limit
+
         subroutine omp_display_env(verbose) bind(c)
           import
           logical (kind=omp_logical_kind), value :: verbose
         end subroutine omp_display_env
 
         function omp_target_alloc(size, device_num) bind(c)
           use, intrinsic :: iso_c_binding, only : c_ptr, c_size_t, c_int
           type(c_ptr) omp_target_alloc

openmp/runtime/src/include/omp_lib.f90.var

!DIR$ ENDIF
 
           function omp_capture_affinity(buffer, format)
             use omp_lib_kinds
             character (len=*) :: format
             character (len=*) :: buffer
             integer (kind=kmp_size_t_kind) :: omp_capture_affinity
           end function omp_capture_affinity
 
+          subroutine omp_set_num_teams(num_teams) bind(c)
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind), value :: num_teams
+          end subroutine omp_set_num_teams
+
+          function omp_get_max_teams() bind(c)
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind) omp_get_max_teams
+          end function omp_get_max_teams
+
+          subroutine omp_set_teams_thread_limit(thread_limit) bind(c)
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind), value :: thread_limit
+          end subroutine omp_set_teams_thread_limit
+
+          function omp_get_teams_thread_limit() bind(c)
+            use omp_lib_kinds
+            integer (kind=omp_integer_kind) omp_get_teams_thread_limit
+          end function omp_get_teams_thread_limit
+
           subroutine omp_display_env(verbose) bind(c)
             use omp_lib_kinds
             logical (kind=omp_logical_kind), value :: verbose
           end subroutine omp_display_env
 
 
           function omp_target_alloc(size, device_num) bind(c)
             use omp_lib_kinds

openmp/runtime/src/kmp.h

 extern const char *__kmp_speculative_statsfile;
 #endif
 
 #endif // KMP_USE_ADAPTIVE_LOCKS
 
 extern int __kmp_display_env; /* TRUE or FALSE */
 extern int __kmp_display_env_verbose; /* TRUE if OMP_DISPLAY_ENV=VERBOSE */
 extern int __kmp_omp_cancellation; /* TRUE or FALSE */
+extern int __kmp_nteams;
+extern int __kmp_teams_thread_limit;
 
 /* ------------------------------------------------------------------------- */
 
 /* the following are protected by the fork/join lock */
 /* write: lock  read: anytime */
 extern kmp_info_t **__kmp_threads; /* Descriptors for the threads */
 /* read/write: lock */
 extern volatile kmp_team_t *__kmp_team_pool;
KMP_EXPORT void *__kmpc_task_reduction_modifier_init(ident_t *loc, int gtid,
                                                      void *data);
 KMP_EXPORT void *__kmpc_taskred_modifier_init(ident_t *loc, int gtid, int is_ws,
                                               int num, void *data);
 KMP_EXPORT void __kmpc_task_reduction_modifier_fini(ident_t *loc, int gtid,
                                                     int is_ws);
 KMP_EXPORT kmp_int32 __kmpc_omp_reg_task_with_affinity(
     ident_t *loc_ref, kmp_int32 gtid, kmp_task_t *new_task, kmp_int32 naffins,
     kmp_task_affinity_info_t *affin_list);
+KMP_EXPORT void __kmp_set_num_teams(int num_teams);
+KMP_EXPORT int __kmp_get_max_teams(void);
+KMP_EXPORT void __kmp_set_teams_thread_limit(int limit);
+KMP_EXPORT int __kmp_get_teams_thread_limit(void);
 
 /* Lock interface routines (fast versions with gtid passed in) */
 KMP_EXPORT void __kmpc_init_lock(ident_t *loc, kmp_int32 gtid,
                                  void **user_lock);
 KMP_EXPORT void __kmpc_init_nest_lock(ident_t *loc, kmp_int32 gtid,
                                       void **user_lock);
 KMP_EXPORT void __kmpc_destroy_lock(ident_t *loc, kmp_int32 gtid,
                                     void **user_lock);

openmp/runtime/src/kmp_ftn_entry.h

 /*
  * kmp_ftn_entry.h -- Fortran entry linkage support for OpenMP.
  */
 
 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FTN_STDCALL
 #error The support file kmp_ftn_entry.h should not be compiled by itself.
 #endif
 
 #ifdef KMP_STUB
 #include "kmp_stub.h"
 #endif
 
 #include "kmp_i18n.h"
 

 
 // This macro helps to reduce code duplication.
 #ifdef PASS_ARGS_BY_VALUE
 #define KMP_DEREF
 #else
 #define KMP_DEREF *
 #endif
 
 void FTN_STDCALL FTN_SET_STACKSIZE(int KMP_DEREF arg) {
 #ifdef KMP_STUB
   __kmps_set_stacksize(KMP_DEREF arg);
 #else
   // __kmp_aux_set_stacksize initializes the library if needed
   __kmp_aux_set_stacksize((size_t)KMP_DEREF arg);
 #endif
 }
 
public:
   const char *get() const { return buf; }
 };
 #endif // KMP_STUB
 
 /*
  * Set the value of the affinity-format-var ICV on the current device to the
  * format specified in the argument.
 */
 void FTN_STDCALL FTN_SET_AFFINITY_FORMAT(char const *format, size_t size) {
 #ifdef KMP_STUB
   return;
 #else
   if (!__kmp_init_serial) {
     __kmp_serial_initialize();
   }
   ConvertedString cformat(format, size);
   // Since the __kmp_affinity_format variable is a C string, do not
#else
   __kmp_entry_thread()->th.th_current_task->td_icvs.default_device =
       KMP_DEREF arg;
 #endif
 }
 
 // Get number of NON-HOST devices.
 // libomptarget, if loaded, provides this function in api.cpp.
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_NUM_DEVICES)(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_GET_NUM_DEVICES)(void) {
 #if KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB)
   return 0;
 #else
   int (*fptr)();
   if ((*(void **)(&fptr) = dlsym(RTLD_DEFAULT, "_Offload_number_of_devices"))) {
     return (*fptr)();
   } else if ((*(void **)(&fptr) = dlsym(RTLD_NEXT, "omp_get_num_devices"))) {
     return (*fptr)();
   } else { // liboffload & libomptarget don't exist
     return 0;
   }
 #endif // KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB)
 }
 
 // This function always returns true when called on host device.
 // Compiler/libomptarget should handle when it is called inside target region.
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
 int FTN_STDCALL KMP_EXPAND_NAME(FTN_IS_INITIAL_DEVICE)(void) {
   return 1; // This is the host
 }
 
 // libomptarget, if loaded, provides this function
 int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
 int FTN_STDCALL FTN_GET_INITIAL_DEVICE(void) {
 #if KMP_MIC || KMP_OS_DARWIN || KMP_OS_WINDOWS || defined(KMP_STUB)
   // same as omp_get_num_devices()
   return 0;
 #else
   int (*fptr)();
   if ((*(void **)(&fptr) = dlsym(RTLD_NEXT, "omp_get_initial_device"))) {
     return (*fptr)();
   } else { // liboffload & libomptarget don't exist
#else
   return __kmp_max_task_priority;
 #endif
 }
 
 // This function will be defined in libomptarget. When libomptarget is not
 // loaded, we assume we are on the host and return KMP_HOST_DEVICE.
 // Compiler/libomptarget will handle this if called inside target.
 int FTN_STDCALL FTN_GET_DEVICE_NUM(void) KMP_WEAK_ATTRIBUTE_EXTERNAL;
 int FTN_STDCALL FTN_GET_DEVICE_NUM(void) { return FTN_GET_INITIAL_DEVICE(); }
 
 // Compiler will ensure that this is only called from host in sequential region
 int FTN_STDCALL FTN_PAUSE_RESOURCE(kmp_pause_status_t kind, int device_num) {
 #ifdef KMP_STUB
   return 1; // just fail
 #else
   if (device_num == FTN_GET_INITIAL_DEVICE())
     return __kmpc_pause_resource(kind);

 }
 
 void FTN_STDCALL FTN_FULFILL_EVENT(kmp_event_t *event) {
 #ifndef KMP_STUB
   __kmp_fulfill_event(event);
 #endif
 }
 
+// nteams-var per-device ICV
+void FTN_STDCALL FTN_SET_NUM_TEAMS(int KMP_DEREF num_teams) {
+#ifdef KMP_STUB
+// Nothing.
+#else
+  if (!__kmp_init_serial) {
+    __kmp_serial_initialize();
+  }
+  __kmp_set_num_teams(KMP_DEREF num_teams);
+#endif
+}
+int FTN_STDCALL FTN_GET_MAX_TEAMS(void) {
+#ifdef KMP_STUB
+  return 1;
+#else
+  if (!__kmp_init_serial) {
+    __kmp_serial_initialize();
+  }
+  return __kmp_get_max_teams();
+#endif
+}
+// teams-thread-limit-var per-device ICV
+void FTN_STDCALL FTN_SET_TEAMS_THREAD_LIMIT(int KMP_DEREF limit) {
+#ifdef KMP_STUB
+// Nothing.
+#else
+  if (!__kmp_init_serial) {
+    __kmp_serial_initialize();
+  }
+  __kmp_set_teams_thread_limit(KMP_DEREF limit);
+#endif
+}
+int FTN_STDCALL FTN_GET_TEAMS_THREAD_LIMIT(void) {
+#ifdef KMP_STUB
+  return 1;
+#else
+  if (!__kmp_init_serial) {
+    __kmp_serial_initialize();
+  }
+  return __kmp_get_teams_thread_limit();
+#endif
+}
+
 // display environment variables when requested
 void FTN_STDCALL FTN_DISPLAY_ENV(int verbose) {
 #ifndef KMP_STUB
   __kmp_omp_display_env(verbose);
 #endif
 }
 
 // GCC compatibility (versioned symbols)
/* These following sections create versioned symbols for the
    Step 2) Set __kmp_api_omp_set_num_threads_10_alias to version:
      omp_set_num_threads@OMP_1.0
    Step 2B) Set __kmp_api_omp_set_num_threads to default version:
      omp_set_num_threads@@VERSION
 */
 
 // OMP_1.0 versioned symbols
 KMP_VERSION_SYMBOL(FTN_SET_NUM_THREADS, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_GET_NUM_THREADS, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_GET_MAX_THREADS, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_GET_THREAD_NUM, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_GET_NUM_PROCS, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_IN_PARALLEL, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_SET_DYNAMIC, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_GET_DYNAMIC, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_SET_NESTED, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_GET_NESTED, 10, "OMP_1.0");
 KMP_VERSION_SYMBOL(FTN_INIT_LOCK, 10, "OMP_1.0");

openmp/runtime/src/kmp_ftn_os.h

 #define FTN_GET_AFFINITY_FORMAT omp_get_affinity_format
 #define FTN_DISPLAY_AFFINITY omp_display_affinity
 #define FTN_CAPTURE_AFFINITY omp_capture_affinity
 #define FTN_PAUSE_RESOURCE omp_pause_resource
 #define FTN_PAUSE_RESOURCE_ALL omp_pause_resource_all
 #define FTN_GET_SUPPORTED_ACTIVE_LEVELS omp_get_supported_active_levels
 #define FTN_DISPLAY_ENV omp_display_env
 #define FTN_FULFILL_EVENT omp_fulfill_event
+#define FTN_SET_NUM_TEAMS omp_set_num_teams
+#define FTN_GET_MAX_TEAMS omp_get_max_teams
+#define FTN_SET_TEAMS_THREAD_LIMIT omp_set_teams_thread_limit
+#define FTN_GET_TEAMS_THREAD_LIMIT omp_get_teams_thread_limit
 
 #endif /* KMP_FTN_PLAIN */
 
 /* ------------------------------------------------------------------------ */
 
 #if KMP_FTN_ENTRIES == KMP_FTN_APPEND
 
 #define FTN_SET_STACKSIZE kmp_set_stacksize_

 #define FTN_GET_AFFINITY_FORMAT omp_get_affinity_format_
 #define FTN_DISPLAY_AFFINITY omp_display_affinity_
 #define FTN_CAPTURE_AFFINITY omp_capture_affinity_
 #define FTN_PAUSE_RESOURCE omp_pause_resource_
 #define FTN_PAUSE_RESOURCE_ALL omp_pause_resource_all_
 #define FTN_GET_SUPPORTED_ACTIVE_LEVELS omp_get_supported_active_levels_
 #define FTN_DISPLAY_ENV omp_display_env_
 #define FTN_FULFILL_EVENT omp_fulfill_event_
+#define FTN_SET_NUM_TEAMS omp_set_num_teams_
+#define FTN_GET_MAX_TEAMS omp_get_max_teams_
+#define FTN_SET_TEAMS_THREAD_LIMIT omp_set_teams_thread_limit_
+#define FTN_GET_TEAMS_THREAD_LIMIT omp_get_teams_thread_limit_
 
 #endif /* KMP_FTN_APPEND */
 
 /* ------------------------------------------------------------------------ */
 
 #if KMP_FTN_ENTRIES == KMP_FTN_UPPER
 
 #define FTN_SET_STACKSIZE KMP_SET_STACKSIZE

 #define FTN_GET_AFFINITY_FORMAT OMP_GET_AFFINITY_FORMAT
 #define FTN_DISPLAY_AFFINITY OMP_DISPLAY_AFFINITY
 #define FTN_CAPTURE_AFFINITY OMP_CAPTURE_AFFINITY
 #define FTN_PAUSE_RESOURCE OMP_PAUSE_RESOURCE
 #define FTN_PAUSE_RESOURCE_ALL OMP_PAUSE_RESOURCE_ALL
 #define FTN_GET_SUPPORTED_ACTIVE_LEVELS OMP_GET_SUPPORTED_ACTIVE_LEVELS
 #define FTN_DISPLAY_ENV OMP_DISPLAY_ENV
 #define FTN_FULFILL_EVENT OMP_FULFILL_EVENT
+#define FTN_SET_NUM_TEAMS OMP_SET_NUM_TEAMS
+#define FTN_GET_MAX_TEAMS OMP_GET_MAX_TEAMS
+#define FTN_SET_TEAMS_THREAD_LIMIT OMP_SET_TEAMS_THREAD_LIMIT
+#define FTN_GET_TEAMS_THREAD_LIMIT OMP_GET_TEAMS_THREAD_LIMIT
 
 #endif /* KMP_FTN_UPPER */
 
 /* ------------------------------------------------------------------------ */
 
 #if KMP_FTN_ENTRIES == KMP_FTN_UAPPEND
 
 #define FTN_SET_STACKSIZE KMP_SET_STACKSIZE_

 #define FTN_GET_AFFINITY_FORMAT OMP_GET_AFFINITY_FORMAT_
 #define FTN_DISPLAY_AFFINITY OMP_DISPLAY_AFFINITY_
 #define FTN_CAPTURE_AFFINITY OMP_CAPTURE_AFFINITY_
 #define FTN_PAUSE_RESOURCE OMP_PAUSE_RESOURCE_
 #define FTN_PAUSE_RESOURCE_ALL OMP_PAUSE_RESOURCE_ALL_
 #define FTN_GET_SUPPORTED_ACTIVE_LEVELS OMP_GET_SUPPORTED_ACTIVE_LEVELS_
 #define FTN_DISPLAY_ENV OMP_DISPLAY_ENV_
 #define FTN_FULFILL_EVENT OMP_FULFILL_EVENT_
+#define FTN_SET_NUM_TEAMS OMP_SET_NUM_TEAMS_
+#define FTN_GET_MAX_TEAMS OMP_GET_MAX_TEAMS_
+#define FTN_SET_TEAMS_THREAD_LIMIT OMP_SET_TEAMS_THREAD_LIMIT_
+#define FTN_GET_TEAMS_THREAD_LIMIT OMP_GET_TEAMS_THREAD_LIMIT_
 
 #endif /* KMP_FTN_UAPPEND */
 
 /* -------------------------- GOMP API NAMES ------------------------ */
 // All GOMP_1.0 symbols
 #define KMP_API_NAME_GOMP_ATOMIC_END GOMP_atomic_end
 #define KMP_API_NAME_GOMP_ATOMIC_START GOMP_atomic_start
 #define KMP_API_NAME_GOMP_BARRIER GOMP_barrier

openmp/runtime/src/kmp_global.cpp

 const char *__kmp_speculative_statsfile = "-";
 #endif
 
 #endif // KMP_USE_ADAPTIVE_LOCKS
 
 int __kmp_display_env = FALSE;
 int __kmp_display_env_verbose = FALSE;
 int __kmp_omp_cancellation = FALSE;
+int __kmp_nteams = 0;
+int __kmp_teams_thread_limit = 0;
 
 #if KMP_HAVE_MWAIT || KMP_HAVE_UMWAIT
 int __kmp_user_level_mwait = FALSE;
 int __kmp_umwait_enabled = FALSE;
 int __kmp_mwait_enabled = FALSE;
 int __kmp_mwait_hints = 0;
 #endif
 

openmp/runtime/src/kmp_runtime.cpp

int __kmp_get_max_active_levels(int gtid) {
   KMP_DEBUG_ASSERT(thread->th.th_current_task);
   KF_TRACE(10, ("__kmp_get_max_active_levels: thread %d, curtask=%p, "
                 "curtask_maxaclevel=%d\n",
                 gtid, thread->th.th_current_task,
                 thread->th.th_current_task->td_icvs.max_active_levels));
   return thread->th.th_current_task->td_icvs.max_active_levels;
 }
 
+// nteams-var per-device ICV
+void __kmp_set_num_teams(int num_teams) {
+  if (num_teams > 0)
+    __kmp_nteams = num_teams;
+}
+int __kmp_get_max_teams(void) { return __kmp_nteams; }
+// teams-thread-limit-var per-device ICV
+void __kmp_set_teams_thread_limit(int limit) {
+  if (limit > 0)
+    __kmp_teams_thread_limit = limit;
+}
+int __kmp_get_teams_thread_limit(void) { return __kmp_teams_thread_limit; }
+
 KMP_BUILD_ASSERT(sizeof(kmp_sched_t) == sizeof(int));
 KMP_BUILD_ASSERT(sizeof(enum sched_type) == sizeof(int));
 
 /* Changes def_sched_var ICV values (run-time schedule kind and chunk) */
 void __kmp_set_schedule(int gtid, kmp_sched_t kind, int chunk) {
   kmp_info_t *thread;
   kmp_sched_t orig_kind;
   //    kmp_team_t *team;

 /* this sets the requested number of teams for the teams region and/or
    the number of threads for the next parallel region encountered  */
 void __kmp_push_num_teams(ident_t *id, int gtid, int num_teams,
                           int num_threads) {
   kmp_info_t *thr = __kmp_threads[gtid];
   KMP_DEBUG_ASSERT(num_teams >= 0);
   KMP_DEBUG_ASSERT(num_threads >= 0);
 
-  if (num_teams == 0)
+  if (num_teams == 0) {
+    if (__kmp_nteams > 0) {
+      num_teams = __kmp_nteams;
+    } else {
       num_teams = 1; // default number of teams is 1.
+    }
+  }
   if (num_teams > __kmp_teams_max_nth) { // if too many teams requested?
     if (!__kmp_reserve_warn) {
       __kmp_reserve_warn = 1;
       __kmp_msg(kmp_ms_warning,
                 KMP_MSG(CantFormThrTeam, num_teams, __kmp_teams_max_nth),
                 KMP_HNT(Unset_ALL_THREADS), __kmp_msg_null);
     }
     num_teams = __kmp_teams_max_nth;
   }
   // Set number of teams (number of threads in the outer "parallel" of the
   // teams)
   thr->th.th_set_nproc = thr->th.th_teams_size.nteams = num_teams;
 
   // Remember the number of threads for inner parallel regions
   if (!TCR_4(__kmp_init_middle))
     __kmp_middle_initialize(); // get internal globals calculated
   KMP_DEBUG_ASSERT(__kmp_avail_proc);
   KMP_DEBUG_ASSERT(__kmp_dflt_team_nth);
   if (num_threads == 0) {
+    if (__kmp_teams_thread_limit > 0) {
+      num_threads = __kmp_teams_thread_limit;
+    } else {
       num_threads = __kmp_avail_proc / num_teams;
+    }
     // adjust num_threads w/o warning as it is not user setting
     // num_threads = min(num_threads, nthreads-var, thread-limit-var)
     // no thread_limit clause specified -  do not change thread-limit-var ICV
     if (num_threads > __kmp_dflt_team_nth) {
       num_threads = __kmp_dflt_team_nth; // honor nthreads-var ICV
     }
     if (num_threads > thr->th.th_current_task->td_icvs.thread_limit) {
       num_threads = thr->th.th_current_task->td_icvs.thread_limit;

openmp/runtime/src/kmp_settings.cpp

 } // __kmp_stg_parse_thread_limit
 
 static void __kmp_stg_print_thread_limit(kmp_str_buf_t *buffer,
                                          char const *name, void *data) {
   __kmp_stg_print_int(buffer, name, __kmp_cg_max_nth);
 } // __kmp_stg_print_thread_limit
 
 // -----------------------------------------------------------------------------
+// OMP_NUM_TEAMS
+static void __kmp_stg_parse_nteams(char const *name, char const *value,
+                                   void *data) {
+  __kmp_stg_parse_int(name, value, 1, __kmp_sys_max_nth, &__kmp_nteams);
+  K_DIAG(1, ("__kmp_nteams == %d\n", __kmp_nteams));
+} // __kmp_stg_parse_nteams
+
+static void __kmp_stg_print_nteams(kmp_str_buf_t *buffer, char const *name,
+                                   void *data) {
+  __kmp_stg_print_int(buffer, name, __kmp_nteams);
+} // __kmp_stg_print_nteams
+
+// -----------------------------------------------------------------------------
+// OMP_TEAMS_THREAD_LIMIT
+static void __kmp_stg_parse_teams_th_limit(char const *name, char const *value,
+                                           void *data) {
+  __kmp_stg_parse_int(name, value, 1, __kmp_sys_max_nth,
+                      &__kmp_teams_thread_limit);
+  K_DIAG(1, ("__kmp_teams_thread_limit == %d\n", __kmp_teams_thread_limit));
+} // __kmp_stg_parse_teams_th_limit
+
+static void __kmp_stg_print_teams_th_limit(kmp_str_buf_t *buffer,
+                                           char const *name, void *data) {
+  __kmp_stg_print_int(buffer, name, __kmp_teams_thread_limit);
+} // __kmp_stg_print_teams_th_limit
+
+// -----------------------------------------------------------------------------
 // KMP_TEAMS_THREAD_LIMIT
 static void __kmp_stg_parse_teams_thread_limit(char const *name,
                                                char const *value, void *data) {
   __kmp_stg_parse_int(name, value, 1, __kmp_sys_max_nth, &__kmp_teams_max_nth);
 } // __kmp_stg_teams_thread_limit
 
 static void __kmp_stg_print_teams_thread_limit(kmp_str_buf_t *buffer,
                                                char const *name, void *data) {
#endif
     {"OMP_MAX_TASK_PRIORITY", __kmp_stg_parse_max_task_priority,
      __kmp_stg_print_max_task_priority, NULL, 0, 0},
     {"KMP_TASKLOOP_MIN_TASKS", __kmp_stg_parse_taskloop_min_tasks,
      __kmp_stg_print_taskloop_min_tasks, NULL, 0, 0},
     {"OMP_THREAD_LIMIT", __kmp_stg_parse_thread_limit,
      __kmp_stg_print_thread_limit, NULL, 0, 0},
     {"KMP_TEAMS_THREAD_LIMIT", __kmp_stg_parse_teams_thread_limit,
      __kmp_stg_print_teams_thread_limit, NULL, 0, 0},
+    {"OMP_NUM_TEAMS", __kmp_stg_parse_nteams, __kmp_stg_print_nteams, NULL, 0,
+     0},
+    {"OMP_TEAMS_THREAD_LIMIT", __kmp_stg_parse_teams_th_limit,
+     __kmp_stg_print_teams_th_limit, NULL, 0, 0},
     {"OMP_WAIT_POLICY", __kmp_stg_parse_wait_policy,
      __kmp_stg_print_wait_policy, NULL, 0, 0},
     {"KMP_DISP_NUM_BUFFERS", __kmp_stg_parse_disp_buffers,
      __kmp_stg_print_disp_buffers, NULL, 0, 0},
 #if KMP_NESTED_HOT_TEAMS
     {"KMP_HOT_TEAMS_MAX_LEVEL", __kmp_stg_parse_hot_teams_level,
      __kmp_stg_print_hot_teams_level, NULL, 0, 0},
     {"KMP_HOT_TEAMS_MODE", __kmp_stg_parse_hot_teams_mode,

openmp/runtime/test/api/omp_nteams_api.c

+// RUN: %libomp-compile-and-run
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+int main(int argc, char** argv)
+{
+  int iset, iget;
+  iset = 5;
+  omp_set_num_teams(iset);
+  iget = omp_get_max_teams();
+  if (iset != iget) {
+    fprintf(stderr, "error: nteams-var set to %d, getter returned %d\n", iset, iget);
+    exit(1);
+  }
+  iset = 7;
+  omp_set_teams_thread_limit(iset);
+  iget = omp_get_teams_thread_limit();
+  if (iset != iget) {
+    fprintf(stderr, "error: teams-thread-limit-var set to %d, getter returned %d\n", iset, iget);
+    exit(1);
+  }
+  printf("passed\n");
+  return 0;
+}

openmp/runtime/test/api/omp_nteams_api_env.c

+// RUN: %libomp-compile && env OMP_NUM_TEAMS=5 OMP_TEAMS_THREAD_LIMIT=7 %libomp-run
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+int main(int argc, char** argv)
+{
+  int iset, iget;
+  iset = 5;
+  iget = omp_get_max_teams();
+  if (iset != iget) {
+    fprintf(stderr, "error: nteams-var set to %d, getter returned %d\n", iset, iget);
+    exit(1);
+  }
+  iset = 7;
+  iget = omp_get_teams_thread_limit();
+  if (iset != iget) {
+    fprintf(stderr, "error: teams-thread-limit-var set to %d, getter returned %d\n", iset, iget);
+    exit(1);
+  }
+  printf("passed\n");
+  return 0;
+}

openmp/runtime/test/api/omp_nteams_api_env2.c

+// RUN: %libomp-compile && env OMP_NUM_TEAMS=5 OMP_TEAMS_THREAD_LIMIT=7 %libomp-run
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+int main(int argc, char** argv)
+{
+  int iset, iget;
+  iset = 4; // should override OMP_NUM_TEAMS=5
+  omp_set_num_teams(iset);
+  iget = omp_get_max_teams();
+  if (iset != iget) {
+    fprintf(stderr, "error: nteams-var set to %d, getter returned %d\n", iset, iget);
+    exit(1);
+  }
+  iset = 6; // should override OMP_TEAMS_THREAD_LIMIT=7
+  omp_set_teams_thread_limit(iset);
+  iget = omp_get_teams_thread_limit();
+  if (iset != iget) {
+    fprintf(stderr, "error: teams-thread-limit-var set to %d, getter returned %d\n", iset, iget);
+    exit(1);
+  }
+  printf("passed\n");
+  return 0;
+}