#include <asm/mmu.h>
 #include <asm/ptrace.h>
 
-static inline void kuap_restore_amr(struct pt_regs *regs)
+static inline void kuap_restore_amr(struct pt_regs *regs, unsigned long amr)
 {
-       if (mmu_has_feature(MMU_FTR_RADIX_KUAP)) {
+       if (mmu_has_feature(MMU_FTR_RADIX_KUAP) && unlikely(regs->kuap != amr)) {
                isync();
                mtspr(SPRN_AMR, regs->kuap);
                /*
        }
 }
 
+static inline unsigned long kuap_get_and_check_amr(void)
+{
+       if (mmu_has_feature(MMU_FTR_RADIX_KUAP)) {
+               unsigned long amr = mfspr(SPRN_AMR);
+               if (IS_ENABLED(CONFIG_PPC_KUAP_DEBUG)) /* kuap_check_amr() */
+                       WARN_ON_ONCE(amr != AMR_KUAP_BLOCKED);
+               return amr;
+       }
+       return 0;
+}
+
 static inline void kuap_check_amr(void)
 {
        if (IS_ENABLED(CONFIG_PPC_KUAP_DEBUG) && mmu_has_feature(MMU_FTR_RADIX_KUAP))
                    "Bug: %s fault blocked by AMR!", is_write ? "Write" : "Read");
 }
 #else /* CONFIG_PPC_KUAP */
-static inline void kuap_restore_amr(struct pt_regs *regs)
+static inline void kuap_restore_amr(struct pt_regs *regs, unsigned long amr)
 {
 }
 
 static inline void kuap_check_amr(void)
 {
 }
+
+static inline unsigned long kuap_get_and_check_amr(void)
+{
+       return 0;
+}
 #endif /* CONFIG_PPC_KUAP */
 
 #endif /* __ASSEMBLY__ */
 
        BUG_ON(!FULL_REGS(regs));
        BUG_ON(regs->softe != IRQS_ENABLED);
 
+       /*
+        * 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.
+        */
        kuap_check_amr();
 
        local_irq_save(flags);
        unsigned long *ti_flagsp = ¤t_thread_info()->flags;
        unsigned long flags;
        unsigned long ret = 0;
+       unsigned long amr;
 
        if (IS_ENABLED(CONFIG_PPC_BOOK3S) && unlikely(!(regs->msr & MSR_RI)))
                unrecoverable_exception(regs);
        BUG_ON(regs->msr & MSR_PR);
        BUG_ON(!FULL_REGS(regs));
 
-       kuap_check_amr();
+       amr = kuap_get_and_check_amr();
 
        if (unlikely(*ti_flagsp & _TIF_EMULATE_STACK_STORE)) {
                clear_bits(_TIF_EMULATE_STACK_STORE, ti_flagsp);
 #endif
 
        /*
-        * We don't need to restore AMR on the way back to userspace for KUAP.
-        * The value of AMR only matters while we're in the kernel.
+        * 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.
         */
-       kuap_restore_amr(regs);
+       kuap_restore_amr(regs, amr);
 
        return ret;
 }