diff --git a/include/linux/task_work.h b/include/linux/task_work.h
index 26b8a47f41fcacf98f5cf5d427e22c6f7c6f1715..8cb9f4a95d7a23de2b934215e92249b93e4e6123 100644
--- a/include/linux/task_work.h
+++ b/include/linux/task_work.h
@@ -14,10 +14,12 @@ init_task_work(struct callback_head *twork, task_work_func_t func)
 }
 
 enum task_work_notify_mode {
-	TWA_NONE,
+	TWA_NONE = 0,
 	TWA_RESUME,
 	TWA_SIGNAL,
 	TWA_SIGNAL_NO_IPI,
+	TWA_FLAGS = 0xff00,
+	TWAF_NO_ALLOC = 0x0100,
 };
 
 static inline bool task_work_pending(struct task_struct *task)
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 462571b26f88d9ba79b5b868145364e0a0960804..0224089e1105c02aaa3fab33e9ff15bc1461cc9f 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -12678,7 +12678,9 @@ void task_tick_mm_cid(struct rq *rq, struct task_struct *curr)
 		return;
 	if (time_before(now, READ_ONCE(curr->mm->mm_cid_next_scan)))
 		return;
-	task_work_add(curr, work, TWA_RESUME);
+
+	/* No page allocation under rq lock */
+	task_work_add(curr, work, TWA_RESUME | TWAF_NO_ALLOC);
 }
 
 void sched_mm_cid_exit_signals(struct task_struct *t)
diff --git a/kernel/task_work.c b/kernel/task_work.c
index 2134ac8057a94e64b42d75607f9753e6e82526aa..45f9963aa2d43dc1b64d5752b55dc749e7811a3f 100644
--- a/kernel/task_work.c
+++ b/kernel/task_work.c
@@ -43,9 +43,18 @@ int task_work_add(struct task_struct *task, struct callback_head *work,
 		  enum task_work_notify_mode notify)
 {
 	struct callback_head *head;
+	int flags = notify & TWA_FLAGS;
+	notify &= ~TWA_FLAGS;
 
-	/* record the work call stack in order to print it in KASAN reports */
-	kasan_record_aux_stack(work);
+	/*
+	 * Record the work call stack in order to print it in KASAN reports.
+	 * Note that stack allocation can fail if TWAF_NO_ALLOC flag
+	 * is set and new page is needed to expand the stack buffer.
+	 */
+	if (flags & TWAF_NO_ALLOC)
+		kasan_record_aux_stack_noalloc(work);
+	else
+		kasan_record_aux_stack(work);
 
 	head = READ_ONCE(task->task_works);
 	do {