pr_debug("restore svm ranges\n");
 
-       /* kfd_process_notifier_release destroys this worker thread. So during
-        * the lifetime of this thread, kfd_process and mm will be valid.
-        */
        p = container_of(svms, struct kfd_process, svms);
-       mm = p->mm;
-       if (!mm)
+
+       /* Keep mm reference when svm_range_validate_and_map ranges */
+       mm = get_task_mm(p->lead_thread);
+       if (!mm) {
+               pr_debug("svms 0x%p process mm gone\n", svms);
                return;
+       }
 
        svm_range_list_lock_and_flush_work(svms, mm);
        mutex_lock(&svms->lock);
 out_reschedule:
        mutex_unlock(&svms->lock);
        mmap_write_unlock(mm);
+       mmput(mm);
 
        /* If validation failed, reschedule another attempt */
        if (evicted_ranges) {
 
        pr_debug("pasid 0x%x svms 0x%p\n", p->pasid, &p->svms);
 
+       cancel_delayed_work_sync(&p->svms.restore_work);
+
        /* Ensure list work is finished before process is destroyed */
        flush_work(&p->svms.deferred_list_work);
 
        atomic_inc(&p->svms.drain_pagefaults);
        svm_range_drain_retry_fault(&p->svms);
 
-
        list_for_each_entry_safe(prange, next, &p->svms.list, list) {
                svm_range_unlink(prange);
                svm_range_remove_notifier(prange);