diff --git a/arch/x86/kvm/mmu.c b/arch/x86/kvm/mmu.c
index 263718d49f730..59d1866398c42 100644
--- a/arch/x86/kvm/mmu.c
+++ b/arch/x86/kvm/mmu.c
@@ -1948,6 +1948,461 @@ static void handle_changed_pte(struct kvm *kvm, int as_id, gfn_t gfn,
 	}
 }
 
+/*
+ * Given a host page table entry and its level, returns a pointer containing
+ * the host virtual address of the child page table referenced by the page table
+ * entry. Returns null if there is no such entry.
+ */
+static u64 *pte_to_child_pt(u64 pte, int level)
+{
+	u64 *pt;
+	/* There's no child entry if this entry isn't present */
+	if (!is_present_direct_pte(pte))
+		return NULL;
+
+	/* There is no child page table if this is a leaf entry. */
+	if (is_last_spte(pte, level))
+		return NULL;
+
+	pt = (u64 *)__va(pte & PT64_BASE_ADDR_MASK);
+	return pt;
+}
+
+enum mmu_lock_mode {
+	MMU_NO_LOCK = 0,
+	MMU_READ_LOCK = 1,
+	MMU_WRITE_LOCK = 2,
+	MMU_LOCK_MAY_RESCHED = 4
+};
+
+/*
+ * A direct walk iterator encapsulates a walk through a direct paging structure.
+ * It handles ensuring that the walk uses RCU to safely access page table
+ * memory.
+ */
+struct direct_walk_iterator {
+	/* Internal */
+	gfn_t walk_start;
+	gfn_t walk_end;
+	gfn_t target_gfn;
+	long tlbs_dirty;
+
+	/* the address space id. */
+	int as_id;
+	u64 *pt_path[PT64_ROOT_4LEVEL];
+	bool walk_in_progress;
+
+	/*
+	 * If set, the next call to direct_walk_iterator_next_pte_raw will
+	 * simply reread the current pte and return. This is useful in cases
+	 * where a thread misses a race to set a pte and wants to retry. This
+	 * should be set with a call to direct_walk_iterator_retry_pte.
+	 */
+	bool retry_pte;
+
+	/*
+	 * If set, the next call to direct_walk_iterator_next_pte_raw will not
+	 * step down to a lower level on its next step, even if it is at a
+	 * present, non-leaf pte. This is useful when, for example, splitting
+	 * pages, since we know that the entries below the now split page don't
+	 * need to be handled again.
+	 */
+	bool skip_step_down;
+
+	enum mmu_lock_mode lock_mode;
+	struct kvm *kvm;
+
+	/* Output */
+
+	/* The iterator's current level within the paging structure */
+	int level;
+	/* A pointer to the current PTE */
+	u64 *ptep;
+	/* The a snapshot of the PTE pointed to by ptep */
+	u64 old_pte;
+	/* The lowest GFN mapped by the current PTE */
+	gfn_t pte_gfn_start;
+	/* The highest GFN mapped by the current PTE, + 1 */
+	gfn_t pte_gfn_end;
+};
+
+static void direct_walk_iterator_start_traversal(
+		struct direct_walk_iterator *iter)
+{
+	int level;
+
+	/*
+	 * Only clear the levels below the root. The root level page table is
+	 * allocated at VM creation time and will never change for the life of
+	 * the VM.
+	 */
+	for (level = PT_PAGE_TABLE_LEVEL; level < PT64_ROOT_4LEVEL; level++)
+		iter->pt_path[level - 1] = NULL;
+	iter->level = 0;
+	iter->ptep = NULL;
+	iter->old_pte = 0;
+	iter->pte_gfn_start = 0;
+	iter->pte_gfn_end = 0;
+	iter->walk_in_progress = false;
+	iter->retry_pte = false;
+	iter->skip_step_down = false;
+}
+
+static bool direct_walk_iterator_flush_needed(struct direct_walk_iterator *iter)
+{
+	long tlbs_dirty;
+
+	if (iter->tlbs_dirty) {
+		tlbs_dirty = xadd(&iter->kvm->tlbs_dirty, iter->tlbs_dirty) +
+				iter->tlbs_dirty;
+		iter->tlbs_dirty = 0;
+	} else {
+		tlbs_dirty = READ_ONCE(iter->kvm->tlbs_dirty);
+	}
+
+	return (iter->lock_mode & MMU_WRITE_LOCK) && tlbs_dirty;
+}
+
+static bool direct_walk_iterator_end_traversal(
+		struct direct_walk_iterator *iter)
+{
+	if (iter->walk_in_progress)
+		rcu_read_unlock();
+	return direct_walk_iterator_flush_needed(iter);
+}
+
+/*
+ * Resets a direct walk iterator to the root of the paging structure and RCU
+ * unlocks. After calling this function, the traversal can be reattempted.
+ */
+static void direct_walk_iterator_reset_traversal(
+		struct direct_walk_iterator *iter)
+{
+	/*
+	 * It's okay it ignore the return value, indicating whether a TLB flush
+	 * is needed here because we are ending and then restarting the
+	 * traversal without releasing the MMU lock. At this point the
+	 * iterator tlbs_dirty will have been flushed to the kvm tlbs_dirty, so
+	 * the next end_traversal will return that a flush is needed, if there's
+	 * not an intervening flush for some other reason.
+	 */
+	direct_walk_iterator_end_traversal(iter);
+	direct_walk_iterator_start_traversal(iter);
+}
+
+/*
+ * Sets a direct walk iterator to seek the gfn range [start, end).
+ * If end is greater than the maximum possible GFN, it will be changed to the
+ * maximum possible gfn + 1. (Note that start/end is and inclusive/exclusive
+ * range, so the last gfn to be interated over would be the largest possible
+ * GFN, in this scenario.)
+ */
+__attribute__((unused))
+static void direct_walk_iterator_setup_walk(struct direct_walk_iterator *iter,
+	struct kvm *kvm, int as_id, gfn_t start, gfn_t end,
+	enum mmu_lock_mode lock_mode)
+{
+	BUG_ON(!kvm->arch.direct_mmu_enabled);
+	BUG_ON((lock_mode & MMU_WRITE_LOCK) && (lock_mode & MMU_READ_LOCK));
+	BUG_ON(as_id < 0);
+	BUG_ON(as_id >= KVM_ADDRESS_SPACE_NUM);
+	BUG_ON(!VALID_PAGE(kvm->arch.direct_root_hpa[as_id]));
+
+	/* End cannot be greater than the maximum possible gfn. */
+	end = min(end, 1ULL << (PT64_ROOT_4LEVEL * PT64_PT_BITS));
+
+	iter->as_id = as_id;
+	iter->pt_path[PT64_ROOT_4LEVEL - 1] =
+			(u64 *)__va(kvm->arch.direct_root_hpa[as_id]);
+
+	iter->walk_start = start;
+	iter->walk_end = end;
+	iter->target_gfn = start;
+
+	iter->lock_mode = lock_mode;
+	iter->kvm = kvm;
+	iter->tlbs_dirty = 0;
+
+	direct_walk_iterator_start_traversal(iter);
+}
+
+__attribute__((unused))
+static void direct_walk_iterator_retry_pte(struct direct_walk_iterator *iter)
+{
+	BUG_ON(!iter->walk_in_progress);
+	iter->retry_pte = true;
+}
+
+__attribute__((unused))
+static void direct_walk_iterator_skip_step_down(
+		struct direct_walk_iterator *iter)
+{
+	BUG_ON(!iter->walk_in_progress);
+	iter->skip_step_down = true;
+}
+
+/*
+ * Steps down one level in the paging structure towards the previously set
+ * target gfn. Returns true if the iterator was able to step down a level,
+ * false otherwise.
+ */
+static bool direct_walk_iterator_try_step_down(
+		struct direct_walk_iterator *iter)
+{
+	u64 *child_pt;
+
+	/*
+	 * Reread the pte before stepping down to avoid traversing into page
+	 * tables that are no longer linked from this entry. This is not
+	 * needed for correctness - just a small optimization.
+	 */
+	iter->old_pte = READ_ONCE(*iter->ptep);
+
+	child_pt = pte_to_child_pt(iter->old_pte, iter->level);
+	if (child_pt == NULL)
+		return false;
+	child_pt = rcu_dereference(child_pt);
+
+	iter->level--;
+	iter->pt_path[iter->level - 1] = child_pt;
+	return true;
+}
+
+/*
+ * Steps to the next entry in the current page table, at the current page table
+ * level. The next entry could map a page of guest memory, another page table,
+ * or it could be non-present or invalid. Returns true if the iterator was able
+ * to step to the next entry in the page table, false otherwise.
+ */
+static bool direct_walk_iterator_try_step_side(
+		struct direct_walk_iterator *iter)
+{
+	/*
+	 * If the current gfn maps past the target gfn range, the next entry in
+	 * the current page table will be outside the target range.
+	 */
+	if (iter->pte_gfn_end >= iter->walk_end)
+		return false;
+
+	/*
+	 * Check if the iterator is already at the end of the current page
+	 * table.
+	 */
+	if (!(iter->pte_gfn_end % KVM_PAGES_PER_HPAGE(iter->level + 1)))
+		return false;
+
+	iter->target_gfn = iter->pte_gfn_end;
+	return true;
+}
+
+/*
+ * Tries to back up a level in the paging structure so that the walk can
+ * continue from the next entry in the parent page table. Returns true on a
+ * successful step up, false otherwise.
+ */
+static bool direct_walk_iterator_try_step_up(struct direct_walk_iterator *iter)
+{
+	if (iter->level == PT64_ROOT_4LEVEL)
+		return false;
+
+	iter->level++;
+	return true;
+}
+
+/*
+ * Step to the next pte in a pre-order traversal of the target gfn range.
+ * To get to the next pte, the iterator either steps down towards the current
+ * target gfn, if at a present, non-leaf pte, or over to a pte mapping a
+ * highter gfn, if there's room in the gfn range. If there is no step within
+ * the target gfn range, returns false.
+ */
+static bool direct_walk_iterator_next_pte_raw(struct direct_walk_iterator *iter)
+{
+	bool retry_pte = iter->retry_pte;
+	bool skip_step_down = iter->skip_step_down;
+
+	iter->retry_pte = false;
+	iter->skip_step_down = false;
+
+	if (iter->target_gfn >= iter->walk_end)
+		return false;
+
+	/* If the walk is just starting, set up initial values. */
+	if (!iter->walk_in_progress) {
+		rcu_read_lock();
+
+		iter->level = PT64_ROOT_4LEVEL;
+		iter->walk_in_progress = true;
+		return true;
+	}
+
+	if (retry_pte)
+		return true;
+
+	if (!skip_step_down && direct_walk_iterator_try_step_down(iter))
+		return true;
+
+	while (!direct_walk_iterator_try_step_side(iter))
+		if (!direct_walk_iterator_try_step_up(iter))
+			return false;
+	return true;
+}
+
+static void direct_walk_iterator_recalculate_output_fields(
+		struct direct_walk_iterator *iter)
+{
+	iter->ptep = iter->pt_path[iter->level - 1] +
+			PT64_INDEX(iter->target_gfn << PAGE_SHIFT, iter->level);
+	iter->old_pte = READ_ONCE(*iter->ptep);
+	iter->pte_gfn_start = ALIGN_DOWN(iter->target_gfn,
+			KVM_PAGES_PER_HPAGE(iter->level));
+	iter->pte_gfn_end = iter->pte_gfn_start +
+			KVM_PAGES_PER_HPAGE(iter->level);
+}
+
+static void direct_walk_iterator_prepare_cond_resched(
+		struct direct_walk_iterator *iter)
+{
+	if (direct_walk_iterator_end_traversal(iter))
+		kvm_flush_remote_tlbs(iter->kvm);
+
+	if (iter->lock_mode & MMU_WRITE_LOCK)
+		write_unlock(&iter->kvm->mmu_lock);
+	else if (iter->lock_mode & MMU_READ_LOCK)
+		read_unlock(&iter->kvm->mmu_lock);
+
+}
+
+static void direct_walk_iterator_finish_cond_resched(
+		struct direct_walk_iterator *iter)
+{
+	if (iter->lock_mode & MMU_WRITE_LOCK)
+		write_lock(&iter->kvm->mmu_lock);
+	else if (iter->lock_mode & MMU_READ_LOCK)
+		read_lock(&iter->kvm->mmu_lock);
+
+	direct_walk_iterator_start_traversal(iter);
+}
+
+static void direct_walk_iterator_cond_resched(struct direct_walk_iterator *iter)
+{
+	if (!(iter->lock_mode & MMU_LOCK_MAY_RESCHED) || !need_resched())
+		return;
+
+	direct_walk_iterator_prepare_cond_resched(iter);
+	cond_resched();
+	direct_walk_iterator_finish_cond_resched(iter);
+}
+
+static bool direct_walk_iterator_next_pte(struct direct_walk_iterator *iter)
+{
+	/*
+	 * This iterator could be iterating over a large number of PTEs, such
+	 * that if this thread did not yield, it would cause scheduler\
+	 * problems. To avoid this, yield if needed. Note the check on
+	 * MMU_LOCK_MAY_RESCHED in direct_walk_iterator_cond_resched. This
+	 * iterator will not yield unless that flag is set in its lock_mode.
+	 */
+	direct_walk_iterator_cond_resched(iter);
+
+	while (true) {
+		if (!direct_walk_iterator_next_pte_raw(iter))
+			return false;
+
+		direct_walk_iterator_recalculate_output_fields(iter);
+		if (iter->old_pte != DISCONNECTED_PTE)
+			break;
+
+		/*
+		 * The iterator has encountered a disconnected pte, so it is in
+		 * a page that has been disconnected from the root. Restart the
+		 * traversal from the root in this case.
+		 */
+		direct_walk_iterator_reset_traversal(iter);
+	}
+
+	trace_kvm_mmu_direct_walk_iterator_step(iter->walk_start,
+			iter->walk_end, iter->pte_gfn_start,
+			iter->level, iter->old_pte);
+
+	return true;
+}
+
+/*
+ * As direct_walk_iterator_next_pte but skips over non-present ptes.
+ * (i.e. ptes that are 0 or invalidated.)
+ */
+static bool direct_walk_iterator_next_present_pte(
+		struct direct_walk_iterator *iter)
+{
+	while (direct_walk_iterator_next_pte(iter))
+		if (is_present_direct_pte(iter->old_pte))
+			return true;
+
+	return false;
+}
+
+/*
+ * As direct_walk_iterator_next_present_pte but skips over non-leaf ptes.
+ */
+__attribute__((unused))
+static bool direct_walk_iterator_next_present_leaf_pte(
+		struct direct_walk_iterator *iter)
+{
+	while (direct_walk_iterator_next_present_pte(iter))
+		if (is_last_spte(iter->old_pte, iter->level))
+			return true;
+
+	return false;
+}
+
+/*
+ * Performs an atomic compare / exchange of ptes.
+ * Returns true if the pte was successfully set to the new value, false if the
+ * there was a race and the compare exchange needs to be retried.
+ */
+static bool cmpxchg_pte(u64 *ptep, u64 old_pte, u64 new_pte, int level, u64 gfn)
+{
+	u64 r;
+
+	r = cmpxchg64(ptep, old_pte, new_pte);
+	if (r == old_pte)
+		trace_kvm_mmu_set_pte_atomic(gfn, level, old_pte, new_pte);
+
+	return r == old_pte;
+}
+
+__attribute__((unused))
+static bool direct_walk_iterator_set_pte(struct direct_walk_iterator *iter,
+					 u64 new_pte)
+{
+	bool r;
+
+	if (!(iter->lock_mode & (MMU_READ_LOCK | MMU_WRITE_LOCK))) {
+		BUG_ON(is_present_direct_pte(iter->old_pte) !=
+				is_present_direct_pte(new_pte));
+		BUG_ON(spte_to_pfn(iter->old_pte) != spte_to_pfn(new_pte));
+		BUG_ON(is_last_spte(iter->old_pte, iter->level) !=
+				is_last_spte(new_pte, iter->level));
+	}
+
+	if (iter->old_pte == new_pte)
+		return true;
+
+	r = cmpxchg_pte(iter->ptep, iter->old_pte, new_pte, iter->level,
+			iter->pte_gfn_start);
+	if (r) {
+		handle_changed_pte(iter->kvm, iter->as_id, iter->pte_gfn_start,
+				   iter->old_pte, new_pte, iter->level, false);
+
+		if (iter->lock_mode & (MMU_WRITE_LOCK | MMU_READ_LOCK))
+			iter->tlbs_dirty++;
+	} else
+		direct_walk_iterator_retry_pte(iter);
+
+	return r;
+}
+
 /**
  * kvm_mmu_write_protect_pt_masked - write protect selected PT level pages
  * @kvm: kvm instance
diff --git a/arch/x86/kvm/mmutrace.h b/arch/x86/kvm/mmutrace.h
index 7ca8831c7d1a2..530723038296a 100644
--- a/arch/x86/kvm/mmutrace.h
+++ b/arch/x86/kvm/mmutrace.h
@@ -166,6 +166,56 @@ TRACE_EVENT(
 		  __entry->created ? "new" : "existing")
 );
 
+TRACE_EVENT(
+	kvm_mmu_direct_walk_iterator_step,
+	TP_PROTO(u64 walk_start, u64 walk_end, u64 base_gfn, int level,
+		u64 pte),
+	TP_ARGS(walk_start, walk_end, base_gfn, level, pte),
+
+	TP_STRUCT__entry(
+		__field(u64, walk_start)
+		__field(u64, walk_end)
+		__field(u64, base_gfn)
+		__field(int, level)
+		__field(u64, pte)
+		),
+
+	TP_fast_assign(
+		__entry->walk_start = walk_start;
+		__entry->walk_end = walk_end;
+		__entry->base_gfn = base_gfn;
+		__entry->level = level;
+		__entry->pte = pte;
+		),
+
+	TP_printk("walk_start=%llx walk_end=%llx base_gfn=%llx lvl=%d pte=%llx",
+		__entry->walk_start, __entry->walk_end, __entry->base_gfn,
+		__entry->level, __entry->pte)
+);
+
+TRACE_EVENT(
+	kvm_mmu_set_pte_atomic,
+	TP_PROTO(u64 gfn, int level, u64 old_pte, u64 new_pte),
+	TP_ARGS(gfn, level, old_pte, new_pte),
+
+	TP_STRUCT__entry(
+		__field(u64, gfn)
+		__field(int, level)
+		__field(u64, old_pte)
+		__field(u64, new_pte)
+		),
+
+	TP_fast_assign(
+		__entry->gfn = gfn;
+		__entry->level = level;
+		__entry->old_pte = old_pte;
+		__entry->new_pte = new_pte;
+		),
+
+	TP_printk("gfn=%llx level=%d old_pte=%llx new_pte=%llx", __entry->gfn,
+		  __entry->level, __entry->old_pte, __entry->new_pte)
+);
+
 DECLARE_EVENT_CLASS(kvm_mmu_page_class,
 
 	TP_PROTO(struct kvm_mmu_page *sp),
