#include <linux/slab.h>
 #include <linux/vmalloc.h>
 #include <linux/kthread.h>
-#include <linux/cgroup.h>
 #include <linux/module.h>
 #include <linux/sort.h>
 #include <linux/sched/mm.h>
 #include <linux/sched/signal.h>
+#include <linux/sched/vhost_task.h>
 #include <linux/interval_tree_generic.h>
 #include <linux/nospec.h>
 #include <linux/kcov.h>
                 * test_and_set_bit() implies a memory barrier.
                 */
                llist_add(&work->node, &dev->worker->work_list);
-               wake_up_process(dev->worker->task);
+               wake_up_process(dev->worker->vtsk->task);
        }
 }
 EXPORT_SYMBOL_GPL(vhost_work_queue);
 static int vhost_worker(void *data)
 {
        struct vhost_worker *worker = data;
-       struct vhost_dev *dev = worker->dev;
        struct vhost_work *work, *work_next;
        struct llist_node *node;
 
-       kthread_use_mm(dev->mm);
-
        for (;;) {
                /* mb paired w/ kthread_stop */
                set_current_state(TASK_INTERRUPTIBLE);
 
-               if (kthread_should_stop()) {
+               if (vhost_task_should_stop(worker->vtsk)) {
                        __set_current_state(TASK_RUNNING);
                        break;
                }
                                schedule();
                }
        }
-       kthread_unuse_mm(dev->mm);
+
        return 0;
 }
 
 }
 EXPORT_SYMBOL_GPL(vhost_dev_check_owner);
 
-struct vhost_attach_cgroups_struct {
-       struct vhost_work work;
-       struct task_struct *owner;
-       int ret;
-};
-
-static void vhost_attach_cgroups_work(struct vhost_work *work)
-{
-       struct vhost_attach_cgroups_struct *s;
-
-       s = container_of(work, struct vhost_attach_cgroups_struct, work);
-       s->ret = cgroup_attach_task_all(s->owner, current);
-}
-
-static int vhost_attach_cgroups(struct vhost_dev *dev)
-{
-       struct vhost_attach_cgroups_struct attach;
-
-       attach.owner = current;
-       vhost_work_init(&attach.work, vhost_attach_cgroups_work);
-       vhost_work_queue(dev, &attach.work);
-       vhost_dev_flush(dev);
-       return attach.ret;
-}
-
 /* Caller should have device mutex */
 bool vhost_dev_has_owner(struct vhost_dev *dev)
 {
 
        dev->worker = NULL;
        WARN_ON(!llist_empty(&worker->work_list));
-       kthread_stop(worker->task);
+       vhost_task_stop(worker->vtsk);
        kfree(worker);
 }
 
 static int vhost_worker_create(struct vhost_dev *dev)
 {
        struct vhost_worker *worker;
-       struct task_struct *task;
+       struct vhost_task *vtsk;
+       char name[TASK_COMM_LEN];
        int ret;
 
        worker = kzalloc(sizeof(*worker), GFP_KERNEL_ACCOUNT);
                return -ENOMEM;
 
        dev->worker = worker;
-       worker->dev = dev;
        worker->kcov_handle = kcov_common_handle();
        init_llist_head(&worker->work_list);
+       snprintf(name, sizeof(name), "vhost-%d", current->pid);
 
-       task = kthread_create(vhost_worker, worker, "vhost-%d", current->pid);
-       if (IS_ERR(task)) {
-               ret = PTR_ERR(task);
+       vtsk = vhost_task_create(vhost_worker, worker, name);
+       if (!vtsk) {
+               ret = -ENOMEM;
                goto free_worker;
        }
 
-       worker->task = task;
-       wake_up_process(task); /* avoid contributing to loadavg */
-
-       ret = vhost_attach_cgroups(dev);
-       if (ret)
-               goto stop_worker;
-
+       worker->vtsk = vtsk;
+       vhost_task_start(vtsk);
        return 0;
 
-stop_worker:
-       kthread_stop(worker->task);
 free_worker:
        kfree(worker);
        dev->worker = NULL;