[llvm] [LLVM][NewPM] Add C API for running the pipeline on a single function. (PR #103773)

[Tim Besard via llvm-commits]

https://github.com/maleadt created https://github.com/llvm/llvm-project/pull/103773

By adding a new entrypoint, `LLVMRunPassesOnFunction`, as suggested in https://discourse.llvm.org/t/newpm-c-api-questions/80598.

Also removes erroneous `LLVMConsumeError`s from the pass builder unit tests as the string conversion already consumes the error.

>From 55cc962d232d543a3012674b3068e9dfd690ad5b Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard at gmail.com>
Date: Wed, 14 Aug 2024 12:51:49 +0200
Subject: [PATCH] [LLVM][NewPM] Add C API for running the pipeline on a single
 function.

Also removes erroneous LLVMConsumeError from the pass builder tests
as the string conversion already consumes the error.
---
 llvm/include/llvm-c/Transforms/PassBuilder.h  | 10 ++++
 llvm/lib/Passes/PassBuilderBindings.cpp       | 52 +++++++++++++++----
 .../PassBuilderBindingsTest.cpp               | 16 +++++-
 3 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/llvm/include/llvm-c/Transforms/PassBuilder.h b/llvm/include/llvm-c/Transforms/PassBuilder.h
index 03a5abaa753312..a91cdb8e4e328f 100644
--- a/llvm/include/llvm-c/Transforms/PassBuilder.h
+++ b/llvm/include/llvm-c/Transforms/PassBuilder.h
@@ -50,6 +50,16 @@ LLVMErrorRef LLVMRunPasses(LLVMModuleRef M, const char *Passes,
                            LLVMTargetMachineRef TM,
                            LLVMPassBuilderOptionsRef Options);
 
+/**
+ * Construct and run a set of passes over a function
+ *
+ * This function behaves the same as LLVMRunPasses, but operates on a single
+ * function instead of an entire module.
+ */
+LLVMErrorRef LLVMRunPassesOnFunction(LLVMValueRef F, const char *Passes,
+                                     LLVMTargetMachineRef TM,
+                                     LLVMPassBuilderOptionsRef Options);
+
 /**
  * Create a new set of options for a PassBuilder
  *
diff --git a/llvm/lib/Passes/PassBuilderBindings.cpp b/llvm/lib/Passes/PassBuilderBindings.cpp
index 4e12dd2226c4d2..e638a40dedc9d6 100644
--- a/llvm/lib/Passes/PassBuilderBindings.cpp
+++ b/llvm/lib/Passes/PassBuilderBindings.cpp
@@ -48,15 +48,24 @@ static TargetMachine *unwrap(LLVMTargetMachineRef P) {
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(LLVMPassBuilderOptions,
                                    LLVMPassBuilderOptionsRef)
 
-LLVMErrorRef LLVMRunPasses(LLVMModuleRef M, const char *Passes,
-                           LLVMTargetMachineRef TM,
-                           LLVMPassBuilderOptionsRef Options) {
+LLVMErrorRef RunPasses(LLVMModuleRef M, LLVMValueRef F, const char *Passes,
+                       LLVMTargetMachineRef TM,
+                       LLVMPassBuilderOptionsRef Options) {
   TargetMachine *Machine = unwrap(TM);
   LLVMPassBuilderOptions *PassOpts = unwrap(Options);
   bool Debug = PassOpts->DebugLogging;
   bool VerifyEach = PassOpts->VerifyEach;
 
-  Module *Mod = unwrap(M);
+  // Determine what to run passes on.
+  Module *Mod;
+  Function *Fun = nullptr;
+  if (F) {
+    Fun = unwrap<Function>(F);
+    Mod = Fun->getParent();
+  } else {
+    Mod = unwrap(M);
+  }
+
   PassInstrumentationCallbacks PIC;
   PassBuilder PB(Machine, PassOpts->PTO, std::nullopt, &PIC);
 
@@ -80,18 +89,39 @@ LLVMErrorRef LLVMRunPasses(LLVMModuleRef M, const char *Passes,
 
   StandardInstrumentations SI(Mod->getContext(), Debug, VerifyEach);
   SI.registerCallbacks(PIC, &MAM);
-  ModulePassManager MPM;
-  if (VerifyEach) {
-    MPM.addPass(VerifierPass());
-  }
-  if (auto Err = PB.parsePassPipeline(MPM, Passes)) {
-    return wrap(std::move(Err));
+
+  // Run the pipeline.
+  if (Fun) {
+    FunctionPassManager FPM;
+    if (VerifyEach)
+      FPM.addPass(VerifierPass());
+    if (auto Err = PB.parsePassPipeline(FPM, Passes))
+      return wrap(std::move(Err));
+    FPM.run(*Fun, FAM);
+  } else {
+    ModulePassManager MPM;
+    if (VerifyEach)
+      MPM.addPass(VerifierPass());
+    if (auto Err = PB.parsePassPipeline(MPM, Passes))
+      return wrap(std::move(Err));
+    MPM.run(*Mod, MAM);
   }
 
-  MPM.run(*Mod, MAM);
   return LLVMErrorSuccess;
 }
 
+LLVMErrorRef LLVMRunPasses(LLVMModuleRef M, const char *Passes,
+                           LLVMTargetMachineRef TM,
+                           LLVMPassBuilderOptionsRef Options) {
+  return RunPasses(M, nullptr, Passes, TM, Options);
+}
+
+LLVMErrorRef LLVMRunPassesOnFunction(LLVMValueRef F, const char *Passes,
+                                     LLVMTargetMachineRef TM,
+                                     LLVMPassBuilderOptionsRef Options) {
+  return RunPasses(nullptr, F, Passes, TM, Options);
+}
+
 LLVMPassBuilderOptionsRef LLVMCreatePassBuilderOptions() {
   return wrap(new LLVMPassBuilderOptions());
 }
diff --git a/llvm/unittests/Passes/PassBuilderBindings/PassBuilderBindingsTest.cpp b/llvm/unittests/Passes/PassBuilderBindings/PassBuilderBindingsTest.cpp
index 2b06033f0c3fa2..4e17b1ad09e2b3 100644
--- a/llvm/unittests/Passes/PassBuilderBindings/PassBuilderBindingsTest.cpp
+++ b/llvm/unittests/Passes/PassBuilderBindings/PassBuilderBindingsTest.cpp
@@ -35,6 +35,9 @@ class PassBuilderCTest : public testing::Test {
     LLVMDisposeMessage(Triple);
     Context = LLVMContextCreate();
     Module = LLVMModuleCreateWithNameInContext("test", Context);
+    LLVMTypeRef FT =
+        LLVMFunctionType(LLVMVoidTypeInContext(Context), nullptr, 0, 0);
+    Function = LLVMAddFunction(Module, "test", FT);
   }
 
   void TearDown() override {
@@ -52,6 +55,7 @@ class PassBuilderCTest : public testing::Test {
 public:
   LLVMTargetMachineRef TM;
   LLVMModuleRef Module;
+  LLVMValueRef Function;
   LLVMContextRef Context;
 };
 
@@ -63,7 +67,6 @@ TEST_F(PassBuilderCTest, Basic) {
   LLVMPassBuilderOptionsSetAAPipeline(Options, "basic-aa");
   if (LLVMErrorRef E = LLVMRunPasses(Module, "default<O2>", TM, Options)) {
     char *Msg = LLVMGetErrorMessage(E);
-    LLVMConsumeError(E);
     LLVMDisposePassBuilderOptions(Options);
     FAIL() << "Failed to run passes: " << Msg;
   }
@@ -80,3 +83,14 @@ TEST_F(PassBuilderCTest, InvalidPassIsError) {
   LLVMConsumeError(E2);
   LLVMDisposePassBuilderOptions(Options);
 }
+
+TEST_F(PassBuilderCTest, Function) {
+  LLVMPassBuilderOptionsRef Options = LLVMCreatePassBuilderOptions();
+  if (LLVMErrorRef E =
+          LLVMRunPassesOnFunction(Function, "no-op-function", TM, Options)) {
+    char *Msg = LLVMGetErrorMessage(E);
+    LLVMDisposePassBuilderOptions(Options);
+    FAIL() << "Failed to run passes on function: " << Msg;
+  }
+  LLVMDisposePassBuilderOptions(Options);
+}