Improve dsq types

Author: tommy-u

scheds/rust/scx_mitosis/src/bpf/dsq.bpf.h

@@ -3,7 +3,7 @@
  * This software may be used and distributed according to the terms of the
  * GNU General Public License version 2.
  *
- * This header defines the 32-bit dispatch queue (DSQ) ID encoding
+ * This header defines the 64-bit dispatch queue (DSQ) ID encoding
  * scheme for scx_mitosis, using type fields to distinguish between
  * per-CPU and cell+L3 domain queues. It includes helper functions to
  * construct, validate, and parse these DSQ IDs for queue management.
@@ -37,94 +37,134 @@
  * Only the low 32 bits are used.
  *
  *   [63     ..     32] [31..0]
- *   [  0s or unused  ] [ VAL ]
+ *   [ 0][   unused   ] [ VAL ]
  *
  *   Mitosis uses VAL as follows:
  *
- *   [31..24] [23..0]
+ *   [31..28] [27..0]
  *   [QTYPE ] [DATA ]
  *
- *   QTYPE encodes the queue type (exactly one bit set):
+ *   QTYPE encodes the queue type:
  *
  *     QTYPE = 0x1 -> Per-CPU Q
- *       [31 .. 24] [23 .. 16] [15    ..      0]
- *       [00000001] [00000000] [      CPU#     ]
+ *       [31..28] [27 ..          ..        0]
+ *       [ 0001 ] [          CPU#            ]
  *       [Q-TYPE:1]
  *
  *     QTYPE = 0x2 -> Cell+L3 Q
- *       [31 .. 24] [23 .. 16] [15      ..    0]
- *       [00000010] [  CELL# ] [      L3ID     ]
+ *       [31..28] [27 .. 16] [15      ..    0]
+ *       [ 0010 ] [  CELL# ] [      L3ID     ]
  *       [Q-TYPE:2]
  *
  */
 
+#if __BYTE_ORDER__ != __ORDER_LITTLE_ENDIAN__
+#error "dsq64 bitfield layout assumes little-endian (bpfel)."
+#endif
+
+/* ---- Bitfield widths (bits) ---- */
+#define CPU_B     28
+#define L3_B      16
+#define CELL_B    12
+#define TYPE_B     4
+#define DATA_B    28
+#define RSVD_B    32
+
+/* Sum checks (in bits) */
+_Static_assert(CPU_B  + TYPE_B          == 32, "CPU layout low half must be 32 bits");
+_Static_assert(L3_B   + CELL_B + TYPE_B == 32, "CELL+L3 layout low half must be 32 bits");
+_Static_assert(DATA_B + TYPE_B          == 32, "Common layout low half must be 32 bits");
+
+typedef union {
+	u64 raw;
+
+	/* Per-CPU user DSQ */
+	struct { u64 cpu: CPU_B;   u64 type: TYPE_B; u64 rsvd: RSVD_B; } cpu_dsq;
+
+	/* Cell+L3 user DSQ */
+	struct { u64 l3: L3_B;     u64 cell: CELL_B; u64 type: TYPE_B; u64 rsvd: RSVD_B; } cell_l3_dsq;
+
+	/* Generic user view */
+	struct { u64 data: DATA_B; u64 type: TYPE_B; u64 rsvd: RSVD_B; } user_dsq;
+
+	/* Built-in DSQ view */
+	struct { u64 value:32; u64 rsvd:30; u64 local_on:1; u64 builtin:1; } builtin_dsq;
+
+	/*
+	 * NOTE: There's an argument for removing both these attributes.
+	 * Packed is making sure the compiler doesn't add padding between
+	 * the fields, but it has the side effect of reducing alignment to 1.
+	 * Right now this is redundant because the bit fields are correctly
+	 * allocated to 64 bits so it's 8-byte aligned. Maybe it's defensive,
+	 * or maybe it's just noise.
+	*/
+} __attribute__((packed, aligned(8))) dsq_id_t;
+
+/*
+ * Invalid DSQ ID Sentinel:
+ * invalid bc bit 63 clear (it's a user DSQ) && dsq_type == 0 (no type)
+ * Good for catching uninitialized DSQ IDs.
+*/
+#define DSQ_INVALID ((dsq_id_t){ .raw = 0 })
+
+_Static_assert(sizeof(((dsq_id_t){0}).cpu_dsq)     == sizeof(u64), "cpu view must be 8 bytes");
+_Static_assert(sizeof(((dsq_id_t){0}).cell_l3_dsq) == sizeof(u64), "cell+l3 view must be 8 bytes");
+_Static_assert(sizeof(((dsq_id_t){0}).user_dsq)    == sizeof(u64), "user common view must be 8 bytes");
+_Static_assert(sizeof(((dsq_id_t){0}).builtin_dsq) == sizeof(u64), "builtin view must be 8 bytes");
+
+/* Compile-time checks (in bytes) */
+_Static_assert(sizeof(dsq_id_t)   == sizeof(u64), "dsq64 must be 8 bytes (64 bits)");
+_Static_assert(_Alignof(dsq_id_t) == 8, "dsq64 must be 8-byte aligned");
+
+/* Range guards */
+_Static_assert(MAX_CPUS  <= (1u << CPU_B),  "MAX_CPUS must fit in field");
+_Static_assert(MAX_L3S   <= (1u << L3_B),   "MAX_L3S must fit in field");
+_Static_assert(MAX_CELLS <= (1u << CELL_B), "MAX_CELLS must fit in field");
+
 /* DSQ type enumeration */
 enum dsq_type {
-	DSQ_UNKNOWN,
+	DSQ_TYPE_NONE,
 	DSQ_TYPE_CPU,
 	DSQ_TYPE_CELL_L3,
 };
 
-/* DSQ ID structure using unions for type-safe access */
-struct dsq_cpu {
-	u32 cpu : 16;
-	u32 unused : 8;
-	u32 type : 8;
-} __attribute__((packed));
-
-struct dsq_cell_l3 {
-	u32 l3 : 16;
-	u32 cell : 8;
-	u32 type : 8;
-} __attribute__((packed));
-
-union dsq_id {
-	u32 raw;
-	struct dsq_cpu cpu;
-	struct dsq_cell_l3 cell_l3;
-	struct {
-		u32 data : 24;
-		u32 type : 8;
-	} common;
-} __attribute__((packed));
-
-/* Static assertions to ensure correct sizes */
-/* Verify that all DSQ structures are exactly 32 bits */
-_Static_assert(sizeof(struct dsq_cpu) == 4, "dsq_cpu must be 32 bits");
-_Static_assert(sizeof(struct dsq_cell_l3) == 4, "dsq_cell_l3 must be 32 bits");
-_Static_assert(sizeof(union dsq_id) == 4, "dsq_id union must be 32 bits");
-
-/* Inline helper functions for DSQ ID manipulation */
+/*
+ * While I considered error propagation, I decided to bail to force errors early.
+*/
+
+static inline bool is_user_dsq(dsq_id_t dsq_id){
+	return !dsq_id.builtin_dsq.builtin;
+}
 
 // Is this a per CPU DSQ?
-static inline bool is_cpu_dsq(u32 dsq_id)
+static inline bool is_cpu_dsq(dsq_id_t dsq_id)
 {
-	union dsq_id id = { .raw = dsq_id };
-	return id.common.type == DSQ_TYPE_CPU;
+	return is_user_dsq(dsq_id) && dsq_id.user_dsq.type == DSQ_TYPE_CPU;
 }
 
 // If this is a per cpu dsq, return the cpu
-static inline u32 get_cpu_from_dsq(u32 dsq_id)
+static inline u32 get_cpu_from_dsq(dsq_id_t dsq_id)
 {
-	union dsq_id id = { .raw = dsq_id };
-	if (id.common.type != DSQ_TYPE_CPU)
-		return DSQ_ERROR;
-	return id.cpu.cpu;
+	if (!is_cpu_dsq(dsq_id))
+		scx_bpf_error("trying to get cpu from non-cpu dsq\n");
+
+	return dsq_id.cpu_dsq.cpu;
 }
 
 /* Helper functions to construct DSQ IDs */
-static inline u32 get_cpu_dsq_id(u32 cpu)
+static inline dsq_id_t get_cpu_dsq_id(u32 cpu)
 {
+	// Check for valid CPU range, 0 indexed so >=.
 	if (cpu >= MAX_CPUS)
-		return DSQ_ERROR;
-	union dsq_id id = { .cpu = { .cpu = cpu, .unused = 0, .type = DSQ_TYPE_CPU } };
-	return id.raw;
+		scx_bpf_error("invalid cpu %d\n", cpu);
+
+	return (dsq_id_t){.cpu_dsq = {.cpu = cpu, .type = DSQ_TYPE_CPU}};
 }
 
-static inline u32 get_cell_l3_dsq_id(u32 cell, u32 l3)
+static inline dsq_id_t get_cell_l3_dsq_id(u32 cell, u32 l3)
 {
 	if (cell >= MAX_CELLS || l3 >= MAX_L3S)
-		return DSQ_ERROR;
-	union dsq_id id = { .cell_l3 = {.l3 = l3, .cell = cell, .type = DSQ_TYPE_CELL_L3 } };
-	return id.raw;
+		scx_bpf_error("cell %d or l3 %d too large\n", cell, l3);
+
+	return (dsq_id_t) { .cell_l3_dsq = { .l3 = l3, .cell = cell, .type = DSQ_TYPE_CELL_L3 } };
 }

scheds/rust/scx_mitosis/src/bpf/mitosis.bpf.c

@@ -118,7 +118,6 @@ static inline struct cgroup *task_cgroup(struct task_struct *p)
 	return cgrp;
 }
 
-
 struct {
 	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
 	__uint(map_flags, BPF_F_NO_PREALLOC);
@@ -347,8 +346,6 @@ static inline int update_task_cpumask(struct task_struct *p,
 
 		/* --- Point to the correct (cell,L3) DSQ and set vtime baseline --- */
 		tctx->dsq = get_cell_l3_dsq_id(tctx->cell, tctx->l3);
-		if (tctx->dsq == DSQ_ERROR)
-			return -EINVAL;
 
 		struct cell *cell = lookup_cell(tctx->cell);
 		if (!cell)
@@ -606,7 +603,7 @@ void BPF_STRUCT_OPS(mitosis_enqueue, struct task_struct *p, u64 enq_flags)
 	if (time_before(vtime, basis_vtime - slice_ns))
 		vtime = basis_vtime - slice_ns;
 
-	scx_bpf_dsq_insert_vtime(p, tctx->dsq, slice_ns, vtime, enq_flags);
+	scx_bpf_dsq_insert_vtime(p, tctx->dsq.raw, slice_ns, vtime, enq_flags);
 
 	/* Kick the CPU if needed */
 	if (!__COMPAT_is_enq_cpu_selected(enq_flags) && cpu >= 0)
@@ -626,14 +623,10 @@ void BPF_STRUCT_OPS(mitosis_dispatch, s32 cpu, struct task_struct *prev)
 	cell = READ_ONCE(cctx->cell);
 
 	/* Start from a valid DSQ */
-	u64 local_dsq = get_cpu_dsq_id(cpu);
-	if (!local_dsq) {
-		scx_bpf_error("get_cpu_dsq_id(%d) returned 0", cpu);
-		return;
-	}
+	dsq_id_t local_dsq = get_cpu_dsq_id(cpu);
 
 	bool found = false;
-	u64 min_vtime_dsq = local_dsq;
+	dsq_id_t min_vtime_dsq = local_dsq;
 	u64 min_vtime = ~0ULL; /* U64_MAX */
 	struct task_struct *p;
 
@@ -644,19 +637,17 @@ void BPF_STRUCT_OPS(mitosis_dispatch, s32 cpu, struct task_struct *prev)
 
 	/* Check the L3 queue */
 	if (l3 != INVALID_L3_ID) {
-		u64 cell_l3_dsq = get_cell_l3_dsq_id(cell, l3);
-		if (cell_l3_dsq != DSQ_ERROR) {
-			bpf_for_each(scx_dsq, p, cell_l3_dsq, 0) {
-				min_vtime = p->scx.dsq_vtime;
-				min_vtime_dsq = cell_l3_dsq;
-				found = true;
-				break;
-			}
+		dsq_id_t cell_l3_dsq = get_cell_l3_dsq_id(cell, l3);
+		bpf_for_each(scx_dsq, p, cell_l3_dsq.raw, 0) {
+			min_vtime = p->scx.dsq_vtime;
+			min_vtime_dsq = cell_l3_dsq;
+			found = true;
+			break;
 		}
 	}
 
 	/* Check the CPU DSQ for a lower vtime */
-	bpf_for_each(scx_dsq, p, local_dsq, 0) {
+	bpf_for_each(scx_dsq, p, local_dsq.raw, 0) {
 		if (!found || time_before(p->scx.dsq_vtime, min_vtime)) {
 			min_vtime = p->scx.dsq_vtime;
 			min_vtime_dsq = local_dsq;
@@ -672,7 +663,7 @@ void BPF_STRUCT_OPS(mitosis_dispatch, s32 cpu, struct task_struct *prev)
 	*/
 	// TODO: The upstream side has "&& min_vtime_dsq != dsq" as part of this condition.
 	// Do we care?
-	if (!scx_bpf_dsq_move_to_local(min_vtime_dsq)) {
+	if (!scx_bpf_dsq_move_to_local(min_vtime_dsq.raw)) {
 #if MITOSIS_ENABLE_STEALING
 		/* Dead-simple work stealing:
 		 * If our local choices are empty, scan sibling (cell,L3) DSQs in the
@@ -684,19 +675,18 @@ void BPF_STRUCT_OPS(mitosis_dispatch, s32 cpu, struct task_struct *prev)
 			// TODO: This math is kinda dumb and confusing.
 			u32 start = ((u32)l3 + 1) % nr_l3;
 			u32 off;
+			// TODO: This might try a bunch of L3s outside of the cell
 			bpf_for (off, 0, nr_l3) {
 				u32 cand = (start + off) % nr_l3;
 				if (cand == (u32)l3)
 					continue;
-				u64 src = get_cell_l3_dsq_id(cell, cand);
-				if (src == DSQ_ERROR)
-					continue;
+				dsq_id_t src = get_cell_l3_dsq_id(cell, cand);
 
 				struct task_struct *q;
 				/* Peek only at the head. */
-				bpf_for_each(scx_dsq, q, src, 0) {
+				bpf_for_each(scx_dsq, q, src.raw, 0) {
 					// TODO maybe this should use if (scx_bpf_dsq_move(BPF_FOR_EACH_ITER, q, SCX_DSQ_LOCAL, 0) == 0)
-					if (scx_bpf_dsq_move_to_local(src)) {
+					if (scx_bpf_dsq_move_to_local(src.raw)) {
 						struct task_ctx *qt = lookup_task_ctx(q);
 						if (qt) {
 							qt->steal_count++;
@@ -722,7 +712,7 @@ void BPF_STRUCT_OPS(mitosis_dispatch, s32 cpu, struct task_struct *prev)
 		}
 		if (!moved)
 #endif
-			scx_bpf_dsq_move_to_local(local_dsq);
+			scx_bpf_dsq_move_to_local(local_dsq.raw);
 	}
 }
 
@@ -1149,7 +1139,7 @@ void BPF_STRUCT_OPS(mitosis_running, struct task_struct *p)
 #endif
 
 	/* Validate task's DSQ before it starts running */
-	if (tctx->dsq == 0) {
+	if (tctx->dsq.raw == DSQ_INVALID.raw) {
 		if (tctx->all_cell_cpus_allowed) {
 			scx_bpf_error(
 				"Task %d has invalid DSQ 0 in running callback (CELL-SCHEDULABLE task, can run on any CPU in cell %d)",
@@ -1437,13 +1427,13 @@ static __always_inline void dump_cell_state(u32 cell_idx)
 	}
 }
 
+// TODO: FIX THIS
 static __always_inline void dump_l3_state(){
-
 }
 
 void BPF_STRUCT_OPS(mitosis_dump, struct scx_dump_ctx *dctx)
 {
-	u64 dsq_id;
+	dsq_id_t dsq_id;
 	int i;
 	struct cell *cell;
 	struct cpu_ctx *cpu_ctx;
@@ -1472,7 +1462,7 @@ void BPF_STRUCT_OPS(mitosis_dump, struct scx_dump_ctx *dctx)
 		dsq_id = get_cpu_dsq_id(i);
 		scx_bpf_dump("CPU[%d] cell=%d vtime=%llu nr_queued=%d\n", i,
 			     cpu_ctx->cell, READ_ONCE(cpu_ctx->vtime_now),
-			     scx_bpf_dsq_nr_queued(dsq_id));
+			     scx_bpf_dsq_nr_queued(dsq_id.raw));
 	}
 
 	dump_l3_state();
@@ -1488,9 +1478,9 @@ void BPF_STRUCT_OPS(mitosis_dump_task, struct scx_dump_ctx *dctx,
 		return;
 
 	scx_bpf_dump(
-		"Task[%d] vtime=%llu basis_vtime=%llu cell=%u dsq=%x all_cell_cpus_allowed=%d\n",
+		"Task[%d] vtime=%llu basis_vtime=%llu cell=%u dsq=%llu all_cell_cpus_allowed=%d\n",
 		p->pid, p->scx.dsq_vtime, tctx->basis_vtime, tctx->cell,
-		tctx->dsq, tctx->all_cell_cpus_allowed);
+		tctx->dsq.raw, tctx->all_cell_cpus_allowed);
 	scx_bpf_dump("Task[%d] CPUS=", p->pid);
 	dump_cpumask(p->cpus_ptr);
 	scx_bpf_dump("\n");
@@ -1522,7 +1512,7 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(mitosis_init)
 		if ((u8_ptr = MEMBER_VPTR(all_cpus, [i / 8]))) {
 			if (*u8_ptr & (1 << (i % 8))) {
 				bpf_cpumask_set_cpu(i, cpumask);
-				ret = scx_bpf_create_dsq(get_cpu_dsq_id(i), ANY_NUMA);
+				ret = scx_bpf_create_dsq(get_cpu_dsq_id(i).raw, ANY_NUMA);
 				if (ret < 0) {
 					bpf_cpumask_release(cpumask);
 					return ret;
@@ -1584,12 +1574,10 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(mitosis_init)
 		u32 l3;
 		bpf_for(l3, 0, nr_l3)
 		{
-			u32 id = get_cell_l3_dsq_id(i, l3);
-			if (id == DSQ_ERROR)
-				return -EINVAL;
-			ret = scx_bpf_create_dsq(id, ANY_NUMA);
+			dsq_id_t id = get_cell_l3_dsq_id(i, l3);
+			ret = scx_bpf_create_dsq(id.raw, ANY_NUMA);
 			if (ret < 0)
-				return ret;
+				scx_bpf_error( "Failed to create DSQ for cell %d, L3 %d: err %d", i, l3, ret);
 		}
 	}