CApi.cpp

Go to the documentation of this file.

//===- CApi.cpp - Enzyme API exported to C for external use -----------===//
//
//                             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
//
// If using this code in an academic setting, please cite the following:
// @incollection{enzymeNeurips,
// title = {Instead of Rewriting Foreign Code for Machine Learning,
//          Automatically Synthesize Fast Gradients},
// author = {Moses, William S. and Churavy, Valentin},
// booktitle = {Advances in Neural Information Processing Systems 33},
// year = {2020},
// note = {To appear in},
// }
//
//===----------------------------------------------------------------------===//
//
// This file defines various utility functions of Enzyme for access via C
//
//===----------------------------------------------------------------------===//
#include "CApi.h"
#include "Utils.h"
#if LLVM_VERSION_MAJOR >= 16
#define private public
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
#undef private
#else
#include "SCEV/ScalarEvolution.h"
#include "SCEV/ScalarEvolutionExpander.h"
#endif
 
#include "DiffeGradientUtils.h"
#include "DifferentialUseAnalysis.h"
#include "EnzymeLogic.h"
#include "GradientUtils.h"
#include "LibraryFuncs.h"
#if LLVM_VERSION_MAJOR >= 16
#include "llvm/Analysis/TargetLibraryInfo.h"
#else
#include "SCEV/TargetLibraryInfo.h"
#endif
#include "TraceInterface.h"
 
// #include "llvm/ADT/Triple.h"
#include "llvm/Analysis/CallGraph.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/IR/DIBuilder.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/Transforms/Utils/Cloning.h"
 
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/IPO/Attributor.h"
 
#define addAttribute addAttributeAtIndex
#define removeAttribute removeAttributeAtIndex
#define getAttribute getAttributeAtIndex
#define hasAttribute hasAttributeAtIndex
 
using namespace llvm;
 
TargetLibraryInfo eunwrap(LLVMTargetLibraryInfoRef P) {
  return TargetLibraryInfo(*reinterpret_cast<TargetLibraryInfoImpl *>(P));
}
 
EnzymeLogic &eunwrap(EnzymeLogicRef LR) { return *(EnzymeLogic *)LR; }
 
TraceInterface *eunwrap(EnzymeTraceInterfaceRef Ref) {
  return (TraceInterface *)Ref;
}
 
TypeAnalysis &eunwrap(EnzymeTypeAnalysisRef TAR) {
  return *(TypeAnalysis *)TAR;
}
AugmentedReturn *eunwrap(EnzymeAugmentedReturnPtr ARP) {
  return (AugmentedReturn *)ARP;
}
EnzymeAugmentedReturnPtr ewrap(const AugmentedReturn &AR) {
  return (EnzymeAugmentedReturnPtr)(&AR);
}
 
ConcreteType eunwrap(CConcreteType CDT, llvm::LLVMContext &ctx) {
  switch (CDT) {
  case DT_Anything:
    return BaseType::Anything;
  case DT_Integer:
    return BaseType::Integer;
  case DT_Pointer:
    return BaseType::Pointer;
  case DT_Half:
    return ConcreteType(llvm::Type::getHalfTy(ctx));
  case DT_Float:
    return ConcreteType(llvm::Type::getFloatTy(ctx));
  case DT_Double:
    return ConcreteType(llvm::Type::getDoubleTy(ctx));
  case DT_X86_FP80:
    return ConcreteType(llvm::Type::getX86_FP80Ty(ctx));
  case DT_BFloat16:
    return ConcreteType(llvm::Type::getBFloatTy(ctx));
  case DT_FP128:
    return ConcreteType(llvm::Type::getFP128Ty(ctx));
  case DT_Unknown:
    return BaseType::Unknown;
  }
  llvm_unreachable("Unknown concrete type to unwrap");
}
 
std::vector<int> eunwrap(IntList IL) {
  std::vector<int> v;
  for (size_t i = 0; i < IL.size; i++) {
    v.push_back((int)IL.data[i]);
  }
  return v;
}
std::set<int64_t> eunwrap64(IntList IL) {
  std::set<int64_t> v;
  for (size_t i = 0; i < IL.size; i++) {
    v.insert((int64_t)IL.data[i]);
  }
  return v;
}
TypeTree eunwrap(CTypeTreeRef CTT) { return *(TypeTree *)CTT; }
 
CConcreteType ewrap(const ConcreteType &CT) {
  if (auto flt = CT.isFloat()) {
    if (flt->isHalfTy())
      return DT_Half;
    if (flt->isFloatTy())
      return DT_Float;
    if (flt->isDoubleTy())
      return DT_Double;
    if (flt->isX86_FP80Ty())
      return DT_X86_FP80;
    if (flt->isBFloatTy())
      return DT_BFloat16;
    if (flt->isFP128Ty())
      return DT_FP128;
  } else {
    switch (CT.SubTypeEnum) {
    case BaseType::Integer:
      return DT_Integer;
    case BaseType::Pointer:
      return DT_Pointer;
    case BaseType::Anything:
      return DT_Anything;
    case BaseType::Unknown:
      return DT_Unknown;
    case BaseType::Float:
      llvm_unreachable("Illegal conversion of concretetype");
    }
  }
  llvm_unreachable("Illegal conversion of concretetype");
}
 
IntList ewrap(const std::vector<int> &offsets) {
  IntList IL;
  IL.size = offsets.size();
  IL.data = new int64_t[IL.size];
  for (size_t i = 0; i < offsets.size(); i++) {
    IL.data[i] = offsets[i];
  }
  return IL;
}
 
CTypeTreeRef ewrap(const TypeTree &TT) {
  return (CTypeTreeRef)(new TypeTree(TT));
}
 
FnTypeInfo eunwrap(CFnTypeInfo CTI, llvm::Function *F) {
  FnTypeInfo FTI(F);
  // auto &ctx = F->getContext();
  FTI.Return = eunwrap(CTI.Return);
 
  size_t argnum = 0;
  for (auto &arg : F->args()) {
    FTI.Arguments[&arg] = eunwrap(CTI.Arguments[argnum]);
    FTI.KnownValues[&arg] = eunwrap64(CTI.KnownValues[argnum]);
    argnum++;
  }
  return FTI;
}
 
extern "C" {
 
void EnzymeSetCLBool(void *ptr, uint8_t val) {
  auto cl = (llvm::cl::opt<bool> *)ptr;
  cl->setValue((bool)val);
}
 
uint8_t EnzymeGetCLBool(void *ptr) {
  auto cl = (llvm::cl::opt<bool> *)ptr;
  return (uint8_t)(bool)cl->getValue();
}
 
void EnzymeSetCLInteger(void *ptr, int64_t val) {
  auto cl = (llvm::cl::opt<int> *)ptr;
  cl->setValue((int)val);
}
 
int64_t EnzymeGetCLInteger(void *ptr) {
  auto cl = (llvm::cl::opt<int> *)ptr;
  return (int64_t)cl->getValue();
}
 
void EnzymeSetCLString(void *ptr, const char *val) {
  if (auto *clopt = static_cast<cl::opt<std::string> *>(ptr))
    clopt->setValue(val);
}
 
EnzymeLogicRef CreateEnzymeLogic(uint8_t PostOpt) {
  return (EnzymeLogicRef)(new EnzymeLogic((bool)PostOpt));
}
 
void EnzymeLogicSetExternalContext(EnzymeLogicRef Ref, void *ExternalContext) {
  eunwrap(Ref).ExternalContext = ExternalContext;
}
 
void *EnzymeLogicGetExternalContext(EnzymeLogicRef Ref) {
  return eunwrap(Ref).ExternalContext;
}
 
EnzymeTraceInterfaceRef FindEnzymeStaticTraceInterface(LLVMModuleRef M) {
  return (EnzymeTraceInterfaceRef)(new StaticTraceInterface(unwrap(M)));
}
 
EnzymeTraceInterfaceRef CreateEnzymeStaticTraceInterface(
    LLVMContextRef C, LLVMValueRef getTraceFunction,
    LLVMValueRef getChoiceFunction, LLVMValueRef insertCallFunction,
    LLVMValueRef insertChoiceFunction, LLVMValueRef insertArgumentFunction,
    LLVMValueRef insertReturnFunction, LLVMValueRef insertFunctionFunction,
    LLVMValueRef insertChoiceGradientFunction,
    LLVMValueRef insertArgumentGradientFunction, LLVMValueRef newTraceFunction,
    LLVMValueRef freeTraceFunction, LLVMValueRef hasCallFunction,
    LLVMValueRef hasChoiceFunction) {
  return (EnzymeTraceInterfaceRef)(new StaticTraceInterface(
      *unwrap(C), cast<Function>(unwrap(getTraceFunction)),
      cast<Function>(unwrap(getChoiceFunction)),
      cast<Function>(unwrap(insertCallFunction)),
      cast<Function>(unwrap(insertChoiceFunction)),
      cast<Function>(unwrap(insertArgumentFunction)),
      cast<Function>(unwrap(insertReturnFunction)),
      cast<Function>(unwrap(insertFunctionFunction)),
      cast<Function>(unwrap(insertChoiceGradientFunction)),
      cast<Function>(unwrap(insertArgumentGradientFunction)),
      cast<Function>(unwrap(newTraceFunction)),
      cast<Function>(unwrap(freeTraceFunction)),
      cast<Function>(unwrap(hasCallFunction)),
      cast<Function>(unwrap(hasChoiceFunction))));
};
 
EnzymeTraceInterfaceRef
CreateEnzymeDynamicTraceInterface(LLVMValueRef interface, LLVMValueRef F) {
  return (EnzymeTraceInterfaceRef)(new DynamicTraceInterface(
      unwrap(interface), cast<Function>(unwrap(F))));
}
 
void ClearEnzymeLogic(EnzymeLogicRef Ref) { eunwrap(Ref).clear(); }
 
void EnzymeLogicErasePreprocessedFunctions(EnzymeLogicRef Ref) {
  auto &Logic = eunwrap(Ref);
  for (const auto &pair : Logic.PPC.cache)
    pair.second->eraseFromParent();
}
 
void FreeEnzymeLogic(EnzymeLogicRef Ref) { delete (EnzymeLogic *)Ref; }
 
void FreeTraceInterface(EnzymeTraceInterfaceRef Ref) {
  delete (TraceInterface *)Ref;
}
 
EnzymeTypeAnalysisRef CreateTypeAnalysis(EnzymeLogicRef Log,
                                         char **customRuleNames,
                                         CustomRuleType *customRules,
                                         size_t numRules) {
  EnzymeLogic &Logic = eunwrap(Log);
  TypeAnalysis *TA = new TypeAnalysis(Logic);
  for (size_t i = 0; i < numRules; i++) {
    CustomRuleType rule = customRules[i];
    TA->CustomRules[customRuleNames[i]] =
        [=](int direction, TypeTree &returnTree, ArrayRef<TypeTree> argTrees,
            ArrayRef<std::set<int64_t>> knownValues, CallBase *call,
            TypeAnalyzer *TA) -> uint8_t {
      CTypeTreeRef creturnTree = (CTypeTreeRef)(&returnTree);
      CTypeTreeRef *cargs = new CTypeTreeRef[argTrees.size()];
      IntList *kvs = new IntList[argTrees.size()];
      for (size_t i = 0; i < argTrees.size(); ++i) {
        cargs[i] = (CTypeTreeRef)(&(argTrees[i]));
        kvs[i].size = knownValues[i].size();
        kvs[i].data = new int64_t[kvs[i].size];
        size_t j = 0;
        for (auto val : knownValues[i]) {
          kvs[i].data[j] = val;
          j++;
        }
      }
      uint8_t result = rule(direction, creturnTree, cargs, kvs, argTrees.size(),
                            wrap(call), TA);
      delete[] cargs;
      for (size_t i = 0; i < argTrees.size(); ++i) {
        delete[] kvs[i].data;
      }
      delete[] kvs;
      return result;
    };
  }
  return (EnzymeTypeAnalysisRef)TA;
}
 
void ClearTypeAnalysis(EnzymeTypeAnalysisRef TAR) { eunwrap(TAR).clear(); }
 
void FreeTypeAnalysis(EnzymeTypeAnalysisRef TAR) {
  TypeAnalysis *TA = (TypeAnalysis *)TAR;
  delete TA;
}
 
EnzymeLogicRef EnzymeTypeAnalysisGetLogic(EnzymeTypeAnalysisRef TAR) {
  return (EnzymeLogicRef) & ((TypeAnalysis *)TAR)->Logic;
}
 
void *EnzymeAnalyzeTypes(EnzymeTypeAnalysisRef TAR, CFnTypeInfo CTI,
                         LLVMValueRef F) {
  FnTypeInfo FTI(eunwrap(CTI, cast<Function>(unwrap(F))));
  return (void *)((TypeAnalysis *)TAR)->analyzeFunction(FTI).analyzer;
}
 
void *EnzymeGradientUtilsTypeAnalyzer(GradientUtils *G) {
  return (void *)&G->TR.analyzer;
}
 
EnzymeTypeAnalysisRef EnzymeGetTypeAnalysisFromTypeAnalyzer(void *TAR) {
  return (EnzymeTypeAnalysisRef) & ((TypeAnalyzer *)TAR)->interprocedural;
}
 
void EnzymeGradientUtilsErase(GradientUtils *G, LLVMValueRef I) {
  return G->erase(cast<Instruction>(unwrap(I)));
}
void EnzymeGradientUtilsEraseWithPlaceholder(GradientUtils *G, LLVMValueRef I,
                                             LLVMValueRef orig, uint8_t erase) {
  return G->eraseWithPlaceholder(cast<Instruction>(unwrap(I)),
                                 cast<Instruction>(unwrap(orig)),
                                 "_replacementABI", erase != 0);
}
 
void EnzymeGradientUtilsReplaceAWithB(GradientUtils *G, LLVMValueRef A,
                                      LLVMValueRef B) {
  return G->replaceAWithB(unwrap(A), unwrap(B));
}
 
void EnzymeRegisterAllocationHandler(char *Name, CustomShadowAlloc AHandle,
                                     CustomShadowFree FHandle) {
  shadowHandlers[Name] = [=](IRBuilder<> &B, CallInst *CI,
                             ArrayRef<Value *> Args,
                             GradientUtils *gutils) -> llvm::Value * {
    SmallVector<LLVMValueRef, 3> refs;
    for (auto a : Args)
      refs.push_back(wrap(a));
    return unwrap(
        AHandle(wrap(&B), wrap(CI), Args.size(), refs.data(), gutils));
  };
  if (FHandle)
    shadowErasers[Name] = [=](IRBuilder<> &B,
                              Value *ToFree) -> llvm::CallInst * {
      return cast_or_null<CallInst>(unwrap(FHandle(wrap(&B), wrap(ToFree))));
    };
}
 
void EnzymeRegisterCallHandler(const char *Name,
                               CustomAugmentedFunctionForward FwdHandle,
                               CustomFunctionReverse RevHandle) {
  auto &pair = customCallHandlers[Name];
  pair.first = [=](IRBuilder<> &B, CallInst *CI, GradientUtils &gutils,
                   Value *&normalReturn, Value *&shadowReturn,
                   Value *&tape) -> bool {
    LLVMValueRef normalR = wrap(normalReturn);
    LLVMValueRef shadowR = wrap(shadowReturn);
    LLVMValueRef tapeR = wrap(tape);
    uint8_t noMod =
        FwdHandle(wrap(&B), wrap(CI), &gutils, &normalR, &shadowR, &tapeR);
    normalReturn = unwrap(normalR);
    shadowReturn = unwrap(shadowR);
    tape = unwrap(tapeR);
    return noMod != 0;
  };
  pair.second = [=](IRBuilder<> &B, CallInst *CI, DiffeGradientUtils &gutils,
                    Value *tape) {
    RevHandle(wrap(&B), wrap(CI), &gutils, wrap(tape));
  };
}
 
void EnzymeRegisterFwdCallHandler(char *Name, CustomFunctionForward FwdHandle) {
  auto &pair = customFwdCallHandlers[Name];
  pair = [=](IRBuilder<> &B, CallInst *CI, GradientUtils &gutils,
             Value *&normalReturn, Value *&shadowReturn) -> bool {
    LLVMValueRef normalR = wrap(normalReturn);
    LLVMValueRef shadowR = wrap(shadowReturn);
    uint8_t noMod = FwdHandle(wrap(&B), wrap(CI), &gutils, &normalR, &shadowR);
    normalReturn = unwrap(normalR);
    shadowReturn = unwrap(shadowR);
    return noMod != 0;
  };
}
 
void EnzymeRegisterDiffUseCallHandler(char *Name,
                                      CustomFunctionDiffUse Handle) {
  auto &pair = customDiffUseHandlers[Name];
  pair = [=](const CallInst *CI, const GradientUtils *gutils, const Value *arg,
             bool isshadow, DerivativeMode mode, bool &useDefault) -> bool {
    uint8_t useDefaultC = 0;
    uint8_t noMod = Handle(wrap(CI), gutils, wrap(arg), isshadow,
                           (CDerivativeMode)(mode), &useDefaultC);
    useDefault = useDefaultC != 0;
    return noMod != 0;
  };
}
 
uint8_t EnzymeGradientUtilsGetRuntimeActivity(GradientUtils *gutils) {
  return gutils->runtimeActivity;
}
 
void *EnzymeGradientUtilsGetExternalContext(GradientUtils *gutils) {
  return gutils->Logic.ExternalContext;
}
 
uint8_t EnzymeGradientUtilsGetStrongZero(GradientUtils *gutils) {
  return gutils->strongZero;
}
 
uint8_t EnzymeGradientUtilsGetAtomicAdd(GradientUtils *gutils) {
  return gutils->AtomicAdd;
}
 
uint64_t EnzymeGradientUtilsGetWidth(GradientUtils *gutils) {
  return gutils->getWidth();
}
 
LLVMTypeRef EnzymeGradientUtilsGetShadowType(GradientUtils *gutils,
                                             LLVMTypeRef T) {
  return wrap(gutils->getShadowType(unwrap(T)));
}
 
LLVMTypeRef EnzymeGetShadowType(uint64_t width, LLVMTypeRef T) {
  return wrap(GradientUtils::getShadowType(unwrap(T), width));
}
 
LLVMValueRef EnzymeGradientUtilsNewFromOriginal(GradientUtils *gutils,
                                                LLVMValueRef val) {
  return wrap(gutils->getNewFromOriginal(unwrap(val)));
}
 
void EnzymeReplaceOriginalToNew(GradientUtils *gutils, LLVMValueRef origC,
                                LLVMValueRef repC) {
  auto orig = cast<Instruction>(unwrap(origC));
  auto rep = cast<Instruction>(unwrap(repC));
  auto found = gutils->originalToNewFn.find(orig);
  assert(found != gutils->originalToNewFn.end());
  auto newCall = found->second;
  gutils->originalToNewFn[orig] = rep;
  gutils->newToOriginalFn.erase(newCall);
  gutils->newToOriginalFn[rep] = orig;
}
 
CDerivativeMode EnzymeGradientUtilsGetMode(GradientUtils *gutils) {
  return (CDerivativeMode)gutils->mode;
}
 
CDIFFE_TYPE
EnzymeGradientUtilsGetDiffeType(GradientUtils *G, LLVMValueRef oval,
                                uint8_t foreignFunction) {
  return (CDIFFE_TYPE)(G->getDiffeType(unwrap(oval), foreignFunction != 0));
}
 
CDIFFE_TYPE
EnzymeGradientUtilsGetReturnDiffeType(GradientUtils *G, LLVMValueRef oval,
                                      uint8_t *needsPrimal,
                                      uint8_t *needsShadow,
                                      CDerivativeMode mode) {
  bool needsPrimalB;
  bool needsShadowB;
  auto res = (CDIFFE_TYPE)(G->getReturnDiffeType(
      unwrap(oval), &needsPrimalB, &needsShadowB, (DerivativeMode)mode));
  if (needsPrimal)
    *needsPrimal = needsPrimalB;
  if (needsShadow)
    *needsShadow = needsShadowB;
  return res;
}
 
void EnzymeGradientUtilsSetDebugLocFromOriginal(GradientUtils *gutils,
                                                LLVMValueRef val,
                                                LLVMValueRef orig) {
  return cast<Instruction>(unwrap(val))
      ->setDebugLoc(gutils->getNewFromOriginal(
          cast<Instruction>(unwrap(orig))->getDebugLoc()));
}
 
LLVMValueRef EnzymeInsertValue(LLVMBuilderRef B, LLVMValueRef val,
                               LLVMValueRef val2, unsigned *sz, int64_t length,
                               const char *name) {
  return wrap(unwrap(B)->CreateInsertValue(
      unwrap(val), unwrap(val2), ArrayRef<unsigned>(sz, sz + length), name));
}
 
LLVMValueRef EnzymeGradientUtilsLookup(GradientUtils *gutils, LLVMValueRef val,
                                       LLVMBuilderRef B) {
  return wrap(gutils->lookupM(unwrap(val), *unwrap(B)));
}
 
LLVMValueRef EnzymeGradientUtilsInvertPointer(GradientUtils *gutils,
                                              LLVMValueRef val,
                                              LLVMBuilderRef B) {
  return wrap(gutils->invertPointerM(unwrap(val), *unwrap(B)));
}
 
LLVMValueRef EnzymeGradientUtilsDiffe(DiffeGradientUtils *gutils,
                                      LLVMValueRef val, LLVMBuilderRef B) {
  return wrap(gutils->diffe(unwrap(val), *unwrap(B)));
}
 
void EnzymeGradientUtilsAddToDiffe(DiffeGradientUtils *gutils, LLVMValueRef val,
                                   LLVMValueRef diffe, LLVMBuilderRef B,
                                   LLVMTypeRef T) {
  gutils->addToDiffe(unwrap(val), unwrap(diffe), *unwrap(B), unwrap(T));
}
 
void EnzymeGradientUtilsAddToInvertedPointerDiffe(
    DiffeGradientUtils *gutils, LLVMValueRef orig, LLVMValueRef origVal,
    LLVMTypeRef addingType, unsigned start, unsigned size, LLVMValueRef origptr,
    LLVMValueRef dif, LLVMBuilderRef BuilderM, unsigned align,
    LLVMValueRef mask) {
  MaybeAlign align2;
  if (align)
    align2 = MaybeAlign(align);
  auto inst = cast_or_null<Instruction>(unwrap(orig));
  gutils->addToInvertedPtrDiffe(inst, unwrap(origVal), unwrap(addingType),
                                start, size, unwrap(origptr), unwrap(dif),
                                *unwrap(BuilderM), align2, unwrap(mask));
}
 
void EnzymeGradientUtilsAddToInvertedPointerDiffeTT(
    DiffeGradientUtils *gutils, LLVMValueRef orig, LLVMValueRef origVal,
    CTypeTreeRef vd, unsigned LoadSize, LLVMValueRef origptr,
    LLVMValueRef prediff, LLVMBuilderRef BuilderM, unsigned align,
    LLVMValueRef premask) {
  MaybeAlign align2;
  if (align)
    align2 = MaybeAlign(align);
  auto inst = cast_or_null<Instruction>(unwrap(orig));
  gutils->addToInvertedPtrDiffe(inst, unwrap(origVal), *(TypeTree *)vd,
                                LoadSize, unwrap(origptr), unwrap(prediff),
                                *unwrap(BuilderM), align2, unwrap(premask));
}
 
void EnzymeGradientUtilsSetDiffe(DiffeGradientUtils *gutils, LLVMValueRef val,
                                 LLVMValueRef diffe, LLVMBuilderRef B) {
  gutils->setDiffe(unwrap(val), unwrap(diffe), *unwrap(B));
}
 
uint8_t EnzymeGradientUtilsIsConstantValue(GradientUtils *gutils,
                                           LLVMValueRef val) {
  return gutils->isConstantValue(unwrap(val));
}
 
uint8_t EnzymeGradientUtilsIsConstantInstruction(GradientUtils *gutils,
                                                 LLVMValueRef val) {
  return gutils->isConstantInstruction(cast<Instruction>(unwrap(val)));
}
 
LLVMBasicBlockRef EnzymeGradientUtilsAllocationBlock(GradientUtils *gutils) {
  return wrap(gutils->inversionAllocs);
}
 
uint8_t EnzymeGradientUtilsGetUncacheableArgs(GradientUtils *gutils,
                                              LLVMValueRef orig, uint8_t *data,
                                              uint64_t size) {
  if (gutils->mode == DerivativeMode::ForwardMode ||
      gutils->mode == DerivativeMode::ForwardModeError)
    return 0;
 
  if (!gutils->overwritten_args_map_ptr)
    return 0;
 
  CallInst *call = cast<CallInst>(unwrap(orig));
 
  assert(gutils->overwritten_args_map_ptr);
  auto found = gutils->overwritten_args_map_ptr->find(call);
  if (found == gutils->overwritten_args_map_ptr->end()) {
    llvm::errs() << " oldFunc " << *gutils->oldFunc << "\n";
    for (auto &pair : *gutils->overwritten_args_map_ptr) {
      llvm::errs() << " + " << *pair.first << "\n";
    }
    llvm::errs() << " could not find call orig in overwritten_args_map_ptr "
                 << *call << "\n";
  }
  assert(found != gutils->overwritten_args_map_ptr->end());
 
  const std::vector<bool> &overwritten_args = found->second.second;
 
  if (size != overwritten_args.size()) {
    llvm::errs() << " orig: " << *call << "\n";
    llvm::errs() << " size: " << size
                 << " overwritten_args.size(): " << overwritten_args.size()
                 << "\n";
  }
  assert(size == overwritten_args.size());
  for (uint64_t i = 0; i < size; i++) {
    data[i] = overwritten_args[i];
  }
  return 1;
}
 
CTypeTreeRef EnzymeGradientUtilsAllocAndGetTypeTree(GradientUtils *gutils,
                                                    LLVMValueRef val) {
  auto v = unwrap(val);
  TypeTree TT = gutils->TR.query(v);
  TypeTree *pTT = new TypeTree(TT);
  return (CTypeTreeRef)pTT;
}
 
void EnzymeGradientUtilsDumpTypeResults(GradientUtils *gutils) {
  gutils->TR.dump();
}
 
void EnzymeGradientUtilsSubTransferHelper(
    GradientUtils *gutils, CDerivativeMode mode, LLVMTypeRef secretty,
    uint64_t intrinsic, uint64_t dstAlign, uint64_t srcAlign, uint64_t offset,
    uint8_t dstConstant, LLVMValueRef shadow_dst, uint8_t srcConstant,
    LLVMValueRef shadow_src, LLVMValueRef length, LLVMValueRef isVolatile,
    LLVMValueRef MTI, uint8_t allowForward, uint8_t shadowsLookedUp) {
  auto orig = unwrap(MTI);
  assert(orig);
  SubTransferHelper(gutils, (DerivativeMode)mode, unwrap(secretty),
                    (Intrinsic::ID)intrinsic, (unsigned)dstAlign,
                    (unsigned)srcAlign, (unsigned)offset, (bool)dstConstant,
                    unwrap(shadow_dst), (bool)srcConstant, unwrap(shadow_src),
                    unwrap(length), unwrap(isVolatile), cast<CallInst>(orig),
                    (bool)allowForward, (bool)shadowsLookedUp);
}
 
LLVMBasicBlockRef EnzymeGradientUtilsAddReverseBlock(GradientUtils *gutils,
                                                     LLVMBasicBlockRef block,
                                                     const char *name,
                                                     uint8_t forkCache,
                                                     uint8_t push) {
  return wrap(gutils->addReverseBlock(cast<BasicBlock>(unwrap(block)), name,
                                      forkCache, push));
}
 
void EnzymeGradientUtilsSetReverseBlock(GradientUtils *gutils,
                                        LLVMBasicBlockRef block) {
  auto endBlock = cast<BasicBlock>(unwrap(block));
  auto found = gutils->reverseBlockToPrimal.find(endBlock);
  assert(found != gutils->reverseBlockToPrimal.end());
  auto &vec = gutils->reverseBlocks[found->second];
  assert(vec.size());
  vec.push_back(endBlock);
}
 
LLVMValueRef EnzymeCreateForwardDiff(
    EnzymeLogicRef Logic, LLVMValueRef request_req, LLVMBuilderRef request_ip,
    LLVMValueRef todiff, CDIFFE_TYPE retType, CDIFFE_TYPE *constant_args,
    size_t constant_args_size, EnzymeTypeAnalysisRef TA, uint8_t returnValue,
    CDerivativeMode mode, uint8_t freeMemory, uint8_t runtimeActivity,
    uint8_t strongZero, unsigned width, LLVMTypeRef additionalArg,
    CFnTypeInfo typeInfo, uint8_t subsequent_calls_may_write,
    uint8_t *_overwritten_args, size_t overwritten_args_size,
    EnzymeAugmentedReturnPtr augmented) {
  SmallVector<DIFFE_TYPE, 4> nconstant_args((DIFFE_TYPE *)constant_args,
                                            (DIFFE_TYPE *)constant_args +
                                                constant_args_size);
  std::vector<bool> overwritten_args;
  assert(overwritten_args_size == cast<Function>(unwrap(todiff))->arg_size());
  for (uint64_t i = 0; i < overwritten_args_size; i++) {
    overwritten_args.push_back(_overwritten_args[i]);
  }
  return wrap(eunwrap(Logic).CreateForwardDiff(
      RequestContext(cast_or_null<Instruction>(unwrap(request_req)),
                     unwrap(request_ip)),
      cast<Function>(unwrap(todiff)), (DIFFE_TYPE)retType, nconstant_args,
      eunwrap(TA), returnValue, (DerivativeMode)mode, freeMemory,
      runtimeActivity, strongZero, width, unwrap(additionalArg),
      eunwrap(typeInfo, cast<Function>(unwrap(todiff))),
      subsequent_calls_may_write, overwritten_args, eunwrap(augmented)));
}
LLVMValueRef EnzymeCreatePrimalAndGradient(
    EnzymeLogicRef Logic, LLVMValueRef request_req, LLVMBuilderRef request_ip,
    LLVMValueRef todiff, CDIFFE_TYPE retType, CDIFFE_TYPE *constant_args,
    size_t constant_args_size, EnzymeTypeAnalysisRef TA, uint8_t returnValue,
    uint8_t dretUsed, CDerivativeMode mode, uint8_t runtimeActivity,
    uint8_t strongZero, unsigned width, uint8_t freeMemory,
    LLVMTypeRef additionalArg, uint8_t forceAnonymousTape, CFnTypeInfo typeInfo,
    uint8_t subsequent_calls_may_write, uint8_t *_overwritten_args,
    size_t overwritten_args_size, EnzymeAugmentedReturnPtr augmented,
    uint8_t AtomicAdd) {
  std::vector<DIFFE_TYPE> nconstant_args((DIFFE_TYPE *)constant_args,
                                         (DIFFE_TYPE *)constant_args +
                                             constant_args_size);
  std::vector<bool> overwritten_args;
  assert(overwritten_args_size == cast<Function>(unwrap(todiff))->arg_size());
  for (uint64_t i = 0; i < overwritten_args_size; i++) {
    overwritten_args.push_back(_overwritten_args[i]);
  }
  return wrap(eunwrap(Logic).CreatePrimalAndGradient(
      RequestContext(cast_or_null<Instruction>(unwrap(request_req)),
                     unwrap(request_ip)),
      (ReverseCacheKey){
          .todiff = cast<Function>(unwrap(todiff)),
          .retType = (DIFFE_TYPE)retType,
          .constant_args = nconstant_args,
          .subsequent_calls_may_write = (bool)subsequent_calls_may_write,
          .overwritten_args = overwritten_args,
          .returnUsed = (bool)returnValue,
          .shadowReturnUsed = (bool)dretUsed,
          .mode = (DerivativeMode)mode,
          .width = width,
          .freeMemory = (bool)freeMemory,
          .AtomicAdd = (bool)AtomicAdd,
          .additionalType = unwrap(additionalArg),
          .forceAnonymousTape = (bool)forceAnonymousTape,
          .typeInfo = eunwrap(typeInfo, cast<Function>(unwrap(todiff))),
          .runtimeActivity = (bool)runtimeActivity,
          .strongZero = (bool)strongZero},
      eunwrap(TA), eunwrap(augmented)));
}
EnzymeAugmentedReturnPtr EnzymeCreateAugmentedPrimal(
    EnzymeLogicRef Logic, LLVMValueRef request_req, LLVMBuilderRef request_ip,
    LLVMValueRef todiff, CDIFFE_TYPE retType, CDIFFE_TYPE *constant_args,
    size_t constant_args_size, EnzymeTypeAnalysisRef TA, uint8_t returnUsed,
    uint8_t shadowReturnUsed, CFnTypeInfo typeInfo,
    uint8_t subsequent_calls_may_write, uint8_t *_overwritten_args,
    size_t overwritten_args_size, uint8_t forceAnonymousTape,
    uint8_t runtimeActivity, uint8_t strongZero, unsigned width,
    uint8_t AtomicAdd) {
 
  SmallVector<DIFFE_TYPE, 4> nconstant_args((DIFFE_TYPE *)constant_args,
                                            (DIFFE_TYPE *)constant_args +
                                                constant_args_size);
  std::vector<bool> overwritten_args;
  assert(overwritten_args_size == cast<Function>(unwrap(todiff))->arg_size());
  for (uint64_t i = 0; i < overwritten_args_size; i++) {
    overwritten_args.push_back(_overwritten_args[i]);
  }
  return ewrap(eunwrap(Logic).CreateAugmentedPrimal(
      RequestContext(cast_or_null<Instruction>(unwrap(request_req)),
                     unwrap(request_ip)),
      cast<Function>(unwrap(todiff)), (DIFFE_TYPE)retType, nconstant_args,
      eunwrap(TA), returnUsed, shadowReturnUsed,
      eunwrap(typeInfo, cast<Function>(unwrap(todiff))),
      subsequent_calls_may_write, overwritten_args, forceAnonymousTape,
      runtimeActivity, strongZero, width, AtomicAdd));
}
 
LLVMValueRef EnzymeCreateBatch(EnzymeLogicRef Logic, LLVMValueRef request_req,
                               LLVMBuilderRef request_ip, LLVMValueRef tobatch,
                               unsigned width, CBATCH_TYPE *arg_types,
                               size_t arg_types_size, CBATCH_TYPE retType) {
 
  return wrap(eunwrap(Logic).CreateBatch(
      RequestContext(cast_or_null<Instruction>(unwrap(request_req)),
                     unwrap(request_ip)),
      cast<Function>(unwrap(tobatch)), width,
      ArrayRef<BATCH_TYPE>((BATCH_TYPE *)arg_types,
                           (BATCH_TYPE *)arg_types + arg_types_size),
      (BATCH_TYPE)retType));
}
 
LLVMValueRef EnzymeCreateTrace(
    EnzymeLogicRef Logic, LLVMValueRef request_req, LLVMBuilderRef request_ip,
    LLVMValueRef totrace, LLVMValueRef *sample_functions,
    size_t sample_functions_size, LLVMValueRef *observe_functions,
    size_t observe_functions_size, const char *active_random_variables[],
    size_t active_random_variables_size, CProbProgMode mode, uint8_t autodiff,
    EnzymeTraceInterfaceRef interface) {
 
  SmallPtrSet<Function *, 4> SampleFunctions;
  for (size_t i = 0; i < sample_functions_size; i++) {
    SampleFunctions.insert(cast<Function>(unwrap(sample_functions[i])));
  }
 
  SmallPtrSet<Function *, 4> ObserveFunctions;
  for (size_t i = 0; i < observe_functions_size; i++) {
    ObserveFunctions.insert(cast<Function>(unwrap(observe_functions[i])));
  }
 
  StringSet<> ActiveRandomVariables;
  for (size_t i = 0; i < active_random_variables_size; i++) {
    ActiveRandomVariables.insert(active_random_variables[i]);
  }
 
  return wrap(eunwrap(Logic).CreateTrace(
      RequestContext(cast_or_null<Instruction>(unwrap(request_req)),
                     unwrap(request_ip)),
      cast<Function>(unwrap(totrace)), SampleFunctions, ObserveFunctions,
      ActiveRandomVariables, (ProbProgMode)mode, (bool)autodiff,
      eunwrap(interface)));
}
 
LLVMValueRef
EnzymeExtractFunctionFromAugmentation(EnzymeAugmentedReturnPtr ret) {
  auto AR = (AugmentedReturn *)ret;
  return wrap(AR->fn);
}
 
LLVMTypeRef
EnzymeExtractUnderlyingTapeTypeFromAugmentation(EnzymeAugmentedReturnPtr ret) {
  auto AR = (AugmentedReturn *)ret;
  return wrap(AR->tapeType);
}
 
LLVMTypeRef
EnzymeExtractTapeTypeFromAugmentation(EnzymeAugmentedReturnPtr ret) {
  auto AR = (AugmentedReturn *)ret;
  auto found = AR->returns.find(AugmentedStruct::Tape);
  if (found == AR->returns.end()) {
    return wrap((Type *)nullptr);
  }
  if (found->second == -1) {
    return wrap(AR->fn->getReturnType());
  }
  return wrap(
      cast<StructType>(AR->fn->getReturnType())->getTypeAtIndex(found->second));
}
void EnzymeExtractReturnInfo(EnzymeAugmentedReturnPtr ret, int64_t *data,
                             uint8_t *existed, size_t len) {
  assert(len == 3);
  auto AR = (AugmentedReturn *)ret;
  AugmentedStruct todo[] = {AugmentedStruct::Tape, AugmentedStruct::Return,
                            AugmentedStruct::DifferentialReturn};
  for (size_t i = 0; i < len; i++) {
    auto found = AR->returns.find(todo[i]);
    if (found != AR->returns.end()) {
      existed[i] = true;
      data[i] = (int64_t)found->second;
    } else {
      existed[i] = false;
    }
  }
}
 
static MDNode *extractMDNode(MetadataAsValue *MAV) {
  Metadata *MD = MAV->getMetadata();
  assert((isa<MDNode>(MD) || isa<ConstantAsMetadata>(MD)) &&
         "Expected a metadata node or a canonicalized constant");
 
  if (MDNode *N = dyn_cast<MDNode>(MD))
    return N;
 
  return MDNode::get(MAV->getContext(), MD);
}
 
CTypeTreeRef EnzymeTypeTreeFromMD(LLVMValueRef Val) {
  TypeTree *Ret = new TypeTree();
  MDNode *N = Val ? extractMDNode(unwrap<MetadataAsValue>(Val)) : nullptr;
  Ret->insertFromMD(N);
  return (CTypeTreeRef)N;
}
 
LLVMValueRef EnzymeTypeTreeToMD(CTypeTreeRef CTR, LLVMContextRef ctx) {
  auto MD = ((TypeTree *)CTR)->toMD(*unwrap(ctx));
  return wrap(MetadataAsValue::get(MD->getContext(), MD));
}
 
CTypeTreeRef EnzymeNewTypeTree() { return (CTypeTreeRef)(new TypeTree()); }
CTypeTreeRef EnzymeNewTypeTreeCT(CConcreteType CT, LLVMContextRef ctx) {
  return (CTypeTreeRef)(new TypeTree(eunwrap(CT, *unwrap(ctx))));
}
CTypeTreeRef EnzymeNewTypeTreeTR(CTypeTreeRef CTR) {
  return (CTypeTreeRef)(new TypeTree(*(TypeTree *)(CTR)));
}
void EnzymeFreeTypeTree(CTypeTreeRef CTT) { delete (TypeTree *)CTT; }
uint8_t EnzymeSetTypeTree(CTypeTreeRef dst, CTypeTreeRef src) {
  return *(TypeTree *)dst = *(TypeTree *)src;
}
uint8_t EnzymeMergeTypeTree(CTypeTreeRef dst, CTypeTreeRef src) {
  return ((TypeTree *)dst)->orIn(*(TypeTree *)src, /*PointerIntSame*/ false);
}
uint8_t EnzymeCheckedMergeTypeTree(CTypeTreeRef dst, CTypeTreeRef src,
                                   uint8_t *legalP) {
  bool legal = true;
  bool res =
      ((TypeTree *)dst)
          ->checkedOrIn(*(TypeTree *)src, /*PointerIntSame*/ false, legal);
  *legalP = legal;
  return res;
}
 
void EnzymeTypeTreeOnlyEq(CTypeTreeRef CTT, int64_t x) {
  // TODO only inst
  *(TypeTree *)CTT = ((TypeTree *)CTT)->Only(x, nullptr);
}
void EnzymeTypeTreeData0Eq(CTypeTreeRef CTT) {
  *(TypeTree *)CTT = ((TypeTree *)CTT)->Data0();
}
 
void EnzymeTypeTreeLookupEq(CTypeTreeRef CTT, int64_t size, const char *dl) {
  *(TypeTree *)CTT = ((TypeTree *)CTT)->Lookup(size, DataLayout(dl));
}
void EnzymeTypeTreeCanonicalizeInPlace(CTypeTreeRef CTT, int64_t size,
                                       const char *dl) {
  ((TypeTree *)CTT)->CanonicalizeInPlace(size, DataLayout(dl));
}
 
CConcreteType EnzymeTypeTreeInner0(CTypeTreeRef CTT) {
  return ewrap(((TypeTree *)CTT)->Inner0());
}
 
void EnzymeTypeTreeShiftIndiciesEq(CTypeTreeRef CTT, const char *datalayout,
                                   int64_t offset, int64_t maxSize,
                                   uint64_t addOffset) {
  DataLayout DL(datalayout);
  *(TypeTree *)CTT =
      ((TypeTree *)CTT)->ShiftIndices(DL, offset, maxSize, addOffset);
}
void EnzymeTypeTreeInsertEq(CTypeTreeRef CTT, const int64_t *indices,
                            size_t len, CConcreteType ct, LLVMContextRef ctx) {
  std::vector<int> seq;
  for (size_t i = 0; i < len; i++) {
    seq.push_back(indices[i]);
  }
  ((TypeTree *)CTT)->insert(seq, eunwrap(ct, *unwrap(ctx)));
}
const char *EnzymeTypeTreeToString(CTypeTreeRef src) {
  std::string tmp = ((TypeTree *)src)->str();
  char *cstr = new char[tmp.length() + 1];
  std::strcpy(cstr, tmp.c_str());
 
  return cstr;
}
 
// TODO deprecated
void EnzymeTypeTreeToStringFree(const char *cstr) { delete[] cstr; }
 
const char *EnzymeTypeAnalyzerToString(void *src) {
  auto TA = (TypeAnalyzer *)src;
  std::string str;
  raw_string_ostream ss(str);
  TA->dump(ss);
  ss.str();
  char *cstr = new char[str.length() + 1];
  std::strcpy(cstr, str.c_str());
  return cstr;
}
 
const char *EnzymeGradientUtilsInvertedPointersToString(GradientUtils *gutils,
                                                        void *src) {
  std::string str;
  raw_string_ostream ss(str);
  for (auto z : gutils->invertedPointers) {
    ss << "available inversion for " << *z.first << " of " << *z.second << "\n";
  }
  ss.str();
  char *cstr = new char[str.length() + 1];
  std::strcpy(cstr, str.c_str());
  return cstr;
}
 
LLVMValueRef EnzymeGradientUtilsCallWithInvertedBundles(
    GradientUtils *gutils, LLVMValueRef func, LLVMTypeRef funcTy,
    LLVMValueRef *args_vr, uint64_t args_size, LLVMValueRef orig_vr,
    CValueType *valTys, uint64_t valTys_size, LLVMBuilderRef B,
    uint8_t lookup) {
  auto orig = cast<CallInst>(unwrap(orig_vr));
 
  ArrayRef<ValueType> ar((ValueType *)valTys, valTys_size);
 
  IRBuilder<> &BR = *unwrap(B);
 
  auto Defs = gutils->getInvertedBundles(orig, ar, BR, lookup != 0);
 
  SmallVector<Value *, 1> args;
  for (size_t i = 0; i < args_size; i++) {
    args.push_back(unwrap(args_vr[i]));
  }
 
  auto callval = unwrap(func);
 
  auto res =
      BR.CreateCall(cast<FunctionType>(unwrap(funcTy)), callval, args, Defs);
  return wrap(res);
}
 
void EnzymeStringFree(const char *cstr) { delete[] cstr; }
 
void EnzymeMoveBefore(LLVMValueRef inst1, LLVMValueRef inst2,
                      LLVMBuilderRef B) {
  Instruction *I1 = cast<Instruction>(unwrap(inst1));
  Instruction *I2 = cast<Instruction>(unwrap(inst2));
  if (I1 != I2) {
    if (B != nullptr) {
      IRBuilder<> &BR = *unwrap(B);
      if (I1->getIterator() == BR.GetInsertPoint()) {
        if (I2->getNextNode() == nullptr)
          BR.SetInsertPoint(I1->getParent());
        else
          BR.SetInsertPoint(I1->getNextNode());
      }
    }
    I1->moveBefore(I2);
  }
}
 
void EnzymeSetStringMD(LLVMValueRef Inst, const char *Kind, LLVMValueRef Val) {
  MDNode *N = Val ? extractMDNode(unwrap<MetadataAsValue>(Val)) : nullptr;
  Value *V = unwrap(Inst);
  if (auto I = dyn_cast<Instruction>(V))
    I->setMetadata(Kind, N);
  else
    cast<GlobalVariable>(V)->setMetadata(Kind, N);
}
 
LLVMValueRef EnzymeGetStringMD(LLVMValueRef Inst, const char *Kind) {
  auto *I = unwrap<Instruction>(Inst);
  assert(I && "Expected instruction");
  if (auto *MD = I->getMetadata(Kind))
    return wrap(MetadataAsValue::get(I->getContext(), MD));
  return nullptr;
}
 
void EnzymeSetMustCache(LLVMValueRef inst1) {
  Instruction *I1 = cast<Instruction>(unwrap(inst1));
  I1->setMetadata("enzyme_mustcache", MDNode::get(I1->getContext(), {}));
}
 
uint8_t EnzymeHasFromStack(LLVMValueRef inst1) {
  Instruction *I1 = cast<Instruction>(unwrap(inst1));
  return hasMetadata(I1, "enzyme_fromstack") != 0;
}
 
void EnzymeCloneFunctionDISubprogramInto(LLVMValueRef NF, LLVMValueRef F) {
  auto &OldFunc = *cast<Function>(unwrap(F));
  auto &NewFunc = *cast<Function>(unwrap(NF));
  auto OldSP = OldFunc.getSubprogram();
  if (!OldSP)
    return;
  DIBuilder DIB(*OldFunc.getParent(), /*AllowUnresolved=*/false,
                OldSP->getUnit());
  auto SPType = DIB.createSubroutineType(DIB.getOrCreateTypeArray({}));
  DISubprogram::DISPFlags SPFlags = DISubprogram::SPFlagDefinition |
                                    DISubprogram::SPFlagOptimized |
                                    DISubprogram::SPFlagLocalToUnit;
  auto NewSP = DIB.createFunction(
      OldSP->getUnit(), NewFunc.getName(), NewFunc.getName(), OldSP->getFile(),
      /*LineNo=*/0, SPType, /*ScopeLine=*/0, DINode::FlagZero, SPFlags);
  NewFunc.setSubprogram(NewSP);
  DIB.finalizeSubprogram(NewSP);
  return;
}
 
void EnzymeReplaceFunctionImplementation(LLVMModuleRef M) {
  ReplaceFunctionImplementation(*unwrap(M));
}
 
void EnzymeDetectReadonlyOrThrow(LLVMModuleRef M) {
  DetectReadonlyOrThrow(*unwrap(M));
}
 
void EnzymeDumpModuleRef(LLVMModuleRef M) {
  llvm::errs() << *unwrap(M) << "\n";
}
 
void EnzymeDumpValueRef(LLVMValueRef M) { llvm::errs() << *unwrap(M) << "\n"; }
 
void EnzymeDumpTypeRef(LLVMTypeRef M) { llvm::errs() << *unwrap(M) << "\n"; }
 
static bool runAttributorOnFunctions(InformationCache &InfoCache,
                                     SetVector<Function *> &Functions,
                                     AnalysisGetter &AG,
                                     CallGraphUpdater &CGUpdater,
                                     bool DeleteFns, bool IsModulePass) {
  if (Functions.empty())
    return false;
 
  // Create an Attributor and initially empty information cache that is filled
  // while we identify default attribute opportunities.
  AttributorConfig AC(CGUpdater);
  AC.RewriteSignatures = false;
  AC.IsModulePass = IsModulePass;
  AC.DeleteFns = DeleteFns;
  Attributor A(Functions, InfoCache, AC);
 
  for (Function *F : Functions) {
    // Populate the Attributor with abstract attribute opportunities in the
    // function and the information cache with IR information.
    A.identifyDefaultAbstractAttributes(*F);
  }
 
  ChangeStatus Changed = A.run();
 
  return Changed == ChangeStatus::CHANGED;
}
 
extern "C" void RunAttributorOnModule(LLVMModuleRef M0) {
  auto &M = *unwrap(M0);
  AnalysisGetter AG;
  SetVector<Function *> Functions;
  for (Function &F : M)
    Functions.insert(&F);
 
  CallGraphUpdater CGUpdater;
  BumpPtrAllocator Allocator;
  InformationCache InfoCache(M, AG, Allocator, /* CGSCC */ nullptr);
  runAttributorOnFunctions(InfoCache, Functions, AG, CGUpdater,
                           /* DeleteFns*/ true,
                           /* IsModulePass */ true);
}
 
struct MyAttributorLegacyPass : public ModulePass {
  static char ID;
 
  MyAttributorLegacyPass() : ModulePass(ID) {}
 
  bool runOnModule(Module &M) override {
    if (skipModule(M))
      return false;
 
    AnalysisGetter AG;
    SetVector<Function *> Functions;
    for (Function &F : M)
      Functions.insert(&F);
 
    CallGraphUpdater CGUpdater;
    BumpPtrAllocator Allocator;
    InformationCache InfoCache(M, AG, Allocator, /* CGSCC */ nullptr);
    return runAttributorOnFunctions(InfoCache, Functions, AG, CGUpdater,
                                    /* DeleteFns*/ true,
                                    /* IsModulePass */ true);
  }
 
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    // FIXME: Think about passes we will preserve and add them here.
    AU.addRequired<TargetLibraryInfoWrapperPass>();
  }
};
extern "C++" char MyAttributorLegacyPass::ID = 0;
void EnzymeAddAttributorLegacyPass(LLVMPassManagerRef PM) {
  unwrap(PM)->add(new MyAttributorLegacyPass());
}
 
LLVMMetadataRef EnzymeMakeNonConstTBAA(LLVMMetadataRef MD) {
  auto M = cast<MDNode>(unwrap(MD));
  if (M->getNumOperands() != 4)
    return MD;
  auto CAM = dyn_cast<ConstantAsMetadata>(M->getOperand(3));
  if (!CAM)
    return MD;
  if (!CAM->getValue()->isOneValue())
    return MD;
  SmallVector<Metadata *, 4> MDs;
  for (auto &M : M->operands())
    MDs.push_back(M);
  MDs[3] =
      ConstantAsMetadata::get(ConstantInt::get(CAM->getValue()->getType(), 0));
  return wrap(MDNode::get(M->getContext(), MDs));
}
void EnzymeCopyMetadata(LLVMValueRef inst1, LLVMValueRef inst2) {
  cast<Instruction>(unwrap(inst1))
      ->copyMetadata(*cast<Instruction>(unwrap(inst2)));
}
void EnzymeCopyAlignment(LLVMValueRef inst1, LLVMValueRef inst2) {
  cast<AllocaInst>(unwrap(inst1))
      ->setAlignment(cast<AllocaInst>(unwrap(inst2))->getAlign());
}
void EnzymeTakeName(LLVMValueRef inst1, LLVMValueRef inst2) {
  unwrap(inst1)->takeName(unwrap(inst2));
}
 
LLVMMetadataRef EnzymeAnonymousAliasScopeDomain(const char *str,
                                                LLVMContextRef ctx) {
  MDBuilder MDB(*unwrap(ctx));
  MDNode *scope = MDB.createAnonymousAliasScopeDomain(str);
  return wrap(scope);
}
LLVMMetadataRef EnzymeAnonymousAliasScope(LLVMMetadataRef domain,
                                          const char *str) {
  auto dom = cast<MDNode>(unwrap(domain));
  MDBuilder MDB(dom->getContext());
  MDNode *scope = MDB.createAnonymousAliasScope(dom, str);
  return wrap(scope);
}
uint8_t EnzymeLowerSparsification(LLVMValueRef F, uint8_t replaceAll) {
  return LowerSparsification(cast<Function>(unwrap(F)), replaceAll != 0);
}
 
void EnzymeAttributeKnownFunctions(LLVMValueRef FC) {
  attributeKnownFunctions(*cast<Function>(unwrap(FC)));
}
 
void EnzymeSetCalledFunction(LLVMValueRef C_CI, LLVMValueRef C_F,
                             uint64_t *argrem, uint64_t num_argrem) {
  auto CI = cast<CallInst>(unwrap(C_CI));
  auto F = cast<Function>(unwrap(C_F));
  auto Attrs = CI->getAttributes();
  AttributeList NewAttrs;
 
  if (CI->getType() == F->getReturnType()) {
    for (auto attr : Attrs.getAttributes(AttributeList::ReturnIndex))
      NewAttrs = NewAttrs.addAttribute(F->getContext(),
                                       AttributeList::ReturnIndex, attr);
  }
  for (auto attr : Attrs.getAttributes(AttributeList::FunctionIndex))
    NewAttrs = NewAttrs.addAttribute(F->getContext(),
                                     AttributeList::FunctionIndex, attr);
 
  size_t argremsz = 0;
  size_t nexti = 0;
  SmallVector<Value *, 1> vals;
  for (size_t i = 0, end = CI->arg_size(); i < end; i++) {
    if (argremsz < num_argrem) {
      if (i == argrem[argremsz]) {
        argremsz++;
        continue;
      }
    }
    for (auto attr : Attrs.getAttributes(AttributeList::FirstArgIndex + i))
      NewAttrs = NewAttrs.addAttribute(
          F->getContext(), AttributeList::FirstArgIndex + nexti, attr);
    vals.push_back(CI->getArgOperand(i));
    nexti++;
  }
  assert(argremsz == num_argrem);
 
  IRBuilder<> B(CI);
  SmallVector<OperandBundleDef, 1> Bundles;
  for (unsigned I = 0, E = CI->getNumOperandBundles(); I != E; ++I)
    Bundles.emplace_back(CI->getOperandBundleAt(I));
  auto NC = B.CreateCall(F, vals, Bundles);
  NC->setAttributes(NewAttrs);
  NC->copyMetadata(*CI);
 
  if (CI->getType() == F->getReturnType())
    CI->replaceAllUsesWith(NC);
 
  if (!NC->getType()->isVoidTy())
    NC->takeName(CI);
  NC->setCallingConv(CI->getCallingConv());
  CI->eraseFromParent();
}
 
// clones a function to now miss the return or args
LLVMValueRef EnzymeCloneFunctionWithoutReturnOrArgs(LLVMValueRef FC,
                                                    uint8_t keepReturnU,
                                                    uint64_t *argrem,
                                                    uint64_t num_argrem) {
  auto F = cast<Function>(unwrap(FC));
  auto FT = F->getFunctionType();
  bool keepReturn = keepReturnU != 0;
 
  size_t argremsz = 0;
  size_t nexti = 0;
  SmallVector<Type *, 1> types;
  auto Attrs = F->getAttributes();
  AttributeList NewAttrs;
  for (size_t i = 0, end = FT->getNumParams(); i < end; i++) {
    if (argremsz < num_argrem) {
      if (i == argrem[argremsz]) {
        argremsz++;
        continue;
      }
    }
    for (auto attr : Attrs.getAttributes(AttributeList::FirstArgIndex + i))
      NewAttrs = NewAttrs.addAttribute(
          F->getContext(), AttributeList::FirstArgIndex + nexti, attr);
    types.push_back(F->getFunctionType()->getParamType(i));
    nexti++;
  }
  if (keepReturn) {
    for (auto attr : Attrs.getAttributes(AttributeList::ReturnIndex))
      NewAttrs = NewAttrs.addAttribute(F->getContext(),
                                       AttributeList::ReturnIndex, attr);
  }
  for (auto attr : Attrs.getAttributes(AttributeList::FunctionIndex))
    NewAttrs = NewAttrs.addAttribute(F->getContext(),
                                     AttributeList::FunctionIndex, attr);
 
  FunctionType *FTy = FunctionType::get(
      keepReturn ? F->getReturnType() : Type::getVoidTy(F->getContext()), types,
      FT->isVarArg());
 
  // Create the new function
  Function *NewF = Function::Create(FTy, F->getLinkage(), F->getAddressSpace(),
                                    F->getName(), F->getParent());
 
  ValueToValueMapTy VMap;
  // Loop over the arguments, copying the names of the mapped arguments over...
  nexti = 0;
  argremsz = 0;
  Function::arg_iterator DestI = NewF->arg_begin();
  for (const Argument &I : F->args()) {
    if (argremsz < num_argrem) {
      if (I.getArgNo() == argrem[argremsz]) {
        VMap[&I] = UndefValue::get(I.getType());
        argremsz++;
        continue;
      }
    }
    DestI->setName(I.getName()); // Copy the name over...
    VMap[&I] = &*DestI++;        // Add mapping to VMap
  }
 
  SmallVector<ReturnInst *, 8> Returns; // Ignore returns cloned.
  CloneFunctionInto(NewF, F, VMap, CloneFunctionChangeType::LocalChangesOnly,
                    Returns, "", nullptr);
 
  if (!keepReturn) {
    for (auto &B : *NewF) {
      if (auto RI = dyn_cast<ReturnInst>(B.getTerminator())) {
        IRBuilder<> B(RI);
        auto NRI = B.CreateRetVoid();
        NRI->copyMetadata(*RI);
        RI->eraseFromParent();
      }
    }
  }
  NewF->setAttributes(NewAttrs);
  if (!keepReturn)
    for (auto &Arg : NewF->args())
      Arg.removeAttr(Attribute::Returned);
  SmallVector<std::pair<unsigned, MDNode *>, 1> MD;
  F->getAllMetadata(MD);
  for (auto pair : MD)
    if (pair.first != LLVMContext::MD_dbg)
      NewF->addMetadata(pair.first, *pair.second);
  NewF->takeName(F);
  NewF->setCallingConv(F->getCallingConv());
  if (!keepReturn)
    NewF->addFnAttr("enzyme_retremove", "");
 
  if (num_argrem) {
    SmallVector<uint64_t, 1> previdx;
    if (Attrs.hasAttribute(AttributeList::FunctionIndex, "enzyme_parmremove")) {
      auto attr =
          Attrs.getAttribute(AttributeList::FunctionIndex, "enzyme_parmremove");
      auto prevstr = attr.getValueAsString();
      SmallVector<StringRef, 1> sub;
      prevstr.split(sub, ",");
      for (auto s : sub) {
        uint64_t ival;
        bool b = s.getAsInteger(10, ival);
        (void)b;
        assert(!b);
        previdx.push_back(ival);
      }
    }
    SmallVector<uint64_t, 1> nextidx;
    for (size_t i = 0; i < num_argrem; i++) {
      auto val = argrem[i];
      nextidx.push_back(val);
    }
 
    size_t prevcnt = 0;
    size_t nextcnt = 0;
    SmallVector<uint64_t, 1> out;
    while (prevcnt < previdx.size() && nextcnt < nextidx.size()) {
      if (previdx[prevcnt] <= nextidx[nextcnt] + prevcnt) {
        out.push_back(previdx[prevcnt]);
        prevcnt++;
      } else {
        out.push_back(nextidx[nextcnt] + prevcnt);
        nextcnt++;
      }
    }
    while (prevcnt < previdx.size()) {
      out.push_back(previdx[prevcnt]);
      prevcnt++;
    }
    while (nextcnt < nextidx.size()) {
      out.push_back(nextidx[nextcnt] + prevcnt);
      nextcnt++;
    }
 
    std::string remstr;
    for (auto arg : out) {
      if (remstr.size())
        remstr += ",";
      remstr += std::to_string(arg);
    }
 
    NewF->addFnAttr("enzyme_parmremove", remstr);
  }
  return wrap(NewF);
}
LLVMTypeRef EnzymeAllocaType(LLVMValueRef V) {
  return wrap(cast<AllocaInst>(unwrap(V))->getAllocatedType());
}
LLVMValueRef EnzymeComputeByteOffsetOfGEP(LLVMBuilderRef B_r, LLVMValueRef V_r,
                                          LLVMTypeRef T_r) {
  IRBuilder<> &B = *unwrap(B_r);
  auto T = cast<IntegerType>(unwrap(T_r));
  auto width = T->getBitWidth();
  auto uw = unwrap(V_r);
  GEPOperator *gep = isa<GetElementPtrInst>(uw)
                         ? cast<GEPOperator>(cast<GetElementPtrInst>(uw))
                         : cast<GEPOperator>(cast<ConstantExpr>(uw));
  auto &DL = B.GetInsertBlock()->getParent()->getParent()->getDataLayout();
 
#if LLVM_VERSION_MAJOR >= 20
  SmallMapVector<Value *, APInt, 4> VariableOffsets;
#else
  MapVector<Value *, APInt> VariableOffsets;
#endif
  APInt Offset(width, 0);
  bool success = collectOffset(gep, DL, width, VariableOffsets, Offset);
  (void)success;
  assert(success);
  Value *start = ConstantInt::get(T, Offset);
  for (auto &pair : VariableOffsets)
    start = B.CreateAdd(
        start, B.CreateMul(pair.first, ConstantInt::get(T, pair.second)));
  return wrap(start);
}
}
 
extern "C" {
 
LLVMValueRef EnzymeBuildExtractValue(LLVMBuilderRef B, LLVMValueRef AggVal,
                                     unsigned *Index, unsigned Size,
                                     const char *Name) {
  return wrap(unwrap(B)->CreateExtractValue(
      unwrap(AggVal), ArrayRef<unsigned>(Index, Size), Name));
}
 
LLVMValueRef EnzymeBuildInsertValue(LLVMBuilderRef B, LLVMValueRef AggVal,
                                    LLVMValueRef EltVal, unsigned *Index,
                                    unsigned Size, const char *Name) {
  return wrap(unwrap(B)->CreateInsertValue(
      unwrap(AggVal), unwrap(EltVal), ArrayRef<unsigned>(Index, Size), Name));
}
}