[MIPS] Reduce race between cpu_wait() and need_resched() checking

If a thread became runnable between need_resched() and the WAIT
instruction, switching to the thread will delay until a next interrupt.
Some CPUs can execute the WAIT instruction with interrupt disabled, so
we can get rid of this race on them (at least UP case).

Original Patch by Atsushi with fixing up for MIPS Technology's cores by
Ralf based on feedback from the RTL designers.

Signed-off-by: Atsushi Nemoto <anemo@mba.ocn.ne.jp>
Signed-off-by: Ralf Baechle <ralf@linux-mips.org>

arch/mips/kernel/cpu-probe.c

@@ -38,15 +38,40 @@ static void r3081_wait(void)
 
 static void r39xx_wait(void)
 {
-	unsigned long cfg = read_c0_conf();
-	write_c0_conf(cfg | TX39_CONF_HALT);
+	local_irq_disable();
+	if (!need_resched())
+		write_c0_conf(read_c0_conf() | TX39_CONF_HALT);
+	local_irq_enable();
 }
 
+/*
+ * There is a race when WAIT instruction executed with interrupt
+ * enabled.
+ * But it is implementation-dependent wheter the pipelie restarts when
+ * a non-enabled interrupt is requested.
+ */
 static void r4k_wait(void)
 {
-	__asm__(".set\tmips3\n\t"
-		"wait\n\t"
-		".set\tmips0");
+	__asm__("	.set	mips3			\n"
+		"	wait				\n"
+		"	.set	mips0			\n");
+}
+
+/*
+ * This variant is preferable as it allows testing need_resched and going to
+ * sleep depending on the outcome atomically.  Unfortunately the "It is
+ * implementation-dependent whether the pipeline restarts when a non-enabled
+ * interrupt is requested" restriction in the MIPS32/MIPS64 architecture makes
+ * using this version a gamble.
+ */
+static void r4k_wait_irqoff(void)
+{
+	local_irq_disable();
+	if (!need_resched())
+		__asm__("	.set	mips3		\n"
+			"	wait			\n"
+			"	.set	mips0		\n");
+	local_irq_enable();
 }
 
 /* The Au1xxx wait is available only if using 32khz counter or
@@ -56,17 +81,17 @@ int allow_au1k_wait;
 static void au1k_wait(void)
 {
 	/* using the wait instruction makes CP0 counter unusable */
-	__asm__(".set mips3\n\t"
-		"cache 0x14, 0(%0)\n\t"
-		"cache 0x14, 32(%0)\n\t"
-		"sync\n\t"
-		"nop\n\t"
-		"wait\n\t"
-		"nop\n\t"
-		"nop\n\t"
-		"nop\n\t"
-		"nop\n\t"
-		".set mips0\n\t"
+	__asm__("	.set	mips3			\n"
+		"	cache	0x14, 0(%0)		\n"
+		"	cache	0x14, 32(%0)		\n"
+		"	sync				\n"
+		"	nop				\n"
+		"	wait				\n"
+		"	nop				\n"
+		"	nop				\n"
+		"	nop				\n"
+		"	nop				\n"
+		"	.set	mips0			\n"
 		: : "r" (au1k_wait));
 }
 
@@ -111,7 +136,6 @@ static inline void check_wait(void)
 	case CPU_NEVADA:
 	case CPU_RM7000:
 	case CPU_RM9000:
-	case CPU_TX49XX:
 	case CPU_4KC:
 	case CPU_4KEC:
 	case CPU_4KSC:
@@ -125,6 +149,10 @@ static inline void check_wait(void)
 		cpu_wait = r4k_wait;
 		printk(" available.\n");
 		break;
+	case CPU_TX49XX:
+		cpu_wait = r4k_wait_irqoff;
+		printk(" available.\n");
+		break;
 	case CPU_AU1000:
 	case CPU_AU1100:
 	case CPU_AU1500: