diff --git a/enzyme/Enzyme/EnzymeLogic.h b/enzyme/Enzyme/EnzymeLogic.h index 4ce25e8ae465..562aa12459a8 100644 --- a/enzyme/Enzyme/EnzymeLogic.h +++ b/enzyme/Enzyme/EnzymeLogic.h @@ -51,6 +51,7 @@ extern "C" { extern llvm::cl::opt EnzymePrint; +extern llvm::cl::opt EnzymeJuliaAddrLoad; } enum class AugmentedStruct { Tape, Return, DifferentialReturn }; diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index 034bf00142a7..6adfc1bd6b78 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp @@ -9325,6 +9325,13 @@ bool GradientUtils::needsCacheWholeAllocation( if (found == knownRecomputeHeuristic.end()) continue; + // If caching a julia base object, this is fine as + // GC will deal with any issues with. + if (auto PT = dyn_cast(cur->getType())) + if (PT->getAddressSpace() == 10) + if (EnzymeJuliaAddrLoad) + continue; + // If caching this user, it cannot be a gep/cast of original if (!found->second) { llvm::errs() << " mod: " << *oldFunc->getParent() << "\n";