ThinLTO: import the functions passed to __enzyme_* calls
Since #3211 a ThinLTO build differentiates in each post-link backend, which
only sees its own module plus what the thin link imported into it. The
importer follows call edges, so a function only passed to __enzyme_autodiff
stayed a declaration there ("failed to find fn to differentiate").
At ThinLTO pre-link, add an internal anchor whose function_entry_count lists
the GUIDs of those functions and of the declarations the module's own
differentiated functions call. The module summary turns them into critical
call edges, which the importer follows with a raised size threshold. At
post-link, before simplification, drop the anchor and give the imported
(available_externally) functions internal copies, so their bodies survive
EliminateAvailableExternally until Enzyme runs.
Functions further down the call graph are still subject to the importer's
size limit, and noinline functions are never imported; the README says so.
Assisted-by: Claude Code (Opus 5.5)
diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp
index 7649ba7..8ada547 100644
--- a/enzyme/Enzyme/Enzyme.cpp
+++ b/enzyme/Enzyme/Enzyme.cpp
@@ -92,6 +92,7 @@
#include "llvm/Transforms/Utils/Mem2Reg.h"
#include "CApi.h"
+#include "ThinLTOImport.h"
using namespace llvm;
#ifdef DEBUG_TYPE
#undef DEBUG_TYPE
@@ -3190,6 +3191,12 @@
MPM.addPass(PreserveNVVMNewPM(/*Begin*/ false,
/*PreserveCustomRuleLinkage*/ true,
/*LTOPreLink*/ true));
+ // A ThinLTO backend only sees the functions imported into its module,
+ // and the importer follows calls, not the function an __enzyme_* call
+ // takes. Ask for those explicitly; the post-link pipeline then keeps
+ // the imported bodies around until Enzyme runs (see below).
+ if (Phase == ThinOrFullLTOPhase::ThinLTOPreLink)
+ MPM.addPass(EnzymeThinLTOImportPass(/*PostLink*/ false));
return;
}
#endif
@@ -3235,6 +3242,21 @@
// that requires a functionpass, and we have a modulepass.
// PB.registerVectorizerStartEPCallback(loadPass);
PB.registerPipelineStartEPCallback(loadNVVM);
+
+#if LLVM_VERSION_MAJOR >= 20
+ // Functions imported by ThinLTO are available_externally, and the
+ // optimization pipeline drops those bodies (EliminateAvailableExternally)
+ // before the OptimizerEarly callback where the post-link run differentiates.
+ // Before the post-link simplification starts, give the imported functions
+ // that __enzyme_* calls differentiate internal copies, which are simplified
+ // along with the rest of the module and stay until Enzyme has run.
+ PB.registerPipelineEarlySimplificationEPCallback(
+ [](ModulePassManager &MPM, OptimizationLevel, ThinOrFullLTOPhase Phase) {
+ if (EnzymeEnable && !EnzymeLTOPreLink &&
+ Phase == ThinOrFullLTOPhase::ThinLTOPostLink)
+ MPM.addPass(EnzymeThinLTOImportPass(/*PostLink*/ true));
+ });
+#endif
PB.registerFullLinkTimeOptimizationEarlyEPCallback(loadNVVM);
auto preLTOPass = [](ModulePassManager &MPM, OptimizationLevel Level) {
diff --git a/enzyme/Enzyme/ThinLTOImport.cpp b/enzyme/Enzyme/ThinLTOImport.cpp
new file mode 100644
index 0000000..7cc320e
--- /dev/null
+++ b/enzyme/Enzyme/ThinLTOImport.cpp
@@ -0,0 +1,193 @@
+//===- ThinLTOImport.cpp - Bring functions to differentiate into ThinLTO --===//
+//
+// Enzyme Project
+//
+// Part of the Enzyme Project, under the Apache License v2.0 with LLVM
+// Exceptions. See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// See ThinLTOImport.h.
+//
+//===----------------------------------------------------------------------===//
+
+#include "ThinLTOImport.h"
+
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Config/llvm-config.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+
+using namespace llvm;
+
+// Internal, kept alive by llvm.compiler.used, and never called: it exists only
+// for the import GUIDs on its !prof.
+static constexpr const char *AnchorName = "enzyme.thinlto.imports";
+static constexpr const char *CopySuffix = ".enzyme.thinlto";
+
+static Function *calledFunction(const CallBase &CB) {
+ return dyn_cast<Function>(CB.getCalledOperand()->stripPointerCasts());
+}
+
+// Calls to the __enzyme_* markers that Enzyme lowers.
+template <typename T> static void forEachEnzymeCall(Module &M, T &&Fn) {
+ for (Function &F : M) {
+ if (!F.isDeclaration() || !F.getName().contains("__enzyme_"))
+ continue;
+ for (User *U : F.users())
+ if (auto *CB = dyn_cast<CallBase>(U))
+ if (calledFunction(*CB) == &F)
+ Fn(*CB);
+ }
+}
+
+static GlobalValue::GUID guidOf(const Function &F) {
+#if LLVM_VERSION_MAJOR >= 24
+ if (auto G = F.getGUIDIfAssigned())
+ return *G;
+ return GlobalValue::getGUIDAssumingExternalLinkage(
+ GlobalValue::getGlobalIdentifier(F.getName(), F.getLinkage(),
+ F.getParent()->getSourceFileName()));
+#else
+ return F.getGUID();
+#endif
+}
+
+bool enzymeThinLTORequestImports(Module &M) {
+ // Every function an __enzyme_* call takes, whether the one to differentiate
+ // or a function pointer passed as an argument, which needs a shadow.
+ SmallVector<Function *, 8> Worklist;
+ forEachEnzymeCall(M, [&](CallBase &CB) {
+ for (Value *A : CB.args())
+ if (auto *G = dyn_cast<Function>(A->stripPointerCasts()))
+ Worklist.push_back(G);
+ });
+
+ // Differentiating a function defined here needs the bodies of the functions
+ // it calls as well, so ask for those it only declares. The importer follows
+ // the call edges of imported functions on its own, with a threshold that
+ // decays per level.
+ DenseSet<GlobalValue::GUID> GUIDs;
+ SmallPtrSet<Function *, 16> Seen;
+ while (!Worklist.empty()) {
+ Function *G = Worklist.pop_back_val();
+ if (!Seen.insert(G).second)
+ continue;
+ if (G->isDeclaration()) {
+ if (!G->isIntrinsic() && !G->hasLocalLinkage() &&
+ !G->getName().starts_with("__enzyme_"))
+ GUIDs.insert(guidOf(*G));
+ continue;
+ }
+ for (Instruction &I : instructions(*G))
+ if (auto *CB = dyn_cast<CallBase>(&I))
+ if (Function *Callee = calledFunction(*CB))
+ Worklist.push_back(Callee);
+ }
+ if (GUIDs.empty())
+ return false;
+
+ LLVMContext &Ctx = M.getContext();
+ Function *Anchor = M.getFunction(AnchorName);
+ if (Anchor) {
+ for (auto G : Anchor->getImportGUIDs())
+ GUIDs.insert(G);
+ } else {
+ Anchor = Function::Create(FunctionType::get(Type::getVoidTy(Ctx), false),
+ GlobalValue::InternalLinkage, AnchorName, M);
+ Anchor->addFnAttr(Attribute::NoInline);
+ Anchor->addFnAttr(Attribute::OptimizeNone);
+ ReturnInst::Create(Ctx, BasicBlock::Create(Ctx, "", Anchor));
+ appendToCompilerUsed(M, {Anchor});
+ }
+ // ModuleSummaryAnalysis turns these GUIDs into call edges of hotness
+ // Critical, which raise the size threshold for importing the callee
+ // (-import-critical-multiplier).
+ Anchor->setMetadata(
+ LLVMContext::MD_prof,
+ MDBuilder(Ctx).createFunctionEntryCount(1, /*Synthetic=*/false, &GUIDs));
+ return true;
+}
+
+bool enzymeThinLTOLocalizeImports(Module &M) {
+ bool Changed = false;
+ if (Function *Anchor = M.getFunction(AnchorName)) {
+ removeFromUsedLists(M, [&](Constant *C) { return C == Anchor; });
+ Anchor->eraseFromParent();
+ Changed = true;
+ }
+
+ // The function to differentiate is the first argument. An imported one has
+ // available_externally linkage, so EliminateAvailableExternally would drop
+ // its body before Enzyme runs.
+ auto IsImported = [](Function *G) {
+ return G && G->hasAvailableExternallyLinkage() && !G->isDeclaration();
+ };
+ SmallVector<Use *, 4> RootUses;
+ forEachEnzymeCall(M, [&](CallBase &CB) {
+ if (CB.arg_size() == 0)
+ return;
+ Use &U = CB.getArgOperandUse(0);
+ if (IsImported(dyn_cast<Function>(U.get()->stripPointerCasts())))
+ RootUses.push_back(&U);
+ });
+ if (RootUses.empty())
+ return Changed;
+
+ // The imported functions they call, transitively, lose their bodies too.
+ SetVector<Function *> Closure;
+ for (Use *U : RootUses)
+ Closure.insert(cast<Function>(U->get()->stripPointerCasts()));
+ for (size_t i = 0; i < Closure.size(); ++i)
+ for (Instruction &I : instructions(*Closure[i]))
+ if (auto *CB = dyn_cast<CallBase>(&I))
+ if (Function *Callee = calledFunction(*CB); IsImported(Callee))
+ Closure.insert(Callee);
+
+ // Internal copies keep their bodies. Only calls and the __enzyme_* operand
+ // are redirected to them: any other use of the function, e.g. its address
+ // stored or compared, still means the original.
+ MapVector<Function *, Function *> Copies;
+ for (Function *G : Closure) {
+ ValueToValueMapTy VMap;
+ Function *C = CloneFunction(G, VMap);
+ C->setName(G->getName() + CopySuffix);
+ C->setLinkage(GlobalValue::InternalLinkage);
+ C->setVisibility(GlobalValue::DefaultVisibility);
+ C->setDLLStorageClass(GlobalValue::DefaultStorageClass);
+ C->setComdat(nullptr);
+ Copies[G] = C;
+ }
+ for (auto &[G, C] : Copies)
+ for (Instruction &I : instructions(*C))
+ if (auto *CB = dyn_cast<CallBase>(&I)) {
+ auto Found = Copies.find(calledFunction(*CB));
+ if (Found != Copies.end() && CB->getCalledOperand() == Found->first)
+ CB->setCalledOperand(Found->second);
+ }
+ for (Use *U : RootUses) {
+ Function *C = Copies[cast<Function>(U->get()->stripPointerCasts())];
+ U->set(ConstantExpr::getPointerCast(C, U->get()->getType()));
+ }
+ return true;
+}
+
+llvm::AnalysisKey EnzymeThinLTOImportPass::Key;
+
+PreservedAnalyses EnzymeThinLTOImportPass::run(Module &M,
+ ModuleAnalysisManager &) {
+ bool Changed = PostLink ? enzymeThinLTOLocalizeImports(M)
+ : enzymeThinLTORequestImports(M);
+ return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
diff --git a/enzyme/Enzyme/ThinLTOImport.h b/enzyme/Enzyme/ThinLTOImport.h
new file mode 100644
index 0000000..3d071a5
--- /dev/null
+++ b/enzyme/Enzyme/ThinLTOImport.h
@@ -0,0 +1,56 @@
+//===- ThinLTOImport.h - Bring functions to differentiate into ThinLTO ----===//
+//
+// Enzyme Project
+//
+// Part of the Enzyme Project, under the Apache License v2.0 with LLVM
+// Exceptions. See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Under ThinLTO, Enzyme differentiates in each post-link backend, which only
+// sees its own module plus the functions the thin link imported into it. The
+// function importer follows call edges of the module summary, but a function
+// passed to __enzyme_autodiff is only referenced, so it is never imported.
+//
+// The pre-link half of this pass asks for those functions: it adds an anchor
+// function whose function_entry_count metadata lists their GUIDs, which the
+// module summary turns into critical call edges (the mechanism sample PGO uses
+// to replay its profiled imports). The post-link half drops the anchor and
+// makes an internal copy of each imported function to differentiate, and of
+// the imported functions it calls, so the copies survive
+// EliminateAvailableExternally until Enzyme runs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ENZYME_THINLTO_IMPORT_H
+#define ENZYME_THINLTO_IMPORT_H
+
+#include "PassUtils.h"
+#include "llvm/IR/PassManager.h"
+
+namespace llvm {
+class Module;
+}
+
+/// At ThinLTO pre-link: request the import of every function passed to an
+/// __enzyme_* call that this module only declares, and of the declarations
+/// that functions it defines and differentiates call.
+bool enzymeThinLTORequestImports(llvm::Module &M);
+
+/// At ThinLTO post-link: drop the anchor the pre-link added, and give every
+/// imported function to differentiate an internal copy.
+bool enzymeThinLTOLocalizeImports(llvm::Module &M);
+
+class EnzymeThinLTOImportPass final
+ : public PassParent<EnzymeThinLTOImportPass> {
+ friend PassParent<EnzymeThinLTOImportPass>;
+ bool PostLink;
+ static llvm::AnalysisKey Key;
+
+public:
+ explicit EnzymeThinLTOImportPass(bool PostLink) : PostLink(PostLink) {}
+ llvm::PreservedAnalyses run(llvm::Module &M, llvm::ModuleAnalysisManager &);
+};
+
+#endif
diff --git a/enzyme/Fortran/README.md b/enzyme/Fortran/README.md
index 4d8d8e1..8b5afd4 100644
--- a/enzyme/Fortran/README.md
+++ b/enzyme/Fortran/README.md
@@ -39,9 +39,13 @@
because their callee may be compiled in another file; lld differentiates the linked
program instead. The config file selects lld and loads `LLDEnzyme-<LLVM version>`
into it, so the same flags work with and without `-flto`. With `-flto=thin`, lld
-differentiates each module on its own, so the function passed to `__enzyme_autodiff`
-has to be defined in the same file as the call (ThinLTO imports only functions that
-are called); use full `-flto` when it is compiled in another file.
+differentiates each module on its own, after ThinLTO imported functions into it.
+The compile step asks ThinLTO to import the function passed to `__enzyme_autodiff`
+and the functions it calls. ThinLTO still applies its size limit to functions further
+down the call graph, and never imports functions marked `noinline` (which includes
+everything compiled at `-O0`). When such a function is missing, Enzyme reports that
+it has no derivative for it; compile it in the same file as the call, raise the limit
+with `-Wl,-mllvm,-import-instr-limit=<N>`, or use full `-flto`.
Without the plugin the derivative has to be produced out of line, by emitting LLVM IR
from flang and running the Enzyme pass over it with `opt`:
diff --git a/enzyme/test/Fortran/ReverseMode/config_file_lto.f90 b/enzyme/test/Fortran/ReverseMode/config_file_lto.f90
index 79d9657..303286d 100644
--- a/enzyme/test/Fortran/ReverseMode/config_file_lto.f90
+++ b/enzyme/test/Fortran/ReverseMode/config_file_lto.f90
@@ -6,11 +6,12 @@
! RUN: %fc -cpp -DCALLEE_ONLY -flto -O2 -c %flangEnzymeConfig %s -o %t.callee.o
! RUN: %fc -cpp -DCALLER_ONLY -flto -O2 -c %flangEnzymeConfig %s -o %t.caller.o
! RUN: %fc -flto -O2 %flangEnzymeConfig %t.caller.o %t.callee.o -o %t1 && %t1 | FileCheck %s
-! With -flto=thin, lld differentiates each module on its own, which then has to
-! define the function: ThinLTO imports functions that are called, not ones that
-! are only passed to __enzyme_autodiff, so here both are in one object file.
-! RUN: %fc -cpp -flto=thin -O2 -c %flangEnzymeConfig %s -o %t.thin.o
-! RUN: %fc -flto=thin -O2 %flangEnzymeConfig %t.thin.o -o %t2 && %t2 | FileCheck %s
+! With -flto=thin, lld differentiates each module on its own after the thin
+! link. Enzyme's pre-link run asks the thin link to import the function passed
+! to __enzyme_autodiff, which ThinLTO would not import for a mere reference.
+! RUN: %fc -cpp -DCALLEE_ONLY -flto=thin -O2 -c %flangEnzymeConfig %s -o %t.callee.thin.o
+! RUN: %fc -cpp -DCALLER_ONLY -flto=thin -O2 -c %flangEnzymeConfig %s -o %t.caller.thin.o
+! RUN: %fc -flto=thin -O2 %flangEnzymeConfig %t.caller.thin.o %t.callee.thin.o -o %t2 && %t2 | FileCheck %s
#ifndef CALLER_ONLY
real function cube(x)
diff --git a/enzyme/test/Integration/ReverseMode/thinlto_import.c b/enzyme/test/Integration/ReverseMode/thinlto_import.c
new file mode 100644
index 0000000..38ee800
--- /dev/null
+++ b/enzyme/test/Integration/ReverseMode/thinlto_import.c
@@ -0,0 +1,34 @@
+// REQUIRES: lld
+// Under ThinLTO, lld differentiates each module after the thin link imported
+// functions into it. The importer follows calls, so a function that a module
+// only passes to __enzyme_autodiff would stay a declaration there. The pre-link
+// run asks for it to be imported, together with the functions it calls.
+
+// The pre-link module requests `f` through an anchor function whose profile
+// metadata lists its GUID, which the module summary turns into a call edge.
+// RUN: if [ %llvmver -ge 20 ]; then %clang -std=c11 -O2 -flto=thin %loadClangEnzyme -c %s -S -emit-llvm -o - | FileCheck %s --check-prefix=PRELINK; fi
+// PRELINK: @llvm.compiler.used = {{.*}}@enzyme.thinlto.imports
+// PRELINK: define internal void @enzyme.thinlto.imports() {{.*}}!prof ![[PROF:[0-9]+]]
+// PRELINK: ![[PROF]] = !{!"function_entry_count", i64 1, i64 {{-?[0-9]+}}}
+
+// `f` and `g` each live in their own object file; d/dx (x*x)*x = 3*x^2 = 12.
+// RUN: if [ %llvmver -ge 20 ]; then %clang -std=c11 -O2 -flto=thin %loadClangEnzyme -c %s -o %t.main.o; fi
+// RUN: if [ %llvmver -ge 20 ]; then %clang -std=c11 -O2 -flto=thin %loadClangEnzyme -DDEFINE_F -c %s -o %t.f.o; fi
+// RUN: if [ %llvmver -ge 20 ]; then %clang -std=c11 -O2 -flto=thin %loadClangEnzyme -DDEFINE_G -c %s -o %t.g.o; fi
+// RUN: if [ %llvmver -ge 20 ]; then %clang -O2 -flto=thin %linkLLDEnzyme %t.main.o %t.f.o %t.g.o -o %t && %t | FileCheck %s; fi
+// CHECK: 12.000000
+
+#if defined(DEFINE_F)
+double g(double);
+double f(double x) { return g(x) * x; }
+#elif defined(DEFINE_G)
+double g(double x) { return x * x; }
+#else
+#include <stdio.h>
+double f(double);
+double __enzyme_autodiff(void *, double);
+int main(void) {
+ printf("%f\n", __enzyme_autodiff((void *)f, 2.0));
+ return 0;
+}
+#endif
diff --git a/enzyme/test/lit.site.cfg.py.in b/enzyme/test/lit.site.cfg.py.in
index cfd59d1..f824199 100644
--- a/enzyme/test/lit.site.cfg.py.in
+++ b/enzyme/test/lit.site.cfg.py.in
@@ -126,6 +126,19 @@
+ config.llvm_ver + config.llvm_shlib_ext)
config.substitutions.append(('%loadLLDEnzymeOpt', loadLLDEnzymeOpt))
+# Link flags that load LLDEnzyme into lld, for clang -flto links; the lld
+# feature needs an ld.lld beside the LLVM tools or on the PATH.
+config.substitutions.append(('%linkLLDEnzyme',
+ ' -fuse-ld=lld -Wl,-mllvm,-load=@ENZYME_BINARY_DIR@/Enzyme/LLDEnzyme-'
+ + config.llvm_ver + config.llvm_shlib_ext
+ + ' -Wl,--load-pass-plugin=@ENZYME_BINARY_DIR@/Enzyme/LLDEnzyme-'
+ + config.llvm_ver + config.llvm_shlib_ext))
+if len("@ENZYME_BINARY_DIR@") != 0:
+ import shutil
+ if (os.path.exists(os.path.join(config.llvm_tools_dir, 'ld.lld'))
+ or shutil.which('ld.lld') is not None):
+ config.available_features.add('lld')
+
config.substitutions.append(('%hasMPFR', has_mpfr))
# Fortran-specific configuration
