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>
11typedef 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
25static void affinity_mask_zero(affinity_mask_t *mask) {
26 CPU_ZERO_S(mask->setsize, mask->set);
27}
28
29static 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(size: sizeof(affinity_mask_t));
33 retval->setsize = setsize;
34 retval->set = set;
35 affinity_mask_zero(mask: retval);
36 return retval;
37}
38
39static void affinity_mask_free(affinity_mask_t *mask) { CPU_FREE(mask->set); }
40
41static void affinity_mask_copy(affinity_mask_t *dest,
42 const affinity_mask_t *src) {
43 memcpy(dest: dest->set, src: src->set, n: dest->setsize);
44}
45
46static void affinity_mask_set(affinity_mask_t *mask, int cpu) {
47 CPU_SET_S(cpu, mask->setsize, mask->set);
48}
49
50static void affinity_mask_clr(affinity_mask_t *mask, int cpu) {
51 CPU_CLR_S(cpu, mask->setsize, mask->set);
52}
53
54static int affinity_mask_isset(const affinity_mask_t *mask, int cpu) {
55 return CPU_ISSET_S(cpu, mask->setsize, mask->set);
56}
57
58static int affinity_mask_count(const affinity_mask_t *mask) {
59 return CPU_COUNT_S(mask->setsize, mask->set);
60}
61
62static 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
67static void get_thread_affinity(affinity_mask_t *mask) {
68 if (sched_getaffinity(pid: 0, cpusetsize: mask->setsize, cpuset: mask->set) != 0) {
69 perror(s: "sched_getaffinity()");
70 exit(EXIT_FAILURE);
71 }
72}
73
74static void set_thread_affinity(const affinity_mask_t *mask) {
75 if (sched_setaffinity(pid: 0, cpusetsize: mask->setsize, cpuset: mask->set) != 0) {
76 perror(s: "sched_setaffinity()");
77 exit(EXIT_FAILURE);
78 }
79}
80
81static 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
93static 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(s: ptr, maxlen: remaining, format: "%c", '{');
102 affinity_update_snprintf_values(ptr: &ptr, remaining: &remaining, n, retval: &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(s: ptr, maxlen: remaining, format: "%c", ',');
110 affinity_update_snprintf_values(ptr: &ptr, remaining: &remaining, n, retval: &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, cpu: end))
116 break;
117 }
118 end--;
119
120 if (end - begin >= 2) {
121 n = snprintf(s: ptr, maxlen: remaining, format: "%d-%d", begin, end);
122 affinity_update_snprintf_values(ptr: &ptr, remaining: &remaining, n, retval: &retval);
123 } else if (end - begin == 1) {
124 n = snprintf(s: ptr, maxlen: remaining, format: "%d,%d", begin, end);
125 affinity_update_snprintf_values(ptr: &ptr, remaining: &remaining, n, retval: &retval);
126 } else if (end - begin == 0) {
127 n = snprintf(s: ptr, maxlen: remaining, format: "%d", begin);
128 affinity_update_snprintf_values(ptr: &ptr, remaining: &remaining, n, retval: &retval);
129 }
130 need_comma = 1;
131 cpu = end;
132 }
133 n = snprintf(s: ptr, maxlen: remaining, format: "%c", '}');
134 affinity_update_snprintf_values(ptr: &ptr, remaining: &remaining, n, retval: &retval);
135 return retval;
136}
137#endif
138

source code of openmp/runtime/test/affinity/libomp_test_affinity.h