138 lines · c
1#ifndef LIBOMP_TEST_AFFINITY_H2#define LIBOMP_TEST_AFFINITY_H3 4#ifndef _GNU_SOURCE5#define _GNU_SOURCE6#endif7#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 is18// a good number to make all tests run correctly.19#define AFFINITY_MAX_CPUS 51220#else21#define AFFINITY_MAX_CPUS (32 * 64)22#endif23 24// Operating system affinity mask API25static 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(sizeof(affinity_mask_t));33 retval->setsize = setsize;34 retval->set = set;35 affinity_mask_zero(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->set, src->set, 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(0, mask->setsize, mask->set) != 0) {69 perror("sched_getaffinity()");70 exit(EXIT_FAILURE);71 }72}73 74static 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 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(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#endif138