Skip to content

Commit 3c0c1a6

Browse files
committed
Merge branch 'codegen-autodiff-test' into autodiff-codegen-test
2 parents 2c2bbd2 + 9aea41c commit 3c0c1a6

File tree

2 files changed

+14
-1
lines changed

2 files changed

+14
-1
lines changed

compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,10 @@ struct LLVMRustSanitizerOptions {
700700
#ifdef ENZYME
701701
extern "C" void registerEnzymeAndPassPipeline(llvm::PassBuilder &PB,
702702
/* augmentPassBuilder */ bool);
703+
704+
extern "C" {
705+
extern llvm::cl::opt<std::string> EnzymeFunctionToAnalyze;
706+
}
703707
#endif
704708

705709
extern "C" LLVMRustResult LLVMRustOptimize(
@@ -1069,6 +1073,15 @@ extern "C" LLVMRustResult LLVMRustOptimize(
10691073
return LLVMRustResult::Failure;
10701074
}
10711075

1076+
// Check if PrintTAFn was used and add type analysis pass if needed
1077+
if (!EnzymeFunctionToAnalyze.empty()) {
1078+
if (auto Err = PB.parsePassPipeline(MPM, "print-type-analysis")) {
1079+
std::string ErrMsg = toString(std::move(Err));
1080+
LLVMRustSetLastError(ErrMsg.c_str());
1081+
return LLVMRustResult::Failure;
1082+
}
1083+
}
1084+
10721085
if (PrintAfterEnzyme) {
10731086
// Handle the Rust flag `-Zautodiff=PrintModAfter`.
10741087
std::string Banner = "Module after EnzymeNewPM";

src/tools/enzyme

0 commit comments

Comments
 (0)