Add LLVM optimizer (#104)

* Add LLVM optimizer

* expr_pass: add bpf_passthrough attributes to make tests pass

Author: varun-r-mallya

pythonbpf/codegen.py

@@ -15,6 +15,7 @@
 )
 from .debuginfo import DW_LANG_C11, DwarfBehaviorEnum, DebugInfoGenerator
 import os
+import shutil
 import subprocess
 import inspect
 from pathlib import Path
@@ -175,8 +176,66 @@ def compile_to_ir(filename: str, output: str, loglevel=logging.INFO):
     return output, structs_sym_tab, maps_sym_tab
 
 
+MIN_OPT_VERSION = 15
+_OPT_VERSIONED_RE = re.compile(r"opt-(\d+)")
+
+
+def _find_opt():
+    """Locate an LLVM opt binary.
+
+    Prefers a plain ``opt`` on PATH. Otherwise falls back to the newest
+    versioned ``opt-N`` (N >= MIN_OPT_VERSION) that distros such as
+    Debian/Ubuntu install, e.g. ``opt-18``. Returns None if none is found.
+    """
+    opt = shutil.which("opt")
+    if opt:
+        return opt
+
+    best_version, best_path = -1, None
+    for directory in os.environ.get("PATH", "").split(os.pathsep):
+        try:
+            entries = os.listdir(directory or ".")
+        except OSError:
+            continue
+        for entry in entries:
+            match = _OPT_VERSIONED_RE.fullmatch(entry)
+            if not match:
+                continue
+            version = int(match.group(1))
+            if version < MIN_OPT_VERSION or version <= best_version:
+                continue
+            path = shutil.which(entry, path=directory or ".")
+            if path:
+                best_version, best_path = version, path
+    return best_path
+
+
+def _run_opt(ll_file):
+    """Run the O2 pipeline over the IR and return it as bitcode, or None if
+    no opt binary is available."""
+
+    opt = _find_opt()
+    if opt is None:
+        logger.warning(
+            f"No LLVM 'opt' (or 'opt-N', N >= {MIN_OPT_VERSION}) found on PATH; "
+            "skipping IR optimization."
+        )
+        return None
+
+    logger.info(f"Optimizing IR with {opt} -O2: {ll_file}")
+    result = subprocess.run(
+        [opt, "-O2", str(ll_file), "-o", "-"],
+        check=True,
+        capture_output=True,
+    )
+    return result.stdout
+
+
 def _run_llc(ll_file, obj_file):
-    """Compile LLVM IR to BPF object file using llc."""
+    """Optimize LLVM IR with opt (when available) and compile it to a BPF
+    object file using llc."""
+
+    bitcode = _run_opt(ll_file)
 
     logger.info(f"Compiling IR to object: {ll_file} -> {obj_file}")
     result = subprocess.run(
@@ -185,20 +244,20 @@ def _run_llc(ll_file, obj_file):
             "-march=bpf",
             "-filetype=obj",
             "-O2",
-            str(ll_file),
+            "-" if bitcode is not None else str(ll_file),
             "-o",
             str(obj_file),
         ],
+        input=bitcode,
         check=True,
         capture_output=True,
-        text=True,
     )
 
     if result.returncode == 0:
         logger.info(f"Object file written to {obj_file}")
         return True
     else:
-        logger.error(f"llc compilation failed: {result.stderr}")
+        logger.error(f"llc compilation failed: {result.stderr.decode()}")
         return False
 
 

pythonbpf/context.py

@@ -72,6 +72,10 @@ def __init__(self, module: ir.Module):
         # Helper management
         self.scratch_pool = ScratchPoolManager()
 
+        # Sequence numbers for llvm.bpf.passthrough barriers; each call needs a
+        # distinct one so the optimizer cannot CSE two barriers together.
+        self.passthrough_seq = 0
+
         # Vmlinux handling (optional, specialized)
         self.vmlinux_handler = None  # Can be VmlinuxHandler instance
 

pythonbpf/expr/expr_pass.py

@@ -407,6 +407,28 @@ def _handle_unary_op(
 # ============================================================================
 
 
+def _bpf_passthrough(builder, compilation_context, val):
+    """Wrap an i1 in llvm.bpf.passthrough, an optimization barrier the BPF
+    backend strips before instruction selection.
+
+    Without it, opt's SimplifyCFG speculates a short-circuit operand's compare
+    into the preceding block, and SelectionDAG then folds the two null checks
+    of `p1 or p2` into `(p1 | p2) != 0`. The verifier rejects that with
+    "pointer |= pointer prohibited". BPFAdjustOpt uses the same intrinsic to
+    keep compares serialized for the verifier.
+    """
+    module = compilation_context.module
+    name = "llvm.bpf.passthrough.i1.i1"
+    fn = module.globals.get(name)
+    if fn is None:
+        fn_type = ir.FunctionType(ir.IntType(1), [ir.IntType(32), ir.IntType(1)])
+        fn = ir.Function(module, fn_type, name=name)
+
+    seq = ir.Constant(ir.IntType(32), compilation_context.passthrough_seq)
+    compilation_context.passthrough_seq += 1
+    return builder.call(fn, [seq, val])
+
+
 def _handle_and_op(func, builder, expr, local_sym_tab, compilation_context):
     """Handle `and` boolean operations."""
 
@@ -432,6 +454,8 @@ def _handle_and_op(func, builder, expr, local_sym_tab, compilation_context):
 
         # Convert to boolean if needed
         operand_bool = convert_to_bool(builder, operand_val)
+        if i > 0:
+            operand_bool = _bpf_passthrough(builder, compilation_context, operand_bool)
         current_block = builder.block
 
         if is_last:
@@ -485,6 +509,8 @@ def _handle_or_op(func, builder, expr, local_sym_tab, compilation_context):
 
         # Convert to boolean if needed
         operand_bool = convert_to_bool(builder, operand_val)
+        if i > 0:
+            operand_bool = _bpf_passthrough(builder, compilation_context, operand_bool)
         current_block = builder.block
 
         if is_last:

tests/failing_tests/xdp/xdp_test_1.py tests/passing_tests/xdp/xdp_test_1.pytests/failing_tests/xdp/xdp_test_1.py renamed to tests/passing_tests/xdp/xdp_test_1.py

@@ -15,14 +15,6 @@
 
 "failing_tests/conditionals/struct_ptr.py" = {reason = "Struct pointer used directly as boolean condition not supported", level = "ir"}
 
-# Compiles cleanly; rejected by the kernel. The `data + 34 < data_end` guard is
-# emitted as a signed compare over values round-tripped through the stack, so the
-# verifier never narrows the packet range (it stays r=0) and the later
-# `iph.saddr` read fails with "invalid access to packet, off=26 size=4" /
-# "R1 offset is outside of the packet". Direct packet access needs bounds checks
-# in a form the verifier can follow.
-"failing_tests/xdp/xdp_test_1.py" = {reason = "XDP direct packet access: the data/data_end guard does not establish a packet range the verifier can use", level = "verifier"}
-
 "failing_tests/license.py" = {reason = "Missing LICENSE global produces IR that llc rejects — should be caught earlier with a clear error message", level = "llc"}
 
 "failing_tests/undeclared_values.py" = {reason = "Undeclared variable used in f-string — should raise SyntaxError (correct behaviour, test documents it)", level = "ir"}