ANDROID: arm64: suspend: Restore the UAO state

Upstream commit d08544127d ("arm64: suspend: Reconfigure PSTATE
after resume from idle") when cherry-picked on LTS linux-4.4.y
removed UAO reset code because UAO is not supported in linux-4.4.y.

But common/android-4.4 has UAO support, added in Change-Id:
I1a6a74a1f33b92d54368bd99387b55cf62930903 ("UPSTREAM: arm64: kernel:
Add support for User Access Override").

This patch pick up the UAO specific changes of upstream commit
dropped in LTS cherry-pick.

Fixes: LTS commit 71710cd35a ("arm64: suspend: Reconfigure PSTATE after resume from idle")
Signed-off-by: Amit Pundir <amit.pundir@linaro.org>

arch/arm64/include/asm/exec.h

@@ -18,6 +18,9 @@
 #ifndef __ASM_EXEC_H
 #define __ASM_EXEC_H
 
+#include <linux/sched.h>
+
 extern unsigned long arch_align_stack(unsigned long sp);
+void uao_thread_switch(struct task_struct *next);
 
 #endif	/* __ASM_EXEC_H */

arch/arm64/kernel/process.c

@@ -49,6 +49,7 @@
 #include <asm/alternative.h>
 #include <asm/compat.h>
 #include <asm/cacheflush.h>
+#include <asm/exec.h>
 #include <asm/fpsimd.h>
 #include <asm/mmu_context.h>
 #include <asm/processor.h>
@@ -379,7 +380,7 @@ static void tls_thread_switch(struct task_struct *next)
 }
 
 /* Restore the UAO state depending on next's addr_limit */
-static void uao_thread_switch(struct task_struct *next)
+void uao_thread_switch(struct task_struct *next)
 {
 	if (IS_ENABLED(CONFIG_ARM64_UAO)) {
 		if (task_thread_info(next)->addr_limit == KERNEL_DS)

arch/arm64/kernel/suspend.c

@@ -5,6 +5,7 @@
 #include <asm/cacheflush.h>
 #include <asm/cpufeature.h>
 #include <asm/debug-monitors.h>
+#include <asm/exec.h>
 #include <asm/pgtable.h>
 #include <asm/memory.h>
 #include <asm/mmu_context.h>
@@ -95,6 +96,7 @@ int cpu_suspend(unsigned long arg, int (*fn)(unsigned long))
 		 */
 		asm(ALTERNATIVE("nop", SET_PSTATE_PAN(1), ARM64_HAS_PAN,
 				CONFIG_ARM64_PAN));
+		uao_thread_switch(current);
 
 		/*
 		 * Restore HW breakpoint registers to sane values