Re: [PATCH] sched/ext: Suppress warning in __this_cpu_write() by disabling preemption

[Andrea Righi <arighi@nvidia.com>]

On Wed, Jul 16, 2025 at 03:36:31PM +0200, Peter Zijlstra wrote:
> On Wed, Jul 16, 2025 at 03:15:12PM +0200, Andrea Righi wrote:
> 
> > The idea is to track the scx callbacks that are invoked with a rq lock held
> > and, in those cases, store the locked rq. However, some callbacks may also
> > be invoked from an unlocked context, where no rq is locked and in this case
> > rq should be NULL.
> > 
> > In the latter case, it's acceptable for preemption to remain enabled, but
> > we still want to explicitly set locked_rq = NULL. If during the execution
> > of the callback we jump on another CPU, it'd still be in an unlocked state,
> > so it's locked_rq is still NULL.
> 
> Right; but doing superfluous NULL stores seems pointless. So better to
> avoid the store entirely, rather than making it more expensive and no
> less pointless, right?

Right, we can definitely avoid rewriting NULL.
The following should do the trick.

Breno, can you give it a try?

Thanks,
-Andrea

 kernel/sched/ext.c | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index e231450768897..c76d6c9e497b4 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -1290,12 +1290,15 @@ static inline void update_locked_rq(struct rq *rq)
 	 */
 	if (rq)
 		lockdep_assert_rq_held(rq);
+
+	WARN_ON_ONCE(preemptible());
 	__this_cpu_write(scx_locked_rq_state, rq);
 }
 
 #define SCX_CALL_OP(sch, mask, op, rq, args...)					\
 do {										\
-	update_locked_rq(rq);							\
+	if (rq)									\
+		update_locked_rq(rq);						\
 	if (mask) {								\
 		scx_kf_allow(mask);						\
 		(sch)->ops.op(args);						\
@@ -1303,14 +1306,16 @@ do {										\
 	} else {								\
 		(sch)->ops.op(args);						\
 	}									\
-	update_locked_rq(NULL);							\
+	if (rq)									\
+		update_locked_rq(NULL);						\
 } while (0)
 
 #define SCX_CALL_OP_RET(sch, mask, op, rq, args...)				\
 ({										\
 	__typeof__((sch)->ops.op(args)) __ret;					\
 										\
-	update_locked_rq(rq);							\
+	if (rq)									\
+		update_locked_rq(rq);						\
 	if (mask) {								\
 		scx_kf_allow(mask);						\
 		__ret = (sch)->ops.op(args);					\
@@ -1318,7 +1323,8 @@ do {										\
 	} else {								\
 		__ret = (sch)->ops.op(args);					\
 	}									\
-	update_locked_rq(NULL);							\
+	if (rq)									\
+		update_locked_rq(NULL);						\
 	__ret;									\
 })