addi    r1,r1,SWITCH_FRAME_SIZE
        blr
 
-#ifdef CONFIG_PPC_BOOK3S
        /*
         * If MSR EE/RI was never enabled, IRQs not reconciled, NVGPRs not
         * touched, no exit work created, then this can be used.
        kuap_check_amr r3, r4
        ld      r5,_MSR(r1)
        andi.   r0,r5,MSR_PR
+#ifdef CONFIG_PPC_BOOK3S
        bne     .Lfast_user_interrupt_return_amr
        kuap_kernel_restore r3, r4
        andi.   r0,r5,MSR_RI
        addi    r3,r1,STACK_FRAME_OVERHEAD
        bl      unrecoverable_exception
        b       . /* should not get here */
+#else
+       bne     .Lfast_user_interrupt_return
+       b       .Lfast_kernel_interrupt_return
+#endif
 
        .balign IFETCH_ALIGN_BYTES
        .globl interrupt_return
        cmpdi   r3,0
        bne-    .Lrestore_nvgprs
 
+#ifdef CONFIG_PPC_BOOK3S
 .Lfast_user_interrupt_return_amr:
        kuap_user_restore r3, r4
+#endif
 .Lfast_user_interrupt_return:
        ld      r11,_NIP(r1)
        ld      r12,_MSR(r1)
 
        RFI_TO_KERNEL
        b       .       /* prevent speculative execution */
-#endif /* CONFIG_PPC_BOOK3S */
 
 #ifdef CONFIG_PPC_RTAS
 /*
 
 #endif
 }
 
+/* temporary hack for context tracking, removed in later patch */
+#include <linux/sched/debug.h>
+asmlinkage __visible void __sched schedule_user(void);
+
 /*
  * This should be called after a syscall returns, with r3 the return value
  * from the syscall. If this function returns non-zero, the system call
        while (unlikely(ti_flags & (_TIF_USER_WORK_MASK & ~_TIF_RESTORE_TM))) {
                local_irq_enable();
                if (ti_flags & _TIF_NEED_RESCHED) {
+#ifdef CONFIG_PPC_BOOK3E_64
+                       schedule_user();
+#else
                        schedule();
+#endif
                } else {
                        /*
                         * SIGPENDING must restore signal handler function
 
        account_cpu_user_exit();
 
-#ifndef CONFIG_PPC_BOOK3E_64 /* BOOK3E not using this */
-       /*
-        * We do this at the end so that we do context switch with KERNEL AMR
-        */
+       /* Restore user access locks last */
        kuap_user_restore(regs);
-#endif
        kuep_unlock();
 
        return ret;
 }
 
-#ifndef CONFIG_PPC_BOOK3E_64 /* BOOK3E not yet using this */
 notrace unsigned long interrupt_exit_user_prepare(struct pt_regs *regs, unsigned long msr)
 {
        unsigned long ti_flags;
        BUG_ON(!(regs->msr & MSR_PR));
        BUG_ON(!FULL_REGS(regs));
        BUG_ON(arch_irq_disabled_regs(regs));
+#ifdef CONFIG_PPC_BOOK3S_64
        CT_WARN_ON(ct_state() == CONTEXT_USER);
+#endif
 
        /*
         * We don't need to restore AMR on the way back to userspace for KUAP.
        while (unlikely(ti_flags & (_TIF_USER_WORK_MASK & ~_TIF_RESTORE_TM))) {
                local_irq_enable(); /* returning to user: may enable */
                if (ti_flags & _TIF_NEED_RESCHED) {
+#ifdef CONFIG_PPC_BOOK3E_64
+                       schedule_user();
+#else
                        schedule();
+#endif
                } else {
                        if (ti_flags & _TIF_SIGPENDING)
                                ret |= _TIF_RESTOREALL;
 
        account_cpu_user_exit();
 
-       /*
-        * We do this at the end so that we do context switch with KERNEL AMR
-        */
+       /* Restore user access locks last */
        kuap_user_restore(regs);
+
        return ret;
 }
 
         * CT_WARN_ON comes here via program_check_exception,
         * so avoid recursion.
         */
-       if (TRAP(regs) != 0x700)
+       if (IS_ENABLED(CONFIG_BOOKS) && TRAP(regs) != 0x700)
                CT_WARN_ON(ct_state() == CONTEXT_USER);
 
        kuap = kuap_get_and_assert_locked();
 #endif
 
        /*
-        * Don't want to mfspr(SPRN_AMR) here, because this comes after mtmsr,
-        * which would cause Read-After-Write stalls. Hence, we take the AMR
-        * value from the check above.
+        * 64s does not want to mfspr(SPRN_AMR) here, because this comes after
+        * mtmsr, which would cause Read-After-Write stalls. Hence, take the
+        * AMR value from the check above.
         */
        kuap_kernel_restore(regs, kuap);
 
        return ret;
 }
-#endif