Re: [PATCH RFC v3 23/35] arm64: mte: Try to reserve tag storage in arch_alloc_page()

[Peter Collingbourne <pcc@google.com>]

On Thu, Jan 25, 2024 at 8:45 AM Alexandru Elisei
<alexandru.elisei@arm.com> wrote:
>
> Reserve tag storage for a page that is being allocated as tagged. This
> is a best effort approach, and failing to reserve tag storage is
> allowed.
>
> When all the associated tagged pages have been freed, return the tag
> storage pages back to the page allocator, where they can be used again for
> data allocations.
>
> Signed-off-by: Alexandru Elisei <alexandru.elisei@arm.com>
> ---
>
> Changes since rfc v2:
>
> * Based on rfc v2 patch #16 ("arm64: mte: Manage tag storage on page
> allocation").
> * Fixed calculation of the number of associated tag storage blocks (Hyesoo
> Yu).
> * Tag storage is reserved in arch_alloc_page() instead of
> arch_prep_new_page().
>
>  arch/arm64/include/asm/mte.h             |  16 +-
>  arch/arm64/include/asm/mte_tag_storage.h |  31 +++
>  arch/arm64/include/asm/page.h            |   5 +
>  arch/arm64/include/asm/pgtable.h         |  19 ++
>  arch/arm64/kernel/mte_tag_storage.c      | 234 +++++++++++++++++++++++
>  arch/arm64/mm/fault.c                    |   7 +
>  fs/proc/page.c                           |   1 +
>  include/linux/kernel-page-flags.h        |   1 +
>  include/linux/page-flags.h               |   1 +
>  include/trace/events/mmflags.h           |   3 +-
>  mm/huge_memory.c                         |   1 +
>  11 files changed, 316 insertions(+), 3 deletions(-)
>
> diff --git a/arch/arm64/include/asm/mte.h b/arch/arm64/include/asm/mte.h
> index 8034695b3dd7..6457b7899207 100644
> --- a/arch/arm64/include/asm/mte.h
> +++ b/arch/arm64/include/asm/mte.h
> @@ -40,12 +40,24 @@ void mte_free_tag_buf(void *buf);
>  #ifdef CONFIG_ARM64_MTE
>
>  /* track which pages have valid allocation tags */
> -#define PG_mte_tagged  PG_arch_2
> +#define PG_mte_tagged          PG_arch_2
>  /* simple lock to avoid multiple threads tagging the same page */
> -#define PG_mte_lock    PG_arch_3
> +#define PG_mte_lock            PG_arch_3
> +/* Track if a tagged page has tag storage reserved */
> +#define PG_tag_storage_reserved        PG_arch_4
> +
> +#ifdef CONFIG_ARM64_MTE_TAG_STORAGE
> +DECLARE_STATIC_KEY_FALSE(tag_storage_enabled_key);
> +extern bool page_tag_storage_reserved(struct page *page);
> +#endif
>
>  static inline void set_page_mte_tagged(struct page *page)
>  {
> +#ifdef CONFIG_ARM64_MTE_TAG_STORAGE
> +       /* Open code mte_tag_storage_enabled() */
> +       WARN_ON_ONCE(static_branch_likely(&tag_storage_enabled_key) &&
> +                    !page_tag_storage_reserved(page));
> +#endif
>         /*
>          * Ensure that the tags written prior to this function are visible
>          * before the page flags update.
> diff --git a/arch/arm64/include/asm/mte_tag_storage.h b/arch/arm64/include/asm/mte_tag_storage.h
> index 7b3f6bff8e6f..09f1318d924e 100644
> --- a/arch/arm64/include/asm/mte_tag_storage.h
> +++ b/arch/arm64/include/asm/mte_tag_storage.h
> @@ -5,6 +5,12 @@
>  #ifndef __ASM_MTE_TAG_STORAGE_H
>  #define __ASM_MTE_TAG_STORAGE_H
>
> +#ifndef __ASSEMBLY__
> +
> +#include <linux/mm_types.h>
> +
> +#include <asm/mte.h>
> +
>  #ifdef CONFIG_ARM64_MTE_TAG_STORAGE
>
>  DECLARE_STATIC_KEY_FALSE(tag_storage_enabled_key);
> @@ -15,6 +21,15 @@ static inline bool tag_storage_enabled(void)
>  }
>
>  void mte_init_tag_storage(void);
> +
> +static inline bool alloc_requires_tag_storage(gfp_t gfp)
> +{
> +       return gfp & __GFP_TAGGED;
> +}
> +int reserve_tag_storage(struct page *page, int order, gfp_t gfp);
> +void free_tag_storage(struct page *page, int order);
> +
> +bool page_tag_storage_reserved(struct page *page);
>  #else
>  static inline bool tag_storage_enabled(void)
>  {
> @@ -23,6 +38,22 @@ static inline bool tag_storage_enabled(void)
>  static inline void mte_init_tag_storage(void)
>  {
>  }
> +static inline bool alloc_requires_tag_storage(struct page *page)

This function should take a gfp_t to match the
CONFIG_ARM64_MTE_TAG_STORAGE case.

Peter

> +{
> +       return false;
> +}
> +static inline int reserve_tag_storage(struct page *page, int order, gfp_t gfp)
> +{
> +       return 0;
> +}
> +static inline void free_tag_storage(struct page *page, int order)
> +{
> +}
> +static inline bool page_tag_storage_reserved(struct page *page)
> +{
> +       return true;
> +}
>  #endif /* CONFIG_ARM64_MTE_TAG_STORAGE */
>
> +#endif /* !__ASSEMBLY__ */
>  #endif /* __ASM_MTE_TAG_STORAGE_H  */
> diff --git a/arch/arm64/include/asm/page.h b/arch/arm64/include/asm/page.h
> index 88bab032a493..3a656492f34a 100644
> --- a/arch/arm64/include/asm/page.h
> +++ b/arch/arm64/include/asm/page.h
> @@ -35,6 +35,11 @@ void copy_highpage(struct page *to, struct page *from);
>  void tag_clear_highpage(struct page *to);
>  #define __HAVE_ARCH_TAG_CLEAR_HIGHPAGE
>
> +#ifdef CONFIG_ARM64_MTE_TAG_STORAGE
> +void arch_alloc_page(struct page *, int order, gfp_t gfp);
> +#define HAVE_ARCH_ALLOC_PAGE
> +#endif
> +
>  #define clear_user_page(page, vaddr, pg)       clear_page(page)
>  #define copy_user_page(to, from, vaddr, pg)    copy_page(to, from)
>
> diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
> index 2499cc4fa4f2..f30466199a9b 100644
> --- a/arch/arm64/include/asm/pgtable.h
> +++ b/arch/arm64/include/asm/pgtable.h
> @@ -10,6 +10,7 @@
>
>  #include <asm/memory.h>
>  #include <asm/mte.h>
> +#include <asm/mte_tag_storage.h>
>  #include <asm/pgtable-hwdef.h>
>  #include <asm/pgtable-prot.h>
>  #include <asm/tlbflush.h>
> @@ -1069,6 +1070,24 @@ static inline void arch_swap_restore(swp_entry_t entry, struct folio *folio)
>                 mte_restore_page_tags_by_swp_entry(entry, &folio->page);
>  }
>
> +#ifdef CONFIG_ARM64_MTE_TAG_STORAGE
> +
> +#define __HAVE_ARCH_FREE_PAGES_PREPARE
> +static inline void arch_free_pages_prepare(struct page *page, int order)
> +{
> +       if (tag_storage_enabled() && page_mte_tagged(page))
> +               free_tag_storage(page, order);
> +}
> +
> +#define __HAVE_ARCH_ALLOC_CMA
> +static inline bool arch_alloc_cma(gfp_t gfp_mask)
> +{
> +       if (tag_storage_enabled() && alloc_requires_tag_storage(gfp_mask))
> +               return false;
> +       return true;
> +}
> +
> +#endif /* CONFIG_ARM64_MTE_TAG_STORAGE */
>  #endif /* CONFIG_ARM64_MTE */
>
>  #define __HAVE_ARCH_CALC_VMA_GFP
> diff --git a/arch/arm64/kernel/mte_tag_storage.c b/arch/arm64/kernel/mte_tag_storage.c
> index d58c68b4a849..762c7c803a70 100644
> --- a/arch/arm64/kernel/mte_tag_storage.c
> +++ b/arch/arm64/kernel/mte_tag_storage.c
> @@ -34,6 +34,31 @@ struct tag_region {
>  static struct tag_region tag_regions[MAX_TAG_REGIONS];
>  static int num_tag_regions;
>
> +/*
> + * A note on locking. Reserving tag storage takes the tag_blocks_lock mutex,
> + * because alloc_contig_range() might sleep.
> + *
> + * Freeing tag storage takes the xa_lock spinlock with interrupts disabled
> + * because pages can be freed from non-preemptible contexts, including from an
> + * interrupt handler.
> + *
> + * Because tag storage can be freed from interrupt contexts, the xarray is
> + * defined with the XA_FLAGS_LOCK_IRQ flag to disable interrupts when calling
> + * xa_store(). This is done to prevent a deadlock with free_tag_storage() being
> + * called from an interrupt raised before xa_store() releases the xa_lock.
> + *
> + * All of the above means that reserve_tag_storage() cannot run concurrently
> + * with itself (no concurrent insertions), but it can run at the same time as
> + * free_tag_storage(). The first thing that reserve_tag_storage() does after
> + * taking the mutex is increase the refcount on all present tag storage blocks
> + * with the xa_lock held, to serialize against freeing the blocks. This is an
> + * optimization to avoid taking and releasing the xa_lock after each iteration
> + * if the refcount operation was moved inside the loop, where it would have had
> + * to be executed for each block.
> + */
> +static DEFINE_XARRAY_FLAGS(tag_blocks_reserved, XA_FLAGS_LOCK_IRQ);
> +static DEFINE_MUTEX(tag_blocks_lock);
> +
>  static u32 __init get_block_size_pages(u32 block_size_bytes)
>  {
>         u32 a = PAGE_SIZE;
> @@ -364,3 +389,212 @@ static int __init mte_enable_tag_storage(void)
>         return -EINVAL;
>  }
>  arch_initcall(mte_enable_tag_storage);
> +
> +static void page_set_tag_storage_reserved(struct page *page, int order)
> +{
> +       int i;
> +
> +       for (i = 0; i < (1 << order); i++)
> +               set_bit(PG_tag_storage_reserved, &(page + i)->flags);
> +}
> +
> +static void block_ref_add(unsigned long block, struct tag_region *region, int order)
> +{
> +       int count;
> +
> +       count = min(1u << order, 32 * region->block_size_pages);
> +       page_ref_add(pfn_to_page(block), count);
> +}
> +
> +static int block_ref_sub_return(unsigned long block, struct tag_region *region, int order)
> +{
> +       int count;
> +
> +       count = min(1u << order, 32 * region->block_size_pages);
> +       return page_ref_sub_return(pfn_to_page(block), count);
> +}
> +
> +static bool tag_storage_block_is_reserved(unsigned long block)
> +{
> +       return xa_load(&tag_blocks_reserved, block) != NULL;
> +}
> +
> +static int tag_storage_reserve_block(unsigned long block, struct tag_region *region, int order)
> +{
> +       int ret;
> +
> +       ret = xa_err(xa_store(&tag_blocks_reserved, block, pfn_to_page(block), GFP_KERNEL));
> +       if (!ret)
> +               block_ref_add(block, region, order);
> +
> +       return ret;
> +}
> +
> +static int order_to_num_blocks(int order, u32 block_size_pages)
> +{
> +       int num_tag_storage_pages = max((1 << order) / 32, 1);
> +
> +       return DIV_ROUND_UP(num_tag_storage_pages, block_size_pages);
> +}
> +
> +static int tag_storage_find_block_in_region(struct page *page, unsigned long *blockp,
> +                                           struct tag_region *region)
> +{
> +       struct range *tag_range = &region->tag_range;
> +       struct range *mem_range = &region->mem_range;
> +       u64 page_pfn = page_to_pfn(page);
> +       u64 block, block_offset;
> +
> +       if (!(mem_range->start <= page_pfn && page_pfn <= mem_range->end))
> +               return -ERANGE;
> +
> +       block_offset = (page_pfn - mem_range->start) / 32;
> +       block = tag_range->start + rounddown(block_offset, region->block_size_pages);
> +
> +       if (block + region->block_size_pages - 1 > tag_range->end) {
> +               pr_err("Block 0x%llx-0x%llx is outside tag region 0x%llx-0x%llx\n",
> +                       PFN_PHYS(block), PFN_PHYS(block + region->block_size_pages + 1) - 1,
> +                       PFN_PHYS(tag_range->start), PFN_PHYS(tag_range->end + 1) - 1);
> +               return -ERANGE;
> +       }
> +       *blockp = block;
> +
> +       return 0;
> +
> +}
> +
> +static int tag_storage_find_block(struct page *page, unsigned long *block,
> +                                 struct tag_region **region)
> +{
> +       int i, ret;
> +
> +       for (i = 0; i < num_tag_regions; i++) {
> +               ret = tag_storage_find_block_in_region(page, block, &tag_regions[i]);
> +               if (ret == 0) {
> +                       *region = &tag_regions[i];
> +                       return 0;
> +               }
> +       }
> +
> +       return -EINVAL;
> +}
> +
> +bool page_tag_storage_reserved(struct page *page)
> +{
> +       return test_bit(PG_tag_storage_reserved, &page->flags);
> +}
> +
> +int reserve_tag_storage(struct page *page, int order, gfp_t gfp)
> +{
> +       unsigned long start_block, end_block;
> +       struct tag_region *region;
> +       unsigned long block;
> +       unsigned long flags;
> +       int ret = 0;
> +
> +       VM_WARN_ON_ONCE(!preemptible());
> +
> +       if (page_tag_storage_reserved(page))
> +               return 0;
> +
> +       /*
> +        * __alloc_contig_migrate_range() ignores gfp when allocating the
> +        * destination page for migration. Regardless, massage gfp flags and
> +        * remove __GFP_TAGGED to avoid recursion in case gfp stops being
> +        * ignored.
> +        */
> +       gfp &= ~__GFP_TAGGED;
> +       if (!(gfp & __GFP_NORETRY))
> +               gfp |= __GFP_RETRY_MAYFAIL;
> +
> +       ret = tag_storage_find_block(page, &start_block, &region);
> +       if (WARN_ONCE(ret, "Missing tag storage block for pfn 0x%lx", page_to_pfn(page)))
> +               return -EINVAL;
> +       end_block = start_block + order_to_num_blocks(order, region->block_size_pages);
> +
> +       mutex_lock(&tag_blocks_lock);
> +
> +       /* Check again, this time with the lock held. */
> +       if (page_tag_storage_reserved(page))
> +               goto out_unlock;
> +
> +       /* Make sure existing entries are not freed from out under out feet. */
> +       xa_lock_irqsave(&tag_blocks_reserved, flags);
> +       for (block = start_block; block < end_block; block += region->block_size_pages) {
> +               if (tag_storage_block_is_reserved(block))
> +                       block_ref_add(block, region, order);
> +       }
> +       xa_unlock_irqrestore(&tag_blocks_reserved, flags);
> +
> +       for (block = start_block; block < end_block; block += region->block_size_pages) {
> +               /* Refcount incremented above. */
> +               if (tag_storage_block_is_reserved(block))
> +                       continue;
> +
> +               ret = cma_alloc_range(region->cma, block, region->block_size_pages, 3, gfp);
> +               /* Should never happen. */
> +               VM_WARN_ON_ONCE(ret == -EEXIST);
> +               if (ret)
> +                       goto out_error;
> +
> +               ret = tag_storage_reserve_block(block, region, order);
> +               if (ret) {
> +                       cma_release(region->cma, pfn_to_page(block), region->block_size_pages);
> +                       goto out_error;
> +               }
> +       }
> +
> +       page_set_tag_storage_reserved(page, order);
> +out_unlock:
> +       mutex_unlock(&tag_blocks_lock);
> +
> +       return 0;
> +
> +out_error:
> +       xa_lock_irqsave(&tag_blocks_reserved, flags);
> +       for (block = start_block; block < end_block; block += region->block_size_pages) {
> +               if (tag_storage_block_is_reserved(block) &&
> +                   block_ref_sub_return(block, region, order) == 1) {
> +                       __xa_erase(&tag_blocks_reserved, block);
> +                       cma_release(region->cma, pfn_to_page(block), region->block_size_pages);
> +               }
> +       }
> +       xa_unlock_irqrestore(&tag_blocks_reserved, flags);
> +
> +       mutex_unlock(&tag_blocks_lock);
> +
> +       return ret;
> +}
> +
> +void free_tag_storage(struct page *page, int order)
> +{
> +       unsigned long block, start_block, end_block;
> +       struct tag_region *region;
> +       unsigned long flags;
> +       int ret;
> +
> +       ret = tag_storage_find_block(page, &start_block, &region);
> +       if (WARN_ONCE(ret, "Missing tag storage block for pfn 0x%lx", page_to_pfn(page)))
> +               return;
> +
> +       end_block = start_block + order_to_num_blocks(order, region->block_size_pages);
> +
> +       xa_lock_irqsave(&tag_blocks_reserved, flags);
> +       for (block = start_block; block < end_block; block += region->block_size_pages) {
> +               if (WARN_ONCE(!tag_storage_block_is_reserved(block),
> +                   "Block 0x%lx is not reserved for pfn 0x%lx", block, page_to_pfn(page)))
> +                       continue;
> +
> +               if (block_ref_sub_return(block, region, order) == 1) {
> +                       __xa_erase(&tag_blocks_reserved, block);
> +                       cma_release(region->cma, pfn_to_page(block), region->block_size_pages);
> +               }
> +       }
> +       xa_unlock_irqrestore(&tag_blocks_reserved, flags);
> +}
> +
> +void arch_alloc_page(struct page *page, int order, gfp_t gfp)
> +{
> +       if (tag_storage_enabled() && alloc_requires_tag_storage(gfp))
> +               reserve_tag_storage(page, order, gfp);
> +}
> diff --git a/arch/arm64/mm/fault.c b/arch/arm64/mm/fault.c
> index c022e473c17c..1ffaeccecda2 100644
> --- a/arch/arm64/mm/fault.c
> +++ b/arch/arm64/mm/fault.c
> @@ -37,6 +37,7 @@
>  #include <asm/esr.h>
>  #include <asm/kprobes.h>
>  #include <asm/mte.h>
> +#include <asm/mte_tag_storage.h>
>  #include <asm/processor.h>
>  #include <asm/sysreg.h>
>  #include <asm/system_misc.h>
> @@ -950,6 +951,12 @@ gfp_t arch_calc_vma_gfp(struct vm_area_struct *vma, gfp_t gfp)
>
>  void tag_clear_highpage(struct page *page)
>  {
> +       if (tag_storage_enabled() && !page_tag_storage_reserved(page)) {
> +               /* Don't zero the tags if tag storage is not reserved */
> +               clear_page(page_address(page));
> +               return;
> +       }
> +
>         /* Newly allocated page, shouldn't have been tagged yet */
>         WARN_ON_ONCE(!try_page_mte_tagging(page));
>         mte_zero_clear_page_tags(page_address(page));
> diff --git a/fs/proc/page.c b/fs/proc/page.c
> index 195b077c0fac..e7eb584a9234 100644
> --- a/fs/proc/page.c
> +++ b/fs/proc/page.c
> @@ -221,6 +221,7 @@ u64 stable_page_flags(struct page *page)
>  #ifdef CONFIG_ARCH_USES_PG_ARCH_X
>         u |= kpf_copy_bit(k, KPF_ARCH_2,        PG_arch_2);
>         u |= kpf_copy_bit(k, KPF_ARCH_3,        PG_arch_3);
> +       u |= kpf_copy_bit(k, KPF_ARCH_4,        PG_arch_4);
>  #endif
>
>         return u;
> diff --git a/include/linux/kernel-page-flags.h b/include/linux/kernel-page-flags.h
> index 859f4b0c1b2b..4a0d719ffdd4 100644
> --- a/include/linux/kernel-page-flags.h
> +++ b/include/linux/kernel-page-flags.h
> @@ -19,5 +19,6 @@
>  #define KPF_SOFTDIRTY          40
>  #define KPF_ARCH_2             41
>  #define KPF_ARCH_3             42
> +#define KPF_ARCH_4             43
>
>  #endif /* LINUX_KERNEL_PAGE_FLAGS_H */
> diff --git a/include/linux/page-flags.h b/include/linux/page-flags.h
> index b7237bce7446..03f03e6d735e 100644
> --- a/include/linux/page-flags.h
> +++ b/include/linux/page-flags.h
> @@ -135,6 +135,7 @@ enum pageflags {
>  #ifdef CONFIG_ARCH_USES_PG_ARCH_X
>         PG_arch_2,
>         PG_arch_3,
> +       PG_arch_4,
>  #endif
>         __NR_PAGEFLAGS,
>
> diff --git a/include/trace/events/mmflags.h b/include/trace/events/mmflags.h
> index 6ca0d5ed46c0..ba962fd10a2c 100644
> --- a/include/trace/events/mmflags.h
> +++ b/include/trace/events/mmflags.h
> @@ -125,7 +125,8 @@ IF_HAVE_PG_HWPOISON(hwpoison)                                               \
>  IF_HAVE_PG_IDLE(idle)                                                  \
>  IF_HAVE_PG_IDLE(young)                                                 \
>  IF_HAVE_PG_ARCH_X(arch_2)                                              \
> -IF_HAVE_PG_ARCH_X(arch_3)
> +IF_HAVE_PG_ARCH_X(arch_3)                                              \
> +IF_HAVE_PG_ARCH_X(arch_4)
>
>  #define show_page_flags(flags)                                         \
>         (flags) ? __print_flags(flags, "|",                             \
> diff --git a/mm/huge_memory.c b/mm/huge_memory.c
> index 2bad63a7ec16..47932539cc50 100644
> --- a/mm/huge_memory.c
> +++ b/mm/huge_memory.c
> @@ -2804,6 +2804,7 @@ static void __split_huge_page_tail(struct folio *folio, int tail,
>  #ifdef CONFIG_ARCH_USES_PG_ARCH_X
>                          (1L << PG_arch_2) |
>                          (1L << PG_arch_3) |
> +                        (1L << PG_arch_4) |
>  #endif
>                          (1L << PG_dirty) |
>                          LRU_GEN_MASK | LRU_REFS_MASK));
> --
> 2.43.0
>