}
 }
 
+static void cpu_init_stage2(void *dummy)
+{
+       __cpu_init_stage2();
+}
+
 static void cpu_init_hyp_mode(void *dummy)
 {
        phys_addr_t boot_pgd_ptr;
 }
 #endif
 
+static void teardown_common_resources(void)
+{
+       free_percpu(kvm_host_cpu_state);
+}
+
+static int init_common_resources(void)
+{
+       kvm_host_cpu_state = alloc_percpu(kvm_cpu_context_t);
+       if (!kvm_host_cpu_state) {
+               kvm_err("Cannot allocate host CPU state\n");
+               return -ENOMEM;
+       }
+
+       return 0;
+}
+
+static int init_subsystems(void)
+{
+       int err;
+
+       /*
+        * Init HYP view of VGIC
+        */
+       err = kvm_vgic_hyp_init();
+       switch (err) {
+       case 0:
+               vgic_present = true;
+               break;
+       case -ENODEV:
+       case -ENXIO:
+               vgic_present = false;
+               break;
+       default:
+               return err;
+       }
+
+       /*
+        * Init HYP architected timer support
+        */
+       err = kvm_timer_hyp_init();
+       if (err)
+               return err;
+
+       kvm_perf_init();
+       kvm_coproc_table_init();
+
+       return 0;
+}
+
+static void teardown_hyp_mode(void)
+{
+       int cpu;
+
+       if (is_kernel_in_hyp_mode())
+               return;
+
+       free_hyp_pgds();
+       for_each_possible_cpu(cpu)
+               free_page(per_cpu(kvm_arm_hyp_stack_page, cpu));
+}
+
+static int init_vhe_mode(void)
+{
+       /*
+        * Execute the init code on each CPU.
+        */
+       on_each_cpu(cpu_init_stage2, NULL, 1);
+
+       /* set size of VMID supported by CPU */
+       kvm_vmid_bits = kvm_get_vmid_bits();
+       kvm_info("%d-bit VMID\n", kvm_vmid_bits);
+
+       kvm_info("VHE mode initialized successfully\n");
+       return 0;
+}
+
 /**
  * Inits Hyp-mode on all online CPUs
  */
                stack_page = __get_free_page(GFP_KERNEL);
                if (!stack_page) {
                        err = -ENOMEM;
-                       goto out_free_stack_pages;
+                       goto out_err;
                }
 
                per_cpu(kvm_arm_hyp_stack_page, cpu) = stack_page;
        err = create_hyp_mappings(__hyp_text_start, __hyp_text_end);
        if (err) {
                kvm_err("Cannot map world-switch code\n");
-               goto out_free_mappings;
+               goto out_err;
        }
 
        err = create_hyp_mappings(__start_rodata, __end_rodata);
        if (err) {
                kvm_err("Cannot map rodata section\n");
-               goto out_free_mappings;
+               goto out_err;
        }
 
        /*
 
                if (err) {
                        kvm_err("Cannot map hyp stack\n");
-                       goto out_free_mappings;
+                       goto out_err;
                }
        }
 
-       /*
-        * Map the host CPU structures
-        */
-       kvm_host_cpu_state = alloc_percpu(kvm_cpu_context_t);
-       if (!kvm_host_cpu_state) {
-               err = -ENOMEM;
-               kvm_err("Cannot allocate host CPU state\n");
-               goto out_free_mappings;
-       }
-
        for_each_possible_cpu(cpu) {
                kvm_cpu_context_t *cpu_ctxt;
 
 
                if (err) {
                        kvm_err("Cannot map host CPU state: %d\n", err);
-                       goto out_free_context;
+                       goto out_err;
                }
        }
 
         */
        on_each_cpu(cpu_init_hyp_mode, NULL, 1);
 
-       /*
-        * Init HYP view of VGIC
-        */
-       err = kvm_vgic_hyp_init();
-       switch (err) {
-       case 0:
-               vgic_present = true;
-               break;
-       case -ENODEV:
-       case -ENXIO:
-               vgic_present = false;
-               break;
-       default:
-               goto out_free_context;
-       }
-
-       /*
-        * Init HYP architected timer support
-        */
-       err = kvm_timer_hyp_init();
-       if (err)
-               goto out_free_context;
-
 #ifndef CONFIG_HOTPLUG_CPU
        free_boot_hyp_pgd();
 #endif
 
-       kvm_perf_init();
+       cpu_notifier_register_begin();
+
+       err = __register_cpu_notifier(&hyp_init_cpu_nb);
+
+       cpu_notifier_register_done();
+
+       if (err) {
+               kvm_err("Cannot register HYP init CPU notifier (%d)\n", err);
+               goto out_err;
+       }
+
+       hyp_cpu_pm_init();
 
        /* set size of VMID supported by CPU */
        kvm_vmid_bits = kvm_get_vmid_bits();
        kvm_info("Hyp mode initialized successfully\n");
 
        return 0;
-out_free_context:
-       free_percpu(kvm_host_cpu_state);
-out_free_mappings:
-       free_hyp_pgds();
-out_free_stack_pages:
-       for_each_possible_cpu(cpu)
-               free_page(per_cpu(kvm_arm_hyp_stack_page, cpu));
+
 out_err:
+       teardown_hyp_mode();
        kvm_err("error initializing Hyp mode: %d\n", err);
        return err;
 }
                }
        }
 
-       cpu_notifier_register_begin();
-
-       err = init_hyp_mode();
+       err = init_common_resources();
        if (err)
-               goto out_err;
+               return err;
 
-       err = __register_cpu_notifier(&hyp_init_cpu_nb);
-       if (err) {
-               kvm_err("Cannot register HYP init CPU notifier (%d)\n", err);
+       if (is_kernel_in_hyp_mode())
+               err = init_vhe_mode();
+       else
+               err = init_hyp_mode();
+       if (err)
                goto out_err;
-       }
-
-       cpu_notifier_register_done();
 
-       hyp_cpu_pm_init();
+       err = init_subsystems();
+       if (err)
+               goto out_hyp;
 
-       kvm_coproc_table_init();
        return 0;
+
+out_hyp:
+       teardown_hyp_mode();
 out_err:
-       cpu_notifier_register_done();
+       teardown_common_resources();
        return err;
 }