#include <linux/uio.h>
 #include <linux/ioprio.h>
 #include <linux/blk-cgroup.h>
+#include <linux/sched/mm.h>
 
 #include "loop.h"
 
 {
        struct loop_cmd *cmd = container_of(iocb, struct loop_cmd, iocb);
 
-       if (cmd->css)
-               css_put(cmd->css);
        cmd->ret = ret;
        lo_rw_aio_do_completion(cmd);
 }
        cmd->iocb.ki_complete = lo_rw_aio_complete;
        cmd->iocb.ki_flags = IOCB_DIRECT;
        cmd->iocb.ki_ioprio = IOPRIO_PRIO_VALUE(IOPRIO_CLASS_NONE, 0);
-       if (cmd->css)
-               kthread_associate_blkcg(cmd->css);
 
        if (rw == WRITE)
                ret = call_write_iter(file, &cmd->iocb, &iter);
                ret = call_read_iter(file, &cmd->iocb, &iter);
 
        lo_rw_aio_do_completion(cmd);
-       kthread_associate_blkcg(NULL);
 
        if (ret != -EIOCBQUEUED)
                cmd->iocb.ki_complete(&cmd->iocb, ret, 0);
        struct list_head cmd_list;
        struct list_head idle_list;
        struct loop_device *lo;
-       struct cgroup_subsys_state *css;
+       struct cgroup_subsys_state *blkcg_css;
        unsigned long last_ran_at;
 };
 
 
        spin_lock_irq(&lo->lo_work_lock);
 
-       if (queue_on_root_worker(cmd->css))
+       if (queue_on_root_worker(cmd->blkcg_css))
                goto queue_work;
 
        node = &lo->worker_tree.rb_node;
        while (*node) {
                parent = *node;
                cur_worker = container_of(*node, struct loop_worker, rb_node);
-               if (cur_worker->css == cmd->css) {
+               if (cur_worker->blkcg_css == cmd->blkcg_css) {
                        worker = cur_worker;
                        break;
-               } else if ((long)cur_worker->css < (long)cmd->css) {
+               } else if ((long)cur_worker->blkcg_css < (long)cmd->blkcg_css) {
                        node = &(*node)->rb_left;
                } else {
                        node = &(*node)->rb_right;
        worker = kzalloc(sizeof(struct loop_worker), GFP_NOWAIT | __GFP_NOWARN);
        /*
         * In the event we cannot allocate a worker, just queue on the
-        * rootcg worker
+        * rootcg worker and issue the I/O as the rootcg
         */
-       if (!worker)
+       if (!worker) {
+               cmd->blkcg_css = NULL;
+               if (cmd->memcg_css)
+                       css_put(cmd->memcg_css);
+               cmd->memcg_css = NULL;
                goto queue_work;
+       }
 
-       worker->css = cmd->css;
-       css_get(worker->css);
+       worker->blkcg_css = cmd->blkcg_css;
+       css_get(worker->blkcg_css);
        INIT_WORK(&worker->work, loop_workfn);
        INIT_LIST_HEAD(&worker->cmd_list);
        INIT_LIST_HEAD(&worker->idle_list);
                                idle_list) {
                list_del(&worker->idle_list);
                rb_erase(&worker->rb_node, &lo->worker_tree);
-               css_put(worker->css);
+               css_put(worker->blkcg_css);
                kfree(worker);
        }
        spin_unlock_irq(&lo->lo_work_lock);
        }
 
        /* always use the first bio's css */
+       cmd->blkcg_css = NULL;
+       cmd->memcg_css = NULL;
 #ifdef CONFIG_BLK_CGROUP
-       if (cmd->use_aio && rq->bio && rq->bio->bi_blkg) {
-               cmd->css = &bio_blkcg(rq->bio)->css;
-               css_get(cmd->css);
-       } else
+       if (rq->bio && rq->bio->bi_blkg) {
+               cmd->blkcg_css = &bio_blkcg(rq->bio)->css;
+#ifdef CONFIG_MEMCG
+               cmd->memcg_css =
+                       cgroup_get_e_css(cmd->blkcg_css->cgroup,
+                                       &memory_cgrp_subsys);
+#endif
+       }
 #endif
-               cmd->css = NULL;
        loop_queue_work(lo, cmd);
 
        return BLK_STS_OK;
        const bool write = op_is_write(req_op(rq));
        struct loop_device *lo = rq->q->queuedata;
        int ret = 0;
+       struct mem_cgroup *old_memcg = NULL;
 
        if (write && (lo->lo_flags & LO_FLAGS_READ_ONLY)) {
                ret = -EIO;
                goto failed;
        }
 
+       if (cmd->blkcg_css)
+               kthread_associate_blkcg(cmd->blkcg_css);
+       if (cmd->memcg_css)
+               old_memcg = set_active_memcg(
+                       mem_cgroup_from_css(cmd->memcg_css));
+
        ret = do_req_filebacked(lo, rq);
+
+       if (cmd->blkcg_css)
+               kthread_associate_blkcg(NULL);
+
+       if (cmd->memcg_css) {
+               set_active_memcg(old_memcg);
+               css_put(cmd->memcg_css);
+       }
  failed:
        /* complete non-aio request */
        if (!cmd->use_aio || ret) {
                        break;
                list_del(&worker->idle_list);
                rb_erase(&worker->rb_node, &lo->worker_tree);
-               css_put(worker->css);
+               css_put(worker->blkcg_css);
                kfree(worker);
        }
        if (!list_empty(&lo->idle_worker_list))