sched_ext: Convert sub-cap kfuncs to __arena cmask arguments

The sub-cap kfuncs take their cmask arguments as __ign pointers. The values
cross the kfunc boundary as unchecked scalars and scx_cmask_ref_init()
rebases them into the arena by hand.

BPF now translates between BPF and kernel arena addresses for __arena
arguments. Tag the cmask arguments __arena so the kfuncs receive kernel
addresses and scx_cmask_ref_init() loses the hand-rolled conversion. The
optional denied_out keeps its NULL not-provided signal via
__arena__nullable. The mandatory masks use plain __arena.

scx_qmap's call sites drop the (void *)(long) casts since the BPF-side
declarations type the cmask arguments __arena and take arena pointers
directly.

The arena argument address translation is currently implemented only on
x86-64. Schedulers calling these kfuncs load only there for now.

Signed-off-by: Tejun Heo <tj@kernel.org>

2 files changed, 15 insertions, 16 deletions

diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
index 6035d007c43f..979d4cabfaf9 100644
--- a/tools/sched_ext/include/scx/common.bpf.h
+++ b/tools/sched_ext/include/scx/common.bpf.h
@@ -26,6 +26,7 @@
 #include <asm-generic/errno.h>
 #include "user_exit_info.bpf.h"
 #include "enum_defs.autogen.h"
+#include "bpf_arena_common.bpf.h"
 
 #define PF_IDLE				0x00000002	/* I am an IDLE thread */
 #define PF_IO_WORKER			0x00000010	/* Task is an IO worker */
@@ -115,12 +116,10 @@ u32 scx_bpf_cidperf_cur(s32 cid) __ksym __weak;
 s32 scx_bpf_cidperf_set(s32 cid, u32 perf) __ksym __weak;
 
 /* sub-scheduler cap control, scx_bpf_sub_caps() cgroup_id 0 == self */
-s32 scx_bpf_sub_grant(u64 cgroup_id, u64 caps, const struct scx_cmask *cmask,
-		      struct scx_cmask *denied) __ksym __weak;
-void scx_bpf_sub_revoke(u64 cgroup_id, u64 caps, const struct scx_cmask *cmask) __ksym __weak;
-s32 scx_bpf_sub_caps(u64 cgroup_id, u64 caps, struct scx_cmask *out) __ksym __weak;
-s32 scx_bpf_sub_kill_bstr(u64 cgroup_id, char *fmt,
-			  unsigned long long *data, u32 data__sz) __ksym __weak;
+s32 scx_bpf_sub_grant(u64 cgroup_id, u64 caps, const struct scx_cmask __arena *cmask__arena, struct scx_cmask __arena *denied_out__arena__nullable) __ksym __weak;
+void scx_bpf_sub_revoke(u64 cgroup_id, u64 caps, const struct scx_cmask __arena *cmask__arena) __ksym __weak;
+s32 scx_bpf_sub_caps(u64 cgroup_id, u64 caps, struct scx_cmask __arena *out__arena) __ksym __weak;
+s32 scx_bpf_sub_kill_bstr(u64 cgroup_id, char *fmt, unsigned long long *data, u32 data__sz) __ksym __weak;
 
 /*
  * Use the following as @it__iter when calling scx_bpf_dsq_move[_vtime]() from
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index dd0434437835..61dc748742a2 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -1385,8 +1385,8 @@ __noinline void compute_partition(void)
 	}
 
 	/* find out the cids we hold */
-	scx_bpf_sub_caps(0, SCX_CAP_ENQ, (void *)(long)&qa.held_excl.mask);
-	scx_bpf_sub_caps(0, SCX_CAP_ENQ_IMMED, (void *)(long)&qa.held_shared.mask);
+	scx_bpf_sub_caps(0, SCX_CAP_ENQ, &qa.held_excl.mask);
+	scx_bpf_sub_caps(0, SCX_CAP_ENQ_IMMED, &qa.held_shared.mask);
 	cmask_andnot(&qa.held_shared.mask, &qa.held_excl.mask);	/* held only as ENQ_IMMED */
 
 	qa.part.nr_shared = 0;
@@ -1613,13 +1613,13 @@ __noinline void apply_partition(void)
 		cmask_andnot(&qa.to_grant_cids.mask, &ssc->prev_granted.mask);
 
 		scx_bpf_sub_revoke(cgid, SCX_CAP_ENQ_IMMED | SCX_CAP_PERF,
-				   (void *)(long)&qa.prev_rr_cids.mask);
+				   &qa.prev_rr_cids.mask);
 		scx_bpf_sub_revoke(cgid, SCX_CAP_ENQ | SCX_CAP_PREEMPT |
 				   SCX_CAP_ENQ_IMMED | SCX_CAP_PERF,
-				   (void *)(long)&qa.to_revoke_cids.mask);
+				   &qa.to_revoke_cids.mask);
 		scx_bpf_sub_grant(cgid, SCX_CAP_ENQ | SCX_CAP_PREEMPT |
 				  SCX_CAP_ENQ_IMMED | SCX_CAP_PERF,
-				  (void *)(long)&qa.to_grant_cids.mask, NULL);
+				  &qa.to_grant_cids.mask, NULL);
 	}
 
 	/* the current holder of the shared pool gets ENQ_IMMED on all of it */
@@ -1636,7 +1636,7 @@ __noinline void apply_partition(void)
 		if (holder_cgid)
 			scx_bpf_sub_grant(holder_cgid,
 					  SCX_CAP_ENQ_IMMED | SCX_CAP_PERF,
-					  (void *)(long)&qa.rr_cids.mask, NULL);
+					  &qa.rr_cids.mask, NULL);
 	}
 }
 
@@ -1728,11 +1728,11 @@ static void rr_advance(void)
 		if (old_cgid)
 			scx_bpf_sub_revoke(old_cgid,
 					   SCX_CAP_ENQ_IMMED | SCX_CAP_PERF,
-					   (void *)(long)&qa.rr_cids.mask);
+					   &qa.rr_cids.mask);
 		if (new_cgid)
 			scx_bpf_sub_grant(new_cgid,
 					  SCX_CAP_ENQ_IMMED | SCX_CAP_PERF,
-					  (void *)(long)&qa.rr_cids.mask, NULL);
+					  &qa.rr_cids.mask, NULL);
 	}
 
 	part_end();
@@ -1840,8 +1840,8 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init)
 	cmask_init(&qa.held_excl.mask, 0, nr_cids);
 	cmask_init(&qa.held_shared.mask, 0, nr_cids);
 
-	scx_bpf_sub_caps(0, SCX_CAP_ENQ, (void *)(long)&qa.held_excl.mask);
-	scx_bpf_sub_caps(0, SCX_CAP_ENQ_IMMED, (void *)(long)&qa.held_shared.mask);
+	scx_bpf_sub_caps(0, SCX_CAP_ENQ, &qa.held_excl.mask);
+	scx_bpf_sub_caps(0, SCX_CAP_ENQ_IMMED, &qa.held_shared.mask);
 	cmask_andnot(&qa.held_shared.mask, &qa.held_excl.mask);
 
 	bpf_for(i, 0, MAX_SUB_SCHEDS) {