Acked-by: Artem Kuzin <artem.kuzin@huawei.com>
Acked-by: Alexander Grubnikov <alexander.grubnikov@huawei.com>
Acked-by: Ilya Hanov <ilya.hanov@huawei-partners.com>
Acked-by: Denis Darvish <darvish.denis@huawei.com>
Signed-off-by: Nikita Panov <panov.nikita@huawei.com>
---
 arch/arm64/include/asm/mmu_context.h | 4 ++++
 arch/arm64/kernel/hibernate.c        | 5 +++++
 arch/arm64/kernel/sleep.S            | 8 ++++++++
 mm/numa_kernel_replication.c         | 5 +++++
 4 files changed, 22 insertions(+)

diff --git a/arch/arm64/include/asm/mmu_context.h b/arch/arm64/include/asm/mmu_context.h
index 803b68758152..f364281d3442 100644
--- a/arch/arm64/include/asm/mmu_context.h
+++ b/arch/arm64/include/asm/mmu_context.h
@@ -136,10 +136,14 @@ static inline void cpu_install_ttbr0(phys_addr_t ttbr0, unsigned long t0sz)
 
 void __cpu_replace_ttbr1(pgd_t *pgdp, bool cnp);
 
+#ifndef CONFIG_KERNEL_REPLICATION
 static inline void cpu_enable_swapper_cnp(void)
 {
 	__cpu_replace_ttbr1(lm_alias(swapper_pg_dir), true);
 }
+#else
+void cpu_enable_swapper_cnp(void);
+#endif
 
 static inline void cpu_replace_ttbr1(pgd_t *pgdp)
 {
diff --git a/arch/arm64/kernel/hibernate.c b/arch/arm64/kernel/hibernate.c
index 7bf117427777..e9f5ddf37c4a 100644
--- a/arch/arm64/kernel/hibernate.c
+++ b/arch/arm64/kernel/hibernate.c
@@ -16,6 +16,7 @@
 #include <linux/sched.h>
 #include <linux/suspend.h>
 #include <linux/utsname.h>
+#include <linux/numa_kernel_replication.h>
 
 #include <asm/barrier.h>
 #include <asm/cacheflush.h>
@@ -113,7 +114,11 @@ int arch_hibernation_header_save(void *addr, unsigned int max_size)
 		return -EOVERFLOW;
 
 	arch_hdr_invariants(&hdr->invariants);
+#ifdef CONFIG_KERNEL_REPLICATION
+	hdr->ttbr1_el1		= virt_to_phys(this_node_pgd(&init_mm));
+#else
 	hdr->ttbr1_el1		= __pa_symbol(swapper_pg_dir);
+#endif /* CONFIG_KERNEL_REPLICATION */
 	hdr->reenter_kernel	= _cpu_resume;
 
 	/* We can't use __hyp_get_vectors() because kvm may still be loaded */
diff --git a/arch/arm64/kernel/sleep.S b/arch/arm64/kernel/sleep.S
index f093cdf71be1..8ad66ed70122 100644
--- a/arch/arm64/kernel/sleep.S
+++ b/arch/arm64/kernel/sleep.S
@@ -5,6 +5,10 @@
 #include <asm/assembler.h>
 #include <asm/smp.h>
 
+#ifdef CONFIG_KERNEL_REPLICATION
+.extern numa_setup_pgd
+#endif
+
 	.text
 /*
  * Implementation of MPIDR_EL1 hash algorithm through shifting
@@ -144,6 +148,10 @@ SYM_FUNC_START(_cpu_resume)
 	bl	kasan_unpoison_task_stack_below
 #endif
 
+#ifdef CONFIG_KERNEL_REPLICATION
+	bl 	numa_setup_pgd
+#endif
+
 	ldp	x19, x20, [x29, #16]
 	ldp	x21, x22, [x29, #32]
 	ldp	x23, x24, [x29, #48]
diff --git a/mm/numa_kernel_replication.c b/mm/numa_kernel_replication.c
index 082aabc6b8bc..b1c5cd2a1fa1 100644
--- a/mm/numa_kernel_replication.c
+++ b/mm/numa_kernel_replication.c
@@ -67,6 +67,11 @@ static int node_to_memory_node[MAX_NUMNODES];
 static bool pgtables_extra;
 static DEFINE_SPINLOCK(debugfs_lock);
 
+void cpu_enable_swapper_cnp(void)
+{
+	__cpu_replace_ttbr1(this_node_pgd(&init_mm), true);
+}
+
 bool is_text_replicated(void)
 {
 	return text_replicated;
-- 
2.34.1