diff --git a/lib/test_hmm.c b/lib/test_hmm.c
index 056f2e4..bd8cd29 100644
--- a/lib/test_hmm.c
+++ b/lib/test_hmm.c
@@ -979,6 +979,8 @@ static int dmirror_migrate_to_device(struct dmirror *dmirror,
 
 	mmap_read_lock(mm);
 	for (addr = start; addr < end; addr = next) {
+		int i, retried = 0;
+
 		vma = vma_lookup(mm, addr);
 		if (!vma || !(vma->vm_flags & VM_READ)) {
 			ret = -EINVAL;
@@ -987,7 +989,7 @@ static int dmirror_migrate_to_device(struct dmirror *dmirror,
 		next = min(end, addr + (ARRAY_SIZE(src_pfns) << PAGE_SHIFT));
 		if (next > vma->vm_end)
 			next = vma->vm_end;
-
+retry:
 		args.vma = vma;
 		args.src = src_pfns;
 		args.dst = dst_pfns;
@@ -1004,6 +1006,16 @@ static int dmirror_migrate_to_device(struct dmirror *dmirror,
 		migrate_vma_pages(&args);
 		dmirror_migrate_finalize_and_map(&args, dmirror);
 		migrate_vma_finalize(&args);
+
+		for (i = 0; i < ((next - addr) >> PAGE_SHIFT); i++) {
+			if (!(src_pfns[i] & MIGRATE_PFN_MIGRATE)
+			    && migrate_pfn_to_page(src_pfns[i])
+			    && retried++ < 3) {
+				wait_on_page_writeback(
+					migrate_pfn_to_page(src_pfns[i]));
+				goto retry;
+			}
+		}
 	}
 	mmap_read_unlock(mm);
 	mmput(mm);
@@ -1404,6 +1416,10 @@ static void dmirror_devmem_free(struct page *page)
 	if (rpage != page)
 		__free_page(rpage);
 
+	/* Page has been freed so reinitialize these fields */
+	ClearPageDirty(page);
+	folio_clear_swapbacked(page_folio(page));
+
 	mdevice = dmirror_page_to_device(page);
 	spin_lock(&mdevice->lock);
 
@@ -1459,9 +1475,18 @@ static vm_fault_t dmirror_devmem_fault(struct vm_fault *vmf)
 	return 0;
 }
 
+static int dmirror_devmem_pagecache(struct page *page, struct page *newpage)
+{
+	set_page_dirty(newpage);
+	copy_highpage(newpage, BACKING_PAGE(page));
+
+	return 0;
+}
+
 static const struct dev_pagemap_ops dmirror_devmem_ops = {
 	.page_free	= dmirror_devmem_free,
 	.migrate_to_ram	= dmirror_devmem_fault,
+	.migrate_to_pagecache = dmirror_devmem_pagecache,
 };
 
 static int dmirror_device_init(struct dmirror_device *mdevice, int id)
diff --git a/tools/testing/selftests/mm/hmm-tests.c b/tools/testing/selftests/mm/hmm-tests.c
index 141bf63..4b77edd 100644
--- a/tools/testing/selftests/mm/hmm-tests.c
+++ b/tools/testing/selftests/mm/hmm-tests.c
@@ -999,6 +999,254 @@ TEST_F(hmm, migrate)
 }
 
 /*
+ * Migrate file memory to device private memory.
+ */
+TEST_F(hmm, migrate_file)
+{
+	struct hmm_buffer *buffer;
+	unsigned long npages;
+	unsigned long size;
+	unsigned long i;
+	int *ptr;
+	int ret;
+
+	npages = ALIGN(HMM_BUFFER_SIZE, self->page_size) >> self->page_shift;
+	ASSERT_NE(npages, 0);
+	size = npages << self->page_shift;
+
+	buffer = malloc(sizeof(*buffer));
+	ASSERT_NE(buffer, NULL);
+
+	buffer->fd = hmm_create_file(size);
+	buffer->size = size;
+	buffer->mirror = malloc(size);
+	ASSERT_NE(buffer->mirror, NULL);
+
+	buffer->ptr = mmap(NULL, size,
+			   PROT_READ | PROT_WRITE,
+			   MAP_SHARED,
+			   buffer->fd, 0);
+	ASSERT_NE(buffer->ptr, MAP_FAILED);
+
+	/* Initialize buffer in system memory. */
+	for (i = 0, ptr = buffer->ptr; i < size / sizeof(*ptr); ++i)
+		ptr[i] = i;
+
+	/*
+	 * TODO: Migration code should try and clean the pages, but it's not
+	 * working.
+	 */
+	fsync(buffer->fd);
+
+	/* Migrate memory to device. */
+	ret = hmm_migrate_sys_to_dev(self->fd, buffer, npages);
+	ASSERT_EQ(ret, 0);
+	ASSERT_EQ(buffer->cpages, npages);
+
+	/* Check what the device read. */
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i);
+
+	hmm_buffer_free(buffer);
+}
+
+TEST_F(hmm, migrate_file_fault)
+{
+	struct hmm_buffer *buffer;
+	unsigned long npages;
+	unsigned long size;
+	unsigned long i;
+	int *ptr;
+	int ret;
+
+	npages = ALIGN(HMM_BUFFER_SIZE, self->page_size) >> self->page_shift;
+	ASSERT_NE(npages, 0);
+	size = npages << self->page_shift;
+
+	buffer = malloc(sizeof(*buffer));
+	ASSERT_NE(buffer, NULL);
+
+	buffer->fd = hmm_create_file(size);
+	buffer->size = size;
+	buffer->mirror = malloc(size);
+	ASSERT_NE(buffer->mirror, NULL);
+
+	buffer->ptr = mmap(NULL, size,
+			   PROT_READ | PROT_WRITE,
+			   MAP_SHARED,
+			   buffer->fd, 0);
+	ASSERT_NE(buffer->ptr, MAP_FAILED);
+
+	/* Initialize buffer in system memory. */
+	for (i = 0, ptr = buffer->ptr; i < size / sizeof(*ptr); ++i)
+		ptr[i] = i;
+
+	/*
+	 * TODO: Migration code should try and clean the pages, but it's not
+	 * working.
+	 */
+	fsync(buffer->fd);
+
+	/* Migrate memory to device. */
+	ret = hmm_migrate_sys_to_dev(self->fd, buffer, npages);
+	ASSERT_EQ(ret, 0);
+	ASSERT_EQ(buffer->cpages, npages);
+
+	/* Check what the device read. */
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i);
+
+	/* Fault half the pages back to system memory and check them. */
+	for (i = 0, ptr = buffer->ptr; i < size / (2 * sizeof(*ptr)); ++i)
+		ASSERT_EQ(ptr[i], i);
+
+	/* Migrate memory to the device again. */
+	ret = hmm_migrate_sys_to_dev(self->fd, buffer, npages);
+	ASSERT_EQ(ret, 0);
+	ASSERT_EQ(buffer->cpages, npages);
+
+	/* Check what the device read. */
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i);
+
+	hmm_buffer_free(buffer);
+}
+
+TEST_F(hmm, migrate_fault_read_buf)
+{
+	struct hmm_buffer *buffer;
+	unsigned long npages;
+	unsigned long size;
+	unsigned long i;
+	int *ptr;
+	int ret;
+
+	npages = ALIGN(HMM_BUFFER_SIZE, self->page_size) >> self->page_shift;
+	ASSERT_NE(npages, 0);
+	size = npages << self->page_shift;
+
+	buffer = malloc(sizeof(*buffer));
+	ASSERT_NE(buffer, NULL);
+
+	buffer->fd = hmm_create_file(size);
+	buffer->size = size;
+	buffer->mirror = malloc(size);
+	ASSERT_NE(buffer->mirror, NULL);
+
+	buffer->ptr = mmap(NULL, size,
+			   PROT_READ | PROT_WRITE,
+			   MAP_SHARED,
+			   buffer->fd, 0);
+	ASSERT_NE(buffer->ptr, MAP_FAILED);
+
+	/* Initialize buffer in system memory. */
+	for (i = 0, ptr = buffer->ptr; i < size / sizeof(*ptr); ++i)
+		ptr[i] = i;
+
+	/*
+	 * TODO: Migration code should try and clean the pages, but it's not
+	 * working.
+	 */
+	fsync(buffer->fd);
+
+	/* Migrate memory to device. */
+	ret = hmm_migrate_sys_to_dev(self->fd, buffer, npages);
+	ASSERT_EQ(ret, 0);
+	ASSERT_EQ(buffer->cpages, npages);
+
+	/* Check what the device read. */
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i);
+
+	/* Use read and check what we read  */
+	read(buffer->fd, buffer->mirror, size);
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i);
+
+	hmm_buffer_free(buffer);
+}
+
+TEST_F(hmm, migrate_fault_write_buf)
+{
+	struct hmm_buffer *buffer;
+	unsigned long npages;
+	unsigned long size;
+	unsigned long i;
+	int *ptr;
+	int ret;
+
+	npages = ALIGN(HMM_BUFFER_SIZE, self->page_size) >> self->page_shift;
+	ASSERT_NE(npages, 0);
+	size = npages << self->page_shift;
+
+	buffer = malloc(sizeof(*buffer));
+	ASSERT_NE(buffer, NULL);
+
+	buffer->fd = hmm_create_file(size);
+	buffer->size = size;
+	buffer->mirror = malloc(size);
+	ASSERT_NE(buffer->mirror, NULL);
+
+	buffer->ptr = mmap(NULL, size,
+			   PROT_READ | PROT_WRITE,
+			   MAP_SHARED,
+			   buffer->fd, 0);
+	ASSERT_NE(buffer->ptr, MAP_FAILED);
+
+	/* Initialize buffer in system memory. */
+	for (i = 0, ptr = buffer->ptr; i < size / sizeof(*ptr); ++i)
+		ptr[i] = i;
+
+	/*
+	 * TODO: Migration code should try and clean the pages, but it's not
+	 * working.
+	 */
+	fsync(buffer->fd);
+
+	/* Migrate memory to device. */
+	ret = hmm_migrate_sys_to_dev(self->fd, buffer, npages);
+	ASSERT_EQ(ret, 0);
+	ASSERT_EQ(buffer->cpages, npages);
+
+	/* Check what the device read and update to write to the device. */
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i]++, i);
+
+	/* Write to the buffer from the device */
+	ret = hmm_dmirror_cmd(self->fd, HMM_DMIRROR_WRITE, buffer, npages);
+	ASSERT_EQ(ret, 0);
+
+	/* Truncate half the file */
+	size >>= 1;
+	ret = truncate("hmm-test-file", size);
+	ASSERT_EQ(ret, 0);
+
+	/* Use read and check what we read  */
+	ret = read(buffer->fd, buffer->mirror, size);
+	ASSERT_EQ(ret, size);
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i + 1);
+
+	/* Should see the same in the mmap */
+	for (i = 0, ptr = buffer->ptr; i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], i + 1);
+
+	/* And check we get zeros in the second half */
+	size <<= 1;
+	ret = truncate("hmm-test-file", size);
+	ASSERT_EQ(ret, 0);
+
+	for (i = 0, ptr = buffer->ptr; i < size / (2*sizeof(*ptr)); ++i)
+		ASSERT_EQ(ptr[i], i + 1);
+
+	for (i = size/(2*sizeof(*ptr)), ptr = buffer->ptr;
+	     i < size / sizeof(*ptr); ++i)
+		ASSERT_EQ(ptr[i], 0);
+
+	hmm_buffer_free(buffer);
+}
+
+/*
  * Migrate anonymous memory to device private memory and fault some of it back
  * to system memory, then try migrating the resulting mix of system and device
  * private memory to the device.
@@ -1040,8 +1288,10 @@ TEST_F(hmm, migrate_fault)
 	ASSERT_EQ(buffer->cpages, npages);
 
 	/* Check what the device read. */
-	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i)
+	for (i = 0, ptr = buffer->mirror; i < size / sizeof(*ptr); ++i) {
 		ASSERT_EQ(ptr[i], i);
+		ptr[i]++;
+	}
 
 	/* Fault half the pages back to system memory and check them. */
 	for (i = 0, ptr = buffer->ptr; i < size / (2 * sizeof(*ptr)); ++i)
