1 files changed, 7 insertions, 9 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index bd9052a5d3a..e37c44d88d4 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -735,7 +735,7 @@ struct mem_cgroup *mem_cgroup_from_task(struct task_struct *p)
                                 struct mem_cgroup, css);
 }
 
-static struct mem_cgroup *try_get_mem_cgroup_from_mm(struct mm_struct *mm)
+struct mem_cgroup *try_get_mem_cgroup_from_mm(struct mm_struct *mm)
 {
         struct mem_cgroup *mem = NULL;
 
@@ -5414,18 +5414,16 @@ static void mem_cgroup_move_task(struct cgroup_subsys *ss,
                                 struct cgroup *old_cont,
                                 struct task_struct *p)
 {
-        struct mm_struct *mm;
+        struct mm_struct *mm = get_task_mm(p);
 
-        if (!mc.to)
-                /* no need to move charge */
-                return;
-
-        mm = get_task_mm(p);
         if (mm) {
-                mem_cgroup_move_charge(mm);
+                if (mc.to)
+                        mem_cgroup_move_charge(mm);
+                put_swap_token(mm);
                 mmput(mm);
         }
-        mem_cgroup_clear_mc();
+        if (mc.to)
+                mem_cgroup_clear_mc();
 }
 #else   /* !CONFIG_MMU */
 static int mem_cgroup_can_attach(struct cgroup_subsys *ss,