[PATCH v5 00/22] kho: make boot time huge page allocation work nicely with KHO

[Pratyush Yadav <pratyush@kernel.org>]

From: "Pratyush Yadav (Google)" <pratyush@kernel.org>

Hi,

Gigantic huge page allocation is somewhat broken currently with KHO.

First, they break scratch size accounting. Since they are allocated
using the memblock alloc APIs, they count towards RSRV_KERN, and this
scratch size when using scratch_scale. This means if huge pages take a
large enough chunk of system memory scratch size will blow up and fail
to allocate.

Second, scratch can not contain preserved memory, and if hugepages are
allocated from scratch, they will fail to be preserved with the upcoming
hugetlb preservation series [0].

Fix this by introducing the concept of extended scratch areas. They are
areas that the kernel discovers on boot by walking the radix tree and
finding free memory ranges. See patch 10 for more details.

Discovering the scratch areas needs some preparatory changes to KHO, the
radix tree APIs, and to memblock. Patches 1-17 do that.

Patch 18 adds the scratch discovery logic.

Patch 19-20 add the dedicated memblock hugetlb allocator.

Patch 21-22 fix the scratch size calculation with using scratch_scale.

This series is based on top of the liveupdate/kho-scratch branch.

[0] https://lore.kernel.org/linux-mm/20251206230222.853493-1-pratyush@kernel.org/T/#u

Changes in v5:

- Fix memblock test compilation because KHO headers are now always
  included. Since some stubs also need to be added, move the change to a
  new patch "memblock: always include KHO headers".
- Reset start and end to default values when opening up search to all
  nodes in memblock_alloc_hugetlb().
- Update search end instead of start in memblock_alloc_hugetlb() in top
  down mode.

Changes in v4:

- Commit message fixups.
- Rename leaktrace to kmemleak_trace in memblock_prep_allocation().
- Move definition of kho_scratch_overlap() in
  tools/testing/memblock/internal.h to the patch that exposed it.
- Rename KHO_EXT_SHIFT to KHO_EXT_BLKSHIFT.
- Add a comment above kho_decode_radix_key() call in kho_ext_walk_leaf()
  which makes it clearer what the key is for and what it decodes to.
- Get rid of kho_init_scratch_migratetype(). Call
  kho_scratch_migratetype() for each pageblock.
- Use Muchun's patch
  https://lore.kernel.org/linux-mm/20260612035903.2468601-6-songmuchun@bytedance.com/
  to simplify migratetype init for deferred init.
- Drop #ifdef CONFIG_KEXEC_HANDOVER around KHO includes in memblock to
  make sure it builds when the config is disabled. This is because
  memblock_alloc_hugetlb() uses kho_scratch_overlap(), and KHO header
  provides a stub when the config is disabled.

Changes in v3:

More notable changes:

- Refactor how key width and table depth is calculated.
- I noticed that extended scratch breaks migrate type init during
  testing. Below two changes fix it:
- Remove memblock_is_kho_scratch_memory() and kho_scratch_migratetype().
  Instead, add kho_init_scratch_migratetype() and rework how migrate
  type for scratch pages is initialized.
- Move migrate type initialization from deferred_free_pages() to
  deferred_init_pages().

Other changes:

- Use kho_get_mem_map() directly in kho_populate() and add a comment on
  why it is okay to use.
- Don't say "export kho_scratch_overlap()".
- Add a comment in kho_extend_scratch() that it should always be called
  after efi_init()::reserve_regions().
- Get rid of memblock_reserved_size_flags(). Use
  memblock_reserved_hugetlb_size() instead.
- Define KHO_EXT_BLKSIZE to SZ_1G and add a comment explaining why the
  size was picked.
- Fix memblock tests build.
- Rename kho_ext_walk_key() to kho_ext_walk_leaf().
- Add some comments around the kho_ext tree callbacks to clarify which
  tree the function is called for and which one it uses.
- Move initialization of kho_in.radix_tree to kho_memory_init_early().
- Do not check if the tree is already initialized in kho_radix_init_tree().

Changes in v2:

Detailed changelog below.

At a high level, the major change in this version is to remove
MEMBLOCK_KHO_SCRATCH_EXT. Keep MEMBLOCK_KHO_SCRATCH as the only memory
type and mark the discovered areas with it. For HugeTLB, add a dedicated
allocation routine and if allocated memory lands in scratch, do a retry.
Also introduce MEMBLOCK_RSRV_HUGETLB to help with accounting of scratch
area sizes.

- Fixup commit message in patch 1 to make namespacing change clearer.
- Use @key in kernel-doc for radix functions.
- Add a runtime check on key width.
- Move all mem retrieval logic to kho_mem_retrieve().
- Add a comment in kho_mem_retrieve() explaining why mem_map won't be NULL.
- Rename callbacks to ->leaf() and ->node().
- Fixup commit messages.
- Clear tree->root in kho_radix_destroy_tree(). This lets the tree be
  re-initialized by calling kho_radix_init_tree()
- Add kho_get_mem_map() earlier in the series.
- Export kho_scratch_overlap() and use it in memblock_is_kho_scratch_memory().
- Get rid of MEMBLOCK_KHO_SCRATCH_EXT.
- Introduce MEMBLOCK_RSRV_HUGETLB.
- Introduce memblock_alloc_hugetlb() for hugetlb bootmem allocations.
- Refactor memblock_reserved_kern_size() to allow calculating size by flags.
- Exclude hugetlb memory from scratch size calculation.
- Collect R-bys.

Regards,
Pratyush Yadav

Muchun Song (1):
  mm/mm_init: simplify deferred_free_pages() migratetype init

Pratyush Yadav (Google) (21):
  kho: generalize radix tree APIs
  kho: make radix max key width more obvious
  kho: disallow wide keys in radix tree
  kho: return virtual address of mem_map from kho_get_mem_map()
  kho: store incoming radix tree in kho_in
  kho: move all memory retrieval logic to kho_mem_retrieve()
  kho: add a struct for radix callbacks
  kho: add callback for table pages
  kho: add data argument to radix walk callback
  kho: allow early-boot usage of the KHO radix tree
  kho: allow destroying KHO radix tree
  kho: add kho_radix_init_tree()
  kho: expose kho_scratch_overlap() to kexec_handover.h
  kho: initialize kho_scratch pointer earlier in boot
  kho: initialize preserved memory map radix tree earlier
  mm/mm_init: don't rely on memblock to get KHO scratch migratetype
  kho: extend scratch
  memblock: always include KHO headers
  memblock: make HugeTLB bootmem allocation work with KHO
  memblock: add memblock_reserved_hugetlb_size()
  kho: exclude hugetlb memory from scratch size calculation

 include/linux/kexec_handover.h                |  25 +
 include/linux/kho/abi/kexec_handover.h        |   9 +-
 include/linux/kho_radix_tree.h                |  44 +-
 include/linux/memblock.h                      |  23 +-
 kernel/liveupdate/Makefile                    |   1 -
 kernel/liveupdate/kexec_handover.c            | 526 ++++++++++++++----
 kernel/liveupdate/kexec_handover_debug.c      |  25 -
 kernel/liveupdate/kexec_handover_internal.h   |   9 -
 mm/hugetlb.c                                  |  22 +-
 mm/memblock.c                                 | 154 +++--
 mm/mm_init.c                                  |  45 +-
 tools/testing/memblock/internal.h             |   5 +
 tools/testing/memblock/linux/kexec_handover.h |   9 +
 .../testing/memblock/linux/kho/abi/memblock.h |  10 +
 tools/testing/memblock/linux/libfdt.h         |   9 +
 15 files changed, 650 insertions(+), 266 deletions(-)
 delete mode 100644 kernel/liveupdate/kexec_handover_debug.c
 create mode 100644 tools/testing/memblock/linux/kexec_handover.h
 create mode 100644 tools/testing/memblock/linux/kho/abi/memblock.h
 create mode 100644 tools/testing/memblock/linux/libfdt.h


base-commit: 08412b8c707fdbccb7bf2116f0554fe09113cd53
-- 
2.55.0.508.g3f0d502094-goog