#include <linux/notifier.h>
 
+#if defined(CONFIG_CPU_CAVIUM_OCTEON)
+
+extern void octeon_cop2_save(struct octeon_cop2_state *);
+extern void octeon_cop2_restore(struct octeon_cop2_state *);
+
+#define cop2_save(r)           octeon_cop2_save(r)
+#define cop2_restore(r)                octeon_cop2_restore(r)
+
+#define cop2_present           1
+#define cop2_lazy_restore      1
+
+#else
+
+#define cop2_present           0
+#define cop2_lazy_restore      0
+#define cop2_save(r)
+#define cop2_restore(r)
+#endif
+
 enum cu2_ops {
        CU2_EXCEPTION,
        CU2_LWC2_OP,
 
        struct mips3264_watch_reg_state mips3264;
 };
 
-#ifdef CONFIG_CPU_CAVIUM_OCTEON
+#if defined(CONFIG_CPU_CAVIUM_OCTEON)
 
 struct octeon_cop2_state {
        /* DMFC2 rt, 0x0201 */
        /* DMFC2 rt, 0x025A; DMFC2 rt, 0x025B - Pass2 */
        unsigned long   cop2_gfm_result[2];
 };
-#define INIT_OCTEON_COP2 {0,}
+#define COP2_INIT                                              \
+       .cp2                    = {0,},
 
 struct octeon_cvmseg_state {
        unsigned long cvmseg[CONFIG_CAVIUM_OCTEON_CVMSEG_SIZE]
                            [cpu_dcache_line_size() / sizeof(unsigned long)];
 };
 
+#else
+#define COP2_INIT
 #endif
 
 typedef struct {
 #define FPAFF_INIT
 #endif /* CONFIG_MIPS_MT_FPAFF */
 
-#ifdef CONFIG_CPU_CAVIUM_OCTEON
-#define OCTEON_INIT                                            \
-       .cp2                    = INIT_OCTEON_COP2,
-#else
-#define OCTEON_INIT
-#endif /* CONFIG_CPU_CAVIUM_OCTEON */
-
 #define INIT_THREAD  {                                         \
        /*                                                      \
         * Saved main processor registers                       \
        .cp0_baduaddr           = 0,                            \
        .error_code             = 0,                            \
        /*                                                      \
-        * Cavium Octeon specifics (null if not Octeon)         \
+        * Platform specific cop2 registers(null if no COP2)    \
         */                                                     \
-       OCTEON_INIT                                             \
+       COP2_INIT                                               \
 }
 
 struct task_struct;
 
 #include <asm/cpu-features.h>
 #include <asm/watch.h>
 #include <asm/dsp.h>
+#include <asm/cop2.h>
 
 struct task_struct;
 
 
 #define switch_to(prev, next, last)                                    \
 do {                                                                   \
-       u32 __usedfpu;                                                  \
+       u32 __usedfpu, __c0_stat;                                       \
        __mips_mt_fpaff_switch_to(prev);                                \
        if (cpu_has_dsp)                                                \
                __save_dsp(prev);                                       \
+       if (cop2_present && (KSTK_STATUS(prev) & ST0_CU2)) {            \
+               if (cop2_lazy_restore)                                  \
+                       KSTK_STATUS(prev) &= ~ST0_CU2;                  \
+               __c0_stat = read_c0_status();                           \
+               write_c0_status(__c0_stat | ST0_CU2);                   \
+               cop2_save(&prev->thread.cp2);                           \
+               write_c0_status(__c0_stat & ~ST0_CU2);                  \
+       }                                                               \
        __clear_software_ll_bit();                                      \
        __usedfpu = test_and_clear_tsk_thread_flag(prev, TIF_USEDFPU);  \
        (last) = resume(prev, next, task_thread_info(next), __usedfpu); \
 
 #define finish_arch_switch(prev)                                       \
 do {                                                                   \
+       u32 __c0_stat;                                                  \
+       if (cop2_present && !cop2_lazy_restore &&                       \
+                       (KSTK_STATUS(current) & ST0_CU2)) {             \
+               __c0_stat = read_c0_status();                           \
+               write_c0_status(__c0_stat | ST0_CU2);                   \
+               cop2_restore(¤t->thread.cp2);                     \
+               write_c0_status(__c0_stat & ~ST0_CU2);                  \
+       }                                                               \
        if (cpu_has_dsp)                                                \
                __restore_dsp(current);                                 \
        if (cpu_has_userlocal)                                          \
 
        cpu_save_nonscratch a0
        LONG_S  ra, THREAD_REG31(a0)
 
-       /* check if we need to save COP2 registers */
-       PTR_L   t2, TASK_THREAD_INFO(a0)
-       LONG_L  t0, ST_OFF(t2)
-       bbit0   t0, 30, 1f
-
-       /* Disable COP2 in the stored process state */
-       li      t1, ST0_CU2
-       xor     t0, t1
-       LONG_S  t0, ST_OFF(t2)
-
-       /* Enable COP2 so we can save it */
-       mfc0    t0, CP0_STATUS
-       or      t0, t1
-       mtc0    t0, CP0_STATUS
-
-       /* Save COP2 */
-       daddu   a0, THREAD_CP2
-       jal octeon_cop2_save
-       dsubu   a0, THREAD_CP2
-
-       /* Disable COP2 now that we are done */
-       mfc0    t0, CP0_STATUS
-       li      t1, ST0_CU2
-       xor     t0, t1
-       mtc0    t0, CP0_STATUS
-
-1:
 #if CONFIG_CAVIUM_OCTEON_CVMSEG_SIZE > 0
        /* Check if we need to store CVMSEG state */
        mfc0    t0, $11,7       /* CvmMemCtl */