4 files changed, 90 insertions, 0 deletions

diff --git a/include/sched.h b/include/sched.h
index 994260d0..6a6b2fcb 100644
--- a/include/sched.h
+++ b/include/sched.h
@@ -10,6 +10,10 @@ extern "C" {
 #define __NEED_pid_t
 #define __NEED_time_t
 
+#ifdef _GNU_SOURCE
+#define __NEED_size_t
+#endif
+
 #include <bits/alltypes.h>
 
 struct sched_param {
@@ -63,6 +67,63 @@ int     sched_yield(void);
 int clone (int (*)(void *), void *, int, void *, ...);
 int unshare(int);
 int setns(int, int);
+
+void *memcpy(void *__restrict, const void *__restrict, size_t);
+int memcmp(const void *, const void *, size_t);
+void *calloc(size_t, size_t);
+void free(void *);
+
+typedef struct cpu_set_t { unsigned long __bits[128/sizeof(long)]; } cpu_set_t;
+int __sched_cpucount(size_t, const cpu_set_t *);
+int sched_getaffinity(pid_t, size_t, cpu_set_t *);
+int sched_setaffinity(pid_t, size_t, const cpu_set_t *);
+
+#define __CPU_op_S(i, size, set, op) ( (i)/8 >= (size) ? 0 : \
+	((set)->__bits[(i)/8/sizeof(long)] op (1UL<<((i)%(8*sizeof(long))))) )
+
+#define CPU_SET_S(i, size, set) __CPU_op_S(i, size, set, |=)
+#define CPU_CLR_S(i, size, set) __CPU_op_S(i, size, set, &=~)
+#define CPU_ISSET_S(i, size, set) __CPU_op_S(i, size, set, &)
+
+#define __CPU_op_func_S(func, op) \
+static __inline void __CPU_##func##_S(size_t __size, cpu_set_t *__dest, \
+	const cpu_set_t *__src1, const cpu_set_t *__src2) \
+{ \
+	size_t __i; \
+	for (__i=0; __i<__size/sizeof(long); __i++) \
+		__dest->__bits[__i] = __src1->__bits[__i] \
+			op __src2->__bits[__i] ; \
+}
+
+__CPU_op_func_S(AND, &)
+__CPU_op_func_S(OR, |)
+__CPU_op_func_S(XOR, ^)
+
+#define CPU_AND_S(a,b,c,d) __CPU_AND_S(a,b,c,d)
+#define CPU_OR_S(a,b,c,d) __CPU_OR_S(a,b,c,d)
+#define CPU_XOR_S(a,b,c,d) __CPU_XOR_S(a,b,c,d)
+
+#define CPU_COUNT_S(size,set) __sched_cpucount(size,set)
+#define CPU_ZERO_S(size,set) memset(set,0,size)
+#define CPU_EQUAL_S(size,set1,set2) (!memcmp(set1,set2,size))
+
+#define CPU_ALLOC_SIZE(n) (sizeof(long) * ( (n)/(8*sizeof(long)) \
+	+ ((n)%(8*sizeof(long)) + 8*sizeof(long)-1)/(8*sizeof(long)) ) )
+#define CPU_ALLOC(n) ((cpu_set_t *)calloc(1,CPU_ALLOC_SIZE(n)))
+#define CPU_FREE(set) free(set)
+
+#define CPU_SETSIZE 128
+
+#define CPU_SET(i, set) CPU_SET_S(i,sizeof(cpu_set_t),set)
+#define CPU_CLR(i, set) CPU_CLR_S(i,sizeof(cpu_set_t),set)
+#define CPU_ISSET(i, set) CPU_ISSET_S(i,sizeof(cpu_set_t),set)
+#define CPU_AND(d,s1,s2) CPU_AND_S(sizeof(cpu_set_t),d,s1,s2)
+#define CPU_OR(d,s1,s2) CPU_OR_S(sizeof(cpu_set_t),d,s1,s2)
+#define CPU_XOR(d,s1,s2) CPU_XOR_S(sizeof(cpu_set_t),d,s1,s2)
+#define CPU_COUNT(set) CPU_COUNT_S(sizeof(cpu_set_t),set)
+#define CPU_ZERO(set) CPU_ZERO_S(sizeof(cpu_set_t),set)
+#define CPU_EQUAL(set) CPU_EQUAL_S(sizeof(cpu_set_t),set)
+
 #endif
 
 #ifdef __cplusplus
diff --git a/src/sched/sched_cpucount.c b/src/sched/sched_cpucount.c
new file mode 100644
index 00000000..94aa259e
--- /dev/null
+++ b/src/sched/sched_cpucount.c
@@ -0,0 +1,11 @@
+#define _GNU_SOURCE
+#include <sched.h>
+
+int __sched_cpucount(size_t size, const cpu_set_t *set)
+{
+	size_t i, j, cnt=0;
+	const unsigned char *p = (const void *)set;
+	for (i=0; i<size; i++) for (j=0; j<8; j++)
+		if (p[i] & (1<<j)) cnt++;
+	return cnt;
+}
diff --git a/src/sched/sched_getaffinity.c b/src/sched/sched_getaffinity.c
new file mode 100644
index 00000000..0aa4c65a
--- /dev/null
+++ b/src/sched/sched_getaffinity.c
@@ -0,0 +1,10 @@
+#define _GNU_SOURCE
+#include <sched.h>
+#include "syscall.h"
+
+int sched_getaffinity(pid_t tid, size_t size, cpu_set_t *set)
+{
+	long ret = __syscall(SYS_sched_getaffinity, tid, size, set);
+	if (ret > 0) ret = 0;
+	return __syscall_ret(ret);
+}
diff --git a/src/sched/sched_setaffinity.c b/src/sched/sched_setaffinity.c
new file mode 100644
index 00000000..4344df17
--- /dev/null
+++ b/src/sched/sched_setaffinity.c
@@ -0,0 +1,8 @@
+#define _GNU_SOURCE
+#include <sched.h>
+#include "syscall.h"
+
+int sched_setaffinity(pid_t tid, size_t size, const cpu_set_t *set)
+{
+	return syscall(SYS_sched_setaffinity, tid, size, set);
+}