Give TypeAnalysis caller context for dispatch-table callees

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<TypeTree> &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<CallBase *, 4> sites;
+  for (auto &F : *M) {
+    if (!analyzable(F))
+      continue;
+    for (auto &BB : F)
+      for (auto &I : BB) {
+        auto CB = dyn_cast<CallBase>(&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<InlineAsm>(CB->getCalledOperand()))
+          continue;
+        SmallVector<Function *, 4> 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<Argument *, TypeTree>(&a, {}));
+      callerInfo.KnownValues.insert(
+          std::pair<Argument *, std::set<int64_t>>(&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<TypeTree, 4> 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<DIFFE_TYPE> 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<Argument *, TypeTree>(&a, TT));
     type_args.KnownValues.insert(
         std::pair<Argument *, std::set<int64_t>>(&a, {}));

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"))

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<bool> 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<Constant>(V))
+    return C;
+  if (depth >= 8)
+    return nullptr;
+  auto I = dyn_cast<Instruction>(V);
+  if (!I)
+    return nullptr;
+  if (!isa<LoadInst>(I) && !isa<GetElementPtrInst>(I) && !isa<CastInst>(I))
+    return nullptr;
+  if (auto LI = dyn_cast<LoadInst>(I))
+    if (!LI->isSimple())
+      return nullptr;
+
+  SmallVector<Constant *, 4> Ops;
+  for (auto &Op : I->operands()) {
+    auto C = foldConstantMemory(Op, DL, depth + 1);
+    if (!C)
+      return nullptr;
+    Ops.push_back(C);
+  }
+
+  if (isa<LoadInst>(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<ConstantExpr>(C))
+    if (CE->isCast())
+      C = CE->getOperand(0)->stripPointerCasts();
+  return dyn_cast<Function>(C);
+}
+
+llvm::Function *getDevirtualizedCallee(llvm::CallBase *CB) {
+  if (!EnzymeDevirtualize)
+    return nullptr;
+  auto called = CB->getCalledOperand();
+  if (!called || isa<Constant>(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<Constant *, 16> todo = {GV.getInitializer()};
+  SmallPtrSet<Constant *, 16> 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<ConstantAggregate>(C))
+      for (unsigned i = 0, e = C->getNumOperands(); i != e; ++i)
+        todo.push_back(cast<Constant>(C->getOperand(i)));
+  }
+  return false;
+}
+
+bool getIndirectCallCandidates(llvm::CallBase &CB,
+                               llvm::SmallVectorImpl<llvm::Function *> &Out) {
+  if (!EnzymeDevirtualize)
+    return false;
+  auto called = CB.getCalledOperand();
+  if (!called || isa<Constant>(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<CastInst>(V)) {
+    if (CI->getOpcode() != Instruction::IntToPtr &&
+        CI->getOpcode() != Instruction::BitCast &&
+        CI->getOpcode() != Instruction::PtrToInt)
+      return false;
+    V = CI->getOperand(0);
+  }
+  auto LI = dyn_cast<LoadInst>(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<Function *, 4> 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,
@@ -2631,26 +2775,26 @@
   using namespace llvm;
   std::map<BasicBlock *, SmallVector<Instruction *, 1>> maybeBlocks;
   BasicBlock *instBlk = inst->getParent();
   for (auto store : stores) {
     BasicBlock *storeBlk = store->getParent();
     if (instBlk == storeBlk) {
       // if store doesn't come before, exit.
 
       if (store != inst) {
         BasicBlock::const_iterator It = storeBlk->begin();
         for (; &*It != store && &*It != inst; ++It)
           /*empty*/;
         // if inst comes first (e.g. before store) in the
         // block, return true
         if (&*It == inst) {
           results.push_back(store);
         }
       }
       maybeBlocks[storeBlk].push_back(store);
     } else {
       maybeBlocks[storeBlk].push_back(store);
     }
   }
 
   if (maybeBlocks.size() == 0)
     return;

enzyme/Enzyme/Utils.h

@@ -1338,6 +1338,26 @@ template <typename T> static inline llvm::Function *getFunctionFromCall(T *op) {
   return called ? const_cast<llvm::Function *>(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<llvm::Function *> &Out);
+
 static inline llvm::StringRef getFuncName(llvm::Function *called) {
   if (called->hasFnAttribute("enzyme_math"))
     return called->getFnAttribute("enzyme_math").getValueAsString();