* Foundation, 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
  */
 
+#include <linux/cpu.h>
 #include <linux/errno.h>
 #include <linux/err.h>
 #include <linux/kvm_host.h>
        }
 }
 
-static void cpu_init_hyp_mode(void *vector)
+static void cpu_init_hyp_mode(void *dummy)
 {
        unsigned long long boot_pgd_ptr;
        unsigned long long pgd_ptr;
        __cpu_init_hyp_mode(boot_pgd_ptr, pgd_ptr, hyp_stack_ptr, vector_ptr);
 }
 
+static int hyp_init_cpu_notify(struct notifier_block *self,
+                              unsigned long action, void *cpu)
+{
+       switch (action) {
+       case CPU_STARTING:
+       case CPU_STARTING_FROZEN:
+               cpu_init_hyp_mode(NULL);
+               break;
+       }
+
+       return NOTIFY_OK;
+}
+
+static struct notifier_block hyp_init_cpu_nb = {
+       .notifier_call = hyp_init_cpu_notify,
+};
+
 /**
  * Inits Hyp-mode on all online CPUs
  */
 static int init_hyp_mode(void)
 {
-       phys_addr_t init_phys_addr;
        int cpu;
        int err = 0;
 
                per_cpu(kvm_arm_hyp_stack_page, cpu) = stack_page;
        }
 
-       /*
-        * Execute the init code on each CPU.
-        *
-        * Note: The stack is not mapped yet, so don't do anything else than
-        * initializing the hypervisor mode on each CPU using a local stack
-        * space for temporary storage.
-        */
-       init_phys_addr = virt_to_phys(__kvm_hyp_init);
-       for_each_online_cpu(cpu) {
-               smp_call_function_single(cpu, cpu_init_hyp_mode,
-                                        (void *)(long)init_phys_addr, 1);
-       }
-
        /*
         * Map the Hyp-code called directly from the host
         */
                }
        }
 
+       /*
+        * Execute the init code on each CPU.
+        */
+       on_each_cpu(cpu_init_hyp_mode, NULL, 1);
+
        /*
         * Init HYP view of VGIC
         */
        if (err)
                goto out_free_mappings;
 
+#ifndef CONFIG_HOTPLUG_CPU
+       free_boot_hyp_pgd();
+#endif
+
        kvm_perf_init();
 
        kvm_info("Hyp mode initialized successfully\n");
        if (err)
                goto out_err;
 
+       err = register_cpu_notifier(&hyp_init_cpu_nb);
+       if (err) {
+               kvm_err("Cannot register HYP init CPU notifier (%d)\n", err);
+               goto out_err;
+       }
+
        kvm_coproc_table_init();
        return 0;
 out_err:
 
        }
 }
 
+/**
+ * free_boot_hyp_pgd - free HYP boot page tables
+ *
+ * Free the HYP boot page tables. The bounce page is also freed.
+ */
+void free_boot_hyp_pgd(void)
+{
+       mutex_lock(&kvm_hyp_pgd_mutex);
+
+       if (boot_hyp_pgd) {
+               unmap_range(boot_hyp_pgd, hyp_idmap_start, PAGE_SIZE);
+               unmap_range(boot_hyp_pgd, TRAMPOLINE_VA, PAGE_SIZE);
+               kfree(boot_hyp_pgd);
+               boot_hyp_pgd = NULL;
+       }
+
+       if (hyp_pgd)
+               unmap_range(hyp_pgd, TRAMPOLINE_VA, PAGE_SIZE);
+
+       kfree(init_bounce_page);
+       init_bounce_page = NULL;
+
+       mutex_unlock(&kvm_hyp_pgd_mutex);
+}
+
 /**
  * free_hyp_pgds - free Hyp-mode page tables
  *
 {
        unsigned long addr;
 
-       mutex_lock(&kvm_hyp_pgd_mutex);
+       free_boot_hyp_pgd();
 
-       if (boot_hyp_pgd) {
-               unmap_range(boot_hyp_pgd, hyp_idmap_start, PAGE_SIZE);
-               unmap_range(boot_hyp_pgd, TRAMPOLINE_VA, PAGE_SIZE);
-               kfree(boot_hyp_pgd);
-       }
+       mutex_lock(&kvm_hyp_pgd_mutex);
 
        if (hyp_pgd) {
                for (addr = PAGE_OFFSET; virt_addr_valid(addr); addr += PGDIR_SIZE)
                for (addr = VMALLOC_START; is_vmalloc_addr((void*)addr); addr += PGDIR_SIZE)
                        unmap_range(hyp_pgd, KERN_TO_HYP(addr), PGDIR_SIZE);
                kfree(hyp_pgd);
+               hyp_pgd = NULL;
        }
 
-       kfree(init_bounce_page);
        mutex_unlock(&kvm_hyp_pgd_mutex);
 }