#ifdef __ASSEMBLY__
 
-.macro kuep_update_sr  gpr1, gpr2              /* 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 */
-       bdnz    101b
-       isync
-.endm
-
-.macro kuep_lock       gpr1, gpr2
-#ifdef CONFIG_PPC_KUEP
-       li      \gpr1, NUM_USER_SEGMENTS
-       li      \gpr2, 0
-       mtctr   \gpr1
-       mfsrin  \gpr1, \gpr2
-       oris    \gpr1, \gpr1, SR_NX@h           /* set Nx */
-       kuep_update_sr \gpr1, \gpr2
-#endif
-.endm
-
-.macro kuep_unlock     gpr1, gpr2
-#ifdef CONFIG_PPC_KUEP
-       li      \gpr1, NUM_USER_SEGMENTS
-       li      \gpr2, 0
-       mtctr   \gpr1
-       mfsrin  \gpr1, \gpr2
-       rlwinm  \gpr1, \gpr1, 0, ~SR_NX         /* Clear Nx */
-       kuep_update_sr \gpr1, \gpr2
-#endif
-.endm
-
 #ifdef CONFIG_PPC_KUAP
 
 .macro kuap_update_sr  gpr1, gpr2, gpr3        /* NEVER use r0 as gpr2 due to addis */
 
        if (!arch_irq_disabled_regs(regs))
                trace_hardirqs_off();
 
-       if (user_mode(regs))
+       if (user_mode(regs)) {
+               kuep_lock();
                account_cpu_user_entry();
+       }
 #endif
        /*
         * Book3E reconciles irq soft mask in asm
        exception_exit(state->ctx_state);
 #endif
 
+       if (user_mode(regs))
+               kuep_unlock();
        /*
         * Book3S exits to user via interrupt_exit_user_prepare(), which does
         * context tracking, which is a cleaner way to handle PREEMPT=y
 
 static inline void setup_kuep(bool disabled) { }
 #endif /* CONFIG_PPC_KUEP */
 
+#if defined(CONFIG_PPC_KUEP) && defined(CONFIG_PPC_BOOK3S_32)
+void kuep_lock(void);
+void kuep_unlock(void);
+#else
+static inline void kuep_lock(void) { }
+static inline void kuep_unlock(void) { }
+#endif
+
 #ifdef CONFIG_PPC_KUAP
 void setup_kuap(bool disabled);
 #else
 
 #if defined(CONFIG_PPC_BOOK3S_32) || defined(CONFIG_E500)
        .globl  prepare_transfer_to_handler
 prepare_transfer_to_handler:
-       andi.   r12,r9,MSR_PR
        addi    r12,r2,THREAD
-       beq     2f
-#ifdef CONFIG_PPC_BOOK3S_32
-       kuep_lock r11, r12
-#endif
-       blr
 
        /* if from kernel, check interrupted DOZE/NAP mode */
-2:
        kuap_save_and_lock r11, r12, r9, r5, r6
        lwz     r12,TI_LOCAL_FLAGS(r2)
        mtcrf   0x01,r12
        .globl  transfer_to_syscall
 transfer_to_syscall:
        SAVE_NVGPRS(r1)
-#ifdef CONFIG_PPC_BOOK3S_32
-       kuep_lock r11, r12
-#endif
 
        /* Calling convention has r9 = orig r0, r10 = regs */
        addi    r10,r1,STACK_FRAME_OVERHEAD
        cmplwi  cr0,r5,0
        bne-    2f
 #endif /* CONFIG_PPC_47x */
-#ifdef CONFIG_PPC_BOOK3S_32
-       kuep_unlock r5, r7
-#endif
        kuap_check r2, r4
        lwz     r4,_LINK(r1)
        lwz     r5,_CCR(r1)
        bne-    .Lrestore_nvgprs
 
 .Lfast_user_interrupt_return:
-#ifdef CONFIG_PPC_BOOK3S_32
-       kuep_unlock     r10, r11
-#endif
        kuap_check r2, r4
        lwz     r11,_NIP(r1)
        lwz     r12,_MSR(r1)
 
 
 .macro prepare_transfer_to_handler
 #ifdef CONFIG_PPC_BOOK3S_32
+       andi.   r12,r9,MSR_PR
+       bne     777f
        bl      prepare_transfer_to_handler
+777:
 #endif
 .endm
 
 
 
 .macro prepare_transfer_to_handler
 #ifdef CONFIG_E500
+       andi.   r12,r9,MSR_PR
+       bne     777f
        bl      prepare_transfer_to_handler
+777:
 #endif
 .endm
 
 
 {
        syscall_fn f;
 
+       kuep_lock();
+
        regs->orig_gpr3 = r3;
 
        if (IS_ENABLED(CONFIG_PPC_IRQ_SOFT_MASK_DEBUG))
         */
        kuap_user_restore(regs);
 #endif
+       kuep_unlock();
+
        return ret;
 }
 
 
 obj-y += mmu.o mmu_context.o
 obj-$(CONFIG_PPC_BOOK3S_603) += nohash_low.o
 obj-$(CONFIG_PPC_BOOK3S_604) += hash_low.o tlb.o
+obj-$(CONFIG_PPC_KUEP) += kuep.o
 
--- /dev/null
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+#include <asm/kup.h>
+#include <asm/reg.h>
+#include <asm/task_size_32.h>
+#include <asm/mmu.h>
+
+#define KUEP_UPDATE_TWO_USER_SEGMENTS(n) do {          \
+       if (TASK_SIZE > ((n) << 28))                    \
+               mtsr(val1, (n) << 28);                  \
+       if (TASK_SIZE > (((n) + 1) << 28))              \
+               mtsr(val2, ((n) + 1) << 28);            \
+       val1 = (val1 + 0x222) & 0xf0ffffff;             \
+       val2 = (val2 + 0x222) & 0xf0ffffff;             \
+} while (0)
+
+static __always_inline void kuep_update(u32 val)
+{
+       int val1 = val;
+       int val2 = (val + 0x111) & 0xf0ffffff;
+
+       KUEP_UPDATE_TWO_USER_SEGMENTS(0);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(2);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(4);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(6);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(8);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(10);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(12);
+       KUEP_UPDATE_TWO_USER_SEGMENTS(14);
+}
+
+void kuep_lock(void)
+{
+       kuep_update(mfsr(0) | SR_NX);
+}
+
+void kuep_unlock(void)
+{
+       kuep_update(mfsr(0) & ~SR_NX);
+}