return err;
 
        pr_debug("Looking for gpu id 0x%x\n", args->gpu_id);
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev) {
-               pr_debug("Could not find gpu id 0x%x\n", args->gpu_id);
-               return -EINVAL;
-       }
 
        mutex_lock(&p->mutex);
 
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               pr_debug("Could not find gpu id 0x%x\n", args->gpu_id);
+               err = -EINVAL;
+               goto err_pdd;
+       }
+       dev = pdd->dev;
+
        pdd = kfd_bind_process_to_device(dev, p);
        if (IS_ERR(pdd)) {
                err = -ESRCH;
 
 err_create_queue:
 err_bind_process:
+err_pdd:
        mutex_unlock(&p->mutex);
        return err;
 }
                                        struct kfd_process *p, void *data)
 {
        struct kfd_ioctl_set_memory_policy_args *args = data;
-       struct kfd_dev *dev;
        int err = 0;
        struct kfd_process_device *pdd;
        enum cache_policy default_policy, alternate_policy;
                return -EINVAL;
        }
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
-
        mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               pr_debug("Could not find gpu id 0x%x\n", args->gpu_id);
+               err = -EINVAL;
+               goto err_pdd;
+       }
 
-       pdd = kfd_bind_process_to_device(dev, p);
+       pdd = kfd_bind_process_to_device(pdd->dev, p);
        if (IS_ERR(pdd)) {
                err = -ESRCH;
                goto out;
                (args->alternate_policy == KFD_IOC_CACHE_POLICY_COHERENT)
                   ? cache_policy_coherent : cache_policy_noncoherent;
 
-       if (!dev->dqm->ops.set_cache_memory_policy(dev->dqm,
+       if (!pdd->dev->dqm->ops.set_cache_memory_policy(pdd->dev->dqm,
                                &pdd->qpd,
                                default_policy,
                                alternate_policy,
                err = -EINVAL;
 
 out:
+err_pdd:
        mutex_unlock(&p->mutex);
 
        return err;
                                        struct kfd_process *p, void *data)
 {
        struct kfd_ioctl_set_trap_handler_args *args = data;
-       struct kfd_dev *dev;
        int err = 0;
        struct kfd_process_device *pdd;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
-
        mutex_lock(&p->mutex);
 
-       pdd = kfd_bind_process_to_device(dev, p);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               err = -EINVAL;
+               goto err_pdd;
+       }
+
+       pdd = kfd_bind_process_to_device(pdd->dev, p);
        if (IS_ERR(pdd)) {
                err = -ESRCH;
                goto out;
        kfd_process_set_trap_handler(&pdd->qpd, args->tba_addr, args->tma_addr);
 
 out:
+err_pdd:
        mutex_unlock(&p->mutex);
 
        return err;
        bool create_ok;
        long status = 0;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               status = -EINVAL;
+               goto err_pdd;
+       }
+       dev = pdd->dev;
 
        if (dev->adev->asic_type == CHIP_CARRIZO) {
                pr_debug("kfd_ioctl_dbg_register not supported on CZ\n");
-               return -EINVAL;
+               status = -EINVAL;
+               goto err_chip_unsupp;
        }
 
-       mutex_lock(&p->mutex);
        mutex_lock(kfd_get_dbgmgr_mutex());
 
        /*
 
 out:
        mutex_unlock(kfd_get_dbgmgr_mutex());
+err_pdd:
+err_chip_unsupp:
        mutex_unlock(&p->mutex);
 
        return status;
                                struct kfd_process *p, void *data)
 {
        struct kfd_ioctl_dbg_unregister_args *args = data;
+       struct kfd_process_device *pdd;
        struct kfd_dev *dev;
        long status;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev || !dev->dbgmgr)
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       mutex_unlock(&p->mutex);
+       if (!pdd || !pdd->dev->dbgmgr)
                return -EINVAL;
 
+       dev = pdd->dev;
+
        if (dev->adev->asic_type == CHIP_CARRIZO) {
                pr_debug("kfd_ioctl_dbg_unregister not supported on CZ\n");
                return -EINVAL;
 {
        struct kfd_ioctl_dbg_address_watch_args *args = data;
        struct kfd_dev *dev;
+       struct kfd_process_device *pdd;
        struct dbg_address_watch_info aw_info;
        unsigned char *args_buff;
        long status;
 
        memset((void *) &aw_info, 0, sizeof(struct dbg_address_watch_info));
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       mutex_unlock(&p->mutex);
+       if (!pdd) {
+               pr_debug("Could not find gpu id 0x%x\n", args->gpu_id);
                return -EINVAL;
+       }
+       dev = pdd->dev;
 
        if (dev->adev->asic_type == CHIP_CARRIZO) {
                pr_debug("kfd_ioctl_dbg_wave_control not supported on CZ\n");
                return -EINVAL;
        }
-
        cmd_from_user = (void __user *) args->content_ptr;
 
        /* Validate arguments */
 {
        struct kfd_ioctl_dbg_wave_control_args *args = data;
        struct kfd_dev *dev;
+       struct kfd_process_device *pdd;
        struct dbg_wave_control_info wac_info;
        unsigned char *args_buff;
        uint32_t computed_buff_size;
                                sizeof(wac_info.dbgWave_msg.MemoryVA) +
                                sizeof(wac_info.trapId);
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       mutex_unlock(&p->mutex);
+       if (!pdd) {
+               pr_debug("Could not find gpu id 0x%x\n", args->gpu_id);
                return -EINVAL;
+       }
+       dev = pdd->dev;
 
        if (dev->adev->asic_type == CHIP_CARRIZO) {
                pr_debug("kfd_ioctl_dbg_wave_control not supported on CZ\n");
                                struct kfd_process *p, void *data)
 {
        struct kfd_ioctl_get_clock_counters_args *args = data;
-       struct kfd_dev *dev;
+       struct kfd_process_device *pdd;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (dev)
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       mutex_unlock(&p->mutex);
+       if (pdd)
                /* Reading GPU clock counter from KGD */
-               args->gpu_clock_counter = amdgpu_amdkfd_get_gpu_clock_counter(dev->adev);
+               args->gpu_clock_counter = amdgpu_amdkfd_get_gpu_clock_counter(pdd->dev->adev);
        else
                /* Node without GPU resource */
                args->gpu_clock_counter = 0;
        struct kfd_dev *dev;
        long err;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
-
        mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               err = -EINVAL;
+               goto err_pdd;
+       }
+       dev = pdd->dev;
 
        pdd = kfd_bind_process_to_device(dev, p);
        if (IS_ERR(pdd)) {
        return 0;
 
 bind_process_to_device_fail:
+err_pdd:
        mutex_unlock(&p->mutex);
        return err;
 }
                struct kfd_process *p, void *data)
 {
        struct kfd_ioctl_get_tile_config_args *args = data;
-       struct kfd_dev *dev;
+       struct kfd_process_device *pdd;
        struct tile_config config;
        int err = 0;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       mutex_unlock(&p->mutex);
+       if (!pdd)
                return -EINVAL;
 
-       amdgpu_amdkfd_get_tile_config(dev->adev, &config);
+       amdgpu_amdkfd_get_tile_config(pdd->dev->adev, &config);
 
        args->gb_addr_config = config.gb_addr_config;
        args->num_banks = config.num_banks;
 {
        struct kfd_ioctl_acquire_vm_args *args = data;
        struct kfd_process_device *pdd;
-       struct kfd_dev *dev;
        struct file *drm_file;
        int ret;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
-
        drm_file = fget(args->drm_fd);
        if (!drm_file)
                return -EINVAL;
 
        mutex_lock(&p->mutex);
-
-       pdd = kfd_get_process_device_data(dev, p);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
        if (!pdd) {
                ret = -EINVAL;
-               goto err_unlock;
+               goto err_pdd;
        }
 
        if (pdd->drm_file) {
                ret = pdd->drm_file == drm_file ? 0 : -EBUSY;
-               goto err_unlock;
+               goto err_drm_file;
        }
 
        ret = kfd_process_device_init_vm(pdd, drm_file);
        if (ret)
                goto err_unlock;
+
        /* On success, the PDD keeps the drm_file reference */
        mutex_unlock(&p->mutex);
 
        return 0;
 
 err_unlock:
+err_pdd:
+err_drm_file:
        mutex_unlock(&p->mutex);
        fput(drm_file);
        return ret;
        }
        mutex_unlock(&p->svms.lock);
 #endif
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
+       mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               err = -EINVAL;
+               goto err_pdd;
+       }
+
+       dev = pdd->dev;
 
        if ((flags & KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC) &&
                (flags & KFD_IOC_ALLOC_MEM_FLAGS_VRAM) &&
                !kfd_dev_is_large_bar(dev)) {
                pr_err("Alloc host visible vram on small bar is not allowed\n");
-               return -EINVAL;
+               err = -EINVAL;
+               goto err_large_bar;
        }
 
-       mutex_lock(&p->mutex);
-
        pdd = kfd_bind_process_to_device(dev, p);
        if (IS_ERR(pdd)) {
                err = PTR_ERR(pdd);
        amdgpu_amdkfd_gpuvm_free_memory_of_gpu(dev->adev, (struct kgd_mem *)mem,
                                               pdd->drm_priv, NULL);
 err_unlock:
+err_pdd:
+err_large_bar:
        mutex_unlock(&p->mutex);
        return err;
 }
        struct kfd_ioctl_free_memory_of_gpu_args *args = data;
        struct kfd_process_device *pdd;
        void *mem;
-       struct kfd_dev *dev;
        int ret;
        uint64_t size = 0;
 
-       dev = kfd_device_by_id(GET_GPU_ID(args->handle));
-       if (!dev)
-               return -EINVAL;
-
        mutex_lock(&p->mutex);
        /*
         * Safeguard to prevent user space from freeing signal BO.
                goto err_unlock;
        }
 
-       pdd = kfd_get_process_device_data(dev, p);
+       pdd = kfd_process_device_data_by_id(p, GET_GPU_ID(args->handle));
        if (!pdd) {
                pr_err("Process device data doesn't exist\n");
                ret = -EINVAL;
-               goto err_unlock;
+               goto err_pdd;
        }
 
        mem = kfd_process_device_translate_handle(
                goto err_unlock;
        }
 
-       ret = amdgpu_amdkfd_gpuvm_free_memory_of_gpu(dev->adev,
+       ret = amdgpu_amdkfd_gpuvm_free_memory_of_gpu(pdd->dev->adev,
                                (struct kgd_mem *)mem, pdd->drm_priv, &size);
 
        /* If freeing the buffer failed, leave the handle in place for
        WRITE_ONCE(pdd->vram_usage, pdd->vram_usage - size);
 
 err_unlock:
+err_pdd:
        mutex_unlock(&p->mutex);
        return ret;
 }
        struct kfd_ioctl_map_memory_to_gpu_args *args = data;
        struct kfd_process_device *pdd, *peer_pdd;
        void *mem;
-       struct kfd_dev *dev, *peer;
+       struct kfd_dev *dev;
        long err = 0;
        int i;
        uint32_t *devices_arr = NULL;
        bool table_freed = false;
 
-       dev = kfd_device_by_id(GET_GPU_ID(args->handle));
-       if (!dev)
-               return -EINVAL;
-
        if (!args->n_devices) {
                pr_debug("Device IDs array empty\n");
                return -EINVAL;
        }
 
        mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, GET_GPU_ID(args->handle));
+       if (!pdd) {
+               err = -EINVAL;
+               goto get_process_device_data_failed;
+       }
+       dev = pdd->dev;
 
        pdd = kfd_bind_process_to_device(dev, p);
        if (IS_ERR(pdd)) {
        }
 
        for (i = args->n_success; i < args->n_devices; i++) {
-               peer = kfd_device_by_id(devices_arr[i]);
-               if (!peer) {
+               peer_pdd = kfd_process_device_data_by_id(p, devices_arr[i]);
+               if (!peer_pdd) {
                        pr_debug("Getting device by id failed for 0x%x\n",
                                 devices_arr[i]);
                        err = -EINVAL;
                        goto get_mem_obj_from_handle_failed;
                }
 
-               peer_pdd = kfd_bind_process_to_device(peer, p);
+               peer_pdd = kfd_bind_process_to_device(peer_pdd->dev, p);
                if (IS_ERR(peer_pdd)) {
                        err = PTR_ERR(peer_pdd);
                        goto get_mem_obj_from_handle_failed;
                }
+
                err = amdgpu_amdkfd_gpuvm_map_memory_to_gpu(
-                       peer->adev, (struct kgd_mem *)mem,
+                       peer_pdd->dev->adev, (struct kgd_mem *)mem,
                        peer_pdd->drm_priv, &table_freed);
                if (err) {
                        pr_err("Failed to map to gpu %d/%d\n",
        /* Flush TLBs after waiting for the page table updates to complete */
        if (table_freed || !kfd_flush_tlb_after_unmap(dev)) {
                for (i = 0; i < args->n_devices; i++) {
-                       peer = kfd_device_by_id(devices_arr[i]);
-                       if (WARN_ON_ONCE(!peer))
-                               continue;
-                       peer_pdd = kfd_get_process_device_data(peer, p);
+                       peer_pdd = kfd_process_device_data_by_id(p, devices_arr[i]);
                        if (WARN_ON_ONCE(!peer_pdd))
                                continue;
                        kfd_flush_tlb(peer_pdd, TLB_FLUSH_LEGACY);
 
        return err;
 
+get_process_device_data_failed:
 bind_process_to_device_failed:
 get_mem_obj_from_handle_failed:
 map_memory_to_gpu_failed:
        struct kfd_ioctl_unmap_memory_from_gpu_args *args = data;
        struct kfd_process_device *pdd, *peer_pdd;
        void *mem;
-       struct kfd_dev *dev, *peer;
        long err = 0;
        uint32_t *devices_arr = NULL, i;
 
-       dev = kfd_device_by_id(GET_GPU_ID(args->handle));
-       if (!dev)
-               return -EINVAL;
-
        if (!args->n_devices) {
                pr_debug("Device IDs array empty\n");
                return -EINVAL;
        }
 
        mutex_lock(&p->mutex);
-
-       pdd = kfd_get_process_device_data(dev, p);
+       pdd = kfd_process_device_data_by_id(p, GET_GPU_ID(args->handle));
        if (!pdd) {
                err = -EINVAL;
                goto bind_process_to_device_failed;
        }
 
        for (i = args->n_success; i < args->n_devices; i++) {
-               peer = kfd_device_by_id(devices_arr[i]);
-               if (!peer) {
-                       err = -EINVAL;
-                       goto get_mem_obj_from_handle_failed;
-               }
-
-               peer_pdd = kfd_get_process_device_data(peer, p);
+               peer_pdd = kfd_process_device_data_by_id(p, devices_arr[i]);
                if (!peer_pdd) {
-                       err = -ENODEV;
+                       err = -EINVAL;
                        goto get_mem_obj_from_handle_failed;
                }
                err = amdgpu_amdkfd_gpuvm_unmap_memory_from_gpu(
-                       peer->adev, (struct kgd_mem *)mem, peer_pdd->drm_priv);
+                       peer_pdd->dev->adev, (struct kgd_mem *)mem, peer_pdd->drm_priv);
                if (err) {
                        pr_err("Failed to unmap from gpu %d/%d\n",
                               i, args->n_devices);
        }
        mutex_unlock(&p->mutex);
 
-       if (kfd_flush_tlb_after_unmap(dev)) {
-               err = amdgpu_amdkfd_gpuvm_sync_memory(dev->adev,
+       if (kfd_flush_tlb_after_unmap(pdd->dev)) {
+               err = amdgpu_amdkfd_gpuvm_sync_memory(pdd->dev->adev,
                                (struct kgd_mem *) mem, true);
                if (err) {
                        pr_debug("Sync memory failed, wait interrupted by user signal\n");
 
                /* Flush TLBs after waiting for the page table updates to complete */
                for (i = 0; i < args->n_devices; i++) {
-                       peer = kfd_device_by_id(devices_arr[i]);
-                       if (WARN_ON_ONCE(!peer))
-                               continue;
-                       peer_pdd = kfd_get_process_device_data(peer, p);
+                       peer_pdd = kfd_process_device_data_by_id(p, devices_arr[i]);
                        if (WARN_ON_ONCE(!peer_pdd))
                                continue;
                        kfd_flush_tlb(peer_pdd, TLB_FLUSH_HEAVYWEIGHT);
        struct kfd_ioctl_import_dmabuf_args *args = data;
        struct kfd_process_device *pdd;
        struct dma_buf *dmabuf;
-       struct kfd_dev *dev;
        int idr_handle;
        uint64_t size;
        void *mem;
        int r;
 
-       dev = kfd_device_by_id(args->gpu_id);
-       if (!dev)
-               return -EINVAL;
-
        dmabuf = dma_buf_get(args->dmabuf_fd);
        if (IS_ERR(dmabuf))
                return PTR_ERR(dmabuf);
 
        mutex_lock(&p->mutex);
+       pdd = kfd_process_device_data_by_id(p, args->gpu_id);
+       if (!pdd) {
+               r = -EINVAL;
+               goto err_unlock;
+       }
 
-       pdd = kfd_bind_process_to_device(dev, p);
+       pdd = kfd_bind_process_to_device(pdd->dev, p);
        if (IS_ERR(pdd)) {
                r = PTR_ERR(pdd);
                goto err_unlock;
        }
 
-       r = amdgpu_amdkfd_gpuvm_import_dmabuf(dev->adev, dmabuf,
+       r = amdgpu_amdkfd_gpuvm_import_dmabuf(pdd->dev->adev, dmabuf,
                                              args->va_addr, pdd->drm_priv,
                                              (struct kgd_mem **)&mem, &size,
                                              NULL);
        return 0;
 
 err_free:
-       amdgpu_amdkfd_gpuvm_free_memory_of_gpu(dev->adev, (struct kgd_mem *)mem,
+       amdgpu_amdkfd_gpuvm_free_memory_of_gpu(pdd->dev->adev, (struct kgd_mem *)mem,
                                               pdd->drm_priv, NULL);
 err_unlock:
        mutex_unlock(&p->mutex);
                                struct kfd_process *p, void *data)
 {
        struct kfd_ioctl_smi_events_args *args = data;
-       struct kfd_dev *dev;
+       struct kfd_process_device *pdd;
 
-       dev = kfd_device_by_id(args->gpuid);
-       if (!dev)
+       mutex_lock(&p->mutex);
+
+       pdd = kfd_process_device_data_by_id(p, args->gpuid);
+       mutex_unlock(&p->mutex);
+       if (!pdd)
                return -EINVAL;
 
-       return kfd_smi_event_open(dev, &args->anon_fd);
+       return kfd_smi_event_open(pdd->dev, &args->anon_fd);
 }
 
 static int kfd_ioctl_set_xnack_mode(struct file *filep,
        return ret;
 }
 
+static int criu_checkpoint_devices(struct kfd_process *p,
+                            uint32_t num_devices,
+                            uint8_t __user *user_addr,
+                            uint8_t __user *user_priv_data,
+                            uint64_t *priv_offset)
+{
+       struct kfd_criu_device_priv_data *device_priv = NULL;
+       struct kfd_criu_device_bucket *device_buckets = NULL;
+       int ret = 0, i;
+
+       device_buckets = kvzalloc(num_devices * sizeof(*device_buckets), GFP_KERNEL);
+       if (!device_buckets) {
+               ret = -ENOMEM;
+               goto exit;
+       }
+
+       device_priv = kvzalloc(num_devices * sizeof(*device_priv), GFP_KERNEL);
+       if (!device_priv) {
+               ret = -ENOMEM;
+               goto exit;
+       }
+
+       for (i = 0; i < num_devices; i++) {
+               struct kfd_process_device *pdd = p->pdds[i];
+
+               device_buckets[i].user_gpu_id = pdd->user_gpu_id;
+               device_buckets[i].actual_gpu_id = pdd->dev->id;
+
+               /*
+                * priv_data does not contain useful information for now and is reserved for
+                * future use, so we do not set its contents.
+                */
+       }
+
+       ret = copy_to_user(user_addr, device_buckets, num_devices * sizeof(*device_buckets));
+       if (ret) {
+               pr_err("Failed to copy device information to user\n");
+               ret = -EFAULT;
+               goto exit;
+       }
+
+       ret = copy_to_user(user_priv_data + *priv_offset,
+                          device_priv,
+                          num_devices * sizeof(*device_priv));
+       if (ret) {
+               pr_err("Failed to copy device information to user\n");
+               ret = -EFAULT;
+       }
+       *priv_offset += num_devices * sizeof(*device_priv);
+
+exit:
+       kvfree(device_buckets);
+       kvfree(device_priv);
+       return ret;
+}
+
 uint32_t get_process_num_bos(struct kfd_process *p)
 {
        uint32_t num_of_bos = 0;
                        bo_bucket = &bo_buckets[bo_index];
                        bo_priv = &bo_privs[bo_index];
 
+                       bo_bucket->gpu_id = pdd->user_gpu_id;
                        bo_bucket->addr = (uint64_t)kgd_mem->va;
                        bo_bucket->size = amdgpu_bo_size(dumper_bo);
-                       bo_bucket->gpu_id = pdd->dev->id;
                        bo_bucket->alloc_flags = (uint32_t)kgd_mem->alloc_flags;
                        bo_priv->idr_handle = id;
 
 
                        for (i = 0; i < p->n_pdds; i++) {
                                if (amdgpu_amdkfd_bo_mapped_to_dev(p->pdds[i]->dev->adev, kgd_mem))
-                                       bo_priv->mapped_gpuids[dev_idx++] = p->pdds[i]->dev->id;
+                                       bo_priv->mapped_gpuids[dev_idx++] = p->pdds[i]->user_gpu_id;
                        }
 
                        pr_debug("bo_size = 0x%llx, bo_addr = 0x%llx bo_offset = 0x%llx\n"
 }
 
 static int criu_get_process_object_info(struct kfd_process *p,
+                                       uint32_t *num_devices,
                                        uint32_t *num_bos,
                                        uint32_t *num_objects,
                                        uint64_t *objs_priv_size)
        uint32_t num_queues, num_events, num_svm_ranges;
        uint64_t queues_priv_data_size;
 
+       *num_devices = p->n_pdds;
        *num_bos = get_process_num_bos(p);
 
        ret = kfd_process_get_queue_info(p, &num_queues, &queues_priv_data_size);
 
        if (objs_priv_size) {
                priv_size = sizeof(struct kfd_criu_process_priv_data);
+               priv_size += *num_devices * sizeof(struct kfd_criu_device_priv_data);
                priv_size += *num_bos * sizeof(struct kfd_criu_bo_priv_data);
                priv_size += queues_priv_data_size;
                priv_size += num_events * sizeof(struct kfd_criu_event_priv_data);
                           struct kfd_ioctl_criu_args *args)
 {
        int ret;
-       uint32_t num_bos, num_objects;
+       uint32_t num_devices, num_bos, num_objects;
        uint64_t priv_size, priv_offset = 0;
 
-       if (!args->bos || !args->priv_data)
+       if (!args->devices || !args->bos || !args->priv_data)
                return -EINVAL;
 
        mutex_lock(&p->mutex);
                goto exit_unlock;
        }
 
-       ret = criu_get_process_object_info(p, &num_bos, &num_objects, &priv_size);
+       ret = criu_get_process_object_info(p, &num_devices, &num_bos, &num_objects, &priv_size);
        if (ret)
                goto exit_unlock;
 
-       if (num_bos != args->num_bos ||
+       if (num_devices != args->num_devices ||
+           num_bos != args->num_bos ||
            num_objects != args->num_objects ||
            priv_size != args->priv_data_size) {
 
        if (ret)
                goto exit_unlock;
 
+       ret = criu_checkpoint_devices(p, num_devices, (uint8_t __user *)args->devices,
+                               (uint8_t __user *)args->priv_data, &priv_offset);
+       if (ret)
+               goto exit_unlock;
+
        ret = criu_checkpoint_bos(p, num_bos, (uint8_t __user *)args->bos,
                            (uint8_t __user *)args->priv_data, &priv_offset);
        if (ret)
        return ret;
 }
 
+static int criu_restore_devices(struct kfd_process *p,
+                               struct kfd_ioctl_criu_args *args,
+                               uint64_t *priv_offset,
+                               uint64_t max_priv_data_size)
+{
+       struct kfd_criu_device_bucket *device_buckets;
+       struct kfd_criu_device_priv_data *device_privs;
+       int ret = 0;
+       uint32_t i;
+
+       if (args->num_devices != p->n_pdds)
+               return -EINVAL;
+
+       if (*priv_offset + (args->num_devices * sizeof(*device_privs)) > max_priv_data_size)
+               return -EINVAL;
+
+       device_buckets = kmalloc_array(args->num_devices, sizeof(*device_buckets), GFP_KERNEL);
+       if (!device_buckets)
+               return -ENOMEM;
+
+       ret = copy_from_user(device_buckets, (void __user *)args->devices,
+                               args->num_devices * sizeof(*device_buckets));
+       if (ret) {
+               pr_err("Failed to copy devices buckets from user\n");
+               ret = -EFAULT;
+               goto exit;
+       }
+
+       for (i = 0; i < args->num_devices; i++) {
+               struct kfd_dev *dev;
+               struct kfd_process_device *pdd;
+               struct file *drm_file;
+
+               /* device private data is not currently used */
+
+               if (!device_buckets[i].user_gpu_id) {
+                       pr_err("Invalid user gpu_id\n");
+                       ret = -EINVAL;
+                       goto exit;
+               }
+
+               dev = kfd_device_by_id(device_buckets[i].actual_gpu_id);
+               if (!dev) {
+                       pr_err("Failed to find device with gpu_id = %x\n",
+                               device_buckets[i].actual_gpu_id);
+                       ret = -EINVAL;
+                       goto exit;
+               }
+
+               pdd = kfd_get_process_device_data(dev, p);
+               if (!pdd) {
+                       pr_err("Failed to get pdd for gpu_id = %x\n",
+                                       device_buckets[i].actual_gpu_id);
+                       ret = -EINVAL;
+                       goto exit;
+               }
+               pdd->user_gpu_id = device_buckets[i].user_gpu_id;
+
+               drm_file = fget(device_buckets[i].drm_fd);
+               if (!drm_file) {
+                       pr_err("Invalid render node file descriptor sent from plugin (%d)\n",
+                               device_buckets[i].drm_fd);
+                       ret = -EINVAL;
+                       goto exit;
+               }
+
+               if (pdd->drm_file) {
+                       ret = -EINVAL;
+                       goto exit;
+               }
+
+               /* create the vm using render nodes for kfd pdd */
+               if (kfd_process_device_init_vm(pdd, drm_file)) {
+                       pr_err("could not init vm for given pdd\n");
+                       /* On success, the PDD keeps the drm_file reference */
+                       fput(drm_file);
+                       ret = -EINVAL;
+                       goto exit;
+               }
+               /*
+                * pdd now already has the vm bound to render node so below api won't create a new
+                * exclusive kfd mapping but use existing one with renderDXXX but is still needed
+                * for iommu v2 binding  and runtime pm.
+                */
+               pdd = kfd_bind_process_to_device(dev, p);
+               if (IS_ERR(pdd)) {
+                       ret = PTR_ERR(pdd);
+                       goto exit;
+               }
+       }
+
+       /*
+        * We are not copying device private data from user as we are not using the data for now,
+        * but we still adjust for its private data.
+        */
+       *priv_offset += args->num_devices * sizeof(*device_privs);
+
+exit:
+       kfree(device_buckets);
+       return ret;
+}
+
 static int criu_restore_bos(struct kfd_process *p,
                            struct kfd_ioctl_criu_args *args,
                            uint64_t *priv_offset,
                bo_bucket = &bo_buckets[i];
                bo_priv = &bo_privs[i];
 
-               dev = kfd_device_by_id(bo_bucket->gpu_id);
-               if (!dev) {
-                       ret = -EINVAL;
-                       pr_err("Failed to get pdd\n");
-                       goto exit;
-               }
-               pdd = kfd_get_process_device_data(dev, p);
-               if (!pdd) {
-                       ret = -EINVAL;
-                       pr_err("Failed to get pdd\n");
-                       goto exit;
-               }
-
                pr_debug("kfd restore ioctl - bo_bucket[%d]:\n", i);
                pr_debug("size = 0x%llx, bo_addr = 0x%llx bo_offset = 0x%llx\n"
                        "gpu_id = 0x%x alloc_flags = 0x%x\n"
                        bo_bucket->alloc_flags,
                        bo_priv->idr_handle);
 
+               pdd = kfd_process_device_data_by_id(p, bo_bucket->gpu_id);
+               if (!pdd) {
+                       pr_err("Failed to get pdd\n");
+                       ret = -ENODEV;
+                       goto exit;
+               }
+               dev = pdd->dev;
+
                if (bo_bucket->alloc_flags & KFD_IOC_ALLOC_MEM_FLAGS_DOORBELL) {
                        pr_debug("restore ioctl: KFD_IOC_ALLOC_MEM_FLAGS_DOORBELL\n");
                        if (bo_bucket->size != kfd_doorbell_process_slice(dev)) {
                        if (!bo_priv->mapped_gpuids[j])
                                break;
 
-                       peer = kfd_device_by_id(bo_priv->mapped_gpuids[j]);
-                       if (!peer) {
-                               pr_debug("Getting device by id failed for 0x%x\n", pdd->dev->id);
+                       peer_pdd = kfd_process_device_data_by_id(p, bo_priv->mapped_gpuids[j]);
+                       if (IS_ERR(peer_pdd)) {
                                ret = -EINVAL;
                                goto exit;
                        }
+                       peer = peer_pdd->dev;
 
                        peer_pdd = kfd_bind_process_to_device(peer, p);
                        if (IS_ERR(peer_pdd)) {
        if (ret)
                goto exit_unlock;
 
+       ret = criu_restore_devices(p, args, &priv_offset, args->priv_data_size);
+       if (ret)
+               goto exit_unlock;
+
        ret = criu_restore_bos(p, args, &priv_offset, args->priv_data_size);
        if (ret)
                goto exit_unlock;
        args->pid = task_pid_nr_ns(p->lead_thread,
                                        task_active_pid_ns(p->lead_thread));
 
-       ret = criu_get_process_object_info(p, &args->num_bos, &args->num_objects,
-                                          &args->priv_data_size);
+       ret = criu_get_process_object_info(p, &args->num_devices, &args->num_bos,
+                                          &args->num_objects, &args->priv_data_size);
        if (ret)
                goto err_unlock;
 
-       dev_dbg(kfd_device, "Num of bos:%u objects:%u priv_data_size:%lld\n",
-                               args->num_bos, args->num_objects, args->priv_data_size);
+       dev_dbg(kfd_device, "Num of devices:%u bos:%u objects:%u priv_data_size:%lld\n",
+                               args->num_devices, args->num_bos, args->num_objects,
+                               args->priv_data_size);
 
 err_unlock:
        if (ret) {