arm: highmem: Fix CPU_DYING callback for kmap

Add CPUHP_AP_KMAP_DYING state support instead of
use cpu notify for CPU_DYING which is deprecated.

Change-Id: Idaaf628ef4c89f361391151db523f5ffefc5bcc1
Signed-off-by: Maria Yu <aiquny@codeaurora.org>
diff --git a/arch/arm/mm/highmem.c b/arch/arm/mm/highmem.c
index 5d73327..0bb7673 100644
--- a/arch/arm/mm/highmem.c
+++ b/arch/arm/mm/highmem.c
@@ -150,7 +150,7 @@
 }
 
 #ifdef CONFIG_ARCH_WANT_KMAP_ATOMIC_FLUSH
-static void kmap_remove_unused_cpu(int cpu)
+int kmap_remove_unused_cpu(unsigned int cpu)
 {
 	int start_idx, idx, type;
 
@@ -167,6 +167,7 @@
 			set_top_pte(vaddr, __pte(0));
 	}
 	pagefault_enable();
+	return 0;
 }
 
 static void kmap_remove_unused(void *unused)
@@ -179,27 +180,4 @@
 	on_each_cpu(kmap_remove_unused, NULL, 1);
 }
 
-static int hotplug_kmap_atomic_callback(struct notifier_block *nfb,
-					unsigned long action, void *hcpu)
-{
-	switch (action & (~CPU_TASKS_FROZEN)) {
-	case CPU_DYING:
-		kmap_remove_unused_cpu((int)hcpu);
-		break;
-	default:
-		break;
-	}
-
-	return NOTIFY_OK;
-}
-
-static struct notifier_block hotplug_kmap_atomic_notifier = {
-	.notifier_call = hotplug_kmap_atomic_callback,
-};
-
-static int __init init_kmap_atomic(void)
-{
-	return register_hotcpu_notifier(&hotplug_kmap_atomic_notifier);
-}
-early_initcall(init_kmap_atomic);
 #endif
diff --git a/include/linux/cpuhotplug.h b/include/linux/cpuhotplug.h
index 7a2ae2f..d921206 100644
--- a/include/linux/cpuhotplug.h
+++ b/include/linux/cpuhotplug.h
@@ -59,6 +59,7 @@
 	CPUHP_AP_OFFLINE,
 	CPUHP_AP_SCHED_STARTING,
 	CPUHP_AP_RCUTREE_DYING,
+	CPUHP_AP_KMAP_DYING,
 	CPUHP_AP_IRQ_GIC_STARTING,
 	CPUHP_AP_IRQ_GICV3_STARTING,
 	CPUHP_AP_IRQ_HIP04_STARTING,
diff --git a/include/linux/highmem.h b/include/linux/highmem.h
index 61aff32..657b565 100644
--- a/include/linux/highmem.h
+++ b/include/linux/highmem.h
@@ -41,6 +41,7 @@
 
 #ifdef CONFIG_ARCH_WANT_KMAP_ATOMIC_FLUSH
 void kmap_atomic_flush_unused(void);
+int kmap_remove_unused_cpu(unsigned int cpu);
 #else
 static inline void kmap_atomic_flush_unused(void) { }
 #endif
@@ -91,6 +92,10 @@
 
 #endif /* CONFIG_HIGHMEM */
 
+#if !defined(CONFIG_HIGHMEM) || !defined(CONFIG_ARCH_WANT_KMAP_ATOMIC_FLUSH)
+static inline int kmap_remove_unused_cpu(unsigned int cpu) { return 0; }
+#endif
+
 #if defined(CONFIG_HIGHMEM) || defined(CONFIG_X86_32)
 
 DECLARE_PER_CPU(int, __kmap_atomic_idx);
diff --git a/kernel/cpu.c b/kernel/cpu.c
index 0d10ef5..915e750 100644
--- a/kernel/cpu.c
+++ b/kernel/cpu.c
@@ -25,6 +25,7 @@
 #include <linux/smpboot.h>
 #include <linux/relay.h>
 #include <linux/slab.h>
+#include <linux/highmem.h>
 
 #include <trace/events/power.h>
 #define CREATE_TRACE_POINTS
@@ -1433,6 +1434,11 @@
 		.startup.single		= NULL,
 		.teardown.single	= rcutree_dying_cpu,
 	},
+	[CPUHP_AP_KMAP_DYING] = {
+		.name			= "KMAP:dying",
+		.startup.single		= NULL,
+		.teardown.single	= kmap_remove_unused_cpu,
+	},
 	/* Entry state on starting. Interrupts enabled from here on. Transient
 	 * state for synchronsization */
 	[CPUHP_AP_ONLINE] = {