*pfn = phys_to_pfn_t(phys, PFN_DEV|PFN_MAP);
 
+       dax_set_mapping(vmf, *pfn, fault_size);
+
        return vmf_insert_mixed(vmf->vma, vmf->address, *pfn);
 }
 
 
        *pfn = phys_to_pfn_t(phys, PFN_DEV|PFN_MAP);
 
+       dax_set_mapping(vmf, *pfn, fault_size);
+
        return vmf_insert_pfn_pmd(vmf, *pfn, vmf->flags & FAULT_FLAG_WRITE);
 }
 
 
        *pfn = phys_to_pfn_t(phys, PFN_DEV|PFN_MAP);
 
+       dax_set_mapping(vmf, *pfn, fault_size);
+
        return vmf_insert_pfn_pud(vmf, *pfn, vmf->flags & FAULT_FLAG_WRITE);
 }
 #else
                enum page_entry_size pe_size)
 {
        struct file *filp = vmf->vma->vm_file;
-       unsigned long fault_size;
        vm_fault_t rc = VM_FAULT_SIGBUS;
        int id;
        pfn_t pfn;
        id = dax_read_lock();
        switch (pe_size) {
        case PE_SIZE_PTE:
-               fault_size = PAGE_SIZE;
                rc = __dev_dax_pte_fault(dev_dax, vmf, &pfn);
                break;
        case PE_SIZE_PMD:
-               fault_size = PMD_SIZE;
                rc = __dev_dax_pmd_fault(dev_dax, vmf, &pfn);
                break;
        case PE_SIZE_PUD:
-               fault_size = PUD_SIZE;
                rc = __dev_dax_pud_fault(dev_dax, vmf, &pfn);
                break;
        default:
                rc = VM_FAULT_SIGBUS;
        }
 
-       if (rc == VM_FAULT_NOPAGE)
-               dax_set_mapping(vmf, pfn, fault_size);
        dax_read_unlock(id);
 
        return rc;