#include <linux/prctl.h>
 
 #include <asm/alternative.h>
-#include <asm/arch_gicv3.h>
 #include <asm/compat.h>
 #include <asm/cpufeature.h>
+#include <asm/cpuidle.h>
 #include <asm/cacheflush.h>
 #include <asm/exec.h>
 #include <asm/fpsimd.h>
 
 void (*arm_pm_restart)(enum reboot_mode reboot_mode, const char *cmd);
 
-static void noinstr __cpu_do_idle(void)
-{
-       dsb(sy);
-       wfi();
-}
-
-static void noinstr __cpu_do_idle_irqprio(void)
-{
-       unsigned long pmr;
-       unsigned long daif_bits;
-
-       daif_bits = read_sysreg(daif);
-       write_sysreg(daif_bits | PSR_I_BIT | PSR_F_BIT, daif);
-
-       /*
-        * Unmask PMR before going idle to make sure interrupts can
-        * be raised.
-        */
-       pmr = gic_read_pmr();
-       gic_write_pmr(GIC_PRIO_IRQON | GIC_PRIO_PSR_I_SET);
-
-       __cpu_do_idle();
-
-       gic_write_pmr(pmr);
-       write_sysreg(daif_bits, daif);
-}
-
 /*
  *     cpu_do_idle()
  *
  */
 void noinstr cpu_do_idle(void)
 {
-       if (system_uses_irq_prio_masking())
-               __cpu_do_idle_irqprio();
-       else
-               __cpu_do_idle();
+       struct arm_cpuidle_irq_context context;
+
+       arm_cpuidle_save_irq_context(&context);
+
+       dsb(sy);
+       wfi();
+
+       arm_cpuidle_restore_irq_context(&context);
 }
 
 /*