diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index 226579c2c9dc..2f65cf5bba11 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp @@ -4696,6 +4696,11 @@ Constant *GradientUtils::GetOrCreateShadowConstant( arg->setMetadata("enzyme_shadow", MDTuple::get(shadow->getContext(), {ConstantAsMetadata::get(shadow)})); + // Mark the shadow so that dispatch-table analysis does not mistake this + // Enzyme-created mirror of a table for a source-level dispatch table. + shadow->setMetadata( + EnzymeShadowGlobalMD, + MDTuple::get(shadow->getContext(), {ConstantAsMetadata::get(arg)})); shadow->setAlignment(arg->getAlign()); shadow->setUnnamedAddr(arg->getUnnamedAddr()); if (arg->hasInitializer()) @@ -4709,6 +4714,94 @@ Constant *GradientUtils::GetOrCreateShadowConstant( llvm_unreachable("unknown constant to create shadow of"); } +/// Meet the argument types over every call in the module which may reach `fn`, +/// for a function whose address only appears inside a dispatch table and which +/// therefore has no direct call site to take type information from. +/// +/// The result constrains the *shadow* of `fn`, not `fn`, and a shadow is only +/// reachable through the shadow table Enzyme writes into this same module -- an +/// external caller invokes the original. So these really are all of the +/// shadow's call sites, despite `fn` itself having external linkage. +/// +/// Erring towards less information is therefore the safe direction, and both +/// approximations do: a candidate set that is too large only adds call sites, +/// which the meet then weakens against, and an indirect call we cannot bound at +/// all makes us give up and return false. +static bool +collectDispatchCallSiteTypes(llvm::Function *fn, TypeAnalysis &TA, + llvm::SmallVectorImpl &out) { + auto M = fn->getParent(); + auto FT = fn->getFunctionType(); + if (FT->isVarArg()) + return false; + + // We run while Enzyme is part-way through emitting other functions into this + // module; TypeAnalysis asserts on those. They are copies of callers we visit + // anyway, so nothing is lost by skipping them. + auto analyzable = [](Function &F) { + if (F.empty()) + return false; + for (auto &BB : F) + if (!BB.hasTerminator()) + return false; + return true; + }; + + SmallVector sites; + for (auto &F : *M) { + if (!analyzable(F)) + continue; + for (auto &BB : F) + for (auto &I : BB) { + auto CB = dyn_cast(&I); + if (!CB) + continue; + // A mismatched signature can neither land in fn nor be read off + // argument by argument. + if (CB->getFunctionType() != FT) + continue; + if (auto direct = getFunctionFromCall(CB)) { + if (direct == fn) + sites.push_back(CB); + continue; + } + if (isa(CB->getCalledOperand())) + continue; + SmallVector cands; + if (!getIndirectCallCandidates(*CB, cands)) + return false; + if (llvm::is_contained(cands, fn)) + sites.push_back(CB); + } + } + if (sites.empty()) + return false; + + out.assign(fn->arg_size(), TypeTree()); + bool first = true; + for (auto CB : sites) { + auto caller = CB->getFunction(); + FnTypeInfo callerInfo(caller); + for (auto &a : caller->args()) { + callerInfo.Arguments.insert(std::pair(&a, {})); + callerInfo.KnownValues.insert( + std::pair>(&a, {})); + } + TypeResults TR = TA.analyzeFunction(callerInfo); + for (size_t i = 0, e = fn->arg_size(); i < e; ++i) { + TypeTree TT; + if (i < CB->arg_size()) + TT = TR.query(CB->getArgOperand(i)); + if (first) + out[i] = TT; + else + out[i].andIn(TT); + } + first = false; + } + return true; +} + Constant *GradientUtils::GetOrCreateShadowFunction( RequestContext context, EnzymeLogic &Logic, TargetLibraryInfo &TLI, TypeAnalysis &TA, Function *fn, DerivativeMode mode, bool runtimeActivity, @@ -4766,14 +4859,25 @@ Constant *GradientUtils::GetOrCreateShadowFunction( type_args.Return.insert({-1, -1}, BaseType::Pointer); } + // Without this, a pointer argument reaches TypeAnalysis below carrying no + // information at all, which for a dispatch table entry (no direct call site + // anywhere) leaves the body as the only source of type information. + SmallVector callSiteTypes; + bool haveCallSiteTypes = + !isRealloc && collectDispatchCallSiteTypes(fn, TA, callSiteTypes); + // conservatively assume that we can only cache existing floating types // (i.e. that all args are overwritten) std::vector types; + size_t argidx = 0; for (auto &a : fn->args()) { overwritten_args.push_back(!a.getType()->isFPOrFPVectorTy()); TypeTree TT; + if (haveCallSiteTypes) + TT = callSiteTypes[argidx]; if (a.getType()->isFPOrFPVectorTy()) TT.insert({-1}, ConcreteType(a.getType()->getScalarType())); + ++argidx; type_args.Arguments.insert(std::pair(&a, TT)); type_args.KnownValues.insert( std::pair>(&a, {})); diff --git a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp index 9b8e6fce620b..1b49526cc157 100644 --- a/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp +++ b/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp @@ -4729,6 +4729,12 @@ void TypeAnalyzer::visitCallBase(CallBase &call) { Function *ci = getFunctionFromCall(&call); + // A callee loaded out of a constant dispatch table (flang lowers a type-bound + // procedure call this way) is the direct call it folds to, so treat it as + // one -- notably to let visitIPOCall below relate caller and callee types. + if (!ci) + ci = getDevirtualizedCallee(&call); + if (ci) { if (ci->getAttributes().hasAttribute(AttributeList::FunctionIndex, "enzyme_ta_norecur")) diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp index 4bf327c76adf..4e01b3d83485 100644 --- a/enzyme/Enzyme/Utils.cpp +++ b/enzyme/Enzyme/Utils.cpp @@ -35,6 +35,7 @@ #endif #include "TypeAnalysis/TBAA.h" +#include "llvm/Analysis/ConstantFolding.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/DerivedTypes.h" #include "llvm/IR/Function.h" @@ -58,6 +59,149 @@ using namespace llvm; +llvm::cl::opt EnzymeDevirtualize( + "enzyme-devirtualize", cl::init(true), cl::Hidden, + cl::desc("Resolve indirect calls whose target is loaded out of a constant " + "dispatch table (e.g. a flang type-bound procedure binding " + "table), so that interprocedural type analysis applies")); + +/// Fold an address computation and/or load into the constant it provably +/// holds, or null. Loads go exclusively through ConstantFoldLoadFromConstPtr, +/// which only succeeds on constant globals with a definitive initializer; that, +/// plus traversing nothing but Load/GEP/Cast, is what makes a non-null result a +/// guarantee rather than a guess. +static Constant *foldConstantMemory(Value *V, const DataLayout &DL, + unsigned depth) { + if (auto C = dyn_cast(V)) + return C; + if (depth >= 8) + return nullptr; + auto I = dyn_cast(V); + if (!I) + return nullptr; + if (!isa(I) && !isa(I) && !isa(I)) + return nullptr; + if (auto LI = dyn_cast(I)) + if (!LI->isSimple()) + return nullptr; + + SmallVector Ops; + for (auto &Op : I->operands()) { + auto C = foldConstantMemory(Op, DL, depth + 1); + if (!C) + return nullptr; + Ops.push_back(C); + } + + if (isa(I)) + return ConstantFoldLoadFromConstPtr(Ops[0], I->getType(), DL); + return ConstantFoldInstOperands(I, Ops, DL); +} + +/// A dispatch table slot as a function: a raw pointer, or flang's +/// `ptrtoint (ptr @f to i64)`. +static Function *asFunctionSlot(Constant *C) { + if (!C) + return nullptr; + C = C->stripPointerCasts(); + if (auto CE = dyn_cast(C)) + if (CE->isCast()) + C = CE->getOperand(0)->stripPointerCasts(); + return dyn_cast(C); +} + +llvm::Function *getDevirtualizedCallee(llvm::CallBase *CB) { + if (!EnzymeDevirtualize) + return nullptr; + auto called = CB->getCalledOperand(); + if (!called || isa(called)) + return nullptr; + auto &DL = CB->getModule()->getDataLayout(); + auto C = foldConstantMemory(called, DL, 0); + auto F = asFunctionSlot(C); + if (!F) + return nullptr; + // Callers read the target through the call's prototype, so a mismatch would + // misattribute arguments. + if (F->getFunctionType() != CB->getFunctionType()) + return nullptr; + return F; +} + +/// A constant aggregate holding at least one function pointer. +static bool isDispatchTable(GlobalVariable &GV) { + if (!GV.isConstant() || !GV.hasDefinitiveInitializer()) + return false; + // An Enzyme shadow mirrors the layout of the table it shadows, so counting it + // as a table of its own would leave every slot with two candidates. + if (GV.getMetadata(EnzymeShadowGlobalMD)) + return false; + SmallVector todo = {GV.getInitializer()}; + SmallPtrSet seen; + unsigned budget = 4096; + while (!todo.empty()) { + auto C = todo.pop_back_val(); + if (!seen.insert(C).second) + continue; + if (budget-- == 0) + return false; + if (asFunctionSlot(C)) + return true; + if (isa(C)) + for (unsigned i = 0, e = C->getNumOperands(); i != e; ++i) + todo.push_back(cast(C->getOperand(i))); + } + return false; +} + +bool getIndirectCallCandidates(llvm::CallBase &CB, + llvm::SmallVectorImpl &Out) { + if (!EnzymeDevirtualize) + return false; + auto called = CB.getCalledOperand(); + if (!called || isa(called)) + return false; + + auto M = CB.getModule(); + auto &DL = M->getDataLayout(); + + // Peel casts off the callee to reach the load of the slot. + Value *V = called; + while (auto CI = dyn_cast(V)) { + if (CI->getOpcode() != Instruction::IntToPtr && + CI->getOpcode() != Instruction::BitCast && + CI->getOpcode() != Instruction::PtrToInt) + return false; + V = CI->getOperand(0); + } + auto LI = dyn_cast(V); + if (!LI || !LI->isSimple()) + return false; + + auto PtrOp = LI->getPointerOperand(); + APInt Off(DL.getIndexTypeSizeInBits(PtrOp->getType()), 0); + PtrOp->stripAndAccumulateConstantOffsets(DL, Off, + /*AllowNonInbounds*/ true); + if (Off.isNegative()) + return false; + + SmallPtrSet found; + for (auto &GV : M->globals()) { + if (!isDispatchTable(GV)) + continue; + auto slot = + ConstantFoldLoadFromConst(GV.getInitializer(), LI->getType(), Off, DL); + auto F = asFunctionSlot(slot); + if (!F) + continue; + if (F->getFunctionType() != CB.getFunctionType()) + continue; + if (found.insert(F).second) + Out.push_back(F); + } + return true; +} + extern "C" { LLVMValueRef (*CustomErrorHandler)(const char *, LLVMValueRef, ErrorType, const void *, LLVMValueRef, diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index 2b19f14cd481..f273f059fce8 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -1338,6 +1338,26 @@ template static inline llvm::Function *getFunctionFromCall(T *op) { return called ? const_cast(called) : nullptr; } +/// Marks a global as an Enzyme-created shadow, so that a shadowed dispatch +/// table is not mistaken for a source-level one. +static constexpr const char *EnzymeShadowGlobalMD = "enzyme_shadow_of"; + +/// The function an indirect call *must* target, proven by constant-folding the +/// callee, or null. Also null for a callee that is already a direct constant -- +/// use getFunctionFromCall for those. +llvm::Function *getDevirtualizedCallee(llvm::CallBase *CB); + +/// Over-approximate the targets of a `call (inttoptr? (load (base + const +/// offset)))` with the function at that offset in every constant dispatch table +/// in the module. `base` need not be known, which is what reaches Fortran +/// type-bound procedure dispatch off a runtime class descriptor. +/// +/// False means the call is not of that shape and must be read as "may call +/// anything". True fills Out, which remains an over-approximation: a table +/// outside this module could hold something else at the same slot. +bool getIndirectCallCandidates(llvm::CallBase &CB, + llvm::SmallVectorImpl &Out); + static inline llvm::StringRef getFuncName(llvm::Function *called) { if (called->hasFnAttribute("enzyme_math")) return called->getFnAttribute("enzyme_math").getValueAsString(); diff --git a/enzyme/test/Fortran/ForwardMode/type_bound_procedure.f90 b/enzyme/test/Fortran/ForwardMode/type_bound_procedure.f90 new file mode 100644 index 000000000000..dd488179c885 --- /dev/null +++ b/enzyme/test/Fortran/ForwardMode/type_bound_procedure.f90 @@ -0,0 +1,83 @@ +! REQUIRES: fortran +! Note: -S -emit-llvm rather than the usual -flto -c, so as not to require an +! LTO-capable system linker. +! RUN: %fc -O1 %loadFortran -S -emit-llvm %s -o %t.ll && %opt %t.ll %loadEnzyme %enzyme -S -o %t2.ll && %fc -O1 %t2.ll -o %t1 && %t1 | FileCheck %s +! RUN: %fc -O2 %loadFortran -S -emit-llvm %s -o %t.ll && %opt %t.ll %loadEnzyme %enzyme -S -o %t2.ll && %fc -O2 %t2.ll -o %t1 && %t1 | FileCheck %s +! RUN: %if flangenzyme %{ %fc -O2 %loadFortran %loadFlangEnzyme %s -o %t2 && %t2 | FileCheck %s %} + +! Differentiation through a type-bound procedure (vtable) dispatch. +! +! `dispatch` takes a polymorphic dummy argument, so flang lowers `s%step` to a +! load out of the binding table @_QMsolver_modEXvXsolver_t followed by an +! indirect call. `step_impl` has no direct call site anywhere in the module -- +! only a `ptrtoint` entry in that table -- so Enzyme has to recover both it and +! its argument types from the dispatch. + +module solver_mod + implicit none + + type :: solver_t + real :: scale + contains + procedure :: step => step_impl + end type solver_t + +contains + + subroutine step_impl(self, n, x, y) + class(solver_t), intent(in) :: self + integer, intent(in) :: n + real, intent(in) :: x(n) + real, intent(inout) :: y + integer :: i + y = 0 + do i = 1, n + y = y + self%scale * x(i) * x(i) + end do + end subroutine step_impl + + ! Polymorphic dummy => dynamic dispatch through the binding table. + subroutine dispatch(s, n, x, y) + class(solver_t), intent(in) :: s + integer, intent(in) :: n + real, intent(in) :: x(n) + real, intent(inout) :: y + call s%step(n, x, y) + end subroutine dispatch + + subroutine run(n, x, y) + integer, intent(in) :: n + real, intent(in) :: x(n) + real, intent(inout) :: y + type(solver_t) :: s + s%scale = 2.0 + call dispatch(s, n, x, y) + end subroutine run + +end module solver_mod + +program main + use enzyme, only: enzyme_const, enzyme_dup, enzyme_fwddiff + use solver_mod + implicit none + + integer :: n + real :: x(3), dx(3), y, dy + + n = 3 + x = [2.0, 3.0, 4.0] + dx = [1.0, 0.0, 0.0] + y = 0 + dy = 0 + + call enzyme_fwddiff(run, enzyme_const, n, & + enzyme_dup, x, dx, enzyme_dup, y, dy) + + ! y = 2*(2^2 + 3^2 + 4^2) = 58 + ! dy = 2*2*x(1)*dx(1) = 8 + print *, int(y) + print *, int(dy) +end program main + +! CHECK: 58 +! CHECK-NEXT: 8 diff --git a/enzyme/test/TypeAnalysis/flangvtable.ll b/enzyme/test/TypeAnalysis/flangvtable.ll new file mode 100644 index 000000000000..8adf843cab66 --- /dev/null +++ b/enzyme/test/TypeAnalysis/flangvtable.ll @@ -0,0 +1,98 @@ +; RUN: %opt < %s %newLoadEnzyme -passes="print-type-analysis" -type-analysis-func=caller -S -o /dev/null | FileCheck %s +; RUN: %opt < %s %newLoadEnzyme -passes="print-type-analysis" -type-analysis-func=runtime_caller -S -o /dev/null | FileCheck %s --check-prefix=RUNTIME + +; Reduced from flang -O1 output for a type-bound procedure ("s%step(...)"), +; which lowers to: load the derived type descriptor, load its binding table, +; index the table by the bound procedure's constant slot, and inttoptr-call the +; i64 found there. A bound procedure therefore has no direct call site anywhere +; -- its address appears only as a `ptrtoint` inside the table -- leaving type +; analysis nothing but the callee body to work from. + +%_QM__fortran_type_infoTbinding = type { { i64 }, { ptr, i64 } } + +@_QMmEXnXinit = linkonce_odr constant [4 x i8] c"init" +@_QMmEXnXstep = linkonce_odr constant [4 x i8] c"step" + +; Entries are { c_funptr, name-descriptor }, 24 bytes each, so slot 1 (`step`) +; sits at byte 24. +@_QMmEXvXsolver_t = linkonce_odr constant [2 x %_QM__fortran_type_infoTbinding] [ + %_QM__fortran_type_infoTbinding { + { i64 } { i64 ptrtoint (ptr @_QMmPinit to i64) }, + { ptr, i64 } { ptr @_QMmEXnXinit, i64 4 } }, + %_QM__fortran_type_infoTbinding { + { i64 } { i64 ptrtoint (ptr @_QMmPstep to i64) }, + { ptr, i64 } { ptr @_QMmEXnXstep, i64 4 } } +], align 64 + +; The derived type descriptor; its first member is the binding table. +@_QMmEXdtXsolver_t = linkonce_odr constant { ptr, i64 } { ptr @_QMmEXvXsolver_t, i64 2 } + +declare void @use(ptr, ptr) + +; That %rwork is a real(8) work array is visible only here. +define void @caller(ptr %rwork, ptr %n) { +entry: + %d = load double, ptr %rwork, align 8 + %d2 = fadd double %d, 1.000000e+00 + store double %d2, ptr %rwork, align 8 + %vt = load ptr, ptr @_QMmEXdtXsolver_t, align 8 + %slot = getelementptr i8, ptr %vt, i64 24 + %fpi = load i64, ptr %slot, align 8 + %fp = inttoptr i64 %fpi to ptr + call void %fp(ptr %rwork, ptr %n) + ret void +} + +; Nothing in this body says what %yh points at. +define void @_QMmPstep(ptr %yh, ptr %ldyh) { +entry: + %len = load i64, ptr %ldyh, align 8 + %cmp = icmp sgt i64 %len, 0 + br i1 %cmp, label %body, label %exit + +body: + call void @use(ptr %yh, ptr %ldyh) + br label %exit + +exit: + ret void +} + +define void @_QMmPinit(ptr %a, ptr %b) { +entry: + ret void +} + +; Dispatch off a runtime class descriptor: which table %box reaches is unknown. +define void @runtime_caller(ptr %box, ptr %rwork, ptr %n) { +entry: + %d = load double, ptr %rwork, align 8 + %d2 = fadd double %d, 1.000000e+00 + store double %d2, ptr %rwork, align 8 + %dt = load ptr, ptr %box, align 8 + %vt = load ptr, ptr %dt, align 8 + %slot = getelementptr i8, ptr %vt, i64 24 + %fpi = load i64, ptr %slot, align 8 + %fp = inttoptr i64 %fpi to ptr + call void %fp(ptr %rwork, ptr %n) + ret void +} + +; @caller's dispatch resolves, so @_QMmPstep is analyzed interprocedurally at +; all, and %rwork's element type reaches its %yh. + +; CHECK: caller - {} | +; CHECK-NEXT: ptr %rwork: {[-1]:Pointer, [-1,0]:Float@double} + +; CHECK: _QMmPstep - {} |{[-1]:Pointer, [-1,0]:Float@double}:{} {[-1]:Pointer}:{} +; CHECK-NEXT: ptr %yh: {[-1]:Pointer, [-1,0]:Float@double} +; CHECK-NEXT: ptr %ldyh: {[-1]:Pointer} + +; @runtime_caller's does not, and nothing is assumed: no callee is analyzed, +; and no type flows back into %n. + +; RUNTIME: runtime_caller - {} | +; RUNTIME-NEXT: ptr %box: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,0]:Pointer} +; RUNTIME-NEXT: ptr %rwork: {[-1]:Pointer, [-1,0]:Float@double} +; RUNTIME-NEXT: ptr %n: {[-1]:Pointer} +; RUNTIME-NOT: _QMmPstep -