Merge pull request #82 from pythonbpf/compilation-context

PythonBPF: Add Compilation Context to allow parallel compilation of multiple jobs

.github/workflows/python-publish.yml

@@ -33,7 +33,7 @@ jobs:
           python -m build
 
       - name: Upload distributions
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           name: release-dists
           path: dist/
@@ -59,7 +59,7 @@ jobs:
 
     steps:
       - name: Retrieve release distributions
-        uses: actions/download-artifact@v7
+        uses: actions/download-artifact@v8
         with:
           name: release-dists
           path: dist/

pythonbpf/allocation_pass.py

@@ -298,15 +298,6 @@ def allocate_temp_pool(builder, max_temps, local_sym_tab):
             logger.debug(f"Allocated temp variable: {temp_name}")
 
 
-def _get_alignment(tmp_type):
-    """Return alignment for a given type."""
-    if isinstance(tmp_type, ir.PointerType):
-        return 8
-    elif isinstance(tmp_type, ir.IntType):
-        return tmp_type.width // 8
-    return 8
-
-
 def _allocate_for_name(builder, var_name, rval, local_sym_tab):
     """Allocate memory for variable-to-variable assignment (b = a)."""
     source_var = rval.id
@@ -329,16 +320,6 @@ def _allocate_for_name(builder, var_name, rval, local_sym_tab):
     )
 
 
-def _allocate_with_type(builder, var_name, ir_type):
-    """Allocate memory for a variable with a specific type."""
-    var = builder.alloca(ir_type, name=var_name)
-    if isinstance(ir_type, ir.IntType):
-        var.align = ir_type.width // 8
-    elif isinstance(ir_type, ir.PointerType):
-        var.align = 8
-    return var
-
-
 def _allocate_for_attribute(
     builder, var_name, rval, local_sym_tab, compilation_context
 ):
@@ -477,3 +458,20 @@ def _allocate_for_attribute(
     logger.info(
         f"Pre-allocated {var_name} from {struct_var}.{field_name} with type {alloc_type}"
     )
+
+
+def _allocate_with_type(builder, var_name, ir_type):
+    """Allocate variable with appropriate alignment for type."""
+    var = builder.alloca(ir_type, name=var_name)
+    var.align = _get_alignment(ir_type)
+    return var
+
+
+def _get_alignment(ir_type):
+    """Get appropriate alignment for IR type."""
+    if isinstance(ir_type, ir.IntType):
+        return ir_type.width // 8
+    elif isinstance(ir_type, ir.ArrayType) and isinstance(ir_type.element, ir.IntType):
+        return ir_type.element.width // 8
+    else:
+        return 8  # Default: pointer size

pythonbpf/assign_pass.py

@@ -45,7 +45,6 @@ def handle_struct_field_assignment(
     if _is_char_array(field_type) and _is_i8_ptr(val_type):
         _copy_string_to_char_array(
             func,
-            compilation_context,
             builder,
             val,
             field_ptr,
@@ -62,7 +61,6 @@ def handle_struct_field_assignment(
 
 def _copy_string_to_char_array(
     func,
-    compilation_context,
     builder,
     src_ptr,
     dst_ptr,

pythonbpf/globals_pass.py

@@ -32,7 +32,7 @@ def populate_global_symbol_table(tree, compilation_context):
     return False
 
 
-def emit_global(module: ir.Module, node, name):
+def _emit_global(module: ir.Module, node, name):
     logger.info(f"global identifier {name} processing")
     # deduce LLVM type from the annotated return
     if not isinstance(node.returns, ast.Name):
@@ -111,14 +111,14 @@ def globals_processing(tree, compilation_context):
                         node.body[0].value, (ast.Constant, ast.Name, ast.Call)
                     )
                 ):
-                    emit_global(compilation_context.module, node, name)
+                    _emit_global(compilation_context.module, node, name)
                 else:
                     raise SyntaxError(f"ERROR: Invalid syntax for {name} global")
 
     return None
 
 
-def emit_llvm_compiler_used(module: ir.Module, names: list[str]):
+def _emit_llvm_compiler_used(module: ir.Module, names: list[str]):
     """
     Emit the @llvm.compiler.used global given a list of function/global names.
     """
@@ -164,4 +164,4 @@ def globals_list_creation(tree, compilation_context):
                 elif isinstance(dec, ast.Name) and dec.id == "map":
                     collected.append(node.name)
 
-    emit_llvm_compiler_used(module, collected)
+    _emit_llvm_compiler_used(module, collected)

pythonbpf/helper/printk_formatter.py

@@ -41,7 +41,7 @@ def handle_fstring_print(
                 fmt_parts,
                 exprs,
                 local_sym_tab,
-                compilation_context.struct_sym_tab,
+                compilation_context.structs_sym_tab,
             )
         else:
             raise NotImplementedError(f"Unsupported f-string value type: {type(value)}")
@@ -55,12 +55,7 @@ def handle_fstring_print(
 
     for expr in exprs[:3]:
         arg_value = _prepare_expr_args(
-            expr,
-            func,
-            compilation_context.module,
-            builder,
-            local_sym_tab,
-            compilation_context.struct_sym_tab,
+            expr, func, compilation_context, builder, local_sym_tab
         )
         args.append(arg_value)
 
@@ -216,19 +211,19 @@ def _create_format_string_global(fmt_str, func, module, builder):
     return builder.bitcast(fmt_gvar, ir.PointerType())
 
 
-def _prepare_expr_args(expr, func, module, builder, local_sym_tab, struct_sym_tab):
+def _prepare_expr_args(expr, func, compilation_context, builder, local_sym_tab):
     """Evaluate and prepare an expression to use as an arg for bpf_printk."""
 
     # Special case: struct field char array needs pointer to first element
     if isinstance(expr, ast.Attribute):
         char_array_ptr, _ = get_char_array_ptr_and_size(
-            expr, builder, local_sym_tab, struct_sym_tab, func
+            expr, builder, local_sym_tab, compilation_context, func
         )
         if char_array_ptr:
             return char_array_ptr
 
     # Regular expression evaluation
-    val, _ = eval_expr(func, module, builder, expr, local_sym_tab, None, struct_sym_tab)
+    val, _ = eval_expr(func, compilation_context, builder, expr, local_sym_tab)
 
     if not val:
         logger.warning("Failed to evaluate expression for bpf_printk, defaulting to 0")

pythonbpf/maps/map_debug_info.py

@@ -6,9 +6,10 @@ from .map_types import BPFMapType
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def create_map_debug_info(module, map_global, map_name, map_params, structs_sym_tab):
+def create_map_debug_info(compilation_context, map_global, map_name, map_params):
     """Generate debug info metadata for BPF maps HASH and PERF_EVENT_ARRAY"""
-    generator = DebugInfoGenerator(module)
+    generator = DebugInfoGenerator(compilation_context.module)
+    structs_sym_tab = compilation_context.structs_sym_tab
     logger.info(f"Creating debug info for map {map_name} with params {map_params}")
     uint_type = generator.get_uint32_type()
     array_type = generator.create_array_type(
@@ -77,11 +78,9 @@ def create_map_debug_info(module, map_global, map_name, map_params, structs_sym_
 # Ideally we should expose a single create_map_debug_info function that handles all map types.
 # We can probably use a registry pattern to register different map types and their debug info generators.
 # map_params["type"] will be used to determine which generator to use.
-def create_ringbuf_debug_info(
-    module, map_global, map_name, map_params, structs_sym_tab
-):
+def create_ringbuf_debug_info(compilation_context, map_global, map_name, map_params):
     """Generate debug information metadata for BPF RINGBUF map"""
-    generator = DebugInfoGenerator(module)
+    generator = DebugInfoGenerator(compilation_context.module)
 
     int_type = generator.get_int32_type()
 

pythonbpf/maps/maps_pass.py

@@ -31,7 +31,7 @@ def is_map(func_node):
     )
 
 
-def create_bpf_map(module, map_name, map_params):
+def create_bpf_map(compilation_context, map_name, map_params):
     """Create a BPF map in the module with given parameters and debug info"""
 
     # Create the anonymous struct type for BPF map
@@ -40,7 +40,9 @@ def create_bpf_map(module, map_name, map_params):
     )
 
     # Create the global variable
-    map_global = ir.GlobalVariable(module, map_struct_type, name=map_name)
+    map_global = ir.GlobalVariable(
+        compilation_context.module, map_struct_type, name=map_name
+    )
     map_global.linkage = "dso_local"
     map_global.global_constant = False
     map_global.initializer = ir.Constant(map_struct_type, None)
@@ -51,11 +53,13 @@ def create_bpf_map(module, map_name, map_params):
     return MapSymbol(type=map_params["type"], sym=map_global, params=map_params)
 
 
-def _parse_map_params(rval, compilation_context, expected_args=None):
+def _parse_map_params(rval, expected_args=None):
     """Parse map parameters from call arguments and keywords."""
 
     params = {}
-    handler = compilation_context.vmlinux_handler
+
+    # TODO: Replace it with compilation_context.vmlinux_handler someday?
+    handler = VmlinuxHandlerRegistry.get_handler()
     # Parse positional arguments
     if expected_args:
         for i, arg_name in enumerate(expected_args):
@@ -82,14 +86,6 @@ def _parse_map_params(rval, compilation_context, expected_args=None):
 def _get_vmlinux_enum(handler, name):
     if handler and handler.is_vmlinux_enum(name):
         return handler.get_vmlinux_enum_value(name)
-
-    # Fallback to VmlinuxHandlerRegistry if handler invalid
-    # This is for backward compatibility or if refactoring isn't complete
-    if (
-        VmlinuxHandlerRegistry.get_handler()
-        and VmlinuxHandlerRegistry.get_handler().is_vmlinux_enum(name)
-    ):
-        return VmlinuxHandlerRegistry.get_handler().get_vmlinux_enum_value(name)
     return None
 
 
@@ -97,9 +93,7 @@ def _get_vmlinux_enum(handler, name):
 def process_ringbuf_map(map_name, rval, compilation_context):
     """Process a BPF_RINGBUF map declaration"""
     logger.info(f"Processing Ringbuf: {map_name}")
-    map_params = _parse_map_params(
-        rval, compilation_context, expected_args=["max_entries"]
-    )
+    map_params = _parse_map_params(rval, expected_args=["max_entries"])
     map_params["type"] = BPFMapType.RINGBUF
 
     # NOTE: constraints borrowed from https://docs.ebpf.io/linux/map-type/BPF_MAP_TYPE_RINGBUF/
@@ -115,13 +109,12 @@ def process_ringbuf_map(map_name, rval, compilation_context):
 
     logger.info(f"Ringbuf map parameters: {map_params}")
 
-    map_global = create_bpf_map(compilation_context.module, map_name, map_params)
+    map_global = create_bpf_map(compilation_context, map_name, map_params)
     create_ringbuf_debug_info(
-        compilation_context.module,
+        compilation_context,
         map_global.sym,
         map_name,
         map_params,
-        compilation_context.structs_sym_tab,
     )
     return map_global
 
@@ -130,20 +123,17 @@ def process_ringbuf_map(map_name, rval, compilation_context):
 def process_hash_map(map_name, rval, compilation_context):
     """Process a BPF_HASH map declaration"""
     logger.info(f"Processing HashMap: {map_name}")
-    map_params = _parse_map_params(
-        rval, compilation_context, expected_args=["key", "value", "max_entries"]
-    )
+    map_params = _parse_map_params(rval, expected_args=["key", "value", "max_entries"])
     map_params["type"] = BPFMapType.HASH
 
     logger.info(f"Map parameters: {map_params}")
-    map_global = create_bpf_map(compilation_context.module, map_name, map_params)
+    map_global = create_bpf_map(compilation_context, map_name, map_params)
     # Generate debug info for BTF
     create_map_debug_info(
-        compilation_context.module,
+        compilation_context,
         map_global.sym,
         map_name,
         map_params,
-        compilation_context.structs_sym_tab,
     )
     return map_global
 
@@ -152,20 +142,17 @@ def process_hash_map(map_name, rval, compilation_context):
 def process_perf_event_map(map_name, rval, compilation_context):
     """Process a BPF_PERF_EVENT_ARRAY map declaration"""
     logger.info(f"Processing PerfEventArray: {map_name}")
-    map_params = _parse_map_params(
-        rval, compilation_context, expected_args=["key_size", "value_size"]
-    )
+    map_params = _parse_map_params(rval, expected_args=["key_size", "value_size"])
     map_params["type"] = BPFMapType.PERF_EVENT_ARRAY
 
     logger.info(f"Map parameters: {map_params}")
-    map_global = create_bpf_map(compilation_context.module, map_name, map_params)
+    map_global = create_bpf_map(compilation_context, map_name, map_params)
     # Generate debug info for BTF
     create_map_debug_info(
-        compilation_context.module,
+        compilation_context,
         map_global.sym,
         map_name,
         map_params,
-        compilation_context.structs_sym_tab,
     )
     return map_global