2 files changed, 24 insertions, 3 deletions
diff --git a/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp
index 7eb11f9e608..9ce4abdb432 100644
--- a/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp
+++ b/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp
@@ -692,9 +692,12 @@ struct LLVMRustSanitizerOptions {
   bool SanitizeKernelAddressRecover;
 };
 
-// This symbol won't be available or used when Enzyme is not enabled
+// 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 registerEnzyme(llvm::PassBuilder &PB);
+extern "C" void registerEnzymeAndPassPipeline(llvm::PassBuilder &PB,
+                                              /* augmentPassBuilder */ bool);
 #endif
 
 extern "C" LLVMRustResult LLVMRustOptimize(
@@ -1023,7 +1026,7 @@ extern "C" LLVMRustResult LLVMRustOptimize(
   // now load "-enzyme" pass:
 #ifdef ENZYME
   if (RunEnzyme) {
-    registerEnzyme(PB);
+    registerEnzymeAndPassPipeline(PB, true);
     if (auto Err = PB.parsePassPipeline(MPM, "enzyme")) {
       std::string ErrMsg = toString(std::move(Err));
       LLVMRustSetLastError(ErrMsg.c_str());
diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
index b8cef6a7e25..aea2a8dd097 100644
--- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
+++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
@@ -1767,6 +1767,24 @@ extern "C" LLVMValueRef LLVMRustBuildMaxNum(LLVMBuilderRef B, LLVMValueRef LHS,
   return wrap(unwrap(B)->CreateMaxNum(unwrap(LHS), unwrap(RHS)));
 }
 
+#if LLVM_VERSION_LT(19, 0)
+enum {
+  LLVMGEPFlagInBounds = (1 << 0),
+  LLVMGEPFlagNUSW = (1 << 1),
+  LLVMGEPFlagNUW = (1 << 2),
+};
+extern "C" LLVMValueRef
+LLVMBuildGEPWithNoWrapFlags(LLVMBuilderRef B, LLVMTypeRef Ty,
+                            LLVMValueRef Pointer, LLVMValueRef *Indices,
+                            unsigned NumIndices, const char *Name,
+                            unsigned NoWrapFlags) {
+  if (NoWrapFlags & LLVMGEPFlagInBounds)
+    return LLVMBuildInBoundsGEP2(B, Ty, Pointer, Indices, NumIndices, Name);
+  else
+    return LLVMBuildGEP2(B, Ty, Pointer, Indices, NumIndices, Name);
+}
+#endif
+
 // Transfers ownership of DiagnosticHandler unique_ptr to the caller.
 extern "C" DiagnosticHandler *
 LLVMRustContextGetDiagnosticHandler(LLVMContextRef C) {