openmp/runtime/test/teams/kmp_num_teams.c

   1 // RUN: %libomp-compile-and-run
   2 // UNSUPPORTED: gcc
   3 // Linking fails for icc 18/19
   4 // UNSUPPORTED: icc-18, icc-19
   5
   6 #include <stdio.h>
   7 #include <stdlib.h>
   8 #include <omp.h>
   9
  10 #define NT 8
  11
  12 #ifdef __cplusplus
  13 extern "C" {
  14 #endif
  15 typedef int kmp_int32;
  16 typedef struct ident {
  17   kmp_int32 reserved_1;
  18   kmp_int32 flags;
  19   kmp_int32 reserved_2;
  20   kmp_int32 reserved_3;
  21   char const *psource;
  22 } ident_t;
  23 extern int __kmpc_global_thread_num(ident_t *);
  24 extern void __kmpc_push_num_teams_51(ident_t *, kmp_int32, kmp_int32, kmp_int32,
  25                                      kmp_int32);
  26 #ifdef __cplusplus
  27 }
  28 #endif
  29
  30 void check_num_teams(int num_teams_lb, int num_teams_ub, int thread_limit) {
  31   int nteams, nthreads;
  32   int a = 0;
  33
  34   int gtid = __kmpc_global_thread_num(NULL);
  35   __kmpc_push_num_teams_51(NULL, gtid, num_teams_lb, num_teams_ub,
  36                            thread_limit);
  37
  38 #pragma omp teams default(shared)
  39   {
  40     int priv_nteams;
  41     int team_num = omp_get_team_num();
  42     if (team_num == 0)
  43       nteams = omp_get_num_teams();
  44     priv_nteams = omp_get_num_teams();
  45 #pragma omp parallel
  46     {
  47       int priv_nthreads;
  48       int thread_num = omp_get_thread_num();
  49       int teams_ub, teams_lb, thr_limit;
  50       if (team_num == 0 && thread_num == 0)
  51         nthreads = omp_get_num_threads();
  52       priv_nthreads = omp_get_num_threads();
  53
  54       teams_ub = (num_teams_ub ? num_teams_ub : priv_nteams);
  55       teams_lb = (num_teams_lb ? num_teams_lb : teams_ub);
  56       thr_limit = (thread_limit ? thread_limit : priv_nthreads);
  57
  58       if (priv_nteams < teams_lb || priv_nteams > teams_ub) {
  59         fprintf(stderr, "error: invalid number of teams=%d\n", priv_nteams);
  60         exit(1);
  61       }
  62       if (priv_nthreads > thr_limit) {
  63         fprintf(stderr, "error: invalid number of threads=%d\n", priv_nthreads);
  64         exit(1);
  65       }
  66 #pragma omp atomic
  67       a++;
  68     }
  69   }
  70   if (a != nteams * nthreads) {
  71     fprintf(stderr, "error: a (%d) != nteams * nthreads (%d)\n", a,
  72             nteams * nthreads);
  73     exit(1);
  74   } else {
  75     printf("#teams %d, #threads %d: Hello!\n", nteams, nthreads);
  76   }
  77 }
  78
  79 int main(int argc, char *argv[]) {
  80   omp_set_num_threads(NT);
  81
  82   check_num_teams(1, 8, 2);
  83   check_num_teams(2, 2, 2);
  84   check_num_teams(2, 2, 0);
  85   check_num_teams(8, 16, 2);
  86   check_num_teams(9, 16, 0);
  87   check_num_teams(9, 16, 2);
  88   check_num_teams(2, 3, 0);
  89   check_num_teams(0, 0, 2);
  90   check_num_teams(0, 4, 0);
  91   check_num_teams(0, 2, 2);
  92
  93   printf("Test Passed\n");
  94   return 0;
  95 }