#include <asm/mmu_context.h>
 #include <asm/types.h>
 #include <asm/stacktrace.h>
+#include <asm/irq.h>
 
 extern void check_wait(void);
 extern asmlinkage void r4k_wait(void);
 extern int fpu_emulator_cop1Handler(struct pt_regs *xcp,
        struct mips_fpu_struct *ctx, int has_fpu);
 
+#ifdef CONFIG_CPU_CAVIUM_OCTEON
+extern asmlinkage void octeon_cop2_restore(struct octeon_cop2_state *task);
+#endif
+
 void (*board_be_init)(void);
 int (*board_be_handler)(struct pt_regs *regs, int is_fixup);
 void (*board_nmi_handler_setup)(void);
        unsigned int opcode;
        unsigned int cpid;
        int status;
+       unsigned long __maybe_unused flags;
 
        die_if_kernel("do_cpu invoked from kernel context!", regs);
 
                return;
 
        case 2:
+#ifdef CONFIG_CPU_CAVIUM_OCTEON
+               prefetch(¤t->thread.cp2);
+               local_irq_save(flags);
+               KSTK_STATUS(current) |= ST0_CU2;
+               status = read_c0_status();
+               write_c0_status(status | ST0_CU2);
+               octeon_cop2_restore(&(current->thread.cp2));
+               write_c0_status(status & ~ST0_CU2);
+               local_irq_restore(flags);
+               return;
+#endif
        case 3:
                break;
        }
                write_c0_hwrena(enable);
        }
 
+#ifdef CONFIG_CPU_CAVIUM_OCTEON
+       write_c0_hwrena(0xc000000f); /* Octeon has register 30 and 31 */
+#endif
+
 #ifdef CONFIG_MIPS_MT_SMTC
        if (!secondaryTC) {
 #endif /* CONFIG_MIPS_MT_SMTC */