pr_debug("    start:       %lx\n", kimage->start);
        pr_debug("    head:        %lx\n", kimage->head);
        pr_debug("    nr_segments: %lu\n", kimage->nr_segments);
+       pr_debug("    kern_reloc: %pa\n", &kimage->arch.kern_reloc);
 
        for (i = 0; i < kimage->nr_segments; i++) {
                pr_debug("      segment[%lu]: %016lx - %016lx, 0x%lx bytes, %lu pages\n",
        /* Empty routine needed to avoid build errors. */
 }
 
+int machine_kexec_post_load(struct kimage *kimage)
+{
+       void *reloc_code = page_to_virt(kimage->control_code_page);
+
+       memcpy(reloc_code, arm64_relocate_new_kernel,
+              arm64_relocate_new_kernel_size);
+       kimage->arch.kern_reloc = __pa(reloc_code);
+
+       /* Flush the reloc_code in preparation for its execution. */
+       __flush_dcache_area(reloc_code, arm64_relocate_new_kernel_size);
+       flush_icache_range((uintptr_t)reloc_code, (uintptr_t)reloc_code +
+                          arm64_relocate_new_kernel_size);
+
+       return 0;
+}
+
 /**
  * machine_kexec_prepare - Prepare for a kexec reboot.
  *
  */
 void machine_kexec(struct kimage *kimage)
 {
-       phys_addr_t reboot_code_buffer_phys;
-       void *reboot_code_buffer;
        bool in_kexec_crash = (kimage == kexec_crash_image);
        bool stuck_cpus = cpus_are_stuck_in_kernel();
 
        WARN(in_kexec_crash && (stuck_cpus || smp_crash_stop_failed()),
                "Some CPUs may be stale, kdump will be unreliable.\n");
 
-       reboot_code_buffer_phys = page_to_phys(kimage->control_code_page);
-       reboot_code_buffer = phys_to_virt(reboot_code_buffer_phys);
-
        kexec_image_info(kimage);
 
-       /*
-        * Copy arm64_relocate_new_kernel to the reboot_code_buffer for use
-        * after the kernel is shut down.
-        */
-       memcpy(reboot_code_buffer, arm64_relocate_new_kernel,
-               arm64_relocate_new_kernel_size);
-
-       /* Flush the reboot_code_buffer in preparation for its execution. */
-       __flush_dcache_area(reboot_code_buffer, arm64_relocate_new_kernel_size);
-
-       /*
-        * Although we've killed off the secondary CPUs, we don't update
-        * the online mask if we're handling a crash kernel and consequently
-        * need to avoid flush_icache_range(), which will attempt to IPI
-        * the offline CPUs. Therefore, we must use the __* variant here.
-        */
-       __flush_icache_range((uintptr_t)reboot_code_buffer,
-                            (uintptr_t)reboot_code_buffer +
-                            arm64_relocate_new_kernel_size);
-
        /* Flush the kimage list and its buffers. */
        kexec_list_flush(kimage);
 
 
        /*
         * cpu_soft_restart will shutdown the MMU, disable data caches, then
-        * transfer control to the reboot_code_buffer which contains a copy of
+        * transfer control to the kern_reloc which contains a copy of
         * the arm64_relocate_new_kernel routine.  arm64_relocate_new_kernel
         * uses physical addressing to relocate the new image to its final
         * position and transfers control to the image entry point when the
         * userspace (kexec-tools).
         * In kexec_file case, the kernel starts directly without purgatory.
         */
-       cpu_soft_restart(reboot_code_buffer_phys, kimage->head, kimage->start,
+       cpu_soft_restart(kimage->arch.kern_reloc, kimage->head, kimage->start,
                         kimage->arch.dtb_mem);
 
        BUG(); /* Should never get here. */