[RFC 2/3] powerpc64/preempt: Add need_resched bit inside preempt_count

[Shrikanth Hegde <sshegde@linux.ibm.com>]

Current implementation doesn't make use of the PREEMPT_NEED_RESCHED flag,
since this can interact badly with load/store architectures which
rely on the preempt_count word being unchanged across an interrupt.

On PPC64 which is 64 bit and the preempt count is only 32 bits wide,
one can simply pack it next to the resched flag and load the whole thing
in one go, so that a dec-and-test operation doesn't need to check tif.

[Adopted from arm64 implementation, including changelog bits]
Signed-off-by: Shrikanth Hegde <sshegde@linux.ibm.com>
---
 arch/powerpc/include/asm/paca.h    | 13 +++++++-
 arch/powerpc/include/asm/preempt.h | 52 ++++++++++++++++++------------
 arch/powerpc/kexec/core_64.c       |  2 +-
 3 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/arch/powerpc/include/asm/paca.h b/arch/powerpc/include/asm/paca.h
index 0aa729c2faf5..a73e214bd6f3 100644
--- a/arch/powerpc/include/asm/paca.h
+++ b/arch/powerpc/include/asm/paca.h
@@ -164,7 +164,18 @@ struct paca_struct {
 	u64 saved_r1;			/* r1 save for RTAS calls or PM or EE=0 */
 	u64 saved_msr;			/* MSR saved here by enter_rtas */
 	u64 exit_save_r1;		/* Syscall/interrupt R1 save */
-	int preempt_count;
+	union {
+		u64		preempt_count;
+		struct {
+#ifdef __BIG_ENDIAN__
+			u32	need_resched;
+			u32	count;
+#else
+			u32	count;
+			u32	need_resched;
+#endif
+		} preempt;
+	};
 #ifdef CONFIG_PPC_BOOK3E_64
 	u16 trap_save;			/* Used when bad stack is encountered */
 #endif
diff --git a/arch/powerpc/include/asm/preempt.h b/arch/powerpc/include/asm/preempt.h
index e2785d5a13e0..d17c8acac7bd 100644
--- a/arch/powerpc/include/asm/preempt.h
+++ b/arch/powerpc/include/asm/preempt.h
@@ -8,19 +8,20 @@
 #define __ASM_POWERPC_PREEMPT_H
 
 #ifdef CONFIG_PPC64
-#include <linux/thread_info.h>
 #include <asm/paca.h>
 
-#define PREEMPT_ENABLED	(0)
+#define PREEMPT_NEED_RESCHED	BIT(32)
+#define PREEMPT_ENABLED	(PREEMPT_NEED_RESCHED)
 
 static __always_inline int preempt_count(void)
 {
-	return READ_ONCE(local_paca->preempt_count);
+	return READ_ONCE(local_paca->preempt.count);
 }
 
-static __always_inline void preempt_count_set(int pc)
+static __always_inline void preempt_count_set(u64 pc)
 {
-	WRITE_ONCE(local_paca->preempt_count, pc);
+	/* Preserve existing value of PREEMPT_NEED_RESCHED */
+	WRITE_ONCE(local_paca->preempt.count, pc);
 }
 
 #define init_task_preempt_count(p) do { } while (0)
@@ -28,64 +29,73 @@ static __always_inline void preempt_count_set(int pc)
 
 static inline void set_preempt_need_resched(void)
 {
+	WRITE_ONCE(local_paca->preempt.need_resched, 0);
 }
 
 static inline void clear_preempt_need_resched(void)
 {
+	WRITE_ONCE(local_paca->preempt.need_resched, 1);
 }
 
 static inline bool test_preempt_need_resched(void)
 {
-	return false;
+	return !READ_ONCE(local_paca->preempt.need_resched);
 }
 
 static __always_inline void __preempt_count_add(int val)
 {
-	int pc = READ_ONCE(local_paca->preempt_count);
+	u32 pc = READ_ONCE(local_paca->preempt.count);
 
 	pc += val;
-	WRITE_ONCE(local_paca->preempt_count, pc);
+	WRITE_ONCE(local_paca->preempt.count, pc);
 }
 
 static __always_inline void __preempt_count_sub(int val)
 {
-	int pc = READ_ONCE(local_paca->preempt_count);
+	u32 pc = READ_ONCE(local_paca->preempt.count);
 
 	pc -= val;
-	WRITE_ONCE(local_paca->preempt_count, pc);
+	WRITE_ONCE(local_paca->preempt.count, pc);
 }
 
 static inline int __preempt_count_add_return(int val)
 {
-	int pc = READ_ONCE(local_paca->preempt_count);
+	u32 pc = READ_ONCE(local_paca->preempt.count);
 
 	pc += val;
-	WRITE_ONCE(local_paca->preempt_count, pc);
+	WRITE_ONCE(local_paca->preempt.count, pc);
 	return pc;
 }
 
 static inline int __preempt_count_sub_return(int val)
 {
-	int pc = READ_ONCE(local_paca->preempt_count);
+	u32 pc = READ_ONCE(local_paca->preempt.count);
 
 	pc -= val;
-	WRITE_ONCE(local_paca->preempt_count, pc);
+	WRITE_ONCE(local_paca->preempt.count, pc);
 	return pc;
 }
 
 static inline bool __preempt_count_dec_and_test(void)
 {
-	int pc = READ_ONCE(local_paca->preempt_count);
-
-	WRITE_ONCE(local_paca->preempt_count, --pc);
-	return !pc && tif_need_resched();
+	u64 pc = READ_ONCE(local_paca->preempt_count);
+
+	/* Update only the count field, leaving need_resched unchanged */
+	WRITE_ONCE(local_paca->preempt.count, --pc);
+	/*
+	 * If we wrote back all zeroes, then we're preemptible and in
+	 * need of a reschedule. Otherwise, we need to reload the
+	 * preempt_count in case the need_resched flag was cleared by an
+	 * interrupt occurring between the non-atomic READ_ONCE/WRITE_ONCE
+	 * pair.
+	 */
+	return !pc || !READ_ONCE(local_paca->preempt_count);
 }
 
 static inline bool should_resched(int preempt_offset)
 {
-	int pc = READ_ONCE(local_paca->preempt_count);
-
-	return unlikely(pc == preempt_offset && tif_need_resched());
+	u64 pc = READ_ONCE(local_paca->preempt_count);
+	return unlikely(pc == preempt_offset);
 }
 
 #ifdef CONFIG_PREEMPTION
diff --git a/arch/powerpc/kexec/core_64.c b/arch/powerpc/kexec/core_64.c
index f3954695e0b5..0f17c75cee28 100644
--- a/arch/powerpc/kexec/core_64.c
+++ b/arch/powerpc/kexec/core_64.c
@@ -352,7 +352,7 @@ void default_machine_kexec(struct kimage *image)
 	 * XXX: the task struct will likely be invalid once we do the copy!
 	 */
 	current_thread_info()->flags = 0;
-	WRITE_ONCE(local_paca->preempt_count, HARDIRQ_OFFSET);
+	WRITE_ONCE(local_paca->preempt.count, HARDIRQ_OFFSET);
 
 	/* We need a static PACA, too; copy this CPU's PACA over and switch to
 	 * it. Also poison per_cpu_offset and NULL lppaca to catch anyone using
-- 
2.52.0