Implement sched-ext scheduler support and add related tests.

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

SPEC.md

@@ -1802,7 +1802,62 @@ fn main() -> i32 {
 }
 ```
 
-#### 3.10.3 Registration Function
+#### 3.10.3 Sched-ext Scheduler Implementation
+
+KernelScript supports sched-ext (extensible scheduler) through the `sched_ext_ops` struct_ops:
+
+```kernelscript
+// Simple FIFO scheduler using sched-ext
+@struct_ops("sched_ext_ops")
+impl simple_fifo_scheduler {
+    
+    // Select CPU for a waking task
+    fn select_cpu(p: *u8, prev_cpu: i32, wake_flags: u64) -> i32 {
+        // Use default CPU selection with direct dispatch if idle core found
+        var direct: bool = false
+        var cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, wake_flags, &direct)
+        
+        if (direct) {
+            // Insert directly into local DSQ, skipping enqueue
+            scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0)
+        }
+        
+        return cpu
+    }
+    
+    // Enqueue task into global FIFO queue
+    fn enqueue(p: *u8, enq_flags: u64) -> void {
+        // Simple FIFO: insert all tasks into global DSQ
+        scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags)
+    }
+    
+    // Dispatch tasks from global queue to local CPU
+    fn dispatch(cpu: i32, prev: *u8) -> void {
+        // Try to consume a task from the global DSQ
+        if (!scx_bpf_consume(SCX_DSQ_GLOBAL)) {
+            // No tasks available, CPU will go idle
+        }
+    }
+    
+    // Initialize scheduler
+    fn init() -> i32 {
+        return 0  // Success
+    }
+    
+    // Scheduler configuration
+    name: "simple_fifo",
+    timeout_ms: 0,  // No timeout
+    flags: 0,       // Default flags
+}
+
+// Register the scheduler
+fn main() -> i32 {
+    var result = register(simple_fifo_scheduler)
+    return result
+}
+```
+
+#### 3.10.4 Registration Function
 
 The `register()` function is type-aware and generates the appropriate registration code:
 

examples/sched_ext_simple.ks

@@ -0,0 +1,150 @@
+// Simple sched-ext scheduler implementation
+// This demonstrates a basic FIFO scheduler using sched_ext_ops
+
+// kfuncs declarations (extracted from BTF)
+extern scx_bpf_select_cpu_dfl(p: *u8, prev_cpu: i32, wake_flags: u64, direct: *bool) -> i32
+extern scx_bpf_dsq_insert(p: *u8, dsq_id: u64, slice: u64, enq_flags: u64) -> void
+extern scx_bpf_consume(dsq_id: u64, cpu: i32, flags: u64) -> i32
+
+// Kernel enums (extracted from BTF)
+enum scx_public_consts {
+        SCX_OPS_NAME_LEN = 128,
+        SCX_SLICE_DFL = 20000000,
+        SCX_SLICE_INF = 18446744073709551615,
+}
+
+enum scx_dsq_id_flags {
+        SCX_DSQ_FLAG_BUILTIN = 9223372036854775808,
+        SCX_DSQ_FLAG_LOCAL_ON = 4611686018427387904,
+        SCX_DSQ_INVALID = 9223372036854775808,
+        SCX_DSQ_GLOBAL = 9223372036854775809,
+        SCX_DSQ_LOCAL = 9223372036854775810,
+        SCX_DSQ_LOCAL_ON = 13835058055282163712,
+        SCX_DSQ_LOCAL_CPU_MASK = 4294967295,
+}
+
+// Define the sched_ext_ops structure (extracted from BTF)
+struct sched_ext_ops {
+    select_cpu: fn(p: *u8, prev_cpu: i32, wake_flags: u64) -> i32,
+    enqueue: fn(p: *u8, enq_flags: u64) -> void,
+    dispatch: fn(cpu: i32, prev: *u8) -> void,
+    runnable: fn(p: *u8, enq_flags: u64) -> void,
+    running: fn(p: *u8) -> void,
+    stopping: fn(p: *u8, runnable: bool) -> void,
+    quiescent: fn(p: *u8, deq_flags: u64) -> void,
+    init_task: fn(p: *u8, args: *u8) -> i32,
+    exit_task: fn(p: *u8, args: *u8) -> void,
+    enable: fn(p: *u8) -> void,
+    init: fn() -> i32,
+    exit: fn(info: *u8) -> void,
+    name: *u8,
+    timeout_ms: u64,
+    flags: u64,
+}
+
+// Simple FIFO scheduler implementation
+@struct_ops("sched_ext_ops")
+impl simple_fifo_scheduler {
+    
+    // Select CPU for a waking task
+    fn select_cpu(p: *u8, prev_cpu: i32, wake_flags: u64) -> i32 {
+        // Use default CPU selection with direct dispatch if idle core found
+        var direct: bool = false
+        var cpu = scx_bpf_select_cpu_dfl(p, prev_cpu, wake_flags, &direct)
+        
+        if (direct) {
+            // Insert directly into local DSQ, skipping enqueue
+            scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0)
+        }
+        
+        return cpu
+    }
+    
+    // Enqueue task into global FIFO queue
+    fn enqueue(p: *u8, enq_flags: u64) -> void {
+        // Simple FIFO: insert all tasks into global DSQ
+        scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, enq_flags)
+    }
+    
+    // Dispatch tasks from global queue to local CPU
+    fn dispatch(cpu: i32, prev: *u8) -> void {
+        // Try to consume a task from the global DSQ
+        if (scx_bpf_consume(SCX_DSQ_GLOBAL, cpu, 0) == 0) {
+            // No tasks available, CPU will go idle
+        }
+    }
+    
+    // Task becomes runnable
+    fn runnable(p: *u8, enq_flags: u64) -> void {
+        // Optional: track runnable tasks
+        // For simple FIFO, we don't need special handling
+    }
+    
+    // Task starts running
+    fn running(p: *u8) -> void {
+        // Optional: track running tasks
+        // For simple FIFO, we don't need special handling
+    }
+    
+    // Task stops running
+    fn stopping(p: *u8, runnable: bool) -> void {
+        // Optional: handle task stopping
+        // For simple FIFO, we don't need special handling
+    }
+    
+    // Task becomes quiescent
+    fn quiescent(p: *u8, deq_flags: u64) -> void {
+        // Optional: handle quiescent tasks
+        // For simple FIFO, we don't need special handling
+    }
+    
+    // Initialize new task
+    fn init_task(p: *u8, args: *u8) -> i32 {
+        // Return 0 for success
+        return 0
+    }
+    
+    // Clean up exiting task
+    fn exit_task(p: *u8, args: *u8) -> void {
+        // Optional cleanup for exiting tasks
+    }
+    
+    // Enable scheduler
+    fn enable(p: *u8) -> void {
+        // Optional: scheduler enable logic
+    }
+    
+    // Initialize scheduler
+    fn init() -> i32 {
+        // Return 0 for successful initialization
+        return 0
+    }
+    
+    // Exit scheduler
+    fn exit(info: *u8) -> void {
+        // Optional cleanup on scheduler exit
+    }
+    
+    // Scheduler name
+    name: "simple_fifo",
+    
+    // Timeout in milliseconds (0 = no timeout)
+    timeout_ms: 0,
+    
+    // Scheduler flags
+    flags: 0,
+}
+
+// Userspace main function
+fn main() -> i32 {
+    // Register the sched-ext scheduler
+    var result = register(simple_fifo_scheduler)
+    
+    if (result == 0) {
+        print("Simple FIFO scheduler registered successfully")
+    } else {
+        print("Failed to register Simple FIFO scheduler")
+    }
+    
+    return result
+}

src/ir_generator.ml

@@ -3246,7 +3246,12 @@ let lower_multi_program ast symbol_table source_name =
                 (match literal with
                 | Ast.StringLit s -> make_ir_value (IRLiteral (StringLit s)) (IRStr (String.length s + 1)) field_expr.expr_pos
                 | Ast.NullLit -> make_ir_value (IRLiteral NullLit) (IRPointer (IRU8, make_bounds_info ~nullable:true ())) field_expr.expr_pos
-                | _ -> failwith "Unsupported literal type in static field")
+                | Ast.IntLit (value, _) -> 
+                    let ir_type = if value < 0 then IRI32 else IRU32 in
+                    make_ir_value (IRLiteral (IntLit (value, None))) ir_type field_expr.expr_pos
+                | Ast.BoolLit b -> make_ir_value (IRLiteral (BoolLit b)) IRBool field_expr.expr_pos
+                | Ast.CharLit c -> make_ir_value (IRLiteral (CharLit c)) IRChar field_expr.expr_pos
+                | _ -> failwith ("Unsupported literal type in static field: " ^ (Ast.string_of_literal literal)))
             | _ -> failwith "Static fields must be literals"
           in
           Some (field_name, field_val)

src/main.ml

@@ -37,7 +37,7 @@ let rec parse_args () =
       printf "    prog_type: xdp | tc/direction | probe/target_function | tracepoint/category/event\n";
       printf "    Examples: tc/ingress, tc/egress, probe/sys_read, probe/vfs_write, probe/tcp_sendmsg\n";
       printf "    tracepoint: tracepoint/syscalls/sys_enter_read, tracepoint/sched/sched_switch\n";
-      printf "    struct_ops: tcp_congestion_ops\n";
+      printf "    struct_ops: tcp_congestion_ops, sched_ext_ops\n";
       printf "    project_name: Name of the project directory to create\n";
       printf "    --btf-vmlinux-path: Path to BTF vmlinux file (default: /sys/kernel/btf/vmlinux)\n";
       printf "    --kfuncs: Extract available kfuncs from BTF and generate .kh header file\n\n";

src/struct_ops_registry.ml

@@ -46,6 +46,12 @@ let known_struct_ops = [
     kernel_version = Some "5.6+";
     common_usage = ["Testing and development"];
   };
+  {
+    name = "sched_ext_ops";
+    description = "Extensible scheduler operations";
+    kernel_version = Some "6.12+";
+    common_usage = ["Custom scheduling policies"; "BPF-based schedulers"; "Scheduler experimentation"];
+  };
 ]
 
 (** Check if a struct_ops type is known *)