lock_page_memcg(page);
        if (!TestSetPageDirty(page)) {
                struct address_space *mapping = page_mapping(page);
-               unsigned long flags;
 
                if (!mapping) {
                        unlock_page_memcg(page);
                        return 1;
                }
-
-               xa_lock_irqsave(&mapping->i_pages, flags);
-               BUG_ON(page_mapping(page) != mapping);
-               WARN_ON_ONCE(!PagePrivate(page) && !PageUptodate(page));
-               account_page_dirtied(page, mapping);
-               __xa_set_mark(&mapping->i_pages, page_index(page),
-                                  PAGECACHE_TAG_DIRTY);
-               xa_unlock_irqrestore(&mapping->i_pages, flags);
+               __set_page_dirty(page, mapping, !PagePrivate(page));
                unlock_page_memcg(page);
 
                if (mapping->host) {