#include <asm/boot.h>
 #include <asm/cpufeature.h>
+#include <asm/fpsimd.h>
 #include <asm/io.h>
 #include <asm/mmu_context.h>
 #include <asm/neon.h>
 
 #define arch_efi_call_virt_setup()                                     \
 ({                                                                     \
-       kernel_neon_begin();                                            \
        efi_virtmap_load();                                             \
+       __efi_fpsimd_begin();                                           \
 })
 
 #define arch_efi_call_virt(p, f, args...)                              \
 
 #define arch_efi_call_virt_teardown()                                  \
 ({                                                                     \
+       __efi_fpsimd_end();                                             \
        efi_virtmap_unload();                                           \
-       kernel_neon_end();                                              \
 })
 
 #define ARCH_EFI_IRQ_FLAGS_MASK (PSR_D_BIT | PSR_A_BIT | PSR_I_BIT | PSR_F_BIT)
 
 #include <linux/cpu_pm.h>
 #include <linux/kernel.h>
 #include <linux/init.h>
+#include <linux/preempt.h>
 #include <linux/sched/signal.h>
 #include <linux/signal.h>
 #include <linux/hardirq.h>
 
 #include <asm/fpsimd.h>
 #include <asm/cputype.h>
+#include <asm/neon.h>
+#include <asm/simd.h>
 
 #define FPEXC_IOF      (1 << 0)
 #define FPEXC_DZF      (1 << 1)
 }
 EXPORT_SYMBOL(kernel_neon_end);
 
+DEFINE_PER_CPU(struct fpsimd_state, efi_fpsimd_state);
+DEFINE_PER_CPU(bool, efi_fpsimd_state_used);
+
+/*
+ * EFI runtime services support functions
+ *
+ * The ABI for EFI runtime services allows EFI to use FPSIMD during the call.
+ * This means that for EFI (and only for EFI), we have to assume that FPSIMD
+ * is always used rather than being an optional accelerator.
+ *
+ * These functions provide the necessary support for ensuring FPSIMD
+ * save/restore in the contexts from which EFI is used.
+ *
+ * Do not use them for any other purpose -- if tempted to do so, you are
+ * either doing something wrong or you need to propose some refactoring.
+ */
+
+/*
+ * __efi_fpsimd_begin(): prepare FPSIMD for making an EFI runtime services call
+ */
+void __efi_fpsimd_begin(void)
+{
+       if (!system_supports_fpsimd())
+               return;
+
+       WARN_ON(preemptible());
+
+       if (may_use_simd())
+               kernel_neon_begin();
+       else {
+               fpsimd_save_state(this_cpu_ptr(&efi_fpsimd_state));
+               __this_cpu_write(efi_fpsimd_state_used, true);
+       }
+}
+
+/*
+ * __efi_fpsimd_end(): clean up FPSIMD after an EFI runtime services call
+ */
+void __efi_fpsimd_end(void)
+{
+       if (!system_supports_fpsimd())
+               return;
+
+       if (__this_cpu_xchg(efi_fpsimd_state_used, false))
+               fpsimd_load_state(this_cpu_ptr(&efi_fpsimd_state));
+       else
+               kernel_neon_end();
+}
+
 #endif /* CONFIG_KERNEL_MODE_NEON */
 
 #ifdef CONFIG_CPU_PM