[RFC] in-kernel rseq

[Peter Zijlstra <peterz@infradead.org>]

Hi,

It has come to my attention that various people are struggling with
preempt_disable()+preempt_enable() costs for various architectures.

Mostly in relation to things like this_cpu_ and or local_. 

The below is a very crude (and broken, more on that below) POC.

So the 'main' advantage of this over preempt_disable()/preempt_enable(),
it on the preempt_enable() side, this elides the whole conditional and
call schedule() nonsense.

Now, on to the broken part, the below 'commit' address should be the
address of the 'STORE' instruction. In case of LL/SC, it should be the
SC, in case of LSE, it should be the LSE instruction.

This means, it needs to be woven into the asm... and I'm not that handy
with arm64 asm.

The pseudo code would be something like:

	current->sched_seq = &_R;
	...

_start:  compute per cpu-addr
	 load addr
	 $OP
_commit: store addr

	...
	current->sched_rseq = NULL;


Then when preemption happens (from interrupt), the instruction pointer
is 'simply' reset to _start and it tries again.

Anyway, this was aimed at arm64, which chose to use atomics for
this_cpu. But if we move sched_rseq() from schedule-tail into interrupt
entry, then this would also work for things like Power.

Anyway, just throwing ideas out there.


diff --git a/arch/arm64/include/asm/percpu.h b/arch/arm64/include/asm/percpu.h
index b57b2bb00967..080a868391b7 100644
--- a/arch/arm64/include/asm/percpu.h
+++ b/arch/arm64/include/asm/percpu.h
@@ -11,6 +11,7 @@
 #include <asm/cmpxchg.h>
 #include <asm/stack_pointer.h>
 #include <asm/sysreg.h>
+#include <generated/rq-offsets.h>
 
 static inline void set_my_cpu_offset(unsigned long off)
 {
@@ -155,9 +156,23 @@ PERCPU_RET_OP(add, add, ldadd)
 
 #define _pcp_protect(op, pcp, ...)					\
 ({									\
-	preempt_disable_notrace();					\
+	__label__ __rseq_begin;						\
+	__label__ __rseq_end;						\
+	static struct sched_rseq _R = {					\
+		.begin = (unsigned long)&&__rseq_begin,			\
+		.commit = (unsigned long)&&__rseq_end,			\
+		.restart = (unsigned long)&&__rseq_begin,		\
+	};								\
+	struct sched_rseq **this_rseq;					\
+	asm ("mrs %0, sp_el0; add %0, %0, %1;" : "=r" (this_rseq) : "i" (TSK_rseq));\
+	*this_rseq = &_R;						\
+__rseq_begin:								\
+	barrier();							\
 	op(raw_cpu_ptr(&(pcp)), __VA_ARGS__);				\
-	preempt_enable_notrace();					\
+	/* XXX broken */						\
+	barrier();							\
+__rseq_end:								\
+	*this_rseq = NULL;						\
 })
 
 #define _pcp_protect_return(op, pcp, args...)				\
diff --git a/include/linux/sched.h b/include/linux/sched.h
index a7b4a980eb2f..7960f3e21104 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -817,6 +817,12 @@ struct kmap_ctrl {
 #endif
 };
 
+struct sched_rseq {
+	unsigned long begin;
+	unsigned long commit;
+	unsigned long restart;
+};
+
 struct task_struct {
 #ifdef CONFIG_THREAD_INFO_IN_TASK
 	/*
@@ -827,6 +833,8 @@ struct task_struct {
 #endif
 	unsigned int			__state;
 
+	struct sched_rseq		*sched_rseq;
+
 	/* saved state for "spinlock sleepers" */
 	unsigned int			saved_state;
 
diff --git a/include/linux/sched/rseq.h b/include/linux/sched/rseq.h
deleted file mode 100644
index e69de29bb2d1..000000000000
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index bfd280ec0f97..d4702f8590f2 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5087,6 +5087,23 @@ prepare_task_switch(struct rq *rq, struct task_struct *prev,
 	prepare_arch_switch(next);
 }
 
+static inline void sched_rseq(struct task_struct *prev)
+{
+	struct sched_rseq *rseq = prev->sched_rseq;
+	struct pt_regs *regs;
+	unsigned long ip;
+
+	if (likely(!rseq))
+		return;
+
+	regs = task_pt_regs(prev);
+	ip = instruction_pointer(regs);
+	if ((ip - rseq->begin) >= (rseq->commit - rseq->begin))
+		return;
+
+	instruction_pointer_set(regs, rseq->restart);
+}
+
 /**
  * finish_task_switch - clean up after a task-switch
  * @prev: the thread we just switched away from.
@@ -5145,6 +5162,7 @@ static struct rq *finish_task_switch(struct task_struct *prev)
 	prev_state = READ_ONCE(prev->__state);
 	vtime_task_switch(prev);
 	perf_event_task_sched_in(prev, current);
+	sched_rseq(prev);
 	finish_task(prev);
 	tick_nohz_task_switch();
 	finish_lock_switch(rq);
diff --git a/kernel/sched/rq-offsets.c b/kernel/sched/rq-offsets.c
index a23747bbe25b..629989a89395 100644
--- a/kernel/sched/rq-offsets.c
+++ b/kernel/sched/rq-offsets.c
@@ -6,7 +6,8 @@
 
 int main(void)
 {
-	DEFINE(RQ_nr_pinned, offsetof(struct rq, nr_pinned));
+	DEFINE(RQ_nr_pinned,	offsetof(struct rq, nr_pinned));
+	DEFINE(TSK_rseq,	offsetof(struct task_struct, sched_rseq));
 
 	return 0;
 }