openmp/runtime/test/affinity/libomp_test_affinity.h

   1 #ifndef LIBOMP_TEST_AFFINITY_H
   2 #define LIBOMP_TEST_AFFINITY_H
   3
   4 #ifndef _GNU_SOURCE
   5 #define _GNU_SOURCE
   6 #endif
   7 #include <sched.h>
   8 #include <stdio.h>
   9 #include <stdlib.h>
  10 #include <string.h>
  11 typedef struct affinity_mask_t {
  12   size_t setsize;
  13   cpu_set_t *set;
  14 } affinity_mask_t;
  15
  16 #ifdef __ve__
  17 // VE's sched_getaffinity put garbage if the buffer is too big.  512 is
  18 // a good number to make all tests run correctly.
  19 #define AFFINITY_MAX_CPUS 512
  20 #else
  21 #define AFFINITY_MAX_CPUS (32 * 64)
  22 #endif
  23
  24 // Operating system affinity mask API
  25 static void affinity_mask_zero(affinity_mask_t *mask) {
  26   CPU_ZERO_S(mask->setsize, mask->set);
  27 }
  28
  29 static affinity_mask_t *affinity_mask_alloc() {
  30   size_t setsize = CPU_ALLOC_SIZE(AFFINITY_MAX_CPUS);
  31   cpu_set_t *set = CPU_ALLOC(AFFINITY_MAX_CPUS);
  32   affinity_mask_t *retval = (affinity_mask_t *)malloc(sizeof(affinity_mask_t));
  33   retval->setsize = setsize;
  34   retval->set = set;
  35   affinity_mask_zero(retval);
  36   return retval;
  37 }
  38
  39 static void affinity_mask_free(affinity_mask_t *mask) { CPU_FREE(mask->set); }
  40
  41 static void affinity_mask_copy(affinity_mask_t *dest,
  42                                const affinity_mask_t *src) {
  43   memcpy(dest->set, src->set, dest->setsize);
  44 }
  45
  46 static void affinity_mask_set(affinity_mask_t *mask, int cpu) {
  47   CPU_SET_S(cpu, mask->setsize, mask->set);
  48 }
  49
  50 static void affinity_mask_clr(affinity_mask_t *mask, int cpu) {
  51   CPU_CLR_S(cpu, mask->setsize, mask->set);
  52 }
  53
  54 static int affinity_mask_isset(const affinity_mask_t *mask, int cpu) {
  55   return CPU_ISSET_S(cpu, mask->setsize, mask->set);
  56 }
  57
  58 static int affinity_mask_count(const affinity_mask_t *mask) {
  59   return CPU_COUNT_S(mask->setsize, mask->set);
  60 }
  61
  62 static int affinity_mask_equal(const affinity_mask_t *mask1,
  63                                const affinity_mask_t *mask2) {
  64   return CPU_EQUAL_S(mask1->setsize, mask1->set, mask2->set);
  65 }
  66
  67 static void get_thread_affinity(affinity_mask_t *mask) {
  68   if (sched_getaffinity(0, mask->setsize, mask->set) != 0) {
  69     perror("sched_getaffinity()");
  70     exit(EXIT_FAILURE);
  71   }
  72 }
  73
  74 static void set_thread_affinity(const affinity_mask_t *mask) {
  75   if (sched_setaffinity(0, mask->setsize, mask->set) != 0) {
  76     perror("sched_setaffinity()");
  77     exit(EXIT_FAILURE);
  78   }
  79 }
  80
  81 static void affinity_update_snprintf_values(char **ptr, size_t *remaining,
  82                                             size_t n, size_t *retval) {
  83   if (n > *remaining && *remaining > 0) {
  84     *ptr += *remaining;
  85     *remaining = 0;
  86   } else {
  87     *ptr += n;
  88     *remaining -= n;
  89   }
  90   *retval += n;
  91 }
  92
  93 static size_t affinity_mask_snprintf(char *buf, size_t bufsize,
  94                                      const affinity_mask_t *mask) {
  95   int cpu, need_comma, begin, end;
  96   size_t n;
  97   char *ptr = buf;
  98   size_t remaining = bufsize;
  99   size_t retval = 0;
 100
 101   n = snprintf(ptr, remaining, "%c", '{');
 102   affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
 103
 104   need_comma = 0;
 105   for (cpu = 0; cpu < AFFINITY_MAX_CPUS; cpu++) {
 106     if (!affinity_mask_isset(mask, cpu))
 107       continue;
 108     if (need_comma) {
 109       n = snprintf(ptr, remaining, "%c", ',');
 110       affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
 111     }
 112     begin = cpu;
 113     // Find end of range (inclusive end)
 114     for (end = begin + 1; end < AFFINITY_MAX_CPUS; ++end) {
 115       if (!affinity_mask_isset(mask, end))
 116         break;
 117     }
 118     end--;
 119
 120     if (end - begin >= 2) {
 121       n = snprintf(ptr, remaining, "%d-%d", begin, end);
 122       affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
 123     } else if (end - begin == 1) {
 124       n = snprintf(ptr, remaining, "%d,%d", begin, end);
 125       affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
 126     } else if (end - begin == 0) {
 127       n = snprintf(ptr, remaining, "%d", begin);
 128       affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
 129     }
 130     need_comma = 1;
 131     cpu = end;
 132   }
 133   n = snprintf(ptr, remaining, "%c", '}');
 134   affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
 135   return retval;
 136 }
 137 #endif