return container_of(page->pgmap, struct nouveau_dmem, pagemap);
 }
 
-struct nouveau_dmem_fault {
-       struct nouveau_drm *drm;
-       struct nouveau_fence *fence;
-       dma_addr_t *dma;
-       unsigned long npages;
-};
-
 struct nouveau_migrate {
        struct vm_area_struct *vma;
        struct nouveau_drm *drm;
        }
 }
 
-static void
-nouveau_dmem_fault_alloc_and_copy(struct vm_area_struct *vma,
-                                 const unsigned long *src_pfns,
-                                 unsigned long *dst_pfns,
-                                 unsigned long start,
-                                 unsigned long end,
-                                 struct nouveau_dmem_fault *fault)
+static vm_fault_t nouveau_dmem_fault_copy_one(struct nouveau_drm *drm,
+               struct vm_fault *vmf, struct migrate_vma *args,
+               dma_addr_t *dma_addr)
 {
-       struct nouveau_drm *drm = fault->drm;
        struct device *dev = drm->dev->dev;
-       unsigned long addr, i, npages = 0;
-       nouveau_migrate_copy_t copy;
-       int ret;
-
-
-       /* First allocate new memory */
-       for (addr = start, i = 0; addr < end; addr += PAGE_SIZE, i++) {
-               struct page *dpage, *spage;
-
-               dst_pfns[i] = 0;
-               spage = migrate_pfn_to_page(src_pfns[i]);
-               if (!spage || !(src_pfns[i] & MIGRATE_PFN_MIGRATE))
-                       continue;
-
-               dpage = alloc_page_vma(GFP_HIGHUSER, vma, addr);
-               if (!dpage) {
-                       dst_pfns[i] = MIGRATE_PFN_ERROR;
-                       continue;
-               }
-               lock_page(dpage);
-
-               dst_pfns[i] = migrate_pfn(page_to_pfn(dpage)) |
-                             MIGRATE_PFN_LOCKED;
-               npages++;
-       }
-
-       /* Allocate storage for DMA addresses, so we can unmap later. */
-       fault->dma = kmalloc(sizeof(*fault->dma) * npages, GFP_KERNEL);
-       if (!fault->dma)
-               goto error;
-
-       /* Copy things over */
-       copy = drm->dmem->migrate.copy_func;
-       for (addr = start, i = 0; addr < end; addr += PAGE_SIZE, i++) {
-               struct page *spage, *dpage;
-
-               dpage = migrate_pfn_to_page(dst_pfns[i]);
-               if (!dpage || dst_pfns[i] == MIGRATE_PFN_ERROR)
-                       continue;
-
-               spage = migrate_pfn_to_page(src_pfns[i]);
-               if (!spage || !(src_pfns[i] & MIGRATE_PFN_MIGRATE)) {
-                       dst_pfns[i] = MIGRATE_PFN_ERROR;
-                       __free_page(dpage);
-                       continue;
-               }
+       struct page *dpage, *spage;
 
-               fault->dma[fault->npages] =
-                       dma_map_page_attrs(dev, dpage, 0, PAGE_SIZE,
-                                          PCI_DMA_BIDIRECTIONAL,
-                                          DMA_ATTR_SKIP_CPU_SYNC);
-               if (dma_mapping_error(dev, fault->dma[fault->npages])) {
-                       dst_pfns[i] = MIGRATE_PFN_ERROR;
-                       __free_page(dpage);
-                       continue;
-               }
-
-               ret = copy(drm, 1, NOUVEAU_APER_HOST,
-                               fault->dma[fault->npages++],
-                               NOUVEAU_APER_VRAM,
-                               nouveau_dmem_page_addr(spage));
-               if (ret) {
-                       dst_pfns[i] = MIGRATE_PFN_ERROR;
-                       __free_page(dpage);
-                       continue;
-               }
-       }
-
-       nouveau_fence_new(drm->dmem->migrate.chan, false, &fault->fence);
-
-       return;
-
-error:
-       for (addr = start, i = 0; addr < end; addr += PAGE_SIZE, ++i) {
-               struct page *page;
-
-               if (!dst_pfns[i] || dst_pfns[i] == MIGRATE_PFN_ERROR)
-                       continue;
+       spage = migrate_pfn_to_page(args->src[0]);
+       if (!spage || !(args->src[0] & MIGRATE_PFN_MIGRATE))
+               return 0;
 
-               page = migrate_pfn_to_page(dst_pfns[i]);
-               dst_pfns[i] = MIGRATE_PFN_ERROR;
-               if (page == NULL)
-                       continue;
+       dpage = alloc_page_vma(GFP_HIGHUSER, vmf->vma, vmf->address);
+       if (!dpage)
+               return VM_FAULT_SIGBUS;
+       lock_page(dpage);
 
-               __free_page(page);
-       }
-}
+       *dma_addr = dma_map_page(dev, dpage, 0, PAGE_SIZE, DMA_BIDIRECTIONAL);
+       if (dma_mapping_error(dev, *dma_addr))
+               goto error_free_page;
 
-static void
-nouveau_dmem_fault_finalize_and_map(struct nouveau_dmem_fault *fault)
-{
-       struct nouveau_drm *drm = fault->drm;
+       if (drm->dmem->migrate.copy_func(drm, 1, NOUVEAU_APER_HOST, *dma_addr,
+                       NOUVEAU_APER_VRAM, nouveau_dmem_page_addr(spage)))
+               goto error_dma_unmap;
 
-       nouveau_dmem_fence_done(&fault->fence);
+       args->dst[0] = migrate_pfn(page_to_pfn(dpage)) | MIGRATE_PFN_LOCKED;
+       return 0;
 
-       while (fault->npages--) {
-               dma_unmap_page(drm->dev->dev, fault->dma[fault->npages],
-                              PAGE_SIZE, PCI_DMA_BIDIRECTIONAL);
-       }
-       kfree(fault->dma);
+error_dma_unmap:
+       dma_unmap_page(dev, *dma_addr, PAGE_SIZE, DMA_BIDIRECTIONAL);
+error_free_page:
+       __free_page(dpage);
+       return VM_FAULT_SIGBUS;
 }
 
 static vm_fault_t nouveau_dmem_migrate_to_ram(struct vm_fault *vmf)
 {
        struct nouveau_dmem *dmem = page_to_dmem(vmf->page);
-       unsigned long src[1] = {0}, dst[1] = {0};
+       struct nouveau_drm *drm = dmem->drm;
+       struct nouveau_fence *fence;
+       unsigned long src = 0, dst = 0;
+       dma_addr_t dma_addr = 0;
+       vm_fault_t ret;
        struct migrate_vma args = {
                .vma            = vmf->vma,
                .start          = vmf->address,
                .end            = vmf->address + PAGE_SIZE,
-               .src            = src,
-               .dst            = dst,
+               .src            = &src,
+               .dst            = &dst,
        };
-       struct nouveau_dmem_fault fault = { .drm = dmem->drm };
 
        /*
         * FIXME what we really want is to find some heuristic to migrate more
        if (!args.cpages)
                return 0;
 
-       nouveau_dmem_fault_alloc_and_copy(args.vma, src, dst, args.start,
-                       args.end, &fault);
-       migrate_vma_pages(&args);
-       nouveau_dmem_fault_finalize_and_map(&fault);
+       ret = nouveau_dmem_fault_copy_one(drm, vmf, &args, &dma_addr);
+       if (ret || dst == 0)
+               goto done;
 
+       nouveau_fence_new(dmem->migrate.chan, false, &fence);
+       migrate_vma_pages(&args);
+       nouveau_dmem_fence_done(&fence);
+       dma_unmap_page(drm->dev->dev, dma_addr, PAGE_SIZE, DMA_BIDIRECTIONAL);
+done:
        migrate_vma_finalize(&args);
-       if (dst[0] == MIGRATE_PFN_ERROR)
-               return VM_FAULT_SIGBUS;
-
-       return 0;
+       return ret;
 }
 
 static const struct dev_pagemap_ops nouveau_dmem_pagemap_ops = {