diff --git a/enzyme/.bazelversion b/enzyme/.bazelversion index f22d756da39d..815da58b7a9e 100644 --- a/enzyme/.bazelversion +++ b/enzyme/.bazelversion @@ -1 +1 @@ -6.5.0 +7.4.1 diff --git a/enzyme/Enzyme/ActivityAnalysis.cpp b/enzyme/Enzyme/ActivityAnalysis.cpp index 1ce93cd723c2..32b9fa46d481 100644 --- a/enzyme/Enzyme/ActivityAnalysis.cpp +++ b/enzyme/Enzyme/ActivityAnalysis.cpp @@ -301,6 +301,8 @@ const StringSet<> KnownInactiveFunctions = { "__ubsan_handle_pointer_overflow", "__ubsan_handle_type_mismatch_v1", "__ubsan_vptr_type_cache", + "llvm.enzyme.lifetime_start", + "llvm.enzyme.lifetime_end", }; const std::set KnownInactiveIntrinsics = { diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index c24c0f3cc8e7..e9eec6a3e9e5 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -2988,6 +2988,18 @@ class AdjointGenerator : public llvm::InstVisitor { break; if (auto MCI = dyn_cast(MS.getOperand(2))) { if (auto II = dyn_cast(cur)) { + if (II->getCalledFunction()->getName() == + "llvm.enzyme.lifetime_start") { + if (getBaseObject(II->getOperand(1)) == root) { + if (auto CI2 = + dyn_cast(II->getOperand(0))) { + if (MCI->getValue().ule(CI2->getValue())) + break; + } + } + cur = cur->getPrevNode(); + continue; + } // If the start of the lifetime for more memory than being // memset, its valid. if (II->getIntrinsicID() == Intrinsic::lifetime_start) { @@ -3709,7 +3721,8 @@ class AdjointGenerator : public llvm::InstVisitor { return; } if (II.getIntrinsicID() == Intrinsic::stackrestore || - II.getIntrinsicID() == Intrinsic::lifetime_end) { + II.getIntrinsicID() == Intrinsic::lifetime_end || + II.getCalledFunction()->getName() == "llvm.enzyme.lifetime_end") { eraseIfUnused(II, /*erase*/ true, /*check*/ false); return; } @@ -6068,17 +6081,33 @@ class AdjointGenerator : public llvm::InstVisitor { void visitCallInst(llvm::CallInst &call) { using namespace llvm; + StringRef funcName = getFuncNameFromCall(&call); + // When compiling Enzyme against standard LLVM, and not Intel's // modified version of LLVM, the intrinsic `llvm.intel.subscript` is // not fully understood by LLVM. One of the results of this is that the // visitor dispatches to visitCallInst, rather than visitIntrinsicInst, when // presented with the intrinsic - hence why we are handling it here. - if (startsWith(getFuncNameFromCall(&call), ("llvm.intel.subscript"))) { + if (startsWith(funcName, ("llvm.intel.subscript"))) { assert(isa(call)); visitIntrinsicInst(cast(call)); return; } + if (funcName == "llvm.enzyme.lifetime_start") { + visitIntrinsicInst(cast(call)); + return; + } + if (funcName == "llvm.enzyme.lifetime_end") { + SmallVector orig_ops(call.getNumOperands()); + for (unsigned i = 0; i < call.getNumOperands(); ++i) { + orig_ops[i] = call.getOperand(i); + } + handleAdjointForIntrinsic(Intrinsic::lifetime_end, call, orig_ops); + eraseIfUnused(call); + return; + } + CallInst *const newCall = cast(gutils->getNewFromOriginal(&call)); IRBuilder<> BuilderZ(newCall); BuilderZ.setFastMathFlags(getFast()); @@ -6107,7 +6136,6 @@ class AdjointGenerator : public llvm::InstVisitor { : overwritten_args_map.find(&call)->second.second; auto called = getFunctionFromCall(&call); - StringRef funcName = getFuncNameFromCall(&call); bool subretused = false; bool shadowReturnUsed = false; diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp index 0768effaf8aa..91dc92f508f4 100644 --- a/enzyme/Enzyme/EnzymeLogic.cpp +++ b/enzyme/Enzyme/EnzymeLogic.cpp @@ -870,6 +870,12 @@ void calculateUnusedValuesInFunction( }, [&](const Instruction *inst) { if (auto II = dyn_cast(inst)) { + if (II->getCalledFunction()->getName() == + "llvm.enzyme.lifetime_start" || + II->getCalledFunction()->getName() == + "llvm.enzyme.lifetime_end") { + return UseReq::Cached; + } if (II->getIntrinsicID() == Intrinsic::lifetime_start || II->getIntrinsicID() == Intrinsic::lifetime_end || II->getIntrinsicID() == Intrinsic::stacksave || @@ -6636,7 +6642,9 @@ llvm::Function *EnzymeLogic::CreateNoFree(RequestContext context, Function *F) { "__assertfail", "__kmpc_global_thread_num", "nlopt_force_stop", - "cudaRuntimeGetVersion" + "cudaRuntimeGetVersion", + "llvm.enzyme.lifetime_start", + "llvm.enzyme.lifetime_end", }; // clang-format on diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp index b74dc8026fa3..68b252e7b5a5 100644 --- a/enzyme/Enzyme/FunctionUtils.cpp +++ b/enzyme/Enzyme/FunctionUtils.cpp @@ -534,10 +534,64 @@ UpgradeAllocasToMallocs(Function *NewF, DerivativeMode mode, } } +#if LLVM_VERSION_MAJOR >= 22 + Function *start_lifetime = nullptr; + Function *end_lifetime = nullptr; +#endif + for (auto AI : ToConvert) { std::string nam = AI->getName().str(); AI->setName(""); +#if LLVM_VERSION_MAJOR >= 22 + for (auto U : llvm::make_early_inc_range(AI->users())) { + if (auto II = dyn_cast(U)) { + if (II->getIntrinsicID() == Intrinsic::lifetime_start) { + if (!start_lifetime) { + start_lifetime = cast( + NewF->getParent() + ->getOrInsertFunction( + "llvm.enzyme.lifetime_start", + FunctionType::get(Type::getVoidTy(NewF->getContext()), + {}, true)) + .getCallee()); + } + IRBuilder<> B(II); + SmallVector args(II->arg_size()); + for (unsigned i = 0; i < II->arg_size(); ++i) { + args[i] = II->getArgOperand(i); + } + auto newI = B.CreateCall(start_lifetime, args); + newI->takeName(II); + newI->setDebugLoc(II->getDebugLoc()); + II->eraseFromParent(); + continue; + } + if (II->getIntrinsicID() == Intrinsic::lifetime_end) { + if (!end_lifetime) { + end_lifetime = cast( + NewF->getParent() + ->getOrInsertFunction( + "llvm.enzyme.lifetime_end", + FunctionType::get(Type::getVoidTy(NewF->getContext()), + {}, true)) + .getCallee()); + } + IRBuilder<> B(II); + SmallVector args(II->arg_size()); + for (unsigned i = 0; i < II->arg_size(); ++i) { + args[i] = II->getArgOperand(i); + } + auto newI = B.CreateCall(end_lifetime, args); + newI->takeName(II); + newI->setDebugLoc(II->getDebugLoc()); + II->eraseFromParent(); + continue; + } + } + } +#endif + // Ensure we insert the malloc after the allocas Instruction *insertBefore = AI; while (isa(insertBefore->getNextNode())) { @@ -884,6 +938,45 @@ void PreProcessCache::LowerAllocAddr(Function *NewF) { #endif RecursivelyReplaceAddressSpace(T, AIV, /*legal*/ true); } + +#if LLVM_VERSION_MAJOR >= 22 + { + auto start_lifetime = + NewF->getParent()->getFunction("llvm.enzyme.lifetime_start"); + auto end_lifetime = + NewF->getParent()->getFunction("llvm.enzyme.lifetime_end"); + + SmallVector Todo; + for (auto &BB : *NewF) { + for (auto &I : BB) { + if (auto CB = dyn_cast(&I)) { + if (!CB->getCalledFunction()) + continue; + if (CB->getCalledFunction() == start_lifetime || + CB->getCalledFunction() == end_lifetime) { + Todo.push_back(CB); + } + } + } + } + + for (auto CB : Todo) { + if (!isa(CB->getArgOperand(1))) { + CB->eraseFromParent(); + continue; + } + IRBuilder<> B(CB); + if (CB->getCalledFunction() == start_lifetime) { + B.CreateLifetimeStart(CB->getArgOperand(1), + cast(CB->getArgOperand(0))); + } else { + B.CreateLifetimeEnd(CB->getArgOperand(1), + cast(CB->getArgOperand(0))); + } + CB->eraseFromParent(); + } + } +#endif } /// Calls to realloc with an appropriate implementation @@ -7300,6 +7393,9 @@ Constraints::InnerTy Constraints::make_compare(const SCEV *v, bool isEqual, ConstraintContext ctx2(ctx.SE, ctx.loopToSolve, noassumption, ctx.DT); for (auto I : ctx.Assumptions) { bool legal = true; + if (I->getParent()->getParent() != + ctx.loopToSolve->getHeader()->getParent()) + continue; auto parsedCond = getSparseConditions(legal, I->getOperand(0), Constraints::none(), nullptr, ctx2); bool dominates = ctx.DT.dominates(I, ctx.loopToSolve->getHeader()); diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index dd9d95c1e6bc..222ce9baf774 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp @@ -8954,32 +8954,38 @@ void GradientUtils::computeForwardingProperties(Instruction *V) { storingOps.insert(store); } } else if (auto II = dyn_cast(cur)) { - switch (II->getIntrinsicID()) { - case Intrinsic::lifetime_start: + if (II->getCalledFunction()->getName() == "llvm.enzyme.lifetime_start") { LifetimeStarts.insert(II); - break; - case Intrinsic::dbg_declare: - case Intrinsic::dbg_value: - case Intrinsic::dbg_label: + } else if (II->getCalledFunction()->getName() == + "llvm.enzyme.lifetime_end") { + } else { + switch (II->getIntrinsicID()) { + case Intrinsic::lifetime_start: + LifetimeStarts.insert(II); + break; + case Intrinsic::dbg_declare: + case Intrinsic::dbg_value: + case Intrinsic::dbg_label: #if LLVM_VERSION_MAJOR <= 16 - case llvm::Intrinsic::dbg_addr: + case llvm::Intrinsic::dbg_addr: #endif - case Intrinsic::lifetime_end: - break; - case Intrinsic::memset: { - stores.insert(II); - storingOps.insert(II); - break; - } - // TODO memtransfer(cpy/move) - case Intrinsic::memcpy: - case Intrinsic::memmove: - default: - promotable = false; - shadowpromotable = false; - EmitWarning("NotPromotable", *cur, " Could not promote allocation ", *V, - " due to unknown intrinsic ", *cur); - break; + case Intrinsic::lifetime_end: + break; + case Intrinsic::memset: { + stores.insert(II); + storingOps.insert(II); + break; + } + // TODO memtransfer(cpy/move) + case Intrinsic::memcpy: + case Intrinsic::memmove: + default: + promotable = false; + shadowpromotable = false; + EmitWarning("NotPromotable", *cur, " Could not promote allocation ", + *V, " due to unknown intrinsic ", *cur); + break; + } } } else if (auto CI = dyn_cast(cur)) { StringRef funcName = getFuncNameFromCall(CI); diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h index fb6cdb93ff95..3c7ce7f89bab 100644 --- a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h +++ b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.h @@ -61,6 +61,16 @@ extern const llvm::StringMap LIBM_FUNCTIONS; static inline bool isMemFreeLibMFunction(llvm::StringRef str, llvm::Intrinsic::ID *ID = nullptr) { llvm::StringRef ogstr = str; + if (ID) { + if (str == "llvm.enzyme.lifetime_start") { + *ID = llvm::Intrinsic::lifetime_start; + return false; + } + if (str == "llvm.enzyme.lifetime_end") { + *ID = llvm::Intrinsic::lifetime_end; + return false; + } + } if (startsWith(str, "__") && endsWith(str, "_finite")) { str = str.substr(2, str.size() - 2 - 7); } else if (startsWith(str, "__fd_") && endsWith(str, "_1")) { diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp index 9c5f4ffbaf93..3bc46b6281bf 100644 --- a/enzyme/Enzyme/Utils.cpp +++ b/enzyme/Enzyme/Utils.cpp @@ -2801,6 +2801,9 @@ getAllLoadedValuesFrom(AllocaInst *ptr0, size_t offset, size_t valSz, } if (auto II = dyn_cast(U)) { + if (II->getCalledFunction()->getName() == "llvm.enzyme.lifetime_start" || + II->getCalledFunction()->getName() == "llvm.enzyme.lifetime_end") + continue; if (II->getIntrinsicID() == Intrinsic::lifetime_start || II->getIntrinsicID() == Intrinsic::lifetime_end) continue; diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index 097bd2966c7c..b267beb9154f 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -644,6 +644,10 @@ static inline llvm::Type *IntToFloatTy(llvm::Type *T) { static inline bool isDebugFunction(llvm::Function *called) { if (!called) return false; + if (called->getName() == "llvm.enzyme.lifetime_start" || + called->getName() == "llvm.enzyme.lifetime_end") { + return true; + } switch (called->getIntrinsicID()) { case llvm::Intrinsic::dbg_declare: case llvm::Intrinsic::dbg_value: @@ -1729,6 +1733,10 @@ static inline bool isNoAlias(const llvm::Value *val) { static inline bool isNoEscapingAllocation(const llvm::Function *F) { if (F->hasFnAttribute("enzyme_no_escaping_allocation")) return true; + if (F->getName() == "llvm.enzyme.lifetime_start" || + F->getName() == "llvm.enzyme.lifetime_end") { + return true; + } using namespace llvm; switch (F->getIntrinsicID()) { case Intrinsic::memset: