xref: /linux/kernel/sched/ext/arena.c (revision 3cd1f76be638b7386201171e7bb4c88095774dd5)
1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3  * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4  *
5  * scx_arena_pool: kernel-side sub-allocator over BPF-arena pages.
6  *
7  * Each chunk added to @sch->arena_pool comes from one
8  * bpf_arena_alloc_pages_sleepable() call and is registered at the
9  * kernel-side mapping address. Callers translate to the BPF-arena form
10  * themselves if needed.
11  *
12  * Allocations grow the pool on demand. Underlying arena pages are released
13  * when the arena map itself is torn down.
14  *
15  * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
16  * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
17  */
18 #include <linux/genalloc.h>
19 
20 #include "internal.h"
21 #include "arena.h"
22 
23 enum scx_arena_consts {
24 	SCX_ARENA_MIN_ORDER		= 3,	/* 8-byte minimum sub-allocation */
25 	SCX_ARENA_GROW_PAGES		= 4,	/* per growth */
26 };
27 
28 s32 scx_arena_pool_init(struct scx_sched *sch)
29 {
30 	if (!sch->arena_map)
31 		return 0;
32 
33 	sch->arena_pool = gen_pool_create(SCX_ARENA_MIN_ORDER, NUMA_NO_NODE);
34 	if (!sch->arena_pool)
35 		return -ENOMEM;
36 	return 0;
37 }
38 
39 static void scx_arena_clear_chunk(struct gen_pool *pool, struct gen_pool_chunk *chunk,
40 				  void *data)
41 {
42 	int order = pool->min_alloc_order;
43 	size_t chunk_sz = chunk->end_addr - chunk->start_addr + 1;
44 	unsigned long end_bit = chunk_sz >> order;
45 	unsigned long b, e;
46 
47 	for_each_set_bitrange(b, e, chunk->bits, end_bit)
48 		gen_pool_free(pool, chunk->start_addr + (b << order),
49 			      (e - b) << order);
50 }
51 
52 /*
53  * Tear down the pool. Outstanding gen_pool allocations are freed via
54  * scx_arena_clear_chunk() so gen_pool_destroy() doesn't BUG. The underlying
55  * arena pages are released when the arena map itself is torn down.
56  */
57 void scx_arena_pool_destroy(struct scx_sched *sch)
58 {
59 	if (!sch->arena_pool)
60 		return;
61 	gen_pool_for_each_chunk(sch->arena_pool, scx_arena_clear_chunk, NULL);
62 	gen_pool_destroy(sch->arena_pool);
63 	sch->arena_pool = NULL;
64 }
65 
66 /*
67  * Grow the pool by @page_cnt pages. bpf_arena_alloc_pages_sleepable() and
68  * gen_pool_add() (which calls vzalloc(GFP_KERNEL)) require a sleepable
69  * context.
70  */
71 static int scx_arena_grow(struct scx_sched *sch, u32 page_cnt)
72 {
73 	u64 kern_vm_start;
74 	u32 uaddr32;
75 	void *p;
76 	int ret;
77 
78 	if (!sch->arena_map || !sch->arena_pool)
79 		return -EINVAL;
80 
81 	p = bpf_arena_alloc_pages_sleepable(sch->arena_map, NULL,
82 					    page_cnt, NUMA_NO_NODE, 0);
83 	if (!p)
84 		return -ENOMEM;
85 
86 	uaddr32 = (u32)(unsigned long)p;
87 	/* arena.o, which defines these, is built only on MMU && 64BIT */
88 #if defined(CONFIG_MMU) && defined(CONFIG_64BIT)
89 	kern_vm_start = bpf_arena_map_kern_vm_start(sch->arena_map);
90 #else
91 	kern_vm_start = 0;
92 #endif
93 
94 	ret = gen_pool_add(sch->arena_pool, kern_vm_start + uaddr32,
95 			   page_cnt * PAGE_SIZE, NUMA_NO_NODE);
96 	if (ret) {
97 		bpf_arena_free_pages_non_sleepable(sch->arena_map, p, page_cnt);
98 		return ret;
99 	}
100 	return 0;
101 }
102 
103 /*
104  * Allocate @size bytes from the arena pool. Returns kernel VA on success, NULL
105  * on failure. May grow the pool via scx_arena_grow() which sleeps. Caller must
106  * be in a GFP_KERNEL context.
107  */
108 void *scx_arena_alloc(struct scx_sched *sch, size_t size)
109 {
110 	unsigned long kern_va;
111 	u32 page_cnt;
112 
113 	might_sleep();
114 
115 	if (!sch->arena_pool)
116 		return NULL;
117 
118 	while (true) {
119 		kern_va = gen_pool_alloc(sch->arena_pool, size);
120 		if (kern_va)
121 			break;
122 		page_cnt = max_t(u32, SCX_ARENA_GROW_PAGES,
123 				 (size + PAGE_SIZE - 1) >> PAGE_SHIFT);
124 		if (scx_arena_grow(sch, page_cnt))
125 			return NULL;
126 	}
127 
128 	return (void *)kern_va;
129 }
130 
131 void scx_arena_free(struct scx_sched *sch, void *kern_va, size_t size)
132 {
133 	if (sch->arena_pool && kern_va)
134 		gen_pool_free(sch->arena_pool, (unsigned long)kern_va, size);
135 }
136