kfree(vwork);
 }
 
-static void vfio_lock_acct(long npage)
+static void vfio_lock_acct(struct task_struct *task, long npage)
 {
        struct vwork *vwork;
        struct mm_struct *mm;
 
-       if (!current->mm || !npage)
+       if (!npage)
+               return;
+
+       mm = get_task_mm(task);
+       if (!mm)
                return; /* process exited or nothing to do */
 
-       if (down_write_trylock(¤t->mm->mmap_sem)) {
-               current->mm->locked_vm += npage;
-               up_write(¤t->mm->mmap_sem);
+       if (down_write_trylock(&mm->mmap_sem)) {
+               mm->locked_vm += npage;
+               up_write(&mm->mmap_sem);
+               mmput(mm);
                return;
        }
 
         * wouldn't need this silliness
         */
        vwork = kmalloc(sizeof(struct vwork), GFP_KERNEL);
-       if (!vwork)
-               return;
-       mm = get_task_mm(current);
-       if (!mm) {
-               kfree(vwork);
+       if (!vwork) {
+               mmput(mm);
                return;
        }
        INIT_WORK(&vwork->work, vfio_lock_acct_bg);
 
        if (unlikely(disable_hugepages)) {
                if (!rsvd)
-                       vfio_lock_acct(1);
+                       vfio_lock_acct(current, 1);
                return 1;
        }
 
        }
 
        if (!rsvd)
-               vfio_lock_acct(i);
+               vfio_lock_acct(current, i);
 
        return i;
 }
                unlocked += put_pfn(pfn++, prot);
 
        if (do_accounting)
-               vfio_lock_acct(-unlocked);
+               vfio_lock_acct(current, -unlocked);
 
        return unlocked;
 }
                cond_resched();
        }
 
-       vfio_lock_acct(-unlocked);
+       vfio_lock_acct(current, -unlocked);
 }
 
 static void vfio_remove_dma(struct vfio_iommu *iommu, struct vfio_dma *dma)