diff mbox series

[1/2] mm: rmap: make hugetlb pages participate in _nr_pages_mapped

Message ID 20230306230004.1387007-2-jthoughton@google.com (mailing list archive)
State New
Headers show
Series mm: rmap: merge HugeTLB mapcount logic with THPs | expand

Commit Message

James Houghton March 6, 2023, 11 p.m. UTC 
For compound mappings (compound=true), _nr_pages_mapped will now be
incremented by COMPOUND_MAPPED when the first compound mapping is
created.

For small mappings, _nr_pages_mapped is incremented by 1 when the
particular small page is mapped for the first time. This is incompatible
with HPageVmemmapOptimize()ed folios, as most of the tail page structs
will be mapped read-only.

Currently HugeTLB always passes compound=true, but in the future,
HugeTLB pages may be mapped with small mappings.

To implement this change:
 1. Replace most of HugeTLB's calls to page_dup_file_rmap() with
    page_add_file_rmap(). The call in copy_hugetlb_page_range() is kept.
 2. Update page_add_file_rmap() and page_remove_rmap() to support
    HugeTLB folios.
 3. Update hugepage_add_anon_rmap() and hugepage_add_new_anon_rmap() to
    also increment _nr_pages_mapped properly.

With these changes, folio_large_is_mapped() no longer needs to check
_entire_mapcount.

HugeTLB doesn't use LRU or mlock, so page_add_file_rmap() and
page_remove_rmap() excludes those pieces. It is also important that
the folio_test_pmd_mappable() check is removed (or changed), as it's
possible to have a HugeTLB page whose order is not >= HPAGE_PMD_ORDER,
like arm64's CONT_PTE_SIZE HugeTLB pages.

This patch limits HugeTLB pages to 16G in size. That limit can be
increased if COMPOUND_MAPPED is raised.

Signed-off-by: James Houghton <jthoughton@google.com>

Comments

Mike Kravetz March 7, 2023, 9:54 p.m. UTC | #1 
On 03/06/23 23:00, James Houghton wrote:
> For compound mappings (compound=true), _nr_pages_mapped will now be
> incremented by COMPOUND_MAPPED when the first compound mapping is
> created.

This sentence makes it sound like incrementing by COMPOUND_MAPPED for
compound pages is introduced by this patch.  Rather, it is just for
hugetlb (now always) compound mappings.   Perhaps change that to read:
For hugetlb mappings ...

> For small mappings, _nr_pages_mapped is incremented by 1 when the
> particular small page is mapped for the first time. This is incompatible
> with HPageVmemmapOptimize()ed folios, as most of the tail page structs
> will be mapped read-only.
> 
> Currently HugeTLB always passes compound=true, but in the future,
> HugeTLB pages may be mapped with small mappings.
> 
> To implement this change:
>  1. Replace most of HugeTLB's calls to page_dup_file_rmap() with
>     page_add_file_rmap(). The call in copy_hugetlb_page_range() is kept.
>  2. Update page_add_file_rmap() and page_remove_rmap() to support
>     HugeTLB folios.
>  3. Update hugepage_add_anon_rmap() and hugepage_add_new_anon_rmap() to
>     also increment _nr_pages_mapped properly.
> 
> With these changes, folio_large_is_mapped() no longer needs to check
> _entire_mapcount.
> 
> HugeTLB doesn't use LRU or mlock, so page_add_file_rmap() and
> page_remove_rmap() excludes those pieces. It is also important that
> the folio_test_pmd_mappable() check is removed (or changed), as it's
> possible to have a HugeTLB page whose order is not >= HPAGE_PMD_ORDER,
> like arm64's CONT_PTE_SIZE HugeTLB pages.
> 
> This patch limits HugeTLB pages to 16G in size. That limit can be
> increased if COMPOUND_MAPPED is raised.
> 
> Signed-off-by: James Houghton <jthoughton@google.com>
> 

Thanks!

This is a step in the direction of having hugetlb use the same mapcount
scheme as elsewhere.  As you mention, with this in place future mapcount
changes should mostly 'just work' for hugetlb.

Because of this,
Acked-by: Mike Kravetz <mike.kravetz@oracle.com>

I have a few nits below, and I'm sure others will chime in later.

> diff --git a/mm/rmap.c b/mm/rmap.c
> index ba901c416785..4a975429b91a 100644
> --- a/mm/rmap.c
> +++ b/mm/rmap.c
> @@ -1316,19 +1316,21 @@ void page_add_file_rmap(struct page *page, struct vm_area_struct *vma,
>  	int nr = 0, nr_pmdmapped = 0;
>  	bool first;
>  
> -	VM_BUG_ON_PAGE(compound && !PageTransHuge(page), page);
> +	VM_BUG_ON_PAGE(compound && !PageTransHuge(page)
> +				&& !folio_test_hugetlb(folio), page);
>  
>  	/* Is page being mapped by PTE? Is this its first map to be added? */
>  	if (likely(!compound)) {
> +		if (unlikely(folio_test_hugetlb(folio)))
> +			VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
> +				       page);
>  		first = atomic_inc_and_test(&page->_mapcount);
>  		nr = first;
>  		if (first && folio_test_large(folio)) {
>  			nr = atomic_inc_return_relaxed(mapped);
>  			nr = (nr < COMPOUND_MAPPED);
>  		}
> -	} else if (folio_test_pmd_mappable(folio)) {
> -		/* That test is redundant: it's for safety or to optimize out */

I 'think' removing this check is OK.  It would seem that the caller
knows if the folio is mappable.  If we want a similar test, we might be
able to use something like:

	arch_hugetlb_valid_size(folio_size(folio))

> -
> +	} else {
>  		first = atomic_inc_and_test(&folio->_entire_mapcount);
>  		if (first) {
>  			nr = atomic_add_return_relaxed(COMPOUND_MAPPED, mapped);
> @@ -1345,6 +1347,9 @@ void page_add_file_rmap(struct page *page, struct vm_area_struct *vma,
>  		}
>  	}
>  
> +	if (folio_test_hugetlb(folio))
> +		return;

IMO, a comment saying hugetlb is special and does not participate in lru
would be appropriate here.

> +
>  	if (nr_pmdmapped)
>  		__lruvec_stat_mod_folio(folio, folio_test_swapbacked(folio) ?
>  			NR_SHMEM_PMDMAPPED : NR_FILE_PMDMAPPED, nr_pmdmapped);
> @@ -1373,24 +1378,18 @@ void page_remove_rmap(struct page *page, struct vm_area_struct *vma,
>  
>  	VM_BUG_ON_PAGE(compound && !PageHead(page), page);
>  
> -	/* Hugetlb pages are not counted in NR_*MAPPED */
> -	if (unlikely(folio_test_hugetlb(folio))) {
> -		/* hugetlb pages are always mapped with pmds */
> -		atomic_dec(&folio->_entire_mapcount);
> -		return;
> -	}
> -
>  	/* Is page being unmapped by PTE? Is this its last map to be removed? */
>  	if (likely(!compound)) {
> +		if (unlikely(folio_test_hugetlb(folio)))
> +			VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
> +				       page);
>  		last = atomic_add_negative(-1, &page->_mapcount);
>  		nr = last;
>  		if (last && folio_test_large(folio)) {
>  			nr = atomic_dec_return_relaxed(mapped);
>  			nr = (nr < COMPOUND_MAPPED);
>  		}
> -	} else if (folio_test_pmd_mappable(folio)) {
> -		/* That test is redundant: it's for safety or to optimize out */
> -
> +	} else {
>  		last = atomic_add_negative(-1, &folio->_entire_mapcount);
>  		if (last) {
>  			nr = atomic_sub_return_relaxed(COMPOUND_MAPPED, mapped);
> @@ -1407,6 +1406,9 @@ void page_remove_rmap(struct page *page, struct vm_area_struct *vma,
>  		}
>  	}
>  
> +	if (folio_test_hugetlb(folio))
> +		return;

Same as above in page_add_file_rmap.

> +
>  	if (nr_pmdmapped) {
>  		if (folio_test_anon(folio))
>  			idx = NR_ANON_THPS;
> @@ -2541,9 +2543,11 @@ void hugepage_add_anon_rmap(struct page *page, struct vm_area_struct *vma,
>  	first = atomic_inc_and_test(&folio->_entire_mapcount);
>  	VM_BUG_ON_PAGE(!first && (flags & RMAP_EXCLUSIVE), page);
>  	VM_BUG_ON_PAGE(!first && PageAnonExclusive(page), page);
> -	if (first)
> +	if (first) {
> +		atomic_add(COMPOUND_MAPPED, &folio->_nr_pages_mapped);
>  		__page_set_anon_rmap(folio, page, vma, address,
>  				     !!(flags & RMAP_EXCLUSIVE));
> +	}
>  }
>  
>  void hugepage_add_new_anon_rmap(struct folio *folio,
> @@ -2552,6 +2556,7 @@ void hugepage_add_new_anon_rmap(struct folio *folio,
>  	BUG_ON(address < vma->vm_start || address >= vma->vm_end);
>  	/* increment count (starts at -1) */
>  	atomic_set(&folio->_entire_mapcount, 0);
> +	atomic_set(&folio->_nr_pages_mapped, COMPOUND_MAPPED);
>  	folio_clear_hugetlb_restore_reserve(folio);
>  	__page_set_anon_rmap(folio, &folio->page, vma, address, 1);
>  }

Should we look at perhaps modifying page_add_anon_rmap and
folio_add_new_anon_rmap as well?
James Houghton March 8, 2023, 12:36 a.m. UTC | #2 
On Tue, Mar 7, 2023 at 1:54 PM Mike Kravetz <mike.kravetz@oracle.com> wrote:
>
> On 03/06/23 23:00, James Houghton wrote:
> > For compound mappings (compound=true), _nr_pages_mapped will now be
> > incremented by COMPOUND_MAPPED when the first compound mapping is
> > created.
>
> This sentence makes it sound like incrementing by COMPOUND_MAPPED for
> compound pages is introduced by this patch.  Rather, it is just for
> hugetlb (now always) compound mappings.   Perhaps change that to read:
> For hugetlb mappings ...

Yes this is kind of confusing. I'll fix it like you suggest.

>
> > For small mappings, _nr_pages_mapped is incremented by 1 when the
> > particular small page is mapped for the first time. This is incompatible
> > with HPageVmemmapOptimize()ed folios, as most of the tail page structs
> > will be mapped read-only.
> >
> > Currently HugeTLB always passes compound=true, but in the future,
> > HugeTLB pages may be mapped with small mappings.
> >
> > To implement this change:
> >  1. Replace most of HugeTLB's calls to page_dup_file_rmap() with
> >     page_add_file_rmap(). The call in copy_hugetlb_page_range() is kept.
> >  2. Update page_add_file_rmap() and page_remove_rmap() to support
> >     HugeTLB folios.
> >  3. Update hugepage_add_anon_rmap() and hugepage_add_new_anon_rmap() to
> >     also increment _nr_pages_mapped properly.
> >
> > With these changes, folio_large_is_mapped() no longer needs to check
> > _entire_mapcount.
> >
> > HugeTLB doesn't use LRU or mlock, so page_add_file_rmap() and
> > page_remove_rmap() excludes those pieces. It is also important that
> > the folio_test_pmd_mappable() check is removed (or changed), as it's
> > possible to have a HugeTLB page whose order is not >= HPAGE_PMD_ORDER,
> > like arm64's CONT_PTE_SIZE HugeTLB pages.
> >
> > This patch limits HugeTLB pages to 16G in size. That limit can be
> > increased if COMPOUND_MAPPED is raised.
> >
> > Signed-off-by: James Houghton <jthoughton@google.com>
> >
>
> Thanks!
>
> This is a step in the direction of having hugetlb use the same mapcount
> scheme as elsewhere.  As you mention, with this in place future mapcount
> changes should mostly 'just work' for hugetlb.
>
> Because of this,
> Acked-by: Mike Kravetz <mike.kravetz@oracle.com>

Thanks!

>
> I have a few nits below, and I'm sure others will chime in later.
>
> > diff --git a/mm/rmap.c b/mm/rmap.c
> > index ba901c416785..4a975429b91a 100644
> > --- a/mm/rmap.c
> > +++ b/mm/rmap.c
> > @@ -1316,19 +1316,21 @@ void page_add_file_rmap(struct page *page, struct vm_area_struct *vma,
> >       int nr = 0, nr_pmdmapped = 0;
> >       bool first;
> >
> > -     VM_BUG_ON_PAGE(compound && !PageTransHuge(page), page);
> > +     VM_BUG_ON_PAGE(compound && !PageTransHuge(page)
> > +                             && !folio_test_hugetlb(folio), page);
> >
> >       /* Is page being mapped by PTE? Is this its first map to be added? */
> >       if (likely(!compound)) {
> > +             if (unlikely(folio_test_hugetlb(folio)))
> > +                     VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
> > +                                    page);
> >               first = atomic_inc_and_test(&page->_mapcount);
> >               nr = first;
> >               if (first && folio_test_large(folio)) {
> >                       nr = atomic_inc_return_relaxed(mapped);
> >                       nr = (nr < COMPOUND_MAPPED);
> >               }
> > -     } else if (folio_test_pmd_mappable(folio)) {
> > -             /* That test is redundant: it's for safety or to optimize out */
>
> I 'think' removing this check is OK.  It would seem that the caller
> knows if the folio is mappable.  If we want a similar test, we might be
> able to use something like:
>
>         arch_hugetlb_valid_size(folio_size(folio))
>

Ack. I think leaving the check(s) removed is fine.

> > -
> > +     } else {
> >               first = atomic_inc_and_test(&folio->_entire_mapcount);
> >               if (first) {
> >                       nr = atomic_add_return_relaxed(COMPOUND_MAPPED, mapped);
> > @@ -1345,6 +1347,9 @@ void page_add_file_rmap(struct page *page, struct vm_area_struct *vma,
> >               }
> >       }
> >
> > +     if (folio_test_hugetlb(folio))
> > +             return;
>
> IMO, a comment saying hugetlb is special and does not participate in lru
> would be appropriate here.

Will do.

>
> > +
> >       if (nr_pmdmapped)
> >               __lruvec_stat_mod_folio(folio, folio_test_swapbacked(folio) ?
> >                       NR_SHMEM_PMDMAPPED : NR_FILE_PMDMAPPED, nr_pmdmapped);
> > @@ -1373,24 +1378,18 @@ void page_remove_rmap(struct page *page, struct vm_area_struct *vma,
> >
> >       VM_BUG_ON_PAGE(compound && !PageHead(page), page);
> >
> > -     /* Hugetlb pages are not counted in NR_*MAPPED */
> > -     if (unlikely(folio_test_hugetlb(folio))) {
> > -             /* hugetlb pages are always mapped with pmds */
> > -             atomic_dec(&folio->_entire_mapcount);
> > -             return;
> > -     }
> > -
> >       /* Is page being unmapped by PTE? Is this its last map to be removed? */
> >       if (likely(!compound)) {
> > +             if (unlikely(folio_test_hugetlb(folio)))
> > +                     VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
> > +                                    page);
> >               last = atomic_add_negative(-1, &page->_mapcount);
> >               nr = last;
> >               if (last && folio_test_large(folio)) {
> >                       nr = atomic_dec_return_relaxed(mapped);
> >                       nr = (nr < COMPOUND_MAPPED);
> >               }
> > -     } else if (folio_test_pmd_mappable(folio)) {
> > -             /* That test is redundant: it's for safety or to optimize out */
> > -
> > +     } else {
> >               last = atomic_add_negative(-1, &folio->_entire_mapcount);
> >               if (last) {
> >                       nr = atomic_sub_return_relaxed(COMPOUND_MAPPED, mapped);
> > @@ -1407,6 +1406,9 @@ void page_remove_rmap(struct page *page, struct vm_area_struct *vma,
> >               }
> >       }
> >
> > +     if (folio_test_hugetlb(folio))
> > +             return;
>
> Same as above in page_add_file_rmap.
>
> > +
> >       if (nr_pmdmapped) {
> >               if (folio_test_anon(folio))
> >                       idx = NR_ANON_THPS;
> > @@ -2541,9 +2543,11 @@ void hugepage_add_anon_rmap(struct page *page, struct vm_area_struct *vma,
> >       first = atomic_inc_and_test(&folio->_entire_mapcount);
> >       VM_BUG_ON_PAGE(!first && (flags & RMAP_EXCLUSIVE), page);
> >       VM_BUG_ON_PAGE(!first && PageAnonExclusive(page), page);
> > -     if (first)
> > +     if (first) {
> > +             atomic_add(COMPOUND_MAPPED, &folio->_nr_pages_mapped);
> >               __page_set_anon_rmap(folio, page, vma, address,
> >                                    !!(flags & RMAP_EXCLUSIVE));
> > +     }
> >  }
> >
> >  void hugepage_add_new_anon_rmap(struct folio *folio,
> > @@ -2552,6 +2556,7 @@ void hugepage_add_new_anon_rmap(struct folio *folio,
> >       BUG_ON(address < vma->vm_start || address >= vma->vm_end);
> >       /* increment count (starts at -1) */
> >       atomic_set(&folio->_entire_mapcount, 0);
> > +     atomic_set(&folio->_nr_pages_mapped, COMPOUND_MAPPED);
> >       folio_clear_hugetlb_restore_reserve(folio);
> >       __page_set_anon_rmap(folio, &folio->page, vma, address, 1);
> >  }
>
> Should we look at perhaps modifying page_add_anon_rmap and
> folio_add_new_anon_rmap as well?

I think I can merge hugepage_add_anon_rmap with page_add_anon_rmap and
hugepage_add_new_anon_rmap with folio_add_new_anon_rmap. With them
merged, it's pretty easy to see what HugeTLB does differently from
generic mm, which is nice. :)
Peter Xu March 8, 2023, 9:56 p.m. UTC | #3 
On Tue, Mar 07, 2023 at 04:36:51PM -0800, James Houghton wrote:
> > >       if (likely(!compound)) {
> > > +             if (unlikely(folio_test_hugetlb(folio)))
> > > +                     VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
> > > +                                    page);

How about moving folio_test_hugetlb() into the BUG_ON()?

                VM_BUG_ON_PAGE(folio_test_hugetlb(folio) &&
                               HPageVmemmapOptimized(&folio->page),
                               page);

Note that BUG_ON() already contains an "unlikely".

> > >               first = atomic_inc_and_test(&page->_mapcount);
> > >               nr = first;
> > >               if (first && folio_test_large(folio)) {
> > >                       nr = atomic_inc_return_relaxed(mapped);
> > >                       nr = (nr < COMPOUND_MAPPED);
> > >               }
> > > -     } else if (folio_test_pmd_mappable(folio)) {
> > > -             /* That test is redundant: it's for safety or to optimize out */
> >
> > I 'think' removing this check is OK.  It would seem that the caller
> > knows if the folio is mappable.  If we want a similar test, we might be
> > able to use something like:
> >
> >         arch_hugetlb_valid_size(folio_size(folio))
> >
> 
> Ack. I think leaving the check(s) removed is fine.

Would it still be good to keep that as another BUG_ON()?
James Houghton March 9, 2023, 7:58 p.m. UTC | #4 
On Wed, Mar 8, 2023 at 1:56 PM Peter Xu <peterx@redhat.com> wrote:
>
> On Tue, Mar 07, 2023 at 04:36:51PM -0800, James Houghton wrote:
> > > >       if (likely(!compound)) {
> > > > +             if (unlikely(folio_test_hugetlb(folio)))
> > > > +                     VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
> > > > +                                    page);
>
> How about moving folio_test_hugetlb() into the BUG_ON()?
>
>                 VM_BUG_ON_PAGE(folio_test_hugetlb(folio) &&
>                                HPageVmemmapOptimized(&folio->page),
>                                page);
>
> Note that BUG_ON() already contains an "unlikely".

Ok I can do that. It's a little cleaner.

> > > >               first = atomic_inc_and_test(&page->_mapcount);
> > > >               nr = first;
> > > >               if (first && folio_test_large(folio)) {
> > > >                       nr = atomic_inc_return_relaxed(mapped);
> > > >                       nr = (nr < COMPOUND_MAPPED);
> > > >               }
> > > > -     } else if (folio_test_pmd_mappable(folio)) {
> > > > -             /* That test is redundant: it's for safety or to optimize out */
> > >
> > > I 'think' removing this check is OK.  It would seem that the caller
> > > knows if the folio is mappable.  If we want a similar test, we might be
> > > able to use something like:
> > >
> > >         arch_hugetlb_valid_size(folio_size(folio))
> > >
> >
> > Ack. I think leaving the check(s) removed is fine.
>
> Would it still be good to keep that as another BUG_ON()?

Sure, that sounds reasonable to me. I'll add it unless someone disagrees.

As you suggested in your other email, I'll also add a BUG_ON() if we
attempt to do a non-compound mapping of a folio that is larger than
COMPOUND_MAPPED / 2. (Maybe a BUG_ON() in alloc_hugetlb_folio() to
check that the size of the folio we're allocating is less than
COMPOUND_MAPPED / 2 makes sense instead. Just an idea.)

- James
diff mbox series

Patch

diff --git a/include/linux/mm.h b/include/linux/mm.h
index ce1590933995..25c898c51129 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -1143,12 +1143,7 @@  static inline int total_mapcount(struct page *page)
 
 static inline bool folio_large_is_mapped(struct folio *folio)
 {
-	/*
-	 * Reading _entire_mapcount below could be omitted if hugetlb
-	 * participated in incrementing nr_pages_mapped when compound mapped.
-	 */
-	return atomic_read(&folio->_nr_pages_mapped) > 0 ||
-		atomic_read(&folio->_entire_mapcount) >= 0;
+	return atomic_read(&folio->_nr_pages_mapped) > 0;
 }
 
 /**
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index 07abcb6eb203..bf3d327cd1b9 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -5911,7 +5911,7 @@  static vm_fault_t hugetlb_no_page(struct mm_struct *mm,
 	if (anon_rmap)
 		hugepage_add_new_anon_rmap(folio, vma, haddr);
 	else
-		page_dup_file_rmap(&folio->page, true);
+		page_add_file_rmap(&folio->page, vma, true);
 	new_pte = make_huge_pte(vma, &folio->page, ((vma->vm_flags & VM_WRITE)
 				&& (vma->vm_flags & VM_SHARED)));
 	/*
@@ -6293,7 +6293,7 @@  int hugetlb_mcopy_atomic_pte(struct mm_struct *dst_mm,
 		goto out_release_unlock;
 
 	if (folio_in_pagecache)
-		page_dup_file_rmap(&folio->page, true);
+		page_add_file_rmap(&folio->page, dst_vma, true);
 	else
 		hugepage_add_new_anon_rmap(folio, dst_vma, dst_addr);
 
diff --git a/mm/internal.h b/mm/internal.h
index fce94775819c..400451a4dd49 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -55,8 +55,7 @@  void page_writeback_init(void);
 /*
  * If a 16GB hugetlb folio were mapped by PTEs of all of its 4kB pages,
  * its nr_pages_mapped would be 0x400000: choose the COMPOUND_MAPPED bit
- * above that range, instead of 2*(PMD_SIZE/PAGE_SIZE).  Hugetlb currently
- * leaves nr_pages_mapped at 0, but avoid surprise if it participates later.
+ * above that range, instead of 2*(PMD_SIZE/PAGE_SIZE).
  */
 #define COMPOUND_MAPPED		0x800000
 #define FOLIO_PAGES_MAPPED	(COMPOUND_MAPPED - 1)
diff --git a/mm/migrate.c b/mm/migrate.c
index 476cd24e8f32..c95c1cbc7a47 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -254,7 +254,7 @@  static bool remove_migration_pte(struct folio *folio,
 				hugepage_add_anon_rmap(new, vma, pvmw.address,
 						       rmap_flags);
 			else
-				page_dup_file_rmap(new, true);
+				page_add_file_rmap(new, vma, true);
 			set_huge_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
 		} else
 #endif
diff --git a/mm/rmap.c b/mm/rmap.c
index ba901c416785..4a975429b91a 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1316,19 +1316,21 @@  void page_add_file_rmap(struct page *page, struct vm_area_struct *vma,
 	int nr = 0, nr_pmdmapped = 0;
 	bool first;
 
-	VM_BUG_ON_PAGE(compound && !PageTransHuge(page), page);
+	VM_BUG_ON_PAGE(compound && !PageTransHuge(page)
+				&& !folio_test_hugetlb(folio), page);
 
 	/* Is page being mapped by PTE? Is this its first map to be added? */
 	if (likely(!compound)) {
+		if (unlikely(folio_test_hugetlb(folio)))
+			VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
+				       page);
 		first = atomic_inc_and_test(&page->_mapcount);
 		nr = first;
 		if (first && folio_test_large(folio)) {
 			nr = atomic_inc_return_relaxed(mapped);
 			nr = (nr < COMPOUND_MAPPED);
 		}
-	} else if (folio_test_pmd_mappable(folio)) {
-		/* That test is redundant: it's for safety or to optimize out */
-
+	} else {
 		first = atomic_inc_and_test(&folio->_entire_mapcount);
 		if (first) {
 			nr = atomic_add_return_relaxed(COMPOUND_MAPPED, mapped);
@@ -1345,6 +1347,9 @@  void page_add_file_rmap(struct page *page, struct vm_area_struct *vma,
 		}
 	}
 
+	if (folio_test_hugetlb(folio))
+		return;
+
 	if (nr_pmdmapped)
 		__lruvec_stat_mod_folio(folio, folio_test_swapbacked(folio) ?
 			NR_SHMEM_PMDMAPPED : NR_FILE_PMDMAPPED, nr_pmdmapped);
@@ -1373,24 +1378,18 @@  void page_remove_rmap(struct page *page, struct vm_area_struct *vma,
 
 	VM_BUG_ON_PAGE(compound && !PageHead(page), page);
 
-	/* Hugetlb pages are not counted in NR_*MAPPED */
-	if (unlikely(folio_test_hugetlb(folio))) {
-		/* hugetlb pages are always mapped with pmds */
-		atomic_dec(&folio->_entire_mapcount);
-		return;
-	}
-
 	/* Is page being unmapped by PTE? Is this its last map to be removed? */
 	if (likely(!compound)) {
+		if (unlikely(folio_test_hugetlb(folio)))
+			VM_BUG_ON_PAGE(HPageVmemmapOptimized(&folio->page),
+				       page);
 		last = atomic_add_negative(-1, &page->_mapcount);
 		nr = last;
 		if (last && folio_test_large(folio)) {
 			nr = atomic_dec_return_relaxed(mapped);
 			nr = (nr < COMPOUND_MAPPED);
 		}
-	} else if (folio_test_pmd_mappable(folio)) {
-		/* That test is redundant: it's for safety or to optimize out */
-
+	} else {
 		last = atomic_add_negative(-1, &folio->_entire_mapcount);
 		if (last) {
 			nr = atomic_sub_return_relaxed(COMPOUND_MAPPED, mapped);
@@ -1407,6 +1406,9 @@  void page_remove_rmap(struct page *page, struct vm_area_struct *vma,
 		}
 	}
 
+	if (folio_test_hugetlb(folio))
+		return;
+
 	if (nr_pmdmapped) {
 		if (folio_test_anon(folio))
 			idx = NR_ANON_THPS;
@@ -2541,9 +2543,11 @@  void hugepage_add_anon_rmap(struct page *page, struct vm_area_struct *vma,
 	first = atomic_inc_and_test(&folio->_entire_mapcount);
 	VM_BUG_ON_PAGE(!first && (flags & RMAP_EXCLUSIVE), page);
 	VM_BUG_ON_PAGE(!first && PageAnonExclusive(page), page);
-	if (first)
+	if (first) {
+		atomic_add(COMPOUND_MAPPED, &folio->_nr_pages_mapped);
 		__page_set_anon_rmap(folio, page, vma, address,
 				     !!(flags & RMAP_EXCLUSIVE));
+	}
 }
 
 void hugepage_add_new_anon_rmap(struct folio *folio,
@@ -2552,6 +2556,7 @@  void hugepage_add_new_anon_rmap(struct folio *folio,
 	BUG_ON(address < vma->vm_start || address >= vma->vm_end);
 	/* increment count (starts at -1) */
 	atomic_set(&folio->_entire_mapcount, 0);
+	atomic_set(&folio->_nr_pages_mapped, COMPOUND_MAPPED);
 	folio_clear_hugetlb_restore_reserve(folio);
 	__page_set_anon_rmap(folio, &folio->page, vma, address, 1);
 }