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 #define AFFINITY_MAX_CPUS (32 * 64)
17
18 // Operating system affinity mask API
affinity_mask_zero(affinity_mask_t * mask)19 static void affinity_mask_zero(affinity_mask_t *mask) {
20 CPU_ZERO_S(mask->setsize, mask->set);
21 }
22
affinity_mask_alloc()23 static affinity_mask_t *affinity_mask_alloc() {
24 size_t setsize = CPU_ALLOC_SIZE(AFFINITY_MAX_CPUS);
25 cpu_set_t *set = CPU_ALLOC(AFFINITY_MAX_CPUS);
26 affinity_mask_t *retval = (affinity_mask_t *)malloc(sizeof(affinity_mask_t));
27 retval->setsize = setsize;
28 retval->set = set;
29 affinity_mask_zero(retval);
30 return retval;
31 }
32
affinity_mask_free(affinity_mask_t * mask)33 static void affinity_mask_free(affinity_mask_t *mask) { CPU_FREE(mask->set); }
34
affinity_mask_copy(affinity_mask_t * dest,const affinity_mask_t * src)35 static void affinity_mask_copy(affinity_mask_t *dest,
36 const affinity_mask_t *src) {
37 memcpy(dest->set, src->set, dest->setsize);
38 }
39
affinity_mask_set(affinity_mask_t * mask,int cpu)40 static void affinity_mask_set(affinity_mask_t *mask, int cpu) {
41 CPU_SET_S(cpu, mask->setsize, mask->set);
42 }
43
affinity_mask_clr(affinity_mask_t * mask,int cpu)44 static void affinity_mask_clr(affinity_mask_t *mask, int cpu) {
45 CPU_CLR_S(cpu, mask->setsize, mask->set);
46 }
47
affinity_mask_isset(const affinity_mask_t * mask,int cpu)48 static int affinity_mask_isset(const affinity_mask_t *mask, int cpu) {
49 return CPU_ISSET_S(cpu, mask->setsize, mask->set);
50 }
51
affinity_mask_count(const affinity_mask_t * mask)52 static int affinity_mask_count(const affinity_mask_t *mask) {
53 return CPU_COUNT_S(mask->setsize, mask->set);
54 }
55
affinity_mask_equal(const affinity_mask_t * mask1,const affinity_mask_t * mask2)56 static int affinity_mask_equal(const affinity_mask_t *mask1,
57 const affinity_mask_t *mask2) {
58 return CPU_EQUAL_S(mask1->setsize, mask1->set, mask2->set);
59 }
60
get_thread_affinity(affinity_mask_t * mask)61 static void get_thread_affinity(affinity_mask_t *mask) {
62 if (sched_getaffinity(0, mask->setsize, mask->set) != 0) {
63 perror("sched_getaffinity()");
64 exit(EXIT_FAILURE);
65 }
66 }
67
set_thread_affinity(const affinity_mask_t * mask)68 static void set_thread_affinity(const affinity_mask_t *mask) {
69 if (sched_setaffinity(0, mask->setsize, mask->set) != 0) {
70 perror("sched_setaffinity()");
71 exit(EXIT_FAILURE);
72 }
73 }
74
affinity_update_snprintf_values(char ** ptr,size_t * remaining,size_t n,size_t * retval)75 static void affinity_update_snprintf_values(char **ptr, size_t *remaining,
76 size_t n, size_t *retval) {
77 if (n > *remaining && *remaining > 0) {
78 *ptr += *remaining;
79 *remaining = 0;
80 } else {
81 *ptr += n;
82 *remaining -= n;
83 }
84 *retval += n;
85 }
86
affinity_mask_snprintf(char * buf,size_t bufsize,const affinity_mask_t * mask)87 static size_t affinity_mask_snprintf(char *buf, size_t bufsize,
88 const affinity_mask_t *mask) {
89 int cpu, need_comma, begin, end;
90 size_t n;
91 char *ptr = buf;
92 size_t remaining = bufsize;
93 size_t retval = 0;
94
95 n = snprintf(ptr, remaining, "%c", '{');
96 affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
97
98 need_comma = 0;
99 for (cpu = 0; cpu < AFFINITY_MAX_CPUS; cpu++) {
100 if (!affinity_mask_isset(mask, cpu))
101 continue;
102 if (need_comma) {
103 n = snprintf(ptr, remaining, "%c", ',');
104 affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
105 }
106 begin = cpu;
107 // Find end of range (inclusive end)
108 for (end = begin + 1; end < AFFINITY_MAX_CPUS; ++end) {
109 if (!affinity_mask_isset(mask, end))
110 break;
111 }
112 end--;
113
114 if (end - begin >= 2) {
115 n = snprintf(ptr, remaining, "%d-%d", begin, end);
116 affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
117 } else if (end - begin == 1) {
118 n = snprintf(ptr, remaining, "%d,%d", begin, end);
119 affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
120 } else if (end - begin == 0) {
121 n = snprintf(ptr, remaining, "%d", begin);
122 affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
123 }
124 need_comma = 1;
125 cpu = end;
126 }
127 n = snprintf(ptr, remaining, "%c", '}');
128 affinity_update_snprintf_values(&ptr, &remaining, n, &retval);
129 return retval;
130 }
131 #endif
132