* After the initialization objcg->memcg is always pointing at
  * a valid memcg, but can be atomically swapped to the parent memcg.
  *
- * The caller must ensure that the returned memcg won't be released:
- * e.g. acquire the rcu_read_lock or css_set_lock.
+ * The caller must ensure that the returned memcg won't be released.
  */
 static inline struct mem_cgroup *obj_cgroup_memcg(struct obj_cgroup *objcg)
 {
+       lockdep_assert_once(rcu_read_lock_held() || lockdep_is_held(&cgroup_mutex));
        return READ_ONCE(objcg->memcg);
 }
 
        return __folio_memcg(folio);
 }
 
+/*
+ * folio_memcg_charged - If a folio is charged to a memory cgroup.
+ * @folio: Pointer to the folio.
+ *
+ * Returns true if folio is charged to a memory cgroup, otherwise returns false.
+ */
+static inline bool folio_memcg_charged(struct folio *folio)
+{
+       if (folio_memcg_kmem(folio))
+               return __folio_objcg(folio) != NULL;
+       return __folio_memcg(folio) != NULL;
+}
+
 /**
  * folio_memcg_rcu - Locklessly get the memory cgroup associated with a folio.
  * @folio: Pointer to the folio.
        unsigned long memcg_data = READ_ONCE(folio->memcg_data);
 
        VM_BUG_ON_FOLIO(folio_test_slab(folio), folio);
-       WARN_ON_ONCE(!rcu_read_lock_held());
 
        if (memcg_data & MEMCG_DATA_KMEM) {
                struct obj_cgroup *objcg;
                return obj_cgroup_memcg(objcg);
        }
 
+       WARN_ON_ONCE(!rcu_read_lock_held());
+
        return (struct mem_cgroup *)(memcg_data & ~OBJEXTS_FLAGS_MASK);
 }
 
 
 
 static void commit_charge(struct folio *folio, struct mem_cgroup *memcg)
 {
-       VM_BUG_ON_FOLIO(folio_memcg(folio), folio);
+       VM_BUG_ON_FOLIO(folio_memcg_charged(folio), folio);
        /*
         * Any of the following ensures page's memcg stability:
         *
 void split_page_memcg(struct page *head, int old_order, int new_order)
 {
        struct folio *folio = page_folio(head);
-       struct mem_cgroup *memcg = folio_memcg(folio);
        int i;
        unsigned int old_nr = 1 << old_order;
        unsigned int new_nr = 1 << new_order;
 
-       if (mem_cgroup_disabled() || !memcg)
+       if (mem_cgroup_disabled() || !folio_memcg_charged(folio))
                return;
 
        for (i = new_nr; i < old_nr; i += new_nr)
        if (folio_memcg_kmem(folio))
                obj_cgroup_get_many(__folio_objcg(folio), old_nr / new_nr - 1);
        else
-               css_get_many(&memcg->css, old_nr / new_nr - 1);
+               css_get_many(&folio_memcg(folio)->css, old_nr / new_nr - 1);
 }
 
 unsigned long mem_cgroup_usage(struct mem_cgroup *memcg, bool swap)
        struct uncharge_gather ug;
 
        /* Don't touch folio->lru of any random page, pre-check: */
-       if (!folio_memcg(folio))
+       if (!folio_memcg_charged(folio))
                return;
 
        uncharge_gather_clear(&ug);
                return;
 
        /* Page cache replacement: new folio already charged? */
-       if (folio_memcg(new))
+       if (folio_memcg_charged(new))
                return;
 
        memcg = folio_memcg(old);