return (struct thread_info *)sp_el0;
 }
 
-/* Access struct thread_info of another thread */
-static inline struct thread_info *get_thread_info(unsigned long thread_stack)
-{
-       return (struct thread_info *)(thread_stack & ~(THREAD_SIZE - 1));
-}
-
 #define thread_saved_pc(tsk)   \
        ((unsigned long)(tsk->thread.cpu_context.pc))
 #define thread_saved_sp(tsk)   \
 
 /* Restore the UAO state depending on next's addr_limit */
 static void uao_thread_switch(struct task_struct *next)
 {
-       unsigned long next_sp = next->thread.cpu_context.sp;
-
-       if (IS_ENABLED(CONFIG_ARM64_UAO) &&
-           get_thread_info(next_sp)->addr_limit == KERNEL_DS)
-               asm(ALTERNATIVE("nop", SET_PSTATE_UAO(1), ARM64_HAS_UAO,
-                               CONFIG_ARM64_UAO));
-       else
-               asm(ALTERNATIVE("nop", SET_PSTATE_UAO(0), ARM64_HAS_UAO,
-                               CONFIG_ARM64_UAO));
+       if (IS_ENABLED(CONFIG_ARM64_UAO)) {
+               if (task_thread_info(next)->addr_limit == KERNEL_DS)
+                       asm(ALTERNATIVE("nop", SET_PSTATE_UAO(1), ARM64_HAS_UAO));
+               else
+                       asm(ALTERNATIVE("nop", SET_PSTATE_UAO(0), ARM64_HAS_UAO));
+       }
 }
 
 /*
 
        }
 
        if (permission_fault(esr) && (addr < USER_DS)) {
-               if (get_thread_info(regs->sp)->addr_limit == KERNEL_DS)
+               if (get_fs() == KERNEL_DS)
                        panic("Accessing user space memory with fs=KERNEL_DS");
 
                if (!search_exception_tables(regs->pc))