xe->usm.num_vm_in_fault_mode--;
        else if (!(vm->flags & XE_VM_FLAG_MIGRATION))
                xe->usm.num_vm_in_non_fault_mode--;
+
+       if (vm->usm.asid) {
+               void *lookup;
+
+               xe_assert(xe, xe->info.has_asid);
+               xe_assert(xe, !(vm->flags & XE_VM_FLAG_MIGRATION));
+
+               lookup = xa_erase(&xe->usm.asid_to_vm, vm->usm.asid);
+               xe_assert(xe, lookup == vm);
+       }
        mutex_unlock(&xe->usm.lock);
 
        for_each_tile(tile, xe, id)
        struct xe_device *xe = vm->xe;
        struct xe_tile *tile;
        u8 id;
-       void *lookup;
 
        /* xe_vm_close_and_put was not called? */
        xe_assert(xe, !vm->size);
 
        mutex_destroy(&vm->snap_mutex);
 
-       if (!(vm->flags & XE_VM_FLAG_MIGRATION)) {
+       if (!(vm->flags & XE_VM_FLAG_MIGRATION))
                xe_device_mem_access_put(xe);
 
-               if (xe->info.has_asid && vm->usm.asid) {
-                       mutex_lock(&xe->usm.lock);
-                       lookup = xa_erase(&xe->usm.asid_to_vm, vm->usm.asid);
-                       xe_assert(xe, lookup == vm);
-                       mutex_unlock(&xe->usm.lock);
-               }
-       }
-
        for_each_tile(tile, xe, id)
                XE_WARN_ON(vm->pt_root[id]);