1:     wfi
        b       1b
 
-.globl psci_arch_init
-psci_arch_init:
-       mov     r6, lr
-
-       bl      psci_get_cpu_id
-       bl      psci_get_cpu_stack_top
-       mov     sp, r0
-
-       bx      r6
-
        .globl psci_text_end
 psci_text_end:
        .popsection
 
 
        .arch_extension sec
 
-       @ r1 = target CPU
-       @ r2 = target PC
-
-.globl psci_arch_init
-psci_arch_init:
-       mov     r6, lr
-
-       bl      psci_get_cpu_id
-       bl      psci_get_cpu_stack_top
-       mov     sp, r0
-
-       bx      r6
-
-       @ r1 = target CPU
-       @ r2 = target PC
-
 .globl psci_cpu_on
 psci_cpu_on:
        push    {lr}
 
        mcr     p15, 0, r5, c12, c0, 1
        isb
 
-       @ Obtain a secure stack, and configure the PSCI backend
+       @ Obtain a secure stack
+       bl      psci_stack_setup
+
+       @ Configure the PSCI backend
+       push    {r0, r1, r2, ip}
        bl      psci_arch_init
+       pop     {r0, r1, r2, ip}
 #endif
 
 #ifdef CONFIG_ARM_ERRATA_773022
 
        bx      lr
 ENDPROC(psci_get_cpu_stack_top)
 
+@ {r0, r1, r2, ip} from _do_nonsec_entry(kernel_entry, 0, machid, r2) in
+@ arch/arm/lib/bootm.c:boot_jump_linux() must remain unchanged across
+@ this function.
+ENTRY(psci_stack_setup)
+       mov     r6, lr
+       mov     r7, r0
+       bl      psci_get_cpu_id         @ CPU ID => r0
+       bl      psci_get_cpu_stack_top  @ stack top => r0
+       mov     sp, r0
+       mov     r0, r7
+       bx      r6
+ENDPROC(psci_stack_setup)
+
+ENTRY(psci_arch_init)
+       mov     pc, lr
+ENDPROC(psci_arch_init)
+.weak psci_arch_init
+
 ENTRY(psci_cpu_entry)
        bl      psci_enable_smp
 
 
 #define        GICD_BASE               (SUNXI_GIC400_BASE +  0x1000)
 #define        GICC_BASE               (SUNXI_GIC400_BASE +  0x2000)
 
-@ {r0, r1, r2, ip} from _do_nonsec_entry(kernel_entry, 0, machid, r2) in
-@ arch/arm/lib/bootm.c:boot_jump_linux() must remain unchanged across
-@ this function.
 ENTRY(psci_arch_init)
-       mov     r6, lr
-       mov     r7, r0
-       bl      psci_get_cpu_id         @ CPU ID => r0
-       bl      psci_get_cpu_stack_top  @ stack top => r0
-       sub     r0, r0, #4              @ Save space for target PC
-       mov     sp, r0
-       mov     r0, r7
-       mov     lr, r6
-
-       push    {r0, r1, r2, ip, lr}
-       bl      sunxi_gic_init
-       pop     {r0, r1, r2, ip, pc}
+       b       sunxi_gic_init
 ENDPROC(psci_arch_init)
 
 ENTRY(psci_text_end)
 
        ldrne   r7, [r5]
        mcrne   p15, 0, r7, c14, c0, 0  @ write CNTFRQ to CPU1..3
 
-       bl      psci_get_cpu_stack_top  @ stack top => r0
-       mov     sp, r0
-
        bx      r6
 ENDPROC(psci_arch_init)