diff --git a/include/linux/pagewalk.h b/include/linux/pagewalk.h
index ed30c86195b0..bddd9759bab9 100644
--- a/include/linux/pagewalk.h
+++ b/include/linux/pagewalk.h
@@ -8,15 +8,15 @@ struct mm_walk;
 
 /**
  * mm_walk_ops - callbacks for walk_page_range
- * @pgd_entry:		if set, called for each non-empty PGD (top-level) entry
- * @p4d_entry:		if set, called for each non-empty P4D entry
- * @pud_entry:		if set, called for each non-empty PUD entry
- * @pmd_entry:		if set, called for each non-empty PMD entry
+ * @pud_entry:		if set, called for each non-empty PUD (2nd-level) entry
+ *			this handler should only handle pud_trans_huge() puds.
+ *			the pmd_entry or pte_entry callbacks will be used for
+ *			regular PUDs.
+ * @pmd_entry:		if set, called for each non-empty PMD (3rd-level) entry
  *			this handler is required to be able to handle
  *			pmd_trans_huge() pmds.  They may simply choose to
  *			split_huge_page() instead of handling it explicitly.
- * @pte_entry:		if set, called for each non-empty PTE (lowest-level)
- *			entry
+ * @pte_entry:		if set, called for each non-empty PTE (4th-level) entry
  * @pte_hole:		if set, called for each hole at all levels
  * @hugetlb_entry:	if set, called for each hugetlb entry
  * @test_walk:		caller specific callback function to determine whether
@@ -24,15 +24,8 @@ struct mm_walk;
  *			"do page table walk over the current vma", returning
  *			a negative value means "abort current page table walk
  *			right now" and returning 1 means "skip the current vma"
- *
- * p?d_entry callbacks are called even if those levels are folded on a
- * particular architecture/configuration.
  */
 struct mm_walk_ops {
-	int (*pgd_entry)(pgd_t *pgd, unsigned long addr,
-			 unsigned long next, struct mm_walk *walk);
-	int (*p4d_entry)(p4d_t *p4d, unsigned long addr,
-			 unsigned long next, struct mm_walk *walk);
 	int (*pud_entry)(pud_t *pud, unsigned long addr,
 			 unsigned long next, struct mm_walk *walk);
 	int (*pmd_entry)(pmd_t *pmd, unsigned long addr,
@@ -48,25 +41,11 @@ struct mm_walk_ops {
 			struct mm_walk *walk);
 };
 
-/*
- * Action for pud_entry / pmd_entry callbacks.
- * ACTION_SUBTREE is the default
- */
-enum page_walk_action {
-	/* Descend to next level, splitting huge pages if needed and possible */
-	ACTION_SUBTREE = 0,
-	/* Continue to next entry at this level (ignoring any subtree) */
-	ACTION_CONTINUE = 1,
-	/* Call again for this entry */
-	ACTION_AGAIN = 2
-};
-
 /**
  * mm_walk - walk_page_range data
  * @ops:	operation to call during the walk
  * @mm:		mm_struct representing the target process of page table walk
  * @vma:	vma currently walked (NULL if walking outside vmas)
- * @action:	next action to perform (see enum page_walk_action)
  * @private:	private data for callbacks' usage
  *
  * (see the comment on walk_page_range() for more details)
@@ -75,7 +54,6 @@ struct mm_walk {
 	const struct mm_walk_ops *ops;
 	struct mm_struct *mm;
 	struct vm_area_struct *vma;
-	enum page_walk_action action;
 	void *private;
 };
 
diff --git a/mm/hmm.c b/mm/hmm.c
index 231de978d647..902f5fa6bf93 100644
--- a/mm/hmm.c
+++ b/mm/hmm.c
@@ -658,32 +658,23 @@ static int hmm_vma_walk_pud(pud_t *pudp, unsigned long start, unsigned long end,
 {
 	struct hmm_vma_walk *hmm_vma_walk = walk->private;
 	struct hmm_range *range = hmm_vma_walk->range;
-	unsigned long addr = start;
+	unsigned long addr = start, next;
+	pmd_t *pmdp;
 	pud_t pud;
-	int ret = 0;
-	spinlock_t *ptl = pud_trans_huge_lock(pudp, walk->vma);
-
-	if (!ptl)
-		return 0;
-
-	/* Normally we don't want to split the huge page */
-	walk->action = ACTION_CONTINUE;
+	int ret;
 
+again:
 	pud = READ_ONCE(*pudp);
-	if (pud_none(pud)) {
-		ret = hmm_vma_walk_hole(start, end, walk);
-		goto out_unlock;
-	}
+	if (pud_none(pud))
+		return hmm_vma_walk_hole(start, end, walk);
 
 	if (pud_huge(pud) && pud_devmap(pud)) {
 		unsigned long i, npages, pfn;
 		uint64_t *pfns, cpu_flags;
 		bool fault, write_fault;
 
-		if (!pud_present(pud)) {
-			ret = hmm_vma_walk_hole(start, end, walk);
-			goto out_unlock;
-		}
+		if (!pud_present(pud))
+			return hmm_vma_walk_hole(start, end, walk);
 
 		i = (addr - range->start) >> PAGE_SHIFT;
 		npages = (end - addr) >> PAGE_SHIFT;
@@ -692,20 +683,16 @@ static int hmm_vma_walk_pud(pud_t *pudp, unsigned long start, unsigned long end,
 		cpu_flags = pud_to_hmm_pfn_flags(range, pud);
 		hmm_range_need_fault(hmm_vma_walk, pfns, npages,
 				     cpu_flags, &fault, &write_fault);
-		if (fault || write_fault) {
-			ret = hmm_vma_walk_hole_(addr, end, fault,
-						 write_fault, walk);
-			goto out_unlock;
-		}
+		if (fault || write_fault)
+			return hmm_vma_walk_hole_(addr, end, fault,
+						write_fault, walk);
 
 		pfn = pud_pfn(pud) + ((addr & ~PUD_MASK) >> PAGE_SHIFT);
 		for (i = 0; i < npages; ++i, ++pfn) {
 			hmm_vma_walk->pgmap = get_dev_pagemap(pfn,
 					      hmm_vma_walk->pgmap);
-			if (unlikely(!hmm_vma_walk->pgmap)) {
-				ret = -EBUSY;
-				goto out_unlock;
-			}
+			if (unlikely(!hmm_vma_walk->pgmap))
+				return -EBUSY;
 			pfns[i] = hmm_device_entry_from_pfn(range, pfn) |
 				  cpu_flags;
 		}
@@ -714,15 +701,22 @@ static int hmm_vma_walk_pud(pud_t *pudp, unsigned long start, unsigned long end,
 			hmm_vma_walk->pgmap = NULL;
 		}
 		hmm_vma_walk->last = end;
-		goto out_unlock;
+		return 0;
 	}
 
-	/* Ask for the PUD to be split */
-	walk->action = ACTION_SUBTREE;
+	split_huge_pud(walk->vma, pudp, addr);
+	if (pud_none(*pudp))
+		goto again;
 
-out_unlock:
-	spin_unlock(ptl);
-	return ret;
+	pmdp = pmd_offset(pudp, addr);
+	do {
+		next = pmd_addr_end(addr, end);
+		ret = hmm_vma_walk_pmd(pmdp, addr, next, walk);
+		if (ret)
+			return ret;
+	} while (pmdp++, addr = next, addr != end);
+
+	return 0;
 }
 #else
 #define hmm_vma_walk_pud	NULL
diff --git a/mm/pagewalk.c b/mm/pagewalk.c
index 0d4c6159a4ea..ec41e7552f37 100644
--- a/mm/pagewalk.c
+++ b/mm/pagewalk.c
@@ -45,9 +45,6 @@ again:
 				break;
 			continue;
 		}
-
-		walk->action = ACTION_SUBTREE;
-
 		/*
 		 * This implies that each ->pmd_entry() handler
 		 * needs to know about pmd_trans_huge() pmds
@@ -57,21 +54,16 @@ again:
 		if (err)
 			break;
 
-		if (walk->action == ACTION_AGAIN)
-			goto again;
-
 		/*
 		 * Check this here so we only break down trans_huge
 		 * pages when we _need_ to
 		 */
-		if (walk->action == ACTION_CONTINUE ||
-		    !(ops->pte_entry))
+		if (!ops->pte_entry)
 			continue;
 
 		split_huge_pmd(walk->vma, pmd, addr);
 		if (pmd_trans_unstable(pmd))
 			goto again;
-
 		err = walk_pte_range(pmd, addr, next, walk);
 		if (err)
 			break;
@@ -100,25 +92,24 @@ static int walk_pud_range(p4d_t *p4d, unsigned long addr, unsigned long end,
 			continue;
 		}
 
-		walk->action = ACTION_SUBTREE;
+		if (ops->pud_entry) {
+			spinlock_t *ptl = pud_trans_huge_lock(pud, walk->vma);
 
-		if (ops->pud_entry)
-			err = ops->pud_entry(pud, addr, next, walk);
-		if (err)
-			break;
-
-		if (walk->action == ACTION_AGAIN)
-			goto again;
-
-		if (walk->action == ACTION_CONTINUE ||
-		    !(ops->pmd_entry || ops->pte_entry))
-			continue;
+			if (ptl) {
+				err = ops->pud_entry(pud, addr, next, walk);
+				spin_unlock(ptl);
+				if (err)
+					break;
+				continue;
+			}
+		}
 
 		split_huge_pud(walk->vma, pud, addr);
 		if (pud_none(*pud))
 			goto again;
 
-		err = walk_pmd_range(pud, addr, next, walk);
+		if (ops->pmd_entry || ops->pte_entry)
+			err = walk_pmd_range(pud, addr, next, walk);
 		if (err)
 			break;
 	} while (pud++, addr = next, addr != end);
@@ -144,12 +135,7 @@ static int walk_p4d_range(pgd_t *pgd, unsigned long addr, unsigned long end,
 				break;
 			continue;
 		}
-		if (ops->p4d_entry) {
-			err = ops->p4d_entry(p4d, addr, next, walk);
-			if (err)
-				break;
-		}
-		if (ops->pud_entry || ops->pmd_entry || ops->pte_entry)
+		if (ops->pmd_entry || ops->pte_entry)
 			err = walk_pud_range(p4d, addr, next, walk);
 		if (err)
 			break;
@@ -176,13 +162,7 @@ static int walk_pgd_range(unsigned long addr, unsigned long end,
 				break;
 			continue;
 		}
-		if (ops->pgd_entry) {
-			err = ops->pgd_entry(pgd, addr, next, walk);
-			if (err)
-				break;
-		}
-		if (ops->p4d_entry || ops->pud_entry || ops->pmd_entry ||
-		    ops->pte_entry)
+		if (ops->pmd_entry || ops->pte_entry)
 			err = walk_p4d_range(pgd, addr, next, walk);
 		if (err)
 			break;