unlock_page_cgroup(pc);
        /*
         * even after unlock, we have memcg->res.usage here and this memcg
-        * will never be freed.
+        * will never be freed, so it's safe to call css_get().
         */
        memcg_check_events(memcg, page);
        if (do_swap_account && ctype == MEM_CGROUP_CHARGE_TYPE_SWAPOUT) {
                mem_cgroup_swap_statistics(memcg, true);
-               mem_cgroup_get(memcg);
+               css_get(&memcg->css);
        }
        /*
         * Migration does not charge the res_counter for the
 
        /*
         * record memcg information,  if swapout && memcg != NULL,
-        * mem_cgroup_get() was called in uncharge().
+        * css_get() was called in uncharge().
         */
        if (do_swap_account && swapout && memcg)
                swap_cgroup_record(ent, css_id(&memcg->css));
                if (!mem_cgroup_is_root(memcg))
                        res_counter_uncharge(&memcg->memsw, PAGE_SIZE);
                mem_cgroup_swap_statistics(memcg, false);
-               mem_cgroup_put(memcg);
+               css_put(&memcg->css);
        }
        rcu_read_unlock();
 }
                 * This function is only called from task migration context now.
                 * It postpones res_counter and refcount handling till the end
                 * of task migration(mem_cgroup_clear_mc()) for performance
-                * improvement. But we cannot postpone mem_cgroup_get(to)
-                * because if the process that has been moved to @to does
-                * swap-in, the refcount of @to might be decreased to 0.
+                * improvement. But we cannot postpone css_get(to)  because if
+                * the process that has been moved to @to does swap-in, the
+                * refcount of @to might be decreased to 0.
+                *
+                * We are in attach() phase, so the cgroup is guaranteed to be
+                * alive, so we can just call css_get().
                 */
-               mem_cgroup_get(to);
+               css_get(&to->css);
                return 0;
        }
        return -EINVAL;
 {
        struct mem_cgroup *from = mc.from;
        struct mem_cgroup *to = mc.to;
+       int i;
 
        /* we must uncharge all the leftover precharges from mc.to */
        if (mc.precharge) {
                if (!mem_cgroup_is_root(mc.from))
                        res_counter_uncharge(&mc.from->memsw,
                                                PAGE_SIZE * mc.moved_swap);
-               __mem_cgroup_put(mc.from, mc.moved_swap);
+
+               for (i = 0; i < mc.moved_swap; i++)
+                       css_put(&mc.from->css);
 
                if (!mem_cgroup_is_root(mc.to)) {
                        /*
                        res_counter_uncharge(&mc.to->res,
                                                PAGE_SIZE * mc.moved_swap);
                }
-               /* we've already done mem_cgroup_get(mc.to) */
+               /* we've already done css_get(mc.to) */
                mc.moved_swap = 0;
        }
        memcg_oom_recover(from);