if (walker->level == PT_PAGE_TABLE_LEVEL) {
                        walker->gfn = (*ptep & PT_BASE_ADDR_MASK)
                                >> PAGE_SHIFT;
+                       if (write_fault && !is_dirty_pte(*ptep)) {
+                               mark_page_dirty(vcpu->kvm, table_gfn);
+                               *ptep |= PT_DIRTY_MASK;
+                       }
                        break;
                }
 
                        walker->gfn = (*ptep & PT_DIR_BASE_ADDR_MASK)
                                >> PAGE_SHIFT;
                        walker->gfn += PT_INDEX(addr, PT_PAGE_TABLE_LEVEL);
+                       if (write_fault && !is_dirty_pte(*ptep)) {
+                               mark_page_dirty(vcpu->kvm, table_gfn);
+                               *ptep |= PT_DIRTY_MASK;
+                       }
                        break;
                }
 
        return 0;
 }
 
-static void FNAME(mark_pagetable_dirty)(struct kvm *kvm,
-                                       struct guest_walker *walker)
-{
-       mark_page_dirty(kvm, walker->table_gfn[walker->level - 1]);
-}
-
 static void FNAME(set_pte_common)(struct kvm_vcpu *vcpu,
                                  u64 *shadow_pte,
                                  gpa_t gaddr,
                 __FUNCTION__, *shadow_pte, (u64)gpte, access_bits,
                 write_fault, user_fault, gfn);
 
-       if (write_fault && !dirty) {
-               pt_element_t *guest_ent, *tmp = NULL;
-
-               if (walker->ptep)
-                       guest_ent = walker->ptep;
-               else {
-                       tmp = kmap_atomic(walker->page, KM_USER0);
-                       guest_ent = &tmp[walker->index];
-               }
-
-               *guest_ent |= PT_DIRTY_MASK;
-               if (!walker->ptep)
-                       kunmap_atomic(tmp, KM_USER0);
-               dirty = 1;
-               FNAME(mark_pagetable_dirty)(vcpu->kvm, walker);
-       }
-
        /*
         * We don't set the accessed bit, since we sometimes want to see
         * whether the guest actually used the pte (in order to detect