diff --git a/mm/memory.c b/mm/memory.c
index 676020552b32acbede4b70dc8b6b24f48c0789e6..9e6ac954c70bd1b36e976d2328a1477d6d4b5704 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -1026,6 +1026,9 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
 		if (pte_none(ptent))
 			continue;
 
+		if (need_resched())
+			break;
+
 		if (pte_present(ptent)) {
 			struct page *page;
 
@@ -1123,8 +1126,11 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
 	if (force_flush) {
 		force_flush = 0;
 		tlb_flush_mmu(tlb);
-		if (addr != end)
-			goto again;
+	}
+
+	if (addr != end) {
+		cond_resched();
+		goto again;
 	}
 
 	return addr;