rcu_read_unlock();
 }
 
-#ifdef CONFIG_TRANSPARENT_HUGEPAGE
-void mem_cgroup_split_huge_fixup(struct page *head);
-#endif
+void split_page_memcg(struct page *head, unsigned int nr);
 
 #else /* CONFIG_MEMCG */
 
        return 0;
 }
 
-static inline void mem_cgroup_split_huge_fixup(struct page *head)
+static inline void split_page_memcg(struct page *head, unsigned int nr)
 {
 }
 
 
        int i;
 
        /* complete memcg works before add pages to LRU */
-       mem_cgroup_split_huge_fixup(head);
+       split_page_memcg(head, nr);
 
        if (PageAnon(head) && PageSwapCache(head)) {
                swp_entry_t entry = { .val = page_private(head) };
 
 
 #endif /* CONFIG_MEMCG_KMEM */
 
-#ifdef CONFIG_TRANSPARENT_HUGEPAGE
 /*
- * Because page_memcg(head) is not set on compound tails, set it now.
+ * Because page_memcg(head) is not set on tails, set it now.
  */
-void mem_cgroup_split_huge_fixup(struct page *head)
+void split_page_memcg(struct page *head, unsigned int nr)
 {
        struct mem_cgroup *memcg = page_memcg(head);
        int i;
 
-       if (mem_cgroup_disabled())
+       if (mem_cgroup_disabled() || !memcg)
                return;
 
-       for (i = 1; i < HPAGE_PMD_NR; i++) {
-               css_get(&memcg->css);
-               head[i].memcg_data = (unsigned long)memcg;
-       }
+       for (i = 1; i < nr; i++)
+               head[i].memcg_data = head->memcg_data;
+       css_get_many(&memcg->css, nr - 1);
 }
-#endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 
 #ifdef CONFIG_MEMCG_SWAP
 /**