fix: Fix TypePrintFn flag passing for autodiff codegen

Signed-off-by: Karan Janthe <karanjanthe@gmail.com>
This commit is contained in:
Karan Janthe 2025-06-30 14:29:59 +00:00
parent 4e97337005
commit 5773d38c7a
2 changed files with 14 additions and 1 deletions

View file

@ -700,6 +700,10 @@ struct LLVMRustSanitizerOptions {
#ifdef ENZYME
extern "C" void registerEnzymeAndPassPipeline(llvm::PassBuilder &PB,
/* augmentPassBuilder */ bool);
extern "C" {
extern llvm::cl::opt<std::string> EnzymeFunctionToAnalyze;
}
#endif
extern "C" LLVMRustResult LLVMRustOptimize(
@ -1069,6 +1073,15 @@ extern "C" LLVMRustResult LLVMRustOptimize(
return LLVMRustResult::Failure;
}
// Check if PrintTAFn was used and add type analysis pass if needed
if (!EnzymeFunctionToAnalyze.empty()) {
if (auto Err = PB.parsePassPipeline(MPM, "print-type-analysis")) {
std::string ErrMsg = toString(std::move(Err));
LLVMRustSetLastError(ErrMsg.c_str());
return LLVMRustResult::Failure;
}
}
if (PrintAfterEnzyme) {
// Handle the Rust flag `-Zautodiff=PrintModAfter`.
std::string Banner = "Module after EnzymeNewPM";

@ -1 +1 @@
Subproject commit b5098d515d5e1bd0f5470553bc0d18da9794ca8b
Subproject commit 0f65f31837a608db469ebde94adbc921be70bd8a