*/
 #define RMAP_COMPOUND          ((__force rmap_t)BIT(1))
 
+/*
+ * Internally, we're using an enum to specify the granularity. We make the
+ * compiler emit specialized code for each granularity.
+ */
+enum rmap_level {
+       RMAP_LEVEL_PTE = 0,
+       RMAP_LEVEL_PMD,
+};
+
+static inline void __folio_rmap_sanity_checks(struct folio *folio,
+               struct page *page, int nr_pages, enum rmap_level level)
+{
+       /* hugetlb folios are handled separately. */
+       VM_WARN_ON_FOLIO(folio_test_hugetlb(folio), folio);
+       VM_WARN_ON_FOLIO(folio_test_large(folio) &&
+                        !folio_test_large_rmappable(folio), folio);
+
+       VM_WARN_ON_ONCE(nr_pages <= 0);
+       VM_WARN_ON_FOLIO(page_folio(page) != folio, folio);
+       VM_WARN_ON_FOLIO(page_folio(page + nr_pages - 1) != folio, folio);
+
+       switch (level) {
+       case RMAP_LEVEL_PTE:
+               break;
+       case RMAP_LEVEL_PMD:
+               /*
+                * We don't support folios larger than a single PMD yet. So
+                * when RMAP_LEVEL_PMD is set, we assume that we are creating
+                * a single "entire" mapping of the folio.
+                */
+               VM_WARN_ON_FOLIO(folio_nr_pages(folio) != HPAGE_PMD_NR, folio);
+               VM_WARN_ON_FOLIO(nr_pages != HPAGE_PMD_NR, folio);
+               break;
+       default:
+               VM_WARN_ON_ONCE(true);
+       }
+}
+
 /*
  * rmap interfaces called when adding or removing pte of page
  */
                unsigned long address);
 void page_add_file_rmap(struct page *, struct vm_area_struct *,
                bool compound);
-void folio_add_file_rmap_range(struct folio *, struct page *, unsigned int nr,
-               struct vm_area_struct *, bool compound);
+void folio_add_file_rmap_ptes(struct folio *, struct page *, int nr_pages,
+               struct vm_area_struct *);
+#define folio_add_file_rmap_pte(folio, page, vma) \
+       folio_add_file_rmap_ptes(folio, page, 1, vma)
+void folio_add_file_rmap_pmd(struct folio *, struct page *,
+               struct vm_area_struct *);
 void page_remove_rmap(struct page *, struct vm_area_struct *,
                bool compound);
 
 
        __lruvec_stat_mod_folio(folio, NR_ANON_MAPPED, nr);
 }
 
-/**
- * folio_add_file_rmap_range - add pte mapping to page range of a folio
- * @folio:     The folio to add the mapping to
- * @page:      The first page to add
- * @nr_pages:  The number of pages which will be mapped
- * @vma:       the vm area in which the mapping is added
- * @compound:  charge the page as compound or small page
- *
- * The page range of folio is defined by [first_page, first_page + nr_pages)
- *
- * The caller needs to hold the pte lock.
- */
-void folio_add_file_rmap_range(struct folio *folio, struct page *page,
-                       unsigned int nr_pages, struct vm_area_struct *vma,
-                       bool compound)
+static __always_inline void __folio_add_file_rmap(struct folio *folio,
+               struct page *page, int nr_pages, struct vm_area_struct *vma,
+               enum rmap_level level)
 {
        atomic_t *mapped = &folio->_nr_pages_mapped;
-       unsigned int nr_pmdmapped = 0, first;
-       int nr = 0;
+       int nr = 0, nr_pmdmapped = 0, first;
 
-       VM_WARN_ON_FOLIO(folio_test_hugetlb(folio), folio);
-       VM_WARN_ON_FOLIO(compound && !folio_test_pmd_mappable(folio), folio);
+       VM_WARN_ON_FOLIO(folio_test_anon(folio), folio);
+       __folio_rmap_sanity_checks(folio, page, nr_pages, level);
 
-       /* Is page being mapped by PTE? Is this its first map to be added? */
-       if (likely(!compound)) {
+       switch (level) {
+       case RMAP_LEVEL_PTE:
                do {
                        first = atomic_inc_and_test(&page->_mapcount);
                        if (first && folio_test_large(folio)) {
                        if (first)
                                nr++;
                } while (page++, --nr_pages > 0);
-       } else if (folio_test_pmd_mappable(folio)) {
-               /* That test is redundant: it's for safety or to optimize out */
-
+               break;
+       case RMAP_LEVEL_PMD:
                first = atomic_inc_and_test(&folio->_entire_mapcount);
                if (first) {
                        nr = atomic_add_return_relaxed(COMPOUND_MAPPED, mapped);
                                nr = 0;
                        }
                }
+               break;
        }
 
        if (nr_pmdmapped)
                mlock_vma_folio(folio, vma);
 }
 
+/**
+ * folio_add_file_rmap_ptes - add PTE mappings to a page range of a folio
+ * @folio:     The folio to add the mappings to
+ * @page:      The first page to add
+ * @nr_pages:  The number of pages that will be mapped using PTEs
+ * @vma:       The vm area in which the mappings are added
+ *
+ * The page range of the folio is defined by [page, page + nr_pages)
+ *
+ * The caller needs to hold the page table lock.
+ */
+void folio_add_file_rmap_ptes(struct folio *folio, struct page *page,
+               int nr_pages, struct vm_area_struct *vma)
+{
+       __folio_add_file_rmap(folio, page, nr_pages, vma, RMAP_LEVEL_PTE);
+}
+
+/**
+ * folio_add_file_rmap_pmd - add a PMD mapping to a page range of a folio
+ * @folio:     The folio to add the mapping to
+ * @page:      The first page to add
+ * @vma:       The vm area in which the mapping is added
+ *
+ * The page range of the folio is defined by [page, page + HPAGE_PMD_NR)
+ *
+ * The caller needs to hold the page table lock.
+ */
+void folio_add_file_rmap_pmd(struct folio *folio, struct page *page,
+               struct vm_area_struct *vma)
+{
+#ifdef CONFIG_TRANSPARENT_HUGEPAGE
+       __folio_add_file_rmap(folio, page, HPAGE_PMD_NR, vma, RMAP_LEVEL_PMD);
+#else
+       WARN_ON_ONCE(true);
+#endif
+}
+
 /**
  * page_add_file_rmap - add pte mapping to a file page
  * @page:      the page to add the mapping to
                bool compound)
 {
        struct folio *folio = page_folio(page);
-       unsigned int nr_pages;
 
        VM_WARN_ON_ONCE_PAGE(compound && !PageTransHuge(page), page);
 
        if (likely(!compound))
-               nr_pages = 1;
+               folio_add_file_rmap_pte(folio, page, vma);
        else
-               nr_pages = folio_nr_pages(folio);
-
-       folio_add_file_rmap_range(folio, page, nr_pages, vma, compound);
+               folio_add_file_rmap_pmd(folio, page, vma);
 }
 
 /**