Add function existence validation to TypeAnalysisPrinter and ActivityAnalysisPrinter (#2616)

* Initial plan

* Add function existence checks to TypeAnalysisPrinter and ActivityAnalysisPrinter

Co-authored-by: wsmoses <[email protected]>

* Fix compilation errors and add tests for function existence checks

Co-authored-by: wsmoses <[email protected]>

* Use std::once_flag for thread-safe validation in ActivityAnalysisPrinter NewPM

Co-authored-by: wsmoses <[email protected]>

* Improve code style with early returns and flatten nested conditionals

Co-authored-by: wsmoses <[email protected]>

* Final validation and testing complete

Co-authored-by: wsmoses <[email protected]>

* Remove CodeQL artifact

* Replace report_fatal_error with EmitFailure to avoid crashing

Use EmitFailure instead of report_fatal_error when function is not found.
This provides a better error message and avoids crashing the compiler.
Updated tests to match the new error message format.

Co-authored-by: wsmoses <[email protected]>

* Add function existence validation to TypeAnalysisPrinter and ActivityAnalysisPrinter

Co-authored-by: wsmoses <[email protected]>

* Refactor ActivityAnalysisPrinter to module pass and simplify function lookup

- Changed ActivityAnalysisPrinterNewPM from function pass to module pass
- Replaced std::call_once with M.getFunction for simpler function lookup
- Changed "specified in" to "specified by" in error messages
- Updated pass registration in Enzyme.cpp to use ModulePassManager
- Removed unnecessary <mutex> include

Co-authored-by: wsmoses <[email protected]>

* Optimize FAM retrieval in ActivityAnalysisPrinter

Move FunctionAnalysisManager retrieval outside the loop to avoid repeated calls

Co-authored-by: wsmoses <[email protected]>

* Remove legacy FunctionPass implementations and add empty function name validation

- Deleted legacy ActivityAnalysisPrinter FunctionPass (no longer needed with module pass)
- Deleted legacy TypeAnalysisPrinter FunctionPass (no longer needed with module pass)
- Added validation for empty function names in both passes
- Changed ActivityAnalysisPrinter to only run on target function instead of all functions
- Both passes now emit diagnostic errors for empty function names

Co-authored-by: wsmoses <[email protected]>

* Optimize TypeAnalysisPrinter to only run on target function

Changed TypeAnalysisPrinter to only run analysis on the target function instead of iterating through all functions, consistent with ActivityAnalysisPrinter

Co-authored-by: wsmoses <[email protected]>

* Restore legacy passes as ModulePasses

- Added back ActivityAnalysisPrinter as a legacy ModulePass (converted from FunctionPass)
- Added back TypeAnalysisPrinter as a legacy ModulePass (converted from FunctionPass)
- Both legacy passes now have the same validation logic as NewPM passes
- Legacy passes check for empty function names and non-existent functions
- Keeps RegisterPass registration for backward compatibility

Co-authored-by: wsmoses <[email protected]>

* Extract getFirstFunctionDefinition helper function to reduce code duplication

- Added getFirstFunctionDefinition() helper function in Utils.h/Utils.cpp
- Returns the first non-declaration function in a module, or nullptr if none exists
- Updated ActivityAnalysisPrinter to use the helper function (4 occurrences replaced)
- Updated TypeAnalysisPrinter to use the helper function (4 occurrences replaced)
- Eliminates code duplication and improves maintainability

Co-authored-by: wsmoses <[email protected]>

* Add EmitFailure overload for Module and simplify error handling

- Added EmitFailure template overload that takes llvm::Module& instead of Function*
- New overload automatically finds first function definition for diagnostic context
- Simplified all error handling code in ActivityAnalysisPrinter and TypeAnalysisPrinter
- Removed repetitive getFirstFunctionDefinition calls and if/else blocks
- Added forward declaration of getFirstFunctionDefinition for template usage
- Reduced code duplication and improved maintainability

Co-authored-by: wsmoses <[email protected]>

* Fix test expectations to match "specified by" error messages

Updated test files to expect "specified by" instead of "specified in" to match the actual error messages produced by the code.

Co-authored-by: wsmoses <[email protected]>

* Apply clang-format to fix formatting issues

- Fixed line wrapping for long comments
- Removed trailing whitespace
- Standardized indentation in arrays and function calls
- Applied clang-format to all modified C++ files

Co-authored-by: wsmoses <[email protected]>

* Fix formatting with clang-format-16 and llvm style

Applied clang-format-16 with llvm style (as required by CI) instead of clang-format-18. This changes indentation from 4 spaces to 2 spaces in array initializers to match the project's formatting requirements.

Co-authored-by: wsmoses <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: wsmoses <[email protected]>

Author: Copilot

enzyme/Enzyme/ActivityAnalysisPrinter.cpp

@@ -47,6 +47,7 @@
 #include "llvm/Analysis/ScalarEvolution.h"
 
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #include "ActivityAnalysis.h"
 #include "ActivityAnalysisPrinter.h"
@@ -182,20 +183,36 @@ bool printActivityAnalysis(llvm::Function &F, TargetLibraryInfo &TLI) {
   return /*changed*/ false;
 }
 
-class ActivityAnalysisPrinter final : public FunctionPass {
+class ActivityAnalysisPrinter final : public ModulePass {
 public:
   static char ID;
-  ActivityAnalysisPrinter() : FunctionPass(ID) {}
+  ActivityAnalysisPrinter() : ModulePass(ID) {}
+
+  bool runOnModule(Module &M) override {
+    // Check if function name is specified
+    if (FunctionToAnalyze.empty()) {
+      EmitFailure("NoFunctionSpecified", M,
+                  "No function specified for -activity-analysis-func");
+      return false;
+    }
 
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<TargetLibraryInfoWrapperPass>();
-  }
+    // Check if the specified function exists
+    Function *TargetFunc = M.getFunction(FunctionToAnalyze);
 
-  bool runOnFunction(Function &F) override {
+    if (!TargetFunc) {
+      EmitFailure("FunctionNotFound", M, "Function '", FunctionToAnalyze,
+                  "' specified by -activity-analysis-func not found in module");
+      return false;
+    }
 
-    auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
+    // Run analysis only on the target function
+    auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(*TargetFunc);
+    return printActivityAnalysis(*TargetFunc, TLI);
+  }
 
-    return printActivityAnalysis(F, TLI);
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetLibraryInfoWrapperPass>();
+    AU.setPreservesAll();
   }
 };
 
@@ -207,10 +224,28 @@ static RegisterPass<ActivityAnalysisPrinter>
     X("print-activity-analysis", "Print Activity Analysis Results");
 
 ActivityAnalysisPrinterNewPM::Result
-ActivityAnalysisPrinterNewPM::run(llvm::Function &F,
-                                  llvm::FunctionAnalysisManager &FAM) {
-  bool changed = false;
-  changed = printActivityAnalysis(F, FAM.getResult<TargetLibraryAnalysis>(F));
+ActivityAnalysisPrinterNewPM::run(llvm::Module &M,
+                                  llvm::ModuleAnalysisManager &MAM) {
+  // Check if function name is specified
+  if (FunctionToAnalyze.empty()) {
+    EmitFailure("NoFunctionSpecified", M,
+                "No function specified for -activity-analysis-func");
+    return PreservedAnalyses::all();
+  }
+
+  // Check if the specified function exists
+  Function *TargetFunc = M.getFunction(FunctionToAnalyze);
+
+  if (!TargetFunc) {
+    EmitFailure("FunctionNotFound", M, "Function '", FunctionToAnalyze,
+                "' specified by -activity-analysis-func not found in module");
+    return PreservedAnalyses::all();
+  }
+
+  // Run analysis only on the target function
+  auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
+  bool changed = printActivityAnalysis(
+      *TargetFunc, FAM.getResult<TargetLibraryAnalysis>(*TargetFunc));
   return changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
 }
 llvm::AnalysisKey ActivityAnalysisPrinterNewPM::Key;

enzyme/Enzyme/ActivityAnalysisPrinter.h

@@ -46,7 +46,7 @@ class ActivityAnalysisPrinterNewPM final
   using Result = llvm::PreservedAnalyses;
   ActivityAnalysisPrinterNewPM() {}
 
-  Result run(llvm::Function &M, llvm::FunctionAnalysisManager &MAM);
+  Result run(llvm::Module &M, llvm::ModuleAnalysisManager &MAM);
 
   static bool isRequired() { return true; }
 };

enzyme/Enzyme/Enzyme.cpp

@@ -3799,15 +3799,15 @@ extern "C" void registerEnzymeAndPassPipeline(llvm::PassBuilder &PB,
           MPM.addPass(TypeAnalysisPrinterNewPM());
           return true;
         }
+        if (Name == "print-activity-analysis") {
+          MPM.addPass(ActivityAnalysisPrinterNewPM());
+          return true;
+        }
         return false;
       });
   PB.registerPipelineParsingCallback(
       [](llvm::StringRef Name, llvm::FunctionPassManager &FPM,
          llvm::ArrayRef<llvm::PassBuilder::PipelineElement>) {
-        if (Name == "print-activity-analysis") {
-          FPM.addPass(ActivityAnalysisPrinterNewPM());
-          return true;
-        }
         if (Name == "jl-inst-simplify") {
           FPM.addPass(JLInstSimplifyNewPM());
           return true;

enzyme/Enzyme/TypeAnalysis/TypeAnalysisPrinter.cpp

@@ -52,6 +52,7 @@
 #include "llvm/Analysis/ScalarEvolution.h"
 
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #include "../EnzymeLogic.h"
 #include "../FunctionUtils.h"
@@ -165,16 +166,36 @@ bool printTypeAnalyses(llvm::Function &F) {
   }
   return /*changed*/ false;
 }
-class TypeAnalysisPrinter final : public FunctionPass {
+
+class TypeAnalysisPrinter final : public ModulePass {
 public:
   static char ID;
-  TypeAnalysisPrinter() : FunctionPass(ID) {}
+  TypeAnalysisPrinter() : ModulePass(ID) {}
+
+  bool runOnModule(Module &M) override {
+    // Check if function name is specified
+    if (EnzymeFunctionToAnalyze.empty()) {
+      EmitFailure("NoFunctionSpecified", M,
+                  "No function specified for -type-analysis-func");
+      return false;
+    }
 
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<TargetLibraryInfoWrapperPass>();
+    // Check if the specified function exists
+    Function *TargetFunc = M.getFunction(EnzymeFunctionToAnalyze);
+
+    if (!TargetFunc) {
+      EmitFailure("FunctionNotFound", M, "Function '", EnzymeFunctionToAnalyze,
+                  "' specified by -type-analysis-func not found in module");
+      return false;
+    }
+
+    // Run analysis only on the target function
+    return printTypeAnalyses(*TargetFunc);
   }
 
-  bool runOnFunction(Function &F) override { return printTypeAnalyses(F); }
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+  }
 };
 
 } // namespace
@@ -187,9 +208,24 @@ static RegisterPass<TypeAnalysisPrinter> X("print-type-analysis",
 TypeAnalysisPrinterNewPM::Result
 TypeAnalysisPrinterNewPM::run(llvm::Module &M,
                               llvm::ModuleAnalysisManager &MAM) {
-  bool changed = false;
-  for (auto &F : M)
-    changed |= printTypeAnalyses(F);
+  // Check if function name is specified
+  if (EnzymeFunctionToAnalyze.empty()) {
+    EmitFailure("NoFunctionSpecified", M,
+                "No function specified for -type-analysis-func");
+    return PreservedAnalyses::all();
+  }
+
+  // Check if the specified function exists
+  Function *TargetFunc = M.getFunction(EnzymeFunctionToAnalyze);
+
+  if (!TargetFunc) {
+    EmitFailure("FunctionNotFound", M, "Function '", EnzymeFunctionToAnalyze,
+                "' specified by -type-analysis-func not found in module");
+    return PreservedAnalyses::all();
+  }
+
+  // Run analysis only on the target function
+  bool changed = printTypeAnalyses(*TargetFunc);
   return changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
 }
 llvm::AnalysisKey TypeAnalysisPrinterNewPM::Key;

enzyme/Enzyme/Utils.cpp

@@ -3168,6 +3168,15 @@ Function *GetFunctionFromValue(Value *fn) {
   return dyn_cast<Function>(GetFunctionValFromValue(fn));
 }
 
+Function *getFirstFunctionDefinition(Module &M) {
+  for (auto &F : M) {
+    if (!F.isDeclaration()) {
+      return &F;
+    }
+  }
+  return nullptr;
+}
+
 #if LLVM_VERSION_MAJOR >= 16
 std::optional<BlasInfo> extractBLAS(llvm::StringRef in)
 #else

enzyme/Enzyme/Utils.h

@@ -195,6 +195,9 @@ class EnzymeFailure final : public llvm::DiagnosticInfoUnsupported {
                 const llvm::Function *CodeRegion);
 };
 
+// Forward declaration needed for EmitFailure template
+llvm::Function *getFirstFunctionDefinition(llvm::Module &M);
+
 template <typename... Args>
 void EmitFailure(llvm::StringRef RemarkName,
                  const llvm::DiagnosticLocation &Loc,
@@ -217,6 +220,21 @@ void EmitFailure(llvm::StringRef RemarkName,
       (EnzymeFailure("Enzyme: " + ss.str(), Loc, CodeRegion)));
 }
 
+template <typename... Args>
+void EmitFailure(llvm::StringRef RemarkName, llvm::Module &M, Args &...args) {
+  // Use the first function definition in the module as context for the
+  // diagnostic
+  if (llvm::Function *FirstFunc = getFirstFunctionDefinition(M)) {
+    EmitFailure(RemarkName, FirstFunc->getSubprogram(), FirstFunc, args...);
+  } else {
+    // Fallback if no functions in module
+    std::string *str = new std::string();
+    llvm::raw_string_ostream ss(*str);
+    (ss << ... << args);
+    llvm::report_fatal_error(llvm::StringRef(*str));
+  }
+}
+
 static inline llvm::Function *isCalledFunction(llvm::Value *val) {
   if (llvm::CallInst *CI = llvm::dyn_cast<llvm::CallInst>(val)) {
     return CI->getCalledFunction();
@@ -1395,6 +1413,8 @@ void ErrorIfRuntimeInactive(llvm::IRBuilder<> &B, llvm::Value *primal,
 
 llvm::Function *GetFunctionFromValue(llvm::Value *fn);
 
+llvm::Function *getFirstFunctionDefinition(llvm::Module &M);
+
 llvm::Value *simplifyLoad(llvm::Value *LI, size_t valSz = 0,
                           size_t preOffset = 0);
 

enzyme/test/ActivityAnalysis/func-not-found.ll

@@ -0,0 +1,8 @@
+; RUN: not %opt < %s %newLoadEnzyme -passes="print-activity-analysis" -activity-analysis-func=nonexistent -S 2>&1 | FileCheck %s
+
+define void @foo(i64* %x) {
+entry:
+  ret void
+}
+
+; CHECK: Enzyme: Function 'nonexistent' specified by -activity-analysis-func not found in module

enzyme/test/TypeAnalysis/func-not-found.ll

@@ -0,0 +1,8 @@
+; RUN: not %opt < %s %newLoadEnzyme -passes="print-type-analysis" -type-analysis-func=nonexistent -S 2>&1 | FileCheck %s
+
+define void @foo(i64* %x) {
+entry:
+  ret void
+}
+
+; CHECK: Enzyme: Function 'nonexistent' specified by -type-analysis-func not found in module