#include <asm/mmu.h>
 #include <asm/mmu_context.h>
 #include <asm/page.h>
+#include <asm/trans_pgd.h>
 
 #include "cpu-reset.h"
 
        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("    dtb_mem: %pa\n", &kimage->arch.dtb_mem);
        pr_debug("    kern_reloc: %pa\n", &kimage->arch.kern_reloc);
+       pr_debug("    el2_vectors: %pa\n", &kimage->arch.el2_vectors);
 
        for (i = 0; i < kimage->nr_segments; i++) {
                pr_debug("      segment[%lu]: %016lx - %016lx, 0x%lx bytes, %lu pages\n",
        }
 }
 
+/* Allocates pages for kexec page table */
+static void *kexec_page_alloc(void *arg)
+{
+       struct kimage *kimage = (struct kimage *)arg;
+       struct page *page = kimage_alloc_control_pages(kimage, 0);
+
+       if (!page)
+               return NULL;
+
+       memset(page_address(page), 0, PAGE_SIZE);
+
+       return page_address(page);
+}
+
 int machine_kexec_post_load(struct kimage *kimage)
 {
        void *reloc_code = page_to_virt(kimage->control_code_page);
+       struct trans_pgd_info info = {
+               .trans_alloc_page       = kexec_page_alloc,
+               .trans_alloc_arg        = kimage,
+       };
 
        /* If in place, relocation is not used, only flush next kernel */
        if (kimage->head & IND_DONE) {
                return 0;
        }
 
+       kimage->arch.el2_vectors = 0;
+       if (is_hyp_nvhe()) {
+               int rc = trans_pgd_copy_el2_vectors(&info,
+                                                   &kimage->arch.el2_vectors);
+               if (rc)
+                       return rc;
+       }
+
        memcpy(reloc_code, arm64_relocate_new_kernel,
               arm64_relocate_new_kernel_size);
        kimage->arch.kern_reloc = __pa(reloc_code);
                restart(is_hyp_nvhe(), kimage->start, kimage->arch.dtb_mem,
                        0, 0);
        } else {
+               if (is_hyp_nvhe())
+                       __hyp_set_vectors(kimage->arch.el2_vectors);
                cpu_soft_restart(kimage->arch.kern_reloc, virt_to_phys(kimage),
                                 0, 0);
        }