diff --git a/include/linux/hugetlb.h b/include/linux/hugetlb.h
index 551834cd5299..8f85ad0d5bdb 100644
--- a/include/linux/hugetlb.h
+++ b/include/linux/hugetlb.h
@@ -12,6 +12,8 @@
 #include <linux/pgtable.h>
 #include <linux/gfp.h>
 #include <linux/userfaultfd_k.h>
+#include <linux/irqflags.h>
+#include <linux/rcupdate.h>
 
 struct ctl_table;
 struct user_struct;
@@ -24,6 +26,71 @@ typedef struct { unsigned long pd; } hugepd_t;
 #define __hugepd(x) ((hugepd_t) { (x) })
 #endif
 
+/**
+ * @hugetlb_walker_[un]lock(): protects software hugetlb pgtable walkers
+ *
+ * The hugetlb walker lock needs to be taken before huge_pte_offset() and
+ * needs to be released after finishing using the pte_t* returned.  It's
+ * used to guarantee we won't access a freed pgtable page.  Normally in
+ * this case we already have the mmap lock held for read.
+ *
+ * When holding the hugetlb walker lock, the thread cannot sleep, nor can
+ * it already in irq context (just to simplify unlock with no-save for
+ * !RCU_TABLE_TREE).  Before the thread yields itself, it should release
+ * the pgtable lock.  After the lock released, pte_t* can become invalid
+ * anytime so it cannot be accessed anymore.
+ *
+ * The only reason for the lock is to protect concurrent pmd unsharing
+ * followed by e.g. munmap() to drop the pgtable page.  For no-pmd-sharing
+ * archs, it's no-op because it's always safe to access hugetlb pgtables
+ * just like generic pgtables.
+ */
+#if !defined(CONFIG_ARCH_WANT_HUGE_PMD_SHARE)
+static inline void hugetlb_walker_lock(void)
+{
+}
+static inline void hugetlb_walker_unlock(void)
+{
+}
+static inline bool __hugetlb_walker_locked(void)
+{
+	return true;
+}
+#elif defined(CONFIG_MMU_GATHER_RCU_TABLE_FREE)
+static inline void hugetlb_walker_lock(void)
+{
+	rcu_read_lock();
+}
+static inline void hugetlb_walker_unlock(void)
+{
+	rcu_read_unlock();
+}
+static inline bool __hugetlb_walker_locked(void)
+{
+	return rcu_read_lock_held();
+}
+#else
+static inline void hugetlb_walker_lock(void)
+{
+	WARN_ON_ONCE(irqs_disabled());
+	local_irq_disable();
+}
+static inline void hugetlb_walker_unlock(void)
+{
+	local_irq_enable();
+}
+static inline bool __hugetlb_walker_locked(void)
+{
+	return irqs_disabled();
+}
+#endif
+
+#ifdef CONFIG_LOCKDEP
+#define  hugetlb_walker_locked()  __hugetlb_walker_locked()
+#else
+#define  hugetlb_walker_locked()  true
+#endif
+
 #ifdef CONFIG_HUGETLB_PAGE
 
 #include <linux/mempolicy.h>
@@ -192,6 +259,38 @@ extern struct list_head huge_boot_pages;
 
 pte_t *huge_pte_alloc(struct mm_struct *mm, struct vm_area_struct *vma,
 			unsigned long addr, unsigned long sz);
+/*
+ * huge_pte_offset(): Walk the hugetlb pgtable until the last level PTE.
+ * Returns the pte_t* if found, or NULL if the address is not mapped.
+ *
+ * NOTE: since this function will walk all the pgtable pages (including not
+ * only high-level pgtable page, but also PUD entry that can be unshared
+ * concurrently for VM_SHARED), the caller of this function should be
+ * responsible of its thread safety.  One can follow this rule:
+ *
+ *  (1) For private mappings: pmd unsharing is not possible, so it'll
+ *      always be safe if we're with the mmap sem for either read or write.
+ *      This is normally always the case, IOW we don't need to do anything
+ *      special.
+ *
+ *  (2) For shared mappings: pmd unsharing is possible (so the PUD-ranged
+ *      pgtable page can go away from under us!  It can be done by a pmd
+ *      unshare with a follow up munmap() on the other process), then we
+ *      need either:
+ *
+ *     (2.1) hugetlb vma lock read or write held, to make sure pmd unshare
+ *           won't happen upon the range (it also makes sure the pte_t we
+ *           read is the right and stable one), or,
+ *
+ *     (2.2) pgtable walker lock, to make sure even pmd unsharing happened,
+ *           the old shared PUD page won't get freed from under us, so even
+ *           the pteval can be obsolete, at least it's still always safe to
+ *           access the pgtable page (e.g., de-referencing pte_t* would not
+ *           cause use-after-free).
+ *
+ * PS: from the regard of (2.2), it's the same logic of fast-gup being safe
+ * for generic mm when walking the pgtables even without mmap lock.
+ */
 pte_t *huge_pte_offset(struct mm_struct *mm,
 		       unsigned long addr, unsigned long sz);
 unsigned long hugetlb_mask_last_page(struct hstate *h);
