Add function.optimize

README.md

@@ -59,6 +59,7 @@ More examples worth browsing:
 
 - [`examples/intrinsic_memcpy.py`](examples/intrinsic_memcpy.py) - call an LLVM intrinsic by name with `Builder.intrinsic(...)`
 - [`examples/optimize_module.py`](examples/optimize_module.py) - optimize a module with a PassBuilder pipeline string
+- [`examples/optimize_function.py`](examples/optimize_function.py) - optimize one function with a function-level PassBuilder pipeline string
 - [`examples/emit_object_assembly.py`](examples/emit_object_assembly.py) - emit host object code and assembly from a module
 - [`examples/jit_add.py`](examples/jit_add.py) - JIT-compile IR, call it through ctypes, and register a Python callback
 - [`examples/instruction_metadata.py`](examples/instruction_metadata.py) - attach custom metadata to an instruction and print the IR
@@ -75,7 +76,7 @@ More examples worth browsing:
 - IR construction, traversal, and transformation helpers: constants, globals, PHI/control flow/memory/cast/cmp instructions, operands, predecessors, RAUW, split blocks, move/clone/erase instructions
 - IR and bitcode parsing/writing, lazy bitcode modules, module cloning/linking, diagnostics, attributes, COMDATs, calling conventions, linkage/visibility/storage controls
 - Metadata/debug-info APIs including named metadata views, module flag views, instruction/global metadata mappings, DIBuilder recipes, and debug-location scopes
-- Target initialization/lookup, data layouts, `Module.emit_object()`, `Module.emit_assembly()`, target-machine emission, and PassBuilder pipeline execution via `Module.optimize()` / `Module.run_passes()`
+- Target lookup, data layouts, `Module.emit_object()`, `Module.emit_assembly()`, target-machine emission, and PassBuilder pipeline execution via `Module.optimize()`, `Function.optimize()`, and `Module.run_passes()`
 - Generic intrinsic calls with `Builder.intrinsic(...)` and in-process JIT execution through the LLVM-C ORC LLJIT API
 - Lifetime/validity guards that turn many use-after-free, disposed-object, null-reference, and wrong-kind mistakes into Python exceptions instead of hard crashes
 - Auto-generated typed `.pyi` stubs for IDEs/type checkers

devdocs/api-reference.md

@@ -48,6 +48,7 @@ builder.intrinsic("llvm.sqrt", [x], overloaded_types=[x.type])
 
 # Explicit PassBuilder pipeline optimization.
 mod.optimize("default<O2>", target_machine=tm)
+func.optimize("mem2reg,instcombine,simplifycfg", target_machine=tm)
 
 # Direct object/assembly emission. Optimize explicitly first when desired.
 obj = mod.emit_object(target_machine=tm)

devdocs/api-ux-cleanup/plan.md

@@ -82,11 +82,11 @@ builder.intrinsic(
 - Missing overload types error.
 - Generated module verifies.
 
-## Phase 2: Explicit module optimization helper
+## Phase 2: Explicit module/function optimization helpers
 
 ### User problem
 
-Running passes is possible, but the current method name is low-level and does not read like a user task. Users should be able to say “optimize this module with this pipeline” directly.
+Running passes is possible, but the current method name is low-level and does not read like a user task. Users should be able to say “optimize this module” or “optimize this function” with a PassBuilder pipeline directly.
 
 ### Target UX
 
@@ -94,6 +94,9 @@ Running passes is possible, but the current method name is low-level and does no
 mod.optimize("default<O2>")
 mod.optimize("default<Os>", target_machine=tm)
 mod.optimize("function(mem2reg),default<O2>", target_machine=tm, options=opts)
+
+func.optimize("mem2reg,instcombine,simplifycfg")
+func.optimize("instcombine,simplifycfg", target_machine=tm, options=opts)
 ```
 
 ### Public API shape
@@ -105,12 +108,20 @@ mod.optimize(
     target_machine: llvm.TargetMachine | None = None,
     options: llvm.PassBuilderOptions | None = None,
 ) -> None
+
+func.optimize(
+    pipeline: str,
+    *,
+    target_machine: llvm.TargetMachine | None = None,
+    options: llvm.PassBuilderOptions | None = None,
+) -> None
 ```
 
 ### Behavior
 
-- `pipeline` is the LLVM PassBuilder pipeline string.
-- The method mutates the module in place.
+- `Module.optimize` accepts module-level LLVM PassBuilder pipeline strings.
+- `Function.optimize` accepts function-level LLVM PassBuilder pipeline strings.
+- The method mutates the module or function in place.
 - The method wraps the existing pass-running implementation.
 - Existing lower-level pass APIs remain available for advanced users.
 - Error messages should include the failed pipeline string when LLVM rejects it.
@@ -119,8 +130,8 @@ mod.optimize(
 
 - `default<O0>` succeeds on a simple module.
 - `default<O2>` succeeds on a simple module.
-- A custom pipeline succeeds where supported.
-- Invalid pipeline raises a Python exception with the pipeline in the message.
+- A custom function pipeline succeeds on one function without changing sibling functions.
+- Invalid pipelines raise Python exceptions with the pipeline in the message.
 - Passing a target machine works.
 
 ## Phase 3: Object and assembly emission convenience

devdocs/api-ux-cleanup/progress.md

@@ -23,11 +23,12 @@ Implemented and tested the active UX work:
 - [x] Keeps existing lower-level intrinsic APIs public.
 - [x] Tests cover non-overloaded, overloaded, and memory intrinsics.
 
-### Phase 2: Module optimization helper
+### Phase 2: Module/function optimization helpers
 
 - [x] Added `Module.optimize(pipeline, *, target_machine=None, options=None)`.
+- [x] Added `Function.optimize(pipeline, *, target_machine=None, options=None)` for function-level PassBuilder pipelines.
 - [x] Uses LLVM PassBuilder pipeline strings directly.
-- [x] Mutates the module in place.
+- [x] Mutates the module or function in place.
 - [x] Includes the failed pipeline string in error messages.
 - [x] Keeps existing lower-level pass APIs public.
 - [x] Fixed optional pass options by creating a default `PassBuilderOptions` internally when none is provided.
@@ -102,6 +103,7 @@ Coverage:
 - [x] `devdocs/api-ux-cleanup/plan.md` remains the task design reference.
 - [x] `examples/intrinsic_memcpy.py` shows `Builder.intrinsic(...)`.
 - [x] `examples/optimize_module.py` shows `Module.optimize(...)`.
+- [x] `examples/optimize_function.py` shows `Function.optimize(...)`.
 - [x] `examples/emit_object_assembly.py` shows `TargetMachine.host()`, `emit_object()`, and `emit_assembly()`.
 - [x] `examples/jit_add.py` shows `JIT.host()`, `add_module()`, `lookup` via `ctypes_function()`, and `add_symbol()`.
 

examples/optimize_function.py

@@ -0,0 +1,57 @@
+"""Optimize one function with an LLVM PassBuilder function pipeline string.
+
+``Function.optimize`` mutates only that function in place. The pipeline string is
+a function-level PassBuilder pipeline, such as ``mem2reg,instcombine,simplifycfg``.
+
+Run from the repository root with:
+
+    uv run python examples/optimize_function.py
+"""
+
+from __future__ import annotations
+
+import textwrap
+
+import llvm
+
+
+INPUT_IR = """
+define i32 @optimize_me(i32 %x) {
+entry:
+  %tmp = alloca i32
+  %zero = add i32 %x, 0
+  store i32 %zero, ptr %tmp
+  %loaded = load i32, ptr %tmp
+  %result = mul i32 %loaded, 1
+  ret i32 %result
+}
+
+define i32 @leave_me_alone(i32 %x) {
+entry:
+  %result = add i32 %x, 0
+  ret i32 %result
+}
+"""
+
+
+def optimize_one_function(
+    ir_text: str,
+    function_name: str = "optimize_me",
+    pipeline: str = "mem2reg,instcombine,simplifycfg",
+) -> str:
+    with llvm.create_context() as ctx:
+        with ctx.parse_ir(textwrap.dedent(ir_text).strip() + "\n") as mod:
+            assert mod.verify(), mod.verification_error
+            fn = mod.get_function(function_name)
+            assert fn is not None
+            fn.optimize(pipeline)
+            assert mod.verify(), mod.verification_error
+            return str(mod)
+
+
+def main() -> None:
+    print(optimize_one_function(INPUT_IR), end="")
+
+
+if __name__ == "__main__":
+    main()

src/llvm-nanobind.cpp

@@ -3983,6 +3983,9 @@ struct LLVMFunctionWrapper : LLVMValueWrapper {
     return !LLVMVerifyFunction(m_ref, LLVMReturnStatusAction);
   }
 
+  void optimize(const std::string &pipeline, LLVMTargetMachineWrapper *tm,
+                LLVMPassBuilderOptionsWrapper *opts);
+
   /// Verify this function and print any errors to stderr.
   /// Wraps LLVMVerifyFunction with LLVMPrintMessageAction.
   bool verify_and_print() const {
@@ -8590,6 +8593,59 @@ inline void LLVMModuleWrapper::run_passes(const std::string &passes,
   ::run_passes(*this, passes, tm, opts);
 }
 
+void run_passes_on_function(LLVMFunctionWrapper &fn, const std::string &passes,
+                            LLVMTargetMachineWrapper *tm,
+                            LLVMPassBuilderOptionsWrapper *opts) {
+  fn.check_valid();
+  if (LLVMIsDeclaration(fn.m_ref)) {
+    throw LLVMAssertionError(
+        "Function.optimize requires a function definition (check "
+        "is_declaration first)");
+  }
+
+  LLVMTargetMachineRef tm_ref = nullptr;
+  if (tm) {
+    tm->check_valid();
+    tm_ref = tm->m_ref;
+  }
+
+  LLVMPassBuilderOptionsRef opts_ref = nullptr;
+  LLVMPassBuilderOptionsRef owned_opts = nullptr;
+  if (opts) {
+    opts->check_valid();
+    opts_ref = opts->m_ref;
+  } else {
+    // LLVMRunPassesOnFunction expects a valid options object on some LLVM
+    // builds.
+    owned_opts = LLVMCreatePassBuilderOptions();
+    if (!owned_opts)
+      throw LLVMError("Failed to create pass builder options");
+    opts_ref = owned_opts;
+  }
+
+  LLVMErrorRef err = LLVMRunPassesOnFunction(fn.m_ref, passes.c_str(), tm_ref,
+                                             opts_ref);
+  if (owned_opts)
+    LLVMDisposePassBuilderOptions(owned_opts);
+  if (err) {
+    char *msg = LLVMGetErrorMessage(err);
+    std::string error_msg = msg ? msg : "Unknown error";
+    LLVMDisposeErrorMessage(msg);
+    throw LLVMError("Failed to run function passes: " + error_msg);
+  }
+}
+
+inline void LLVMFunctionWrapper::optimize(
+    const std::string &pipeline, LLVMTargetMachineWrapper *tm,
+    LLVMPassBuilderOptionsWrapper *opts) {
+  try {
+    ::run_passes_on_function(*this, pipeline, tm, opts);
+  } catch (const LLVMError &e) {
+    throw LLVMError("Failed to optimize function with pipeline '" + pipeline +
+                    "': " + e.what());
+  }
+}
+
 inline void LLVMModuleWrapper::optimize(const std::string &pipeline,
                                         LLVMTargetMachineWrapper *tm,
                                         LLVMPassBuilderOptionsWrapper *opts) {
@@ -14470,6 +14526,19 @@ Valid when:
            R"(Verify function and print errors to stderr.
 
 <sub>C API: LLVMVerifyFunction</sub>)")
+      .def("optimize", &LLVMFunctionWrapper::optimize, "pipeline"_a,
+           "target_machine"_a.none() = nullptr,
+           "options"_a.none() = nullptr,
+           R"doc(Optimize this function with an LLVM PassBuilder function pipeline string.
+
+Args:
+    pipeline: Function pass pipeline string (e.g., 'instcombine,simplifycfg').
+    target_machine: Optional target machine for target-specific passes.
+    options: Optional PassBuilderOptions.
+
+This mutates the function in place.
+
+<sub>C API: LLVMRunPassesOnFunction</sub>)doc")
       // =====================================================================
       // Intrinsic functions
       // =====================================================================

tests/regressions/test_api_ux_cleanup.py

@@ -4,6 +4,7 @@ Regression tests for the API UX cleanup helpers.
 These tests cover the high-level APIs from devdocs/api-ux-cleanup:
 - Builder.intrinsic(...)
 - Module.optimize(...)
+- Function.optimize(...)
 - TargetMachine.host()
 - Module.emit_object()/emit_assembly()
 - JIT through LLVM-C ORC LLJIT
@@ -153,6 +154,58 @@ def test_module_optimize_success_and_failure():
         _close_module(ctx_mgr, mod_mgr)
 
 
+def test_function_optimize_success_and_failure():
+    with llvm.create_context() as ctx:
+        with ctx.create_module("function_optimize_helper") as mod:
+            i32 = ctx.types.i32
+            fn = mod.add_function("answer", ctx.types.function(i32, [i32]))
+            x = fn.get_param(0)
+            entry = fn.append_basic_block("entry")
+            with entry.create_builder() as builder:
+                result = builder.add(x, i32.constant(0), "result")
+                builder.ret(result)
+
+            untouched = mod.add_function("untouched", ctx.types.function(i32, [i32]))
+            y = untouched.get_param(0)
+            untouched_entry = untouched.append_basic_block("entry")
+            with untouched_entry.create_builder() as builder:
+                result = builder.add(y, i32.constant(0), "result")
+                builder.ret(result)
+
+            decl = mod.add_function("decl", ctx.types.function(i32, []))
+
+            assert mod.verify(), mod.verification_error
+            fn.optimize("instcombine,simplifycfg")
+            assert mod.verify(), mod.verification_error
+
+            ir = str(mod)
+            optimized_body = ir.split("define i32 @answer", 1)[1].split(
+                "define i32 @untouched", 1
+            )[0]
+            untouched_body = ir.split("define i32 @untouched", 1)[1].split(
+                "declare i32 @decl", 1
+            )[0]
+            assert "ret i32 %0" in optimized_body
+            assert " add i32 " not in optimized_body
+            assert " add i32 " in untouched_body
+
+            try:
+                fn.optimize("not-a-real-function-pass")
+            except llvm.LLVMError as exc:
+                message = str(exc)
+                assert "not-a-real-function-pass" in message
+                assert "Failed to optimize function" in message
+            else:
+                raise AssertionError("expected invalid function pipeline error")
+
+            try:
+                decl.optimize("instcombine")
+            except llvm.LLVMAssertionError as exc:
+                assert "function definition" in str(exc)
+            else:
+                raise AssertionError("expected declaration optimization error")
+
+
 def test_module_optimize_with_target_machine():
     tm = _host_target_machine_or_skip()
     if tm is None:

tests/test_examples.py

@@ -103,6 +103,19 @@ def test_optimize_module_example() -> None:
     assert "alloca" not in output
 
 
+def test_optimize_function_example() -> None:
+    output = run_example("examples/optimize_function.py")
+    assert "define i32 @optimize_me" in output
+    assert "define i32 @leave_me_alone" in output
+    assert "ret i32 %x" in output
+    assert "%result = add i32 %x, 0" in output
+    optimized = output.split("define i32 @optimize_me", 1)[1].split(
+        "define i32 @leave_me_alone", 1
+    )[0]
+    assert "alloca" not in optimized
+    assert " add i32 " not in optimized
+
+
 def test_emit_object_assembly_example() -> None:
     output = run_example("examples/emit_object_assembly.py")
     if output.startswith("skipped:"):