Lower kfunc declarations into the IR

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

Author: congwang-mk

src/ebpf_c_codegen.ml

@@ -284,6 +284,16 @@ let escape_c_string s =
 
 let ebpf_type_from_ir_type = Codegen_common.ir_type_to_c Codegen_common.EbpfKernel
 
+(** Type conversion for kfunc signatures. Unlike ebpf_type_from_ir_type, this keeps
+    [IRBool] as C [bool] rather than collapsing it to [__u8], so the emitted prototype
+    matches the kernel's actual kfunc signature (kfunc/extern type matching is strict). *)
+let rec kfunc_signature_type_to_c = function
+  | Ir.IRU8 -> "__u8" | Ir.IRU16 -> "__u16" | Ir.IRU32 -> "__u32" | Ir.IRU64 -> "__u64"
+  | Ir.IRI8 -> "__s8" | Ir.IRI16 -> "__s16" | Ir.IRI32 -> "__s32" | Ir.IRI64 -> "__s64"
+  | Ir.IRBool -> "bool" | Ir.IRChar -> "char" | Ir.IRVoid -> "void"
+  | Ir.IRPointer (inner_type, _) -> sprintf "%s*" (kfunc_signature_type_to_c inner_type)
+  | other -> ebpf_type_from_ir_type other
+
 (** Generate proper C declaration for eBPF, handling function pointers correctly *)
 let generate_ebpf_c_declaration = Codegen_common.c_declaration Codegen_common.EbpfKernel
 
@@ -3036,6 +3046,27 @@ let generate_declarations_in_source_order_unified ctx ir_multi_prog ~_btf_path _
         (* Generate struct_ops instance *)
         emit_line ctx (sprintf "/* eBPF struct_ops instance: %s */" struct_ops_instance.ir_instance_name);
         emit_blank_line ctx
+
+    | Ir.IRDeclKfuncDecl kfunc_decl ->
+        (* Emit a function-signature declaration so eBPF programs can call the kfunc.
+           Kernel-provided kfuncs need `extern ... __ksym;`; locally-defined @kfunc
+           functions just need a forward prototype. *)
+        let params_str = match kfunc_decl.Ir.ikfunc_params with
+          | [] -> "void"
+          | params -> String.concat ", " (List.map (fun (name, ir_type) ->
+              sprintf "%s %s" (kfunc_signature_type_to_c ir_type) name
+            ) params)
+        in
+        let return_type_str = kfunc_signature_type_to_c kfunc_decl.Ir.ikfunc_return_type in
+        if kfunc_decl.Ir.ikfunc_is_extern then
+          emit_line ctx (sprintf "extern %s %s(%s) __ksym;"
+            return_type_str kfunc_decl.Ir.ikfunc_name params_str)
+        else (
+          emit_line ctx "/* kfunc declaration */";
+          emit_line ctx (sprintf "%s %s(%s);"
+            return_type_str kfunc_decl.Ir.ikfunc_name params_str)
+        );
+        emit_blank_line ctx
   ) ir_multi_prog.Ir.source_declarations;
 
   (* Emit callbacks at the end if no functions were found (fallback) *)
@@ -3241,7 +3272,7 @@ let collect_callback_dependencies ir_multi_prog =
 
 (** Compile multi-program IR to eBPF C code with automatic tail call detection *)
 let compile_multi_to_c_with_tail_calls
-    ?(kfunc_declarations=[]) ?(tail_call_analysis=None) ?(btf_path=None)
+    ?(tail_call_analysis=None) ?(btf_path=None)
     (ir_multi_prog : Ir.ir_multi_program) =
 
   let ctx = create_c_context () in
@@ -3260,30 +3291,10 @@ let compile_multi_to_c_with_tail_calls
   let uses_dynptr = check_dynptr_usage ir_multi_prog in
   if uses_dynptr then generate_dynptr_macros ctx;
 
-  (* Generate kfunc declarations *)
-  let rec ast_type_to_c_type ast_type =
-    match ast_type with
-    | Ast.U8 -> "__u8" | Ast.U16 -> "__u16" | Ast.U32 -> "__u32" | Ast.U64 -> "__u64"
-    | Ast.I8 -> "__s8" | Ast.I16 -> "__s16" | Ast.I32 -> "__s32" | Ast.I64 -> "__s64"
-    | Ast.Bool -> "bool" | Ast.Char -> "char" | Ast.Void -> "void"
-    | Ast.Pointer inner_type -> sprintf "%s*" (ast_type_to_c_type inner_type)
-    | _ -> "void"
-  in
-  List.iter (fun kfunc ->
-    let params_str = String.concat ", " (List.map (fun (name, param_type) ->
-      let c_type = ast_type_to_c_type param_type in
-      sprintf "%s %s" c_type name
-    ) kfunc.Ast.func_params) in
-    let return_type_str = match Ast.get_return_type kfunc.Ast.func_return_type with
-      | Some ret_type -> ast_type_to_c_type ret_type
-      | None -> "void"
-    in
-    emit_line ctx (sprintf "/* kfunc declaration */");
-    emit_line ctx (sprintf "%s %s(%s);" return_type_str kfunc.Ast.func_name params_str);
-  ) kfunc_declarations;
-  
-  if kfunc_declarations <> [] then emit_blank_line ctx;
-  
+  (* Kfunc declarations (both kernel-provided `extern` kfuncs and locally-defined
+     @kfunc prototypes) are carried in the IR as IRDeclKfuncDecl and emitted in
+     source order by generate_declarations_in_source_order_unified. *)
+
   (* Generate string type definitions *)
   generate_string_typedefs ctx ir_multi_prog;
 
@@ -3316,16 +3327,16 @@ let compile_multi_to_c_with_tail_calls
   (final_output, final_tail_call_analysis)
 
 (** Multi-program compilation entry point that returns both code and tail call analysis *)
-let compile_multi_to_c ?(kfunc_declarations=[]) ?(tail_call_analysis=None) ?(btf_path=None) ir_multi_program =
+let compile_multi_to_c ?(tail_call_analysis=None) ?(btf_path=None) ir_multi_program =
   compile_multi_to_c_with_tail_calls
-    ~kfunc_declarations ~tail_call_analysis ~btf_path ir_multi_program
+    ~tail_call_analysis ~btf_path ir_multi_program
 
 (** Alias for backward compatibility with existing code *)
 let compile_multi_to_c_with_analysis = compile_multi_to_c
 
 (** Generate complete C program from multiple IR programs - main interface *)
-let generate_c_multi_program ?(kfunc_declarations=[]) ?(btf_path=None) ir_multi_prog =
-  let (c_code, _) = compile_multi_to_c ~kfunc_declarations ~btf_path ir_multi_prog in
+let generate_c_multi_program ?(btf_path=None) ir_multi_prog =
+  let (c_code, _) = compile_multi_to_c ~btf_path ir_multi_prog in
   c_code
 
 (** Generate complete C program from IR *)

src/ir.ml

@@ -399,6 +399,18 @@ and ir_declaration_desc =
   | IRDeclProgramDef of ir_program
   | IRDeclStructOpsDef of ir_struct_ops_declaration
   | IRDeclStructOpsInstance of ir_struct_ops_instance
+  | IRDeclKfuncDecl of ir_kfunc_declaration
+
+(** Kfunc declaration - a function signature to declare at the top of the eBPF C file.
+    Covers both kernel-provided kfuncs (extern ... __ksym) and locally-defined @kfunc
+    prototypes. The full @kfunc body is emitted separately into the kernel module. *)
+and ir_kfunc_declaration = {
+  ikfunc_name: string;
+  ikfunc_params: (string * ir_type) list;
+  ikfunc_return_type: ir_type;
+  ikfunc_is_extern: bool; (* true: kernel-provided (extern __ksym); false: local @kfunc prototype *)
+  ikfunc_pos: ir_position;
+}
 
 (** Utility functions for creating IR nodes *)
 
@@ -536,6 +548,16 @@ let make_ir_struct_ops_def_decl struct_ops_def order =
 let make_ir_struct_ops_instance_decl struct_ops_instance order =
   make_ir_source_declaration (IRDeclStructOpsInstance struct_ops_instance) order struct_ops_instance.ir_instance_pos
 
+let make_ir_kfunc_decl name params return_type is_extern order pos =
+  make_ir_source_declaration
+    (IRDeclKfuncDecl {
+      ikfunc_name = name;
+      ikfunc_params = params;
+      ikfunc_return_type = return_type;
+      ikfunc_is_extern = is_extern;
+      ikfunc_pos = pos;
+    }) order pos
+
 let make_ir_multi_program source_name ?(source_declarations = [])
                           ?userspace_program ?(ring_buffer_registry = create_empty_ring_buffer_registry ())
                           pos =

src/ir_generator.ml

@@ -3375,8 +3375,32 @@ let lower_multi_program ast symbol_table source_name =
                   add_source_declaration (IRDeclStructOpsDef ir_struct_ops_decl) struct_def.struct_pos
               | None -> ())
          | Ast.AttributedFunction attr_func ->
-             (* Insert program entry functions at source position *)
              let func_name = attr_func.attr_function.func_name in
+             let is_kfunc = List.exists (function
+               | Ast.SimpleAttribute "kfunc" -> true
+               | _ -> false
+             ) attr_func.attr_list in
+             if is_kfunc then (
+               (* Locally-defined @kfunc: emit a prototype declaration so eBPF programs
+                  can call it. The body itself goes into the separate kernel module. *)
+               let func = attr_func.attr_function in
+               let ir_params = List.map (fun (name, param_type) ->
+                 (name, ast_type_to_ir_type param_type)
+               ) func.func_params in
+               let ir_return_type = match Ast.get_return_type func.func_return_type with
+                 | Some ret_type -> ast_type_to_ir_type ret_type
+                 | None -> IRVoid
+               in
+               add_source_declaration
+                 (IRDeclKfuncDecl {
+                    ikfunc_name = func.func_name;
+                    ikfunc_params = ir_params;
+                    ikfunc_return_type = ir_return_type;
+                    ikfunc_is_extern = false;
+                    ikfunc_pos = func.func_pos;
+                  }) func.func_pos
+             ) else
+             (* Insert program entry functions at source position *)
              (match Hashtbl.find_opt program_table func_name with
               | Some (_prog_def, ir_program) ->
                   add_source_declaration (IRDeclProgramDef ir_program) ir_program.ir_pos
@@ -3406,6 +3430,23 @@ let lower_multi_program ast symbol_table source_name =
                     | None -> ())
                | Ast.ImplStaticField _ -> ()
              ) impl_block.impl_items
+         | Ast.ExternKfuncDecl extern_decl ->
+             (* Kernel-provided kfunc: emit an `extern ... __ksym;` declaration *)
+             let ir_params = List.map (fun (name, param_type) ->
+               (name, ast_type_to_ir_type param_type)
+             ) extern_decl.Ast.extern_params in
+             let ir_return_type = match extern_decl.Ast.extern_return_type with
+               | Some ret_type -> ast_type_to_ir_type ret_type
+               | None -> IRVoid
+             in
+             add_source_declaration
+               (IRDeclKfuncDecl {
+                  ikfunc_name = extern_decl.Ast.extern_name;
+                  ikfunc_params = ir_params;
+                  ikfunc_return_type = ir_return_type;
+                  ikfunc_is_extern = true;
+                  ikfunc_pos = extern_decl.Ast.extern_pos;
+                }) extern_decl.Ast.extern_pos
          | _ -> () )
   );
 

src/main.ml

@@ -888,15 +888,6 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
     let _resource_plan = Multi_program_ir_optimizer.plan_system_resources (Ir.get_programs ir_with_ring_buffer_analysis) ir_with_ring_buffer_analysis in
     let _optimization_strategies = Multi_program_ir_optimizer.generate_optimization_strategies multi_prog_analysis in
 
-    (* Extract kfunc declarations from AST for eBPF C generation *)
-    let kfunc_declarations = List.filter_map (function
-      | Ast.AttributedFunction attr_func ->
-          (match attr_func.attr_list with
-           | SimpleAttribute "kfunc" :: _ -> Some attr_func.attr_function
-           | _ -> None)
-      | _ -> None
-    ) annotated_ast in
-
     (* Perform tail call analysis on AST *)
     let tail_call_analysis = Tail_call_analyzer.analyze_tail_calls annotated_ast in
 
@@ -915,9 +906,9 @@ let compile_source input_file output_dir _verbose generate_makefile btf_vmlinux_
       { ir_with_ring_buffer_analysis with source_declarations = updated_source_declarations }
     in
 
-    (* Generate eBPF C code (with updated IR and kfunc declarations) *)
+    (* Generate eBPF C code (kfunc declarations are carried in the IR) *)
     let (ebpf_c_code, _final_tail_call_analysis) = Ebpf_c_codegen.compile_multi_to_c_with_analysis
-      ~kfunc_declarations ~tail_call_analysis:(Some tail_call_analysis) ~btf_path:btf_vmlinux_path updated_optimized_ir in
+      ~tail_call_analysis:(Some tail_call_analysis) ~btf_path:btf_vmlinux_path updated_optimized_ir in
 
     (* Analyze kfunc dependencies for automatic kernel module loading *)
     let ir_functions = List.map (fun prog -> prog.Ir.entry_function) (Ir.get_programs ir_with_ring_buffer_analysis) in

src/userspace_codegen.ml

@@ -1274,6 +1274,10 @@ let generate_declarations_in_source_order_userspace ir_multi_prog =
     | Ir.IRDeclStructOpsInstance _struct_ops_instance ->
         (* Skip struct_ops instances in userspace - they're handled separately *)
         ()
+
+    | Ir.IRDeclKfuncDecl _kfunc_decl ->
+        (* Skip kfunc declarations in userspace - they're eBPF-side only *)
+        ()
   ) ir_multi_prog.Ir.source_declarations;
 
   (* Return the declarations in the correct order (reverse since we prepended) *)

tests/test_definition_order.ml

@@ -107,6 +107,7 @@ let get_declaration_name = function
   | IRDeclProgramDef program -> program.entry_function.func_name
   | IRDeclStructOpsDef struct_ops -> struct_ops.ir_struct_ops_name
   | IRDeclStructOpsInstance instance -> instance.ir_instance_name
+  | IRDeclKfuncDecl kfunc_decl -> kfunc_decl.ikfunc_name
 
 (** Test type alias order preservation *)
 let test_type_alias_order () =

tests/test_extern.ml

@@ -155,6 +155,42 @@ let test_extern_kfunc_string_representation () =
   check bool "Contains bpf_ktime_get_ns extern" true contains_bpf_ktime;
   check bool "Contains bpf_trace_printk extern" true contains_bpf_trace
 
+(** Test extern kfunc declarations are emitted into generated eBPF C with __ksym *)
+let test_extern_kfunc_ebpf_codegen () =
+  let program = {|
+    extern bpf_ktime_get_ns() -> u64
+    extern scx_bpf_dsq_insert(p: *u8, dsq_id: u64, slice: u64, enq_flags: u64) -> void
+
+    @xdp
+    fn test_program(ctx: *xdp_md) -> xdp_action {
+        var ts = bpf_ktime_get_ns()
+        scx_bpf_dsq_insert(null, 0, ts, 0)
+        return 2
+    }
+
+    fn main() -> i32 {
+        return 0
+    }
+  |} in
+
+  let ast = Parse.parse_string program in
+  let symbol_table = Symbol_table.build_symbol_table ast in
+  let (typed_ast, _) = Type_checker.type_check_and_annotate_ast ast in
+  let ir = Ir_generator.generate_ir typed_ast symbol_table "test" in
+
+  (* Extern kfunc declarations are lowered into the IR; codegen needs no side-channel *)
+  let (generated_code, _) =
+    Ebpf_c_codegen.compile_multi_to_c_with_analysis ir in
+
+  let contains substr =
+    try ignore (Str.search_forward (Str.regexp_string substr) generated_code 0); true
+    with Not_found -> false
+  in
+  check bool "Contains bpf_ktime_get_ns __ksym extern" true
+    (contains "extern __u64 bpf_ktime_get_ns(void) __ksym;");
+  check bool "Contains scx_bpf_dsq_insert __ksym extern" true
+    (contains "extern void scx_bpf_dsq_insert(__u8* p, __u64 dsq_id, __u64 slice, __u64 enq_flags) __ksym;")
+
 (** Test extern keyword cannot be used in function definitions *)
 let test_extern_in_function_definition_fails () =
   let program = {|
@@ -233,6 +269,7 @@ let tests = [
   "extern kfunc type checking", `Quick, test_extern_kfunc_type_checking;
   "extern kfunc userspace restriction", `Quick, test_extern_kfunc_userspace_restriction;
   "extern kfunc string representation", `Quick, test_extern_kfunc_string_representation;
+  "extern kfunc ebpf codegen", `Quick, test_extern_kfunc_ebpf_codegen;
   "extern in function definition fails", `Quick, test_extern_in_function_definition_fails;
   "extern with body fails", `Quick, test_extern_with_body_fails;
   "extern mixed keywords fails", `Quick, test_extern_mixed_keywords_fails;

tests/test_kfunc_attribute.ml

@@ -148,18 +148,9 @@ let test_ebpf_kfunc_declarations () =
   (* Use the full multi-program type checker for proper expression typing *)
   let (typed_ast, _) = Type_checker.type_check_and_annotate_ast ast in
   let ir = Ir_generator.generate_ir typed_ast symbol_table "test" in
-  
-  (* Extract kfunc declarations *)
-  let kfunc_declarations = List.filter_map (function
-    | Ast.AttributedFunction attr_func ->
-        (match attr_func.attr_list with
-         | SimpleAttribute "kfunc" :: _ -> Some attr_func.attr_function
-         | _ -> None)
-    | _ -> None
-  ) typed_ast in
-  
-  (* Generate eBPF C code *)
-  let (generated_code, _) = Ebpf_c_codegen.compile_multi_to_c_with_analysis ~kfunc_declarations ir in
+
+  (* @kfunc declarations are lowered into the IR; codegen needs no side-channel *)
+  let (generated_code, _) = Ebpf_c_codegen.compile_multi_to_c_with_analysis ir in
 
   (* Check that kfunc declarations are generated *)
   check bool "Contains kfunc declaration" true