[RFC PATCH v0.1 0/9] UMCG early preview/RFC patchset

@@ -368,6 +368,17 @@
 444	common	landlock_create_ruleset	sys_landlock_create_ruleset
 445	common	landlock_add_rule	sys_landlock_add_rule
 446	common	landlock_restrict_self	sys_landlock_restrict_self
+447	common	umcg_api_version	sys_umcg_api_version
+448	common	umcg_register_task	sys_umcg_register_task
+449	common	umcg_unregister_task	sys_umcg_unregister_task
+450	common	umcg_wait		sys_umcg_wait
+451	common	umcg_wake		sys_umcg_wake
+452	common	umcg_swap		sys_umcg_swap
+453	common	umcg_create_group	sys_umcg_create_group
+454	common	umcg_destroy_group	sys_umcg_destroy_group
+455	common	umcg_poll_worker	sys_umcg_poll_worker
+456	common	umcg_run_worker		sys_umcg_run_worker
+457	common	umcg_preempt_worker	sys_umcg_preempt_worker
 
 #
 # Due to a historical design error, certain syscalls are numbered differently

@@ -873,8 +873,31 @@ __SYSCALL(__NR_landlock_add_rule, sys_landlock_add_rule)
 #define __NR_landlock_restrict_self 446
 __SYSCALL(__NR_landlock_restrict_self, sys_landlock_restrict_self)
 
+#define __NR_umcg_api_version 447
+__SYSCALL(__NR_umcg_api_version, sys_umcg_api_version)
+#define __NR_umcg_register_task 448
+__SYSCALL(__NR_umcg_register_task, sys_umcg_register_task)
+#define __NR_umcg_unregister_task 449
+__SYSCALL(__NR_umcg_unregister_task, sys_umcg_unregister_task)
+#define __NR_umcg_wait 450
+__SYSCALL(__NR_umcg_wait, sys_umcg_wait)
+#define __NR_umcg_wake 451
+__SYSCALL(__NR_umcg_wake, sys_umcg_wake)
+#define __NR_umcg_swap 452
+__SYSCALL(__NR_umcg_swap, sys_umcg_swap)
+#define __NR_umcg_create_group 453
+__SYSCALL(__NR_umcg_create_group, sys_umcg_create_group)
+#define __NR_umcg_destroy_group 454
+__SYSCALL(__NR_umcg_destroy_group, sys_umcg_destroy_group)
+#define __NR_umcg_poll_worker 455
+__SYSCALL(__NR_umcg_poll_worker, sys_umcg_poll_worker)
+#define __NR_umcg_run_worker 456
+__SYSCALL(__NR_umcg_run_worker, sys_umcg_run_worker)
+#define __NR_umcg_preempt_worker 457
+__SYSCALL(__NR_umcg_preempt_worker, sys_umcg_preempt_worker)
+
 #undef __NR_syscalls
-#define __NR_syscalls 447
+#define __NR_syscalls 458
 
 /*
  * 32 bit systems traditionally used different

@@ -1661,6 +1661,16 @@ config MEMBARRIER
 
 	  If unsure, say Y.
 
+config UMCG
+	bool "Enable User Managed Concurrency Groups API"
+	default n
+	help
+	  Enable UMCG core wait/wake/swap operations as well as UMCG
+	  group/server/worker API. The core API is useful for fast IPC
+	  and context switching, while the group/server/worker API, together
+	  with the core API, form the basis for an in-process M:N userspace
+	  scheduling framework implemented in lib/umcg.
+
 config KALLSYMS
 	bool "Load all symbols for debugging/ksymoops" if EXPERT
 	default y

@@ -272,6 +272,19 @@ COND_SYSCALL(landlock_create_ruleset);
 COND_SYSCALL(landlock_add_rule);
 COND_SYSCALL(landlock_restrict_self);
 
+/* kernel/sched/umcg.c */
+COND_SYSCALL(umcg_api_version);
+COND_SYSCALL(umcg_register_task);
+COND_SYSCALL(umcg_unregister_task);
+COND_SYSCALL(umcg_wait);
+COND_SYSCALL(umcg_wake);
+COND_SYSCALL(umcg_swap);
+COND_SYSCALL(umcg_create_group);
+COND_SYSCALL(umcg_destroy_group);
+COND_SYSCALL(umcg_poll_worker);
+COND_SYSCALL(umcg_run_worker);
+COND_SYSCALL(umcg_preempt_worker);
+
 /* arch/example/kernel/sys_example.c */
 
 /* mm/fadvise.c */

@@ -71,6 +71,7 @@ struct open_how;
 struct mount_attr;
 struct landlock_ruleset_attr;
 enum landlock_rule_type;
+struct umcg_task;
 
 #include <linux/types.h>
 #include <linux/aio_abi.h>

@@ -1050,6 +1051,14 @@ asmlinkage long sys_landlock_create_ruleset(const struct landlock_ruleset_attr _
 asmlinkage long sys_landlock_add_rule(int ruleset_fd, enum landlock_rule_type rule_type,
 		const void __user *rule_attr, __u32 flags);
 asmlinkage long sys_landlock_restrict_self(int ruleset_fd, __u32 flags);
+asmlinkage long umcg_api_version(u32 api_version, u32 flags);
+asmlinkage long umcg_register_task(u32 api_version, u32 flags, u32 group_id,
+					struct umcg_task __user *umcg_task);
+asmlinkage long umcg_unregister_task(u32 flags);
+asmlinkage long umcg_wait(u32 flags, const struct __kernel_timespec __user *timeout);
+asmlinkage long umcg_wake(u32 flags, u32 next_tid);
+asmlinkage long umcg_swap(u32 wake_flags, u32 next_tid, u32 wait_flags,
+				const struct __kernel_timespec __user *timeout);
 
 /*
  * Architecture-specific system calls

@@ -0,0 +1,70 @@
+/* SPDX-License-Identifier: GPL-2.0+ WITH Linux-syscall-note */
+#ifndef _UAPI_LINUX_UMCG_H
+#define _UAPI_LINUX_UMCG_H
+
+#include <linux/limits.h>
+#include <linux/types.h>
+
+/*
+ * UMCG task states, the first 8 bits.
+ */
+#define UMCG_TASK_NONE			0
+/* UMCG server states. */
+#define UMCG_TASK_POLLING		1
+#define UMCG_TASK_SERVING		2
+#define UMCG_TASK_PROCESSING		3
+/* UMCG worker states. */
+#define UMCG_TASK_RUNNABLE		4
+#define UMCG_TASK_RUNNING		5
+#define UMCG_TASK_BLOCKED		6
+#define UMCG_TASK_UNBLOCKED		7
+
+/* UMCG task state flags, bits 8-15 */
+#define UMCG_TF_WAKEUP_QUEUED		(1 << 8)
+
+/*
+ * Unused at the moment flags reserved for features to be introduced
+ * in the near future.
+ */
+#define UMCG_TF_PREEMPT_DISABLED	(1 << 9)
+#define UMCG_TF_PREEMPTED		(1 << 10)
+
+#define UMCG_NOID	UINT_MAX
+
+/**
+ * struct umcg_task - controls the state of UMCG-enabled tasks.
+ *
+ * While at the moment only one field is present (@state), in future
+ * versions additional fields will be added, e.g. for the userspace to
+ * provide performance-improving hints and for the kernel to export sched
+ * stats.
+ *
+ * The struct is aligned at 32 bytes to ensure that even with future additions
+ * it fits into a single cache line.
+ */
+struct umcg_task {
+	/**
+	 * @state: the current state of the UMCG task described by this struct.
+	 *
+	 * UMCG task state:
+	 *   bits  0 -  7: task state;
+	 *   bits  8 - 15: state flags;
+	 *   bits 16 - 23: reserved; must be zeroes;
+	 *   bits 24 - 31: for userspace use.
+	 */
+	uint32_t	state;
+} __attribute((packed, aligned(4 * sizeof(uint64_t))));
+
+/**
+ * enum umcg_register_flag - flags for sys_umcg_register
+ * @UMCG_REGISTER_CORE_TASK:  Register a UMCG core task (not part of a group);
+ * @UMCG_REGISTER_WORKER:     Register a UMCG worker task;
+ * @UMCG_REGISTER_SERVER:     Register a UMCG server task.
+ */
+enum umcg_register_flag {
+	UMCG_REGISTER_CORE_TASK		= 0,
+	UMCG_REGISTER_WORKER		= 1,
+	UMCG_REGISTER_SERVER		= 2
+};
+
+#endif /* _UAPI_LINUX_UMCG_H */

@@ -37,3 +37,4 @@ obj-$(CONFIG_MEMBARRIER) += membarrier.o
 obj-$(CONFIG_CPU_ISOLATION) += isolation.o
 obj-$(CONFIG_PSI) += psi.o
 obj-$(CONFIG_SCHED_CORE) += core_sched.o
+obj-$(CONFIG_UMCG) += umcg.o

@@ -0,0 +1,143 @@
+// SPDX-License-Identifier: GPL-2.0-only
+
+/*
+ * User Managed Concurrency Groups (UMCG).
+ */
+
+#include <linux/syscalls.h>
+#include <linux/types.h>
+#include <linux/uaccess.h>
+#include <linux/umcg.h>
+
+/**
+ * sys_umcg_api_version - query UMCG API versions that are supported.
+ * @api_version:          Requested API version.
+ * @flags:                Reserved.
+ *
+ * Return:
+ * 0                    - the @api_version is supported;
+ * > 0                  - the maximum supported version of UMCG API if
+ *                        the requested version is not supported.
+ * -EINVAL              - @flags is not zero.
+ *
+ * NOTE: the kernel may drop support for older/deprecated API versions,
+ * so a return of X does not indicate that any version less than X is
+ * supported.
+ */
+SYSCALL_DEFINE2(umcg_api_version, u32, api_version, u32, flags)
+{
+	return -ENOSYS;
+}
+
+/**
+ * sys_umcg_register_task - register the current task as a UMCG task.
+ * @api_version:       The expected/desired API version of the syscall.
+ * @flags:             One of enum umcg_register_flag.
+ * @group_id:          UMCG Group ID. UMCG_NOID for Core tasks.
+ * @umcg_task:         The control struct for the current task.
+ *                     umcg_task->state must be UMCG_TASK_NONE.
+ *
+ * Register the current task as a UMCG task. If this is a core UMCG task,
+ * the syscall marks it as RUNNING and returns immediately.
+ *
+ * If this is a UMCG worker, the syscall marks it UNBLOCKED, and proceeds
+ * with the normal UNBLOCKED worker logic.
+ *
+ * If this is a UMCG server, the syscall immediately returns.
+ *
+ * Return:
+ * 0            - Ok;
+ * -EOPNOTSUPP  - the API version is not supported;
+ * -EINVAL      - one of the parameters is wrong;
+ * -EFAULT      - failed to access @umcg_task.
+ */
+SYSCALL_DEFINE4(umcg_register_task, u32, api_version, u32, flags, u32, group_id,
+		struct umcg_task __user *, umcg_task)
+{
+	return -ENOSYS;
+}
+
+/**
+ * sys_umcg_unregister_task - unregister the current task as a UMCG task.
+ * @flags: reserved.
+ *
+ * Return:
+ * 0       - Ok;
+ * -EINVAL - the current task is not a UMCG task.
+ */
+SYSCALL_DEFINE1(umcg_unregister_task, u32, flags)
+{
+	return -ENOSYS;
+}
+
+/**
+ * sys_umcg_wait - block the current task (if all condtions are met).
+ * @flags:         Reserved.
+ * @timeout:       The absolute timeout of the wait. Not supported yet.
+ *                 Must be NULL.
+ *
+ * Sleep until woken, interrupted, or @timeout expires.
+ *
+ * Return:
+ * 0           - Ok;
+ * -EFAULT     - failed to read struct umcg_task assigned to this task
+ *               via sys_umcg_register();
+ * -EAGAIN     - try again;
+ * -EINTR      - signal pending;
+ * -EOPNOTSUPP - @timeout != NULL (not supported yet).
+ * -EINVAL     - a parameter or a member of struct umcg_task has a wrong value.
+ */
+SYSCALL_DEFINE2(umcg_wait, u32, flags,
+		const struct __kernel_timespec __user *, timeout)
+{
+	return -ENOSYS;
+}
+
+/**
+ * sys_umcg_wake - wake @next_tid task blocked in sys_umcg_wait.
+ * @flags:         Reserved.
+ * @next_tid:      The ID of the task to wake.
+ *
+ * Wake @next identified by @next_tid. @next must be either a UMCG core
+ * task or a UMCG worker task.
+ *
+ * Return:
+ * 0           - Ok;
+ * -EFAULT     - failed to read struct umcg_task assigned to next;
+ * -ESRCH      - @next_tid did not identify a task;
+ * -EAGAIN     - try again;
+ * -EINVAL     - a parameter or a member of next->umcg_task has a wrong value.
+ */
+SYSCALL_DEFINE2(umcg_wake, u32, flags, u32, next_tid)
+{
+	return -ENOSYS;
+}
+
+/**
+ * sys_umcg_swap - wake next, put current to sleep.
+ * @wake_flags:    Reserved.
+ * @next_tid:      The ID of the task to wake.
+ * @wait_flags:    Reserved.
+ * @timeout:       The absolute timeout of the wait. Not supported yet.
+ *
+ * sys_umcg_swap() is semantically equivalent to this code fragment:
+ *
+ *     ret = sys_umcg_wake(wake_flags, next_tid);
+ *     if (ret)
+ *             return ret;
+ *     return sys_umcg_wait(wait_flags, timeout);
+ *
+ * The function attempts to wake @next on the current CPU.
+ *
+ * The current and the next tasks must both be either UMCG core tasks,
+ * or two UMCG workers belonging to the same UMCG group. In the latter
+ * case the UMCG server task that is "running" the current task will
+ * be transferred to the next task.
+ *
+ * Return: see the code snippet above.
+ */
+SYSCALL_DEFINE4(umcg_swap, u32, wake_flags, u32, next_tid, u32, wait_flags,
+		const struct __kernel_timespec __user *, timeout)
+{
+	return -ENOSYS;
+}

@@ -3683,8 +3683,7 @@ static void ttwu_queue(struct task_struct *p, int cpu, int wake_flags)
  * Return: %true if @p->state changes (an actual wakeup was done),
  *	   %false otherwise.
  */
-static int
-try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
+int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
 {
 	unsigned long flags;
 	int cpu, success = 0;

@@ -6764,6 +6764,10 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int wake_flags)
 	if (wake_flags & WF_TTWU) {
 		record_wakee(p);
 
+		if ((wake_flags & WF_CURRENT_CPU) &&
+		    cpumask_test_cpu(cpu, p->cpus_ptr))
+			return cpu;
+
 		if (sched_energy_enabled()) {
 			new_cpu = find_energy_efficient_cpu(p, prev_cpu);
 			if (new_cpu >= 0)

@@ -2027,13 +2027,14 @@ static inline int task_on_rq_migrating(struct task_struct *p)
 }
 
 /* Wake flags. The first three directly map to some SD flag value */
-#define WF_EXEC     0x02 /* Wakeup after exec; maps to SD_BALANCE_EXEC */
-#define WF_FORK     0x04 /* Wakeup after fork; maps to SD_BALANCE_FORK */
-#define WF_TTWU     0x08 /* Wakeup;            maps to SD_BALANCE_WAKE */
+#define WF_EXEC         0x02 /* Wakeup after exec; maps to SD_BALANCE_EXEC */
+#define WF_FORK         0x04 /* Wakeup after fork; maps to SD_BALANCE_FORK */
+#define WF_TTWU         0x08 /* Wakeup;            maps to SD_BALANCE_WAKE */
 
-#define WF_SYNC     0x10 /* Waker goes to sleep after wakeup */
-#define WF_MIGRATED 0x20 /* Internal use, task got migrated */
-#define WF_ON_CPU   0x40 /* Wakee is on_cpu */
+#define WF_SYNC         0x10 /* Waker goes to sleep after wakeup */
+#define WF_MIGRATED     0x20 /* Internal use, task got migrated */
+#define WF_ON_CPU       0x40 /* Wakee is on_cpu */
+#define WF_CURRENT_CPU  0x80 /* Prefer to move the wakee to the current CPU. */
 
 #ifdef CONFIG_SMP
 static_assert(WF_EXEC == SD_BALANCE_EXEC);

@@ -3018,6 +3019,8 @@ static inline bool is_per_cpu_kthread(struct task_struct *p)
 extern void swake_up_all_locked(struct swait_queue_head *q);
 extern void __prepare_to_swait(struct swait_queue_head *q, struct swait_queue *wait);
 
+extern int try_to_wake_up(struct task_struct *tsk, unsigned int state, int wake_flags);
+
 #ifdef CONFIG_PREEMPT_DYNAMIC
 extern int preempt_dynamic_mode;
 extern int sched_dynamic_mode(const char *str);

@@ -0,0 +1,4 @@
+PDX-License-Identifier: GPL-2.0-only
+libumcg.a
+libumcg.o
+

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: GPL-2.0
+
+CFLAGS += -g -I../../../usr/include/ -I../../include/
+
+libumcg.a: libumcg.o
+	ar rc libumcg.a libumcg.o
+
+libumcg.o: libumcg.c
+
+clean :
+	rm libumcg.a libumcg.o

@@ -0,0 +1,350 @@
+// SPDX-License-Identifier: GPL-2.0
+#define _GNU_SOURCE
+#include "libumcg.h"
+
+#include <errno.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdatomic.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <threads.h>
+
+/* UMCG API version supported by this library. */
+static const uint32_t umcg_api_version = 1;
+
+struct umcg_group {
+	uint32_t group_id;
+};
+
+/**
+ * struct umcg_task_tls - per thread struct used to identify/manage UMCG tasks
+ *
+ * Each UMCG task requires an instance of struct umcg_task passed to
+ * sys_umcg_register. This struct contains it, as well as several additional
+ * fields.
+ */
+struct umcg_task_tls {
+	struct umcg_task	umcg_task;
+	umcg_tid		self;
+	intptr_t		tag;
+	pid_t			tid;
+
+} __attribute((aligned(4 * sizeof(uint64_t))));
+
+static thread_local struct umcg_task_tls *umcg_task_tls;
+
+umcg_tid umcg_get_utid(void)
+{
+	return (umcg_tid)&umcg_task_tls;
+}
+
+static umcg_tid umcg_task_to_utid(struct umcg_task *ut)
+{
+	if (!ut)
+		return UMCG_NONE;
+
+	return ((struct umcg_task_tls *)ut)->self;
+}
+
+static struct umcg_task_tls *utid_to_utls(umcg_tid utid)
+{
+	if (!utid || !*(struct umcg_task_tls **)utid) {
+		fprintf(stderr, "utid_to_utls: NULL\n");
+		/* Kill the process rather than corrupt memory. */
+		raise(SIGKILL);
+		return NULL;
+	}
+	return *(struct umcg_task_tls **)utid;
+}
+
+void umcg_set_task_tag(umcg_tid utid, intptr_t tag)
+{
+	utid_to_utls(utid)->tag = tag;
+}
+
+intptr_t umcg_get_task_tag(umcg_tid utid)
+{
+	return utid_to_utls(utid)->tag;
+}
+
+umcg_tid umcg_register_core_task(intptr_t tag)
+{
+	int ret;
+
+	if (umcg_task_tls != NULL) {
+		errno = EINVAL;
+		return UMCG_NONE;
+	}
+
+	umcg_task_tls = malloc(sizeof(struct umcg_task_tls));
+	if (!umcg_task_tls) {
+		errno = ENOMEM;
+		return UMCG_NONE;
+	}
+
+	umcg_task_tls->umcg_task.state = UMCG_TASK_NONE;
+	umcg_task_tls->self = (umcg_tid)&umcg_task_tls;
+	umcg_task_tls->tag = tag;
+	umcg_task_tls->tid = gettid();
+
+	ret = sys_umcg_register_task(umcg_api_version, UMCG_REGISTER_CORE_TASK,
+			UMCG_NOID, &umcg_task_tls->umcg_task);
+	if (ret) {
+		free(umcg_task_tls);
+		umcg_task_tls = NULL;
+		errno = ret;
+		return UMCG_NONE;
+	}
+
+	return umcg_task_tls->self;
+}
+
+int umcg_unregister_task(void)
+{
+	int ret;
+
+	if (!umcg_task_tls) {
+		errno = EINVAL;
+		return -1;
+	}
+
+	ret = sys_umcg_unregister_task(0);
+	if (ret) {
+		errno = ret;
+		return -1;
+	}
+
+	free(umcg_task_tls);
+	atomic_store_explicit(&umcg_task_tls, NULL, memory_order_seq_cst);
+	return 0;
+}
+
+/* Helper return codes. */
+enum umcg_prepare_op_result {
+	UMCG_OP_DONE,
+	UMCG_OP_SYS,
+	UMCG_OP_AGAIN,
+	UMCG_OP_ERROR
+};
+
+static enum umcg_prepare_op_result umcg_prepare_wait(void)
+{
+	struct umcg_task *ut;
+	uint32_t umcg_state;
+	int ret;
+
+	if (!umcg_task_tls) {
+		errno = EINVAL;
+		return UMCG_OP_ERROR;
+	}
+
+	ut = &umcg_task_tls->umcg_task;
+
+	umcg_state = UMCG_TASK_RUNNING;
+	if (atomic_compare_exchange_strong_explicit(&ut->state,
+			&umcg_state, UMCG_TASK_RUNNABLE,
+			memory_order_seq_cst, memory_order_seq_cst))
+		return UMCG_OP_SYS;
+
+	if (umcg_state != (UMCG_TASK_RUNNING | UMCG_TF_WAKEUP_QUEUED)) {
+		fprintf(stderr, "libumcg: unexpected state before wait: %u\n",
+				umcg_state);
+		errno = EINVAL;
+		return UMCG_OP_ERROR;
+	}
+
+	if (atomic_compare_exchange_strong_explicit(&ut->state,
+			&umcg_state, UMCG_TASK_RUNNING,
+			memory_order_seq_cst, memory_order_seq_cst)) {
+		return UMCG_OP_DONE;
+	}
+
+	/* Raced with another wait/wake? This is not supported. */
+	fprintf(stderr, "libumcg: failed to remove the wakeup flag: %u\n",
+			umcg_state);
+	errno = EINVAL;
+	return UMCG_OP_ERROR;
+}
+
+static int umcg_do_wait(const struct timespec *timeout)
+{
+	uint32_t umcg_state;
+	int ret;
+
+	do {
+		ret = sys_umcg_wait(0, timeout);
+		if (ret != 0 && errno != EAGAIN)
+			return ret;
+
+		umcg_state = atomic_load_explicit(
+				&umcg_task_tls->umcg_task.state,
+				memory_order_acquire);
+	} while (umcg_state == UMCG_TASK_RUNNABLE);
+
+	return 0;
+}
+
+int umcg_wait(const struct timespec *timeout)
+{
+	switch (umcg_prepare_wait()) {
+	case UMCG_OP_DONE:
+		return 0;
+	case UMCG_OP_SYS:
+		break;
+	case UMCG_OP_ERROR:
+		return -1;
+	default:
+		fprintf(stderr, "Unknown pre_op result.\n");
+		exit(1);
+		return -1;
+	}
+
+	return umcg_do_wait(timeout);
+}
+
+static enum umcg_prepare_op_result umcg_prepare_wake(struct umcg_task_tls *utls)
+{
+	struct umcg_task *ut = &utls->umcg_task;
+	uint32_t umcg_state, next_state;
+
+	next_state = UMCG_TASK_RUNNING;
+	umcg_state = UMCG_TASK_RUNNABLE;
+	if (atomic_compare_exchange_strong_explicit(&ut->state,
+			&umcg_state, next_state,
+			memory_order_seq_cst, memory_order_seq_cst))
+		return UMCG_OP_SYS;
+
+	if (umcg_state != UMCG_TASK_RUNNING) {
+		if (umcg_state == (UMCG_TASK_RUNNING | UMCG_TF_WAKEUP_QUEUED)) {
+			/*
+			 * With ping-pong mutual swapping using wake/wait
+			 * without synchronization this can happen.
+			 */
+			return UMCG_OP_AGAIN;
+		}
+		fprintf(stderr, "libumcg: unexpected state in umcg_wake(): %u\n",
+				umcg_state);
+		errno = EINVAL;
+		return UMCG_OP_ERROR;
+	}
+
+	if (atomic_compare_exchange_strong_explicit(&ut->state,
+			&umcg_state, UMCG_TASK_RUNNING | UMCG_TF_WAKEUP_QUEUED,
+			memory_order_seq_cst, memory_order_seq_cst)) {
+		return UMCG_OP_DONE;
+	}
+
+	if (umcg_state != UMCG_TASK_RUNNABLE) {
+		fprintf(stderr, "libumcg: unexpected state in umcg_wake (1): %u\n",
+				umcg_state);
+		errno = EINVAL;
+		return UMCG_OP_ERROR;
+	}
+
+	return UMCG_OP_AGAIN;
+}
+
+static int umcg_do_wake_or_swap(struct umcg_task_tls *next_utls,
+				uint64_t prev_wait_counter, bool should_wait,
+				const struct timespec *timeout)
+{
+	int ret;
+
+again:
+
+	if (should_wait)
+		ret = sys_umcg_swap(0, next_utls->tid, 0, timeout);
+	else
+		ret = sys_umcg_wake(0, next_utls->tid);
+
+	if (ret && errno == EAGAIN)
+		goto again;
+
+	return ret;
+}
+
+int umcg_wake(umcg_tid next)
+{
+	struct umcg_task_tls *utls = *(struct umcg_task_tls **)next;
+	uint64_t prev_wait_counter;
+
+	if (!utls) {
+		errno = EINVAL;
+		return -1;
+	}
+
+again:
+	switch (umcg_prepare_wake(utls)) {
+	case UMCG_OP_DONE:
+		return 0;
+	case UMCG_OP_SYS:
+		break;
+	case UMCG_OP_ERROR:
+		return -1;
+	case UMCG_OP_AGAIN:
+		goto again;
+	default:
+		fprintf(stderr, "libumcg: unknown pre_op result.\n");
+		exit(1);
+		return -1;
+	}
+
+	return umcg_do_wake_or_swap(utls, prev_wait_counter, false, NULL);
+}
+
+int umcg_swap(umcg_tid next, const struct timespec *timeout)
+{
+	struct umcg_task_tls *utls = *(struct umcg_task_tls **)next;
+	bool should_wake, should_wait;
+	uint64_t prev_wait_counter;
+	int ret;
+
+	if (!utls) {
+		errno = EINVAL;
+		return -1;
+	}
+
+again:
+	switch (umcg_prepare_wake(utls)) {
+	case UMCG_OP_DONE:
+		should_wake = false;
+		break;
+	case UMCG_OP_SYS:
+		should_wake = true;
+		break;
+	case UMCG_OP_ERROR:
+		return -1;
+	case UMCG_OP_AGAIN:
+		goto again;
+	default:
+		fprintf(stderr, "lubumcg: unknown pre_op result.\n");
+		exit(1);
+		return -1;
+	}
+
+	switch (umcg_prepare_wait()) {
+	case UMCG_OP_DONE:
+		should_wait = false;
+		break;
+	case UMCG_OP_SYS:
+		should_wait = true;
+		break;
+	case UMCG_OP_ERROR:
+		return -1;
+	default:
+		fprintf(stderr, "lubumcg: unknown pre_op result.\n");
+		exit(1);
+		return -1;
+	}
+
+	if (should_wake)
+		return umcg_do_wake_or_swap(utls, prev_wait_counter,
+				should_wait, timeout);
+
+	if (should_wait)
+		return umcg_do_wait(timeout);
+
+	return 0;
+}

@@ -0,0 +1,154 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+
+#ifndef __LIBUMCG_H
+#define __LIBUMCG_H
+
+#define _GNU_SOURCE
+#include <errno.h>
+#include <limits.h>
+#include <unistd.h>
+#include <linux/types.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <syscall.h>
+#include <time.h>
+
+#include <linux/umcg.h>
+
+static int sys_umcg_api_version(uint32_t requested_api_version, uint32_t flags)
+{
+	return syscall(__NR_umcg_api_version, requested_api_version, flags);
+}
+
+static int sys_umcg_register_task(uint32_t api_version, uint32_t flags,
+		uint32_t group_id, struct umcg_task *umcg_task)
+{
+	return syscall(__NR_umcg_register_task, api_version, flags, group_id,
+			umcg_task);
+}
+
+static int sys_umcg_unregister_task(uint32_t flags)
+{
+	return syscall(__NR_umcg_unregister_task, flags);
+}
+
+static int sys_umcg_wait(uint32_t flags, const struct timespec *timeout)
+{
+	return syscall(__NR_umcg_wait, flags, timeout);
+}
+
+static int sys_umcg_wake(uint32_t flags, uint32_t next_tid)
+{
+	return syscall(__NR_umcg_wake, flags, next_tid);
+}
+
+static int sys_umcg_swap(uint32_t wake_flags, uint32_t next_tid,
+		uint32_t wait_flags, const struct timespec *timeout)
+{
+	return syscall(__NR_umcg_swap, wake_flags, next_tid,
+			wait_flags, timeout);
+}
+
+typedef intptr_t umcg_tid; /* UMCG thread ID. */
+
+#define UMCG_NONE	(0)
+
+/**
+ * umcg_get_utid - return the UMCG ID of the current thread.
+ *
+ * The function always succeeds, and the returned ID is guaranteed to be
+ * stable over the life of the thread (and multiple
+ * umcg_register/umcg_unregister calls).
+ *
+ * The ID is NOT guaranteed to be unique over the life of the process.
+ */
+umcg_tid umcg_get_utid(void);
+
+/**
+ * umcg_set_task_tag - add an arbitrary tag to a registered UMCG task.
+ *
+ * Note: non-thread-safe: the user is responsible for proper memory fencing.
+ */
+void umcg_set_task_tag(umcg_tid utid, intptr_t tag);
+
+/*
+ * umcg_get_task_tag - get the task tag. Returns zero if none set.
+ *
+ * Note: non-thread-safe: the user is responsible for proper memory fencing.
+ */
+intptr_t umcg_get_task_tag(umcg_tid utid);
+
+/**
+ * umcg_register_core_task - register the current thread as a UMCG core task
+ *
+ * Return:
+ * UMCG_NONE     - an error occurred. Check errno.
+ * != UMCG_NONE  - the ID of the thread to be used with UMCG API (guaranteed
+ *                 to match the value returned by umcg_get_utid).
+ */
+umcg_tid umcg_register_core_task(intptr_t tag);
+
+/**
+ * umcg_unregister_task - unregister the current thread
+ *
+ * Return:
+ * 0              - OK
+ * -1             - the current thread is not a UMCG thread
+ */
+int umcg_unregister_task(void);
+
+/**
+ * umcg_wait - block the current thread
+ * @timeout:   absolute timeout (not supported at the moment)
+ *
+ * Blocks the current thread, which must have been registered via umcg_register,
+ * until it is waken via umcg_wake or swapped into via umcg_swap. If the current
+ * thread has a wakeup queued (see umcg_wake), returns zero immediately,
+ * consuming the wakeup.
+ *
+ * Return:
+ * 0         - OK, the thread was waken;
+ * -1        - did not wake normally;
+ *               errno:
+ *                 EINTR: interrupted
+ *                 EINVAL: some other error occurred
+ */
+int umcg_wait(const struct timespec *timeout);
+
+/**
+ * umcg_wake - wake @next
+ * @next:      ID of the thread to wake (IDs are returned by umcg_register).
+ *
+ * If @next is blocked via umcg_wait, or umcg_swap, wake it. If @next is
+ * running, queue the wakeup, so that a future block of @next will consume
+ * the wakeup but will not block.
+ *
+ * umcg_wake is non-blocking, but may retry a few times to make sure @next
+ * has indeed woken.
+ *
+ * umcg_wake can queue at most one wakeup; if @next has a wakeup queued,
+ * an error is returned.
+ *
+ * Return:
+ * 0         - OK, @next has woken, or a wakeup has been queued;
+ * -1        - an error occurred.
+ */
+int umcg_wake(umcg_tid next);
+
+/**
+ * umcg_swap - wake @next, put the current thread to sleep
+ * @next:      ID of the thread to wake
+ * @timeout:   absolute timeout (not supported at the moment)
+ *
+ * umcg_swap is semantically equivalent to
+ *
+ *     int ret = umcg_wake(next);
+ *     if (ret)
+ *             return ret;
+ *     return umcg_wait(timeout);
+ *
+ * but may do a synchronous context switch into @next on the current CPU.
+ */
+int umcg_swap(umcg_tid next, const struct timespec *timeout);
+
+#endif  /* __LIBUMCG_H */

@@ -66,6 +66,7 @@ struct sighand_struct;
 struct signal_struct;
 struct task_delay_info;
 struct task_group;
+struct umcg_task_data;
 
 /*
  * Task state bitmask. NOTE! These bits are also

@@ -778,6 +779,10 @@ struct task_struct {
 	struct mm_struct		*mm;
 	struct mm_struct		*active_mm;
 
+#ifdef CONFIG_UMCG
+	struct umcg_task_data __rcu	*umcg_task_data;
+#endif
+
 	/* Per-thread vma caching: */
 	struct vmacache			vmacache;