| 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(size: sizeof(affinity_mask_t)); |
| 33 | retval->setsize = setsize; |
| 34 | retval->set = set; |
| 35 | affinity_mask_zero(mask: 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: dest->set, src: src->set, n: 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(pid: 0, cpusetsize: mask->setsize, cpuset: mask->set) != 0) { |
| 69 | perror(s: "sched_getaffinity()" ); |
| 70 | exit(EXIT_FAILURE); |
| 71 | } |
| 72 | } |
| 73 | |
| 74 | static 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 | |
| 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(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 | |