diff --git a/arch/arm64/include/asm/irqflags.h b/arch/arm64/include/asm/irqflags.h
index f51653fb90e4..e0f5f6b73edd 100644
--- a/arch/arm64/include/asm/irqflags.h
+++ b/arch/arm64/include/asm/irqflags.h
@@ -21,43 +21,77 @@
  * exceptions should be unmasked.
  */
 
-/*
- * CPU interrupt mask handling.
- */
-static inline void arch_local_irq_enable(void)
+static __always_inline bool __irqflags_uses_pmr(void)
 {
-	if (system_has_prio_mask_debugging()) {
-		u32 pmr = read_sysreg_s(SYS_ICC_PMR_EL1);
+	return IS_ENABLED(CONFIG_ARM64_PSEUDO_NMI) &&
+	       alternative_has_feature_unlikely(ARM64_HAS_GIC_PRIO_MASKING);
+}
 
+static __always_inline void __daif_local_irq_enable(void)
+{
+	barrier();
+	asm volatile("msr daifclr, #3");
+	barrier();
+}
+
+static __always_inline void __pmr_local_irq_enable(void)
+{
+	if (IS_ENABLED(CONFIG_ARM64_DEBUG_PRIORITY_MASKING)) {
+		u32 pmr = read_sysreg_s(SYS_ICC_PMR_EL1);
 		WARN_ON_ONCE(pmr != GIC_PRIO_IRQON && pmr != GIC_PRIO_IRQOFF);
 	}
 
-	asm volatile(ALTERNATIVE(
-		"msr	daifclr, #3		// arch_local_irq_enable",
-		__msr_s(SYS_ICC_PMR_EL1, "%0"),
-		ARM64_HAS_GIC_PRIO_MASKING)
-		:
-		: "r" ((unsigned long) GIC_PRIO_IRQON)
-		: "memory");
-
+	barrier();
+	write_sysreg_s(GIC_PRIO_IRQON, SYS_ICC_PMR_EL1);
 	pmr_sync();
+	barrier();
+}
+
+static inline void arch_local_irq_enable(void)
+{
+	if (__irqflags_uses_pmr()) {
+		__pmr_local_irq_enable();
+	} else {
+		__daif_local_irq_enable();
+	}
+}
+
+static __always_inline void __daif_local_irq_disable(void)
+{
+	barrier();
+	asm volatile("msr daifset, #3");
+	barrier();
+}
+
+static __always_inline void __pmr_local_irq_disable(void)
+{
+	if (IS_ENABLED(CONFIG_ARM64_DEBUG_PRIORITY_MASKING)) {
+		u32 pmr = read_sysreg_s(SYS_ICC_PMR_EL1);
+		WARN_ON_ONCE(pmr != GIC_PRIO_IRQON && pmr != GIC_PRIO_IRQOFF);
+	}
+
+	barrier();
+	write_sysreg_s(GIC_PRIO_IRQOFF, SYS_ICC_PMR_EL1);
+	barrier();
 }
 
 static inline void arch_local_irq_disable(void)
 {
-	if (system_has_prio_mask_debugging()) {
-		u32 pmr = read_sysreg_s(SYS_ICC_PMR_EL1);
-
-		WARN_ON_ONCE(pmr != GIC_PRIO_IRQON && pmr != GIC_PRIO_IRQOFF);
+	if (__irqflags_uses_pmr()) {
+		__pmr_local_irq_disable();
+	} else {
+		__daif_local_irq_disable();
 	}
+}
 
-	asm volatile(ALTERNATIVE(
-		"msr	daifset, #3		// arch_local_irq_disable",
-		__msr_s(SYS_ICC_PMR_EL1, "%0"),
-		ARM64_HAS_GIC_PRIO_MASKING)
-		:
-		: "r" ((unsigned long) GIC_PRIO_IRQOFF)
-		: "memory");
+static __always_inline unsigned long __daif_local_save_flags(void)
+{
+	return read_sysreg(daif);
+}
+
+static __always_inline unsigned long __pmr_local_save_flags(void)
+{
+	return read_sysreg_s(SYS_ICC_PMR_EL1);
 }
 
 /*
@@ -65,69 +99,108 @@ static inline void arch_local_irq_disable(void)
  */
 static inline unsigned long arch_local_save_flags(void)
 {
-	unsigned long flags;
+	if (__irqflags_uses_pmr()) {
+		return __pmr_local_save_flags();
+	} else {
+		return __daif_local_save_flags();
+	}
+}
 
-	asm volatile(ALTERNATIVE(
-		"mrs	%0, daif",
-		__mrs_s("%0", SYS_ICC_PMR_EL1),
-		ARM64_HAS_GIC_PRIO_MASKING)
-		: "=&r" (flags)
-		:
-		: "memory");
+static __always_inline bool __daif_irqs_disabled_flags(unsigned long flags)
+{
+	return flags & PSR_I_BIT;
+}
+
+static __always_inline bool __pmr_irqs_disabled_flags(unsigned long flags)
+{
+	return flags != GIC_PRIO_IRQON;
+}
+
+static inline bool arch_irqs_disabled_flags(unsigned long flags)
+{
+	if (__irqflags_uses_pmr()) {
+		return __pmr_irqs_disabled_flags(flags);
+	} else {
+		return __daif_irqs_disabled_flags(flags);
+	}
+}
+
+static __always_inline bool __daif_irqs_disabled(void)
+{
+	return __daif_irqs_disabled_flags(__daif_local_save_flags());
+}
+
+static __always_inline bool __pmr_irqs_disabled(void)
+{
+	return __pmr_irqs_disabled_flags(__pmr_local_save_flags());
+}
+
+static inline bool arch_irqs_disabled(void)
+{
+	if (__irqflags_uses_pmr()) {
+		return __pmr_irqs_disabled();
+	} else {
+		return __daif_irqs_disabled();
+	}
+}
+
+static __always_inline unsigned long __daif_local_irq_save(void)
+{
+	unsigned long flags = __daif_local_save_flags();
+
+	__daif_local_irq_disable();
 
 	return flags;
 }
 
-static inline int arch_irqs_disabled_flags(unsigned long flags)
+static __always_inline unsigned long __pmr_local_irq_save(void)
 {
-	int res;
-
-	asm volatile(ALTERNATIVE(
-		"and	%w0, %w1, #" __stringify(PSR_I_BIT),
-		"eor	%w0, %w1, #" __stringify(GIC_PRIO_IRQON),
-		ARM64_HAS_GIC_PRIO_MASKING)
-		: "=&r" (res)
-		: "r" ((int) flags)
-		: "memory");
-
-	return res;
-}
-
-static inline int arch_irqs_disabled(void)
-{
-	return arch_irqs_disabled_flags(arch_local_save_flags());
-}
-
-static inline unsigned long arch_local_irq_save(void)
-{
-	unsigned long flags;
-
-	flags = arch_local_save_flags();
+	unsigned long flags = __pmr_local_save_flags();
 
 	/*
 	 * There are too many states with IRQs disabled, just keep the current
 	 * state if interrupts are already disabled/masked.
 	 */
-	if (!arch_irqs_disabled_flags(flags))
-		arch_local_irq_disable();
+	if (!__pmr_irqs_disabled_flags(flags))
+		__pmr_local_irq_disable();
 
 	return flags;
 }
 
+static inline unsigned long arch_local_irq_save(void)
+{
+	if (__irqflags_uses_pmr()) {
+		return __pmr_local_irq_save();
+	} else {
+		return __daif_local_irq_save();
+	}
+}
+
+static __always_inline void __daif_local_irq_restore(unsigned long flags)
+{
+	barrier();
+	write_sysreg(flags, daif);
+	barrier();
+}
+
+static __always_inline void __pmr_local_irq_restore(unsigned long flags)
+{
+	barrier();
+	write_sysreg_s(flags, SYS_ICC_PMR_EL1);
+	pmr_sync();
+	barrier();
+}
+
 /*
  * restore saved IRQ state
  */
 static inline void arch_local_irq_restore(unsigned long flags)
 {
-	asm volatile(ALTERNATIVE(
-		"msr	daif, %0",
-		__msr_s(SYS_ICC_PMR_EL1, "%0"),
-		ARM64_HAS_GIC_PRIO_MASKING)
-		:
-		: "r" (flags)
-		: "memory");
-
-	pmr_sync();
+	if (__irqflags_uses_pmr()) {
+		__pmr_local_irq_restore(flags);
+	} else {
+		__daif_local_irq_restore(flags);
+	}
 }
 
 #endif /* __ASM_IRQFLAGS_H */