#include <linux/mman.h>
 #include <linux/ptrace.h>
 #include <linux/dma-buf.h>
-#include <linux/fdtable.h>
 #include <linux/processor.h>
 #include "kfd_priv.h"
 #include "kfd_device_queue_manager.h"
 }
 
 static int criu_get_prime_handle(struct kgd_mem *mem,
-                                int flags, u32 *shared_fd)
+                                int flags, u32 *shared_fd,
+                                struct file **file)
 {
        struct dma_buf *dmabuf;
        int ret;
                return ret;
        }
 
-       ret = dma_buf_fd(dmabuf, flags);
+       ret = get_unused_fd_flags(flags);
        if (ret < 0) {
                pr_err("dmabuf create fd failed, ret:%d\n", ret);
                goto out_free_dmabuf;
        }
 
        *shared_fd = ret;
+       *file = dmabuf->file;
        return 0;
 
 out_free_dmabuf:
        return ret;
 }
 
+static void commit_files(struct file **files,
+                        struct kfd_criu_bo_bucket *bo_buckets,
+                        unsigned int count,
+                        int err)
+{
+       while (count--) {
+               struct file *file = files[count];
+
+               if (!file)
+                       continue;
+               if (err) {
+                       fput(file);
+                       put_unused_fd(bo_buckets[count].dmabuf_fd);
+               } else {
+                       fd_install(bo_buckets[count].dmabuf_fd, file);
+               }
+       }
+}
+
 static int criu_checkpoint_bos(struct kfd_process *p,
                               uint32_t num_bos,
                               uint8_t __user *user_bos,
 {
        struct kfd_criu_bo_bucket *bo_buckets;
        struct kfd_criu_bo_priv_data *bo_privs;
+       struct file **files = NULL;
        int ret = 0, pdd_index, bo_index = 0, id;
        void *mem;
 
                goto exit;
        }
 
+       files = kvzalloc(num_bos * sizeof(struct file *), GFP_KERNEL);
+       if (!files) {
+               ret = -ENOMEM;
+               goto exit;
+       }
+
        for (pdd_index = 0; pdd_index < p->n_pdds; pdd_index++) {
                struct kfd_process_device *pdd = p->pdds[pdd_index];
                struct amdgpu_bo *dumper_bo;
                                ret = criu_get_prime_handle(kgd_mem,
                                                bo_bucket->alloc_flags &
                                                KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE ? DRM_RDWR : 0,
-                                               &bo_bucket->dmabuf_fd);
+                                               &bo_bucket->dmabuf_fd, &files[bo_index]);
                                if (ret)
                                        goto exit;
                        } else {
        *priv_offset += num_bos * sizeof(*bo_privs);
 
 exit:
-       while (ret && bo_index--) {
-               if (bo_buckets[bo_index].alloc_flags
-                   & (KFD_IOC_ALLOC_MEM_FLAGS_VRAM | KFD_IOC_ALLOC_MEM_FLAGS_GTT))
-                       close_fd(bo_buckets[bo_index].dmabuf_fd);
-       }
-
+       commit_files(files, bo_buckets, bo_index, ret);
+       kvfree(files);
        kvfree(bo_buckets);
        kvfree(bo_privs);
        return ret;
 
 static int criu_restore_bo(struct kfd_process *p,
                           struct kfd_criu_bo_bucket *bo_bucket,
-                          struct kfd_criu_bo_priv_data *bo_priv)
+                          struct kfd_criu_bo_priv_data *bo_priv,
+                          struct file **file)
 {
        struct kfd_process_device *pdd;
        struct kgd_mem *kgd_mem;
        if (bo_bucket->alloc_flags
            & (KFD_IOC_ALLOC_MEM_FLAGS_VRAM | KFD_IOC_ALLOC_MEM_FLAGS_GTT)) {
                ret = criu_get_prime_handle(kgd_mem, DRM_RDWR,
-                                           &bo_bucket->dmabuf_fd);
+                                           &bo_bucket->dmabuf_fd, file);
                if (ret)
                        return ret;
        } else {
 {
        struct kfd_criu_bo_bucket *bo_buckets = NULL;
        struct kfd_criu_bo_priv_data *bo_privs = NULL;
+       struct file **files = NULL;
        int ret = 0;
        uint32_t i = 0;
 
        if (!bo_buckets)
                return -ENOMEM;
 
+       files = kvzalloc(args->num_bos * sizeof(struct file *), GFP_KERNEL);
+       if (!files) {
+               ret = -ENOMEM;
+               goto exit;
+       }
+
        ret = copy_from_user(bo_buckets, (void __user *)args->bos,
                             args->num_bos * sizeof(*bo_buckets));
        if (ret) {
 
        /* Create and map new BOs */
        for (; i < args->num_bos; i++) {
-               ret = criu_restore_bo(p, &bo_buckets[i], &bo_privs[i]);
+               ret = criu_restore_bo(p, &bo_buckets[i], &bo_privs[i], &files[i]);
                if (ret) {
                        pr_debug("Failed to restore BO[%d] ret%d\n", i, ret);
                        goto exit;
                ret = -EFAULT;
 
 exit:
-       while (ret && i--) {
-               if (bo_buckets[i].alloc_flags
-                  & (KFD_IOC_ALLOC_MEM_FLAGS_VRAM | KFD_IOC_ALLOC_MEM_FLAGS_GTT))
-                       close_fd(bo_buckets[i].dmabuf_fd);
-       }
+       commit_files(files, bo_buckets, i, ret);
+       kvfree(files);
        kvfree(bo_buckets);
        kvfree(bo_privs);
        return ret;