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
 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));
 }
 
        case CPU_NEVADA:
        case CPU_RM7000:
        case CPU_RM9000:
-       case CPU_TX49XX:
        case CPU_4KC:
        case CPU_4KEC:
        case CPU_4KSC:
                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: