diff --git a/mm/process_reclaim.c b/mm/process_reclaim.c
index f483f7276597..271fb38e0cb0 100644
--- a/mm/process_reclaim.c
+++ b/mm/process_reclaim.c
@@ -88,14 +88,17 @@ static int test_task_flag(struct task_struct *p, int flag)
 {
 	struct task_struct *t = p;
 
-	do {
+	rcu_read_lock();
+	for_each_thread(p, t) {
 		task_lock(t);
 		if (test_tsk_thread_flag(t, flag)) {
 			task_unlock(t);
+			rcu_read_unlock();
 			return 1;
 		}
 		task_unlock(t);
-	} while_each_thread(p, t);
+	}
+	rcu_read_unlock();
 
 	return 0;
 }
@@ -125,10 +128,6 @@ static void swap_fn(struct work_struct *work)
 		if (tsk->flags & PF_KTHREAD)
 			continue;
 
-		/* if task no longer has any memory ignore it */
-		if (test_task_flag(tsk, TIF_MM_RELEASED))
-			continue;
-
 		if (test_task_flag(tsk, TIF_MEMDIE))
 			continue;
 
@@ -165,20 +164,20 @@ static void swap_fn(struct work_struct *work)
 		}
 	}
 
-	for (i = 0; i < si; i++) {
-		get_task_struct(selected[i].p);
+	for (i = 0; i < si; i++)
 		total_sz += selected[i].tasksize;
-	}
-
-	rcu_read_unlock();
 
 	/* Skip reclaim if total size is too less */
 	if (total_sz < SWAP_CLUSTER_MAX) {
-		for (i = 0; i < si; i++)
-			put_task_struct(selected[i].p);
+		rcu_read_unlock();
 		return;
 	}
 
+	for (i = 0; i < si; i++)
+		get_task_struct(selected[i].p);
+
+	rcu_read_unlock();
+
 	while (si--) {
 		nr_to_reclaim =
 			(selected[si].tasksize * per_swap_size) / total_sz;
@@ -186,12 +185,6 @@ static void swap_fn(struct work_struct *work)
 		if (!nr_to_reclaim)
 			nr_to_reclaim = 1;
 
-		if ((test_task_flag(selected[si].p, TIF_MM_RELEASED))
-			|| (test_task_flag(selected[si].p, TIF_MEMDIE))) {
-			put_task_struct(selected[si].p);
-			continue;
-		}
-
 		rp = reclaim_task_anon(selected[si].p, nr_to_reclaim);
 
 		trace_process_reclaim(selected[si].tasksize,