xref: /llvm-project/openmp/runtime/test/affinity/libomp_test_affinity.h (revision e8679b93daa0db1e9b1838186f6309f59ad5cd0b)
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
affinity_mask_zero(affinity_mask_t * mask)25 static void affinity_mask_zero(affinity_mask_t *mask) {
26   CPU_ZERO_S(mask->setsize, mask->set);
27 }
28 
affinity_mask_alloc()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 
affinity_mask_free(affinity_mask_t * mask)39 static void affinity_mask_free(affinity_mask_t *mask) { CPU_FREE(mask->set); }
40 
affinity_mask_copy(affinity_mask_t * dest,const affinity_mask_t * src)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 
affinity_mask_set(affinity_mask_t * mask,int cpu)46 static void affinity_mask_set(affinity_mask_t *mask, int cpu) {
47   CPU_SET_S(cpu, mask->setsize, mask->set);
48 }
49 
affinity_mask_clr(affinity_mask_t * mask,int cpu)50 static void affinity_mask_clr(affinity_mask_t *mask, int cpu) {
51   CPU_CLR_S(cpu, mask->setsize, mask->set);
52 }
53 
affinity_mask_isset(const affinity_mask_t * mask,int cpu)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 
affinity_mask_count(const affinity_mask_t * mask)58 static int affinity_mask_count(const affinity_mask_t *mask) {
59   return CPU_COUNT_S(mask->setsize, mask->set);
60 }
61 
affinity_mask_equal(const affinity_mask_t * mask1,const affinity_mask_t * mask2)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 
get_thread_affinity(affinity_mask_t * mask)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 
set_thread_affinity(const affinity_mask_t * mask)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 
affinity_update_snprintf_values(char ** ptr,size_t * remaining,size_t n,size_t * retval)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 
affinity_mask_snprintf(char * buf,size_t bufsize,const affinity_mask_t * mask)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
138