user0/rust
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: dlopen Enzyme

Browse source
This commit is contained in:
sgasho 2025-11-24 20:50:08 +09:00
parent ee447067e1
commit ddd5aad8a3
13 changed files with 536 additions and 246 deletions
Download patch file Download diff file Expand all files Collapse all files

33
compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp
View file

@ -550,17 +550,8 @@ struct LLVMRustSanitizerOptions {
bool SanitizeKernelAddressRecover;
};
// This symbol won't be available or used when Enzyme is not enabled.
// Always set AugmentPassBuilder to true, since it registers optimizations which
// will improve the performance for Enzyme.
#ifdef ENZYME
extern "C" void registerEnzymeAndPassPipeline(llvm::PassBuilder &PB,
/* augmentPassBuilder */ bool);
extern "C" {
extern llvm::cl::opt<std::string> EnzymeFunctionToAnalyze;
}
#endif
extern "C" typedef void (*registerEnzymeAndPassPipelineFn)(
llvm::PassBuilder &PB, bool augment);
extern "C" LLVMRustResult LLVMRustOptimize(
LLVMModuleRef ModuleRef, LLVMTargetMachineRef TMRef,
@ -569,8 +560,8 @@ extern "C" LLVMRustResult LLVMRustOptimize(
bool LintIR, LLVMRustThinLTOBuffer **ThinLTOBufferRef, bool EmitThinLTO,
bool EmitThinLTOSummary, bool MergeFunctions, bool UnrollLoops,
bool SLPVectorize, bool LoopVectorize, bool DisableSimplifyLibCalls,
bool EmitLifetimeMarkers, bool RunEnzyme, bool PrintBeforeEnzyme,
bool PrintAfterEnzyme, bool PrintPasses,
bool EmitLifetimeMarkers, registerEnzymeAndPassPipelineFn EnzymePtr,
bool PrintBeforeEnzyme, bool PrintAfterEnzyme, bool PrintPasses,
LLVMRustSanitizerOptions *SanitizerOptions, const char *PGOGenPath,
const char *PGOUsePath, bool InstrumentCoverage,
const char *InstrProfileOutput, const char *PGOSampleUsePath,
@ -907,8 +898,8 @@ extern "C" LLVMRustResult LLVMRustOptimize(
}
// now load "-enzyme" pass:
#ifdef ENZYME
if (RunEnzyme) {
// With dlopen, ENZYME macro may not be defined, so check EnzymePtr directly
if (EnzymePtr) {
if (PrintBeforeEnzyme) {
// Handle the Rust flag `-Zautodiff=PrintModBefore`.
@ -916,29 +907,19 @@ extern "C" LLVMRustResult LLVMRustOptimize(
MPM.addPass(PrintModulePass(outs(), Banner, true, false));
}
registerEnzymeAndPassPipeline(PB, false);
EnzymePtr(PB, false);
if (auto Err = PB.parsePassPipeline(MPM, "enzyme")) {
std::string ErrMsg = toString(std::move(Err));
LLVMRustSetLastError(ErrMsg.c_str());
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";
MPM.addPass(PrintModulePass(outs(), Banner, true, false));
}
}
#endif
if (PrintPasses) {
// Print all passes from the PM:
std::string Pipeline;

12
compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
View file

@ -1791,18 +1791,6 @@ extern "C" void LLVMRustSetNoSanitizeHWAddress(LLVMValueRef Global) {
GV.setSanitizerMetadata(MD);
}
#ifdef ENZYME
extern "C" {
extern llvm::cl::opt<unsigned> EnzymeMaxTypeDepth;
}
extern "C" size_t LLVMRustEnzymeGetMaxTypeDepth() { return EnzymeMaxTypeDepth; }
#else
extern "C" size_t LLVMRustEnzymeGetMaxTypeDepth() {
return 6; // Default fallback depth
}
#endif
// Statically assert that the fixed metadata kind IDs declared in
// `metadata_kind.rs` match the ones actually used by LLVM.
#define FIXED_MD_KIND(VARIANT, VALUE) \