#else /* __s390x__ */
-# define TASK_SIZE (test_thread_flag(TIF_31BIT) ? \
+# define TASK_SIZE_OF(tsk) (test_tsk_thread_flag(tsk, TIF_31BIT) ? \
(0x80000000UL) : (0x40000000000UL))
+# define TASK_SIZE TASK_SIZE_OF(current)
# define TASK_UNMAPPED_BASE (TASK_SIZE / 2)
# define DEFAULT_TASK_SIZE (0x40000000000UL)