#include <asm/bug.h>
 #include <asm/book3s/32/mmu-hash.h>
 
-#ifdef __ASSEMBLY__
-
-#ifdef CONFIG_PPC_KUAP
-
-.macro kuap_update_sr  gpr1, gpr2, gpr3        /* NEVER use r0 as gpr2 due to addis */
-101:   mtsrin  \gpr1, \gpr2
-       addi    \gpr1, \gpr1, 0x111             /* next VSID */
-       rlwinm  \gpr1, \gpr1, 0, 0xf0ffffff     /* clear VSID overflow */
-       addis   \gpr2, \gpr2, 0x1000            /* address of next segment */
-       cmplw   \gpr2, \gpr3
-       blt-    101b
-       isync
-.endm
-
-.macro kuap_save_and_lock      sp, thread, gpr1, gpr2, gpr3
-       lwz     \gpr2, KUAP(\thread)
-       rlwinm. \gpr3, \gpr2, 28, 0xf0000000
-       stw     \gpr2, STACK_REGS_KUAP(\sp)
-       beq+    102f
-       li      \gpr1, 0
-       stw     \gpr1, KUAP(\thread)
-       mfsrin  \gpr1, \gpr2
-       oris    \gpr1, \gpr1, SR_KS@h   /* set Ks */
-       kuap_update_sr  \gpr1, \gpr2, \gpr3
-102:
-.endm
-
-.macro kuap_restore    sp, current, gpr1, gpr2, gpr3
-       lwz     \gpr2, STACK_REGS_KUAP(\sp)
-       rlwinm. \gpr3, \gpr2, 28, 0xf0000000
-       stw     \gpr2, THREAD + KUAP(\current)
-       beq+    102f
-       mfsrin  \gpr1, \gpr2
-       rlwinm  \gpr1, \gpr1, 0, ~SR_KS /* Clear Ks */
-       kuap_update_sr  \gpr1, \gpr2, \gpr3
-102:
-.endm
-
-.macro kuap_check      current, gpr
-#ifdef CONFIG_PPC_KUAP_DEBUG
-       lwz     \gpr, THREAD + KUAP(\current)
-999:   twnei   \gpr, 0
-       EMIT_BUG_ENTRY 999b, __FILE__, __LINE__, (BUGFLAG_WARNING | BUGFLAG_ONCE)
-#endif
-.endm
-
-#endif /* CONFIG_PPC_KUAP */
-
-#else /* !__ASSEMBLY__ */
+#ifndef __ASSEMBLY__
 
 #ifdef CONFIG_PPC_KUAP
 
 
        if (user_mode(regs)) {
                kuep_lock();
                account_cpu_user_entry();
+       } else {
+               kuap_save_and_lock(regs);
        }
 #endif
        /*
 
 
 #ifdef __ASSEMBLY__
 #ifndef CONFIG_PPC_KUAP
-.macro kuap_save_and_lock      sp, thread, gpr1, gpr2, gpr3
-.endm
-
-.macro kuap_restore    sp, current, gpr1, gpr2, gpr3
-.endm
-
-.macro kuap_check      current, gpr
-.endm
-
 .macro kuap_check_amr  gpr1, gpr2
 .endm
 
 
 
 #ifdef CONFIG_PPC_KUAP
 
-#ifdef __ASSEMBLY__
-
-.macro kuap_save_and_lock      sp, thread, gpr1, gpr2, gpr3
-       lis     \gpr2, MD_APG_KUAP@h    /* only APG0 and APG1 are used */
-       mfspr   \gpr1, SPRN_MD_AP
-       mtspr   SPRN_MD_AP, \gpr2
-       stw     \gpr1, STACK_REGS_KUAP(\sp)
-.endm
-
-.macro kuap_restore    sp, current, gpr1, gpr2, gpr3
-       lwz     \gpr1, STACK_REGS_KUAP(\sp)
-       mtspr   SPRN_MD_AP, \gpr1
-.endm
-
-.macro kuap_check      current, gpr
-#ifdef CONFIG_PPC_KUAP_DEBUG
-       mfspr   \gpr, SPRN_MD_AP
-       rlwinm  \gpr, \gpr, 16, 0xffff
-999:   twnei   \gpr, MD_APG_KUAP@h
-       EMIT_BUG_ENTRY 999b, __FILE__, __LINE__, (BUGFLAG_WARNING | BUGFLAG_ONCE)
-#endif
-.endm
-
-#else /* !__ASSEMBLY__ */
+#ifndef __ASSEMBLY__
 
 #include <asm/reg.h>
 
 
 #if defined(CONFIG_PPC_BOOK3S_32) || defined(CONFIG_E500)
        .globl  prepare_transfer_to_handler
 prepare_transfer_to_handler:
-       addi    r12,r2,THREAD
-
        /* if from kernel, check interrupted DOZE/NAP mode */
-       kuap_save_and_lock r11, r12, r9, r5, r6
        lwz     r12,TI_LOCAL_FLAGS(r2)
        mtcrf   0x01,r12
        bt-     31-TLF_NAPPING,4f
        lwz     r9,_MSR(r11)            /* if sleeping, clear MSR.EE */
        rlwinm  r9,r9,0,~MSR_EE
        lwz     r12,_LINK(r11)          /* and return to address in LR */
-       kuap_restore r11, r2, r3, r4, r5
        lwz     r2, GPR2(r11)
        b       fast_exception_return
 _ASM_NOKPROBE_SYMBOL(prepare_transfer_to_handler)
        cmplwi  cr0,r5,0
        bne-    2f
 #endif /* CONFIG_PPC_47x */
-       kuap_check r2, r4
        lwz     r4,_LINK(r1)
        lwz     r5,_CCR(r1)
        mtlr    r4
        stw     r10,_CCR(r1)
        stw     r1,KSP(r3)      /* Set old stack pointer */
 
-       kuap_check r2, r0
 #ifdef CONFIG_SMP
        /* We need a sync somewhere here to make sure that if the
         * previous task gets rescheduled on another CPU, it sees all
        bne-    .Lrestore_nvgprs
 
 .Lfast_user_interrupt_return:
-       kuap_check r2, r4
        lwz     r11,_NIP(r1)
        lwz     r12,_MSR(r1)
        mtspr   SPRN_SRR0,r11
 
 .Lfast_kernel_interrupt_return:
        cmpwi   cr1,r3,0
-       kuap_restore r1, r2, r3, r4, r5
        lwz     r11,_NIP(r1)
        lwz     r12,_MSR(r1)
        mtspr   SPRN_SRR0,r11
 
        syscall_fn f;
 
        kuep_lock();
+#ifdef CONFIG_PPC32
+       kuap_save_and_lock(regs);
+#endif
 
        regs->orig_gpr3 = r3;
 
                        isync();
        } else
 #endif
-#ifdef CONFIG_PPC64
                kuap_assert_locked();
-#endif
 
        booke_restore_dbcr0();
 
 
        CT_WARN_ON(ct_state() == CONTEXT_USER);
 
-#ifdef CONFIG_PPC64
        kuap_assert_locked();
-#endif
 
        regs->result = r3;
 
 
        account_cpu_user_exit();
 
-#ifdef CONFIG_PPC_BOOK3S_64 /* BOOK3E and ppc32 not using this */
+#ifndef CONFIG_PPC_BOOK3E_64 /* BOOK3E not using this */
        /*
         * We do this at the end so that we do context switch with KERNEL AMR
         */
         * We don't need to restore AMR on the way back to userspace for KUAP.
         * AMR can only have been unlocked if we interrupted the kernel.
         */
-#ifdef CONFIG_PPC64
        kuap_assert_locked();
-#endif
 
        local_irq_save(flags);
 
        /*
         * We do this at the end so that we do context switch with KERNEL AMR
         */
-#ifdef CONFIG_PPC64
        kuap_user_restore(regs);
-#endif
        return ret;
 }
 
 {
        unsigned long flags;
        unsigned long ret = 0;
-#ifdef CONFIG_PPC64
        unsigned long kuap;
-#endif
 
        if (!IS_ENABLED(CONFIG_BOOKE) && !IS_ENABLED(CONFIG_40x) &&
            unlikely(!(regs->msr & MSR_RI)))
        if (TRAP(regs) != 0x700)
                CT_WARN_ON(ct_state() == CONTEXT_USER);
 
-#ifdef CONFIG_PPC64
        kuap = kuap_get_and_assert_locked();
-#endif
 
        if (unlikely(current_thread_info()->flags & _TIF_EMULATE_STACK_STORE)) {
                clear_bits(_TIF_EMULATE_STACK_STORE, ¤t_thread_info()->flags);
         * which would cause Read-After-Write stalls. Hence, we take the AMR
         * value from the check above.
         */
-#ifdef CONFIG_PPC64
        kuap_kernel_restore(regs, kuap);
-#endif
 
        return ret;
 }
 
         */
        restore_sprs(old_thread, new_thread);
 
+#ifdef CONFIG_PPC32
+       kuap_assert_locked();
+#endif
        last = _switch(old_thread, new_thread);
 
 #ifdef CONFIG_PPC_BOOK3S_64