enzyme/Enzyme/GradientUtils.cpp

//===- GradientUtils.cpp - Helper class and utilities for AD     ---------===//
//
//                             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 define two helper classes GradientUtils and subclass
// DiffeGradientUtils. These classes contain utilities for managing the cache,
// recomputing statements, and in the case of DiffeGradientUtils, managing
// adjoint values and shadow pointers.
//
//===----------------------------------------------------------------------===//

#include <algorithm>
#include <functional>
#include <map>
#include <string>

#include "GradientUtils.h"
#include "MustExitScalarEvolution.h"
#include "Utils.h"

#include "DifferentialUseAnalysis.h"
#include "LibraryFuncs.h"
#include "TypeAnalysis/TBAA.h"

#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/DebugInfoMetadata.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"

#include "llvm/Support/AMDGPUMetadata.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/TimeProfiler.h"

#if LLVM_VERSION_MAJOR >= 14
#define addAttribute addAttributeAtIndex
#define hasAttribute hasAttributeAtIndex
#endif

using namespace llvm;

StringMap<std::function<Value *(IRBuilder<> &, CallInst *, ArrayRef<Value *>,
                                GradientUtils *)>>
    shadowHandlers;
StringMap<std::function<CallInst *(IRBuilder<> &, Value *)>> shadowErasers;

StringMap<
    std::pair<std::function<bool(IRBuilder<> &, CallInst *, GradientUtils &,
                                 Value *&, Value *&, Value *&)>,
              std::function<void(IRBuilder<> &, CallInst *,
                                 DiffeGradientUtils &, Value *)>>>
    customCallHandlers;

StringMap<std::function<bool(IRBuilder<> &, CallInst *, GradientUtils &,
                             Value *&, Value *&)>>
    customFwdCallHandlers;

extern "C" {
llvm::cl::opt<bool>
    EnzymeNewCache("enzyme-new-cache", cl::init(true), cl::Hidden,
                   cl::desc("Use new cache decision algorithm"));

llvm::cl::opt<bool> EnzymeMinCutCache("enzyme-mincut-cache", cl::init(true),
                                      cl::Hidden,
                                      cl::desc("Use Enzyme Mincut algorithm"));

llvm::cl::opt<bool> EnzymeLoopInvariantCache(
    "enzyme-loop-invariant-cache", cl::init(true), cl::Hidden,
    cl::desc("Attempt to hoist cache outside of loop"));

llvm::cl::opt<bool> EnzymeInactiveDynamic(
    "enzyme-inactive-dynamic", cl::init(true), cl::Hidden,
    cl::desc("Force wholy inactive dynamic loops to have 0 iter reverse pass"));

llvm::cl::opt<bool>
    EnzymeSharedForward("enzyme-shared-forward", cl::init(false), cl::Hidden,
                        cl::desc("Forward Shared Memory from definitions"));

llvm::cl::opt<bool>
    EnzymeRegisterReduce("enzyme-register-reduce", cl::init(false), cl::Hidden,
                         cl::desc("Reduce the amount of register reduce"));
llvm::cl::opt<bool>
    EnzymeSpeculatePHIs("enzyme-speculate-phis", cl::init(false), cl::Hidden,
                        cl::desc("Speculatively execute phi computations"));
llvm::cl::opt<bool> EnzymeFreeInternalAllocations(
    "enzyme-free-internal-allocations", cl::init(true), cl::Hidden,
    cl::desc("Always free internal allocations (disable if allocation needs "
             "access outside)"));

llvm::cl::opt<bool>
    EnzymeRematerialize("enzyme-rematerialize", cl::init(true), cl::Hidden,
                        cl::desc("Rematerialize allocations/shadows in the "
                                 "reverse rather than caching"));

llvm::cl::opt<bool>
    EnzymeVectorSplitPhi("enzyme-vector-split-phi", cl::init(true), cl::Hidden,
                         cl::desc("Split phis according to vector size"));

llvm::cl::opt<bool>
    EnzymePrintDiffUse("enzyme-print-diffuse", cl::init(false), cl::Hidden,
                       cl::desc("Print differential use analysis"));
}

SmallVector<unsigned int, 9> MD_ToCopy = {
    LLVMContext::MD_dbg,
    LLVMContext::MD_tbaa,
    LLVMContext::MD_tbaa_struct,
    LLVMContext::MD_range,
    LLVMContext::MD_nonnull,
    LLVMContext::MD_dereferenceable,
    LLVMContext::MD_dereferenceable_or_null};

static bool isPotentialLastLoopValue(llvm::Value *val,
                                     const llvm::BasicBlock *loc,
                                     const llvm::LoopInfo &LI) {
  if (llvm::Instruction *inst = llvm::dyn_cast<llvm::Instruction>(val)) {
    const llvm::Loop *InstLoop = LI.getLoopFor(inst->getParent());
    if (InstLoop == nullptr) {
      return false;
    }
    for (const llvm::Loop *L = LI.getLoopFor(loc); L; L = L->getParentLoop()) {
      if (L == InstLoop)
        return false;
    }
    return true;
  }
  return false;
}

GradientUtils::GradientUtils(
    EnzymeLogic &Logic, Function *newFunc_, Function *oldFunc_,
    TargetLibraryInfo &TLI_, TypeAnalysis &TA_, TypeResults TR_,
    ValueToValueMapTy &invertedPointers_,
    const SmallPtrSetImpl<Value *> &constantvalues_,
    const SmallPtrSetImpl<Value *> &activevals_, DIFFE_TYPE ReturnActivity,
    bool shadowReturnUsed_, ArrayRef<DIFFE_TYPE> ArgDiffeTypes_,
    llvm::ValueMap<const llvm::Value *, AssertingReplacingVH> &originalToNewFn_,
    DerivativeMode mode, bool runtimeActivity, bool strongZero, unsigned width,
    bool omp)
    : CacheUtility(TLI_, newFunc_), Logic(Logic), mode(mode), oldFunc(oldFunc_),
      invertedPointers(),
      OrigDT(oldFunc_->empty()
                 ? ((DominatorTree *)nullptr)
                 : &Logic.PPC.FAM.getResult<llvm::DominatorTreeAnalysis>(
                       *oldFunc_)),
      OrigPDT(oldFunc_->empty()
                  ? ((PostDominatorTree *)nullptr)
                  : &Logic.PPC.FAM.getResult<llvm::PostDominatorTreeAnalysis>(
                        *oldFunc_)),
      OrigLI(oldFunc_->empty()
                 ? ((LoopInfo *)nullptr)
                 : &Logic.PPC.FAM.getResult<llvm::LoopAnalysis>(*oldFunc_)),
      OrigSE(oldFunc_->empty()
                 ? ((ScalarEvolution *)nullptr)
                 : &Logic.PPC.FAM.getResult<llvm::ScalarEvolutionAnalysis>(
                       *oldFunc_)),
      notForAnalysis(getGuaranteedUnreachable(oldFunc_)),
      ATA(oldFunc_->empty()
              ? nullptr
              : new ActivityAnalyzer(
                    Logic.PPC, Logic.PPC.getAAResultsFromFunction(oldFunc_),
                    notForAnalysis, TLI_, constantvalues_, activevals_,
                    ReturnActivity)),
      overwritten_args_map_ptr(nullptr), unnecessaryValuesP(nullptr),
      tid(nullptr), numThreads(nullptr),
      OrigAA(oldFunc_->empty() ? ((AAResults *)nullptr)
                               : &Logic.PPC.getAAResultsFromFunction(oldFunc_)),
      TA(TA_), TR(TR_), omp(omp), runtimeActivity(runtimeActivity),
      strongZero(strongZero), width(width), shadowReturnUsed(shadowReturnUsed_),
      ArgDiffeTypes(ArgDiffeTypes_) {
  if (oldFunc_->empty())
    return;
  if (oldFunc_->getSubprogram()) {
    assert(originalToNewFn_.hasMD());
  }

  for (BasicBlock &BB : *oldFunc) {
    for (Instruction &I : BB) {
      if (auto CI = dyn_cast<CallInst>(&I)) {
        originalCalls.push_back(CI);
      }
    }
  }

  originalToNewFn.getMDMap() = originalToNewFn_.getMDMap();

  if (oldFunc_->getSubprogram()) {
    assert(originalToNewFn.hasMD());
  }
  for (auto pair : invertedPointers_) {
    invertedPointers.insert(std::make_pair(
        (const Value *)pair.first, InvertedPointerVH(this, pair.second)));
  }
  originalToNewFn.insert(originalToNewFn_.begin(), originalToNewFn_.end());
  for (BasicBlock &oBB : *oldFunc) {
    for (Instruction &oI : oBB) {
      newToOriginalFn[originalToNewFn[&oI]] = &oI;
    }
    newToOriginalFn[originalToNewFn[&oBB]] = &oBB;
  }
  for (Argument &oArg : oldFunc->args()) {
    newToOriginalFn[originalToNewFn[&oArg]] = &oArg;
  }
  for (BasicBlock &BB : *newFunc) {
    originalBlocks.push_back(&BB);
  }
  tape = nullptr;
  tapeidx = 0;
  assert(originalBlocks.size() > 0);

  SmallVector<BasicBlock *, 4> ReturningBlocks;
  for (BasicBlock &BB : *oldFunc) {
    if (isa<ReturnInst>(BB.getTerminator()))
      ReturningBlocks.push_back(&BB);
  }
  for (BasicBlock &BB : *oldFunc) {
    bool legal = true;
    for (auto BRet : ReturningBlocks) {
      if (!(BRet == &BB || OrigDT->dominates(&BB, BRet))) {
        legal = false;
        break;
      }
    }
    if (legal)
      BlocksDominatingAllReturns.insert(&BB);
  }
}

// Whether a particular value is neded in rooting the reverse pass
bool GradientUtils::usedInRooting(const llvm::CallBase *orig,
                                  ArrayRef<ValueType> types,
                                  const llvm::Value *val, bool shadow) const {
  SmallVector<OperandBundleDef, 2> OrigDefs;
  orig->getOperandBundlesAsDefs(OrigDefs);
  SmallVector<OperandBundleDef, 2> Defs;
  for (auto bund : OrigDefs) {
    // Only handle jl_roots & gc-transition tags (for now).
    StringRef tag = bund.getTag();
    if (tag == "jl_roots") {
      bool anyPrimal = false;
      bool anyShadow = false;
      for (auto ty : types) {
        if (ty == ValueType::Primal || ty == ValueType::Both)
          anyPrimal = true;
        if (ty == ValueType::Shadow || ty == ValueType::Both)
          anyShadow = true;
      }

      // In the future we can reduce the number of roots
      // we preserve by identifying which operands they
      // correspond to. For now, fall back and preserve all
      // primals and shadows
      // assert(bund.inputs().size() == types.size());
      for (auto inp : bund.inputs()) {
        if (inp != val)
          continue;

        if (anyPrimal && !shadow)
          return true;
        if (anyShadow && shadow)
          return true;
      }
    } else if (tag == "gc-transition") {
      if (shadow)
        continue;
      for (auto inp : bund.inputs()) {
        if (inp != val)
          continue;
        return true;
      }
    } else {
      errs() << "unsupported tag " << bund.getTag() << " for " << *orig << "\n";
      llvm_unreachable("unsupported tag");
    }
  }
  return false;
}

SmallVector<OperandBundleDef, 2>
GradientUtils::getInvertedBundles(CallInst *orig, ArrayRef<ValueType> types,
                                  IRBuilder<> &Builder2, bool lookup,
                                  const ValueToValueMapTy &available) {
  assert(!(lookup && (mode == DerivativeMode::ForwardMode ||
                      mode == DerivativeMode::ForwardModeError)));

  SmallVector<OperandBundleDef, 2> OrigDefs;
  orig->getOperandBundlesAsDefs(OrigDefs);
  SmallVector<OperandBundleDef, 2> Defs;
  for (auto bund : OrigDefs) {
    // Only handle jl_roots & gc-transition tags (for now).
    StringRef tag = bund.getTag();
    if (tag == "jl_roots") {
      // In the future we can reduce the number of roots
      // we preserve by identifying which operands they
      // correspond to. For now, fall back and preserve all
      // primals and shadows
      // assert(bund.inputs().size() == types.size());
      bool anyPrimal = false;
      bool anyShadow = false;
      for (auto ty : types) {
        if (ty == ValueType::Primal || ty == ValueType::Both)
          anyPrimal = true;
        if (ty == ValueType::Shadow || ty == ValueType::Both)
          anyShadow = true;
      }

      SmallVector<Value *, 2> bunds;
      for (auto inp : bund.inputs()) {
        if (anyPrimal) {
          Value *newv = getNewFromOriginal(inp);
          if (lookup)
            newv = lookupM(newv, Builder2, available);
          bunds.push_back(newv);
        }
        if (anyShadow && !isConstantValue(inp)) {
          Value *shadow = invertPointerM(inp, Builder2);
          if (lookup)
            shadow = lookupM(shadow, Builder2);
          bunds.push_back(shadow);
        }
      }
      Defs.push_back(OperandBundleDef(tag.str(), bunds));
    } else if (tag == "gc-transition") {
      SmallVector<Value *, 2> bunds;
      for (auto inp : bund.inputs()) {
        Value *newv = getNewFromOriginal(inp);
        if (lookup)
          newv = lookupM(newv, Builder2, available);
        bunds.push_back(newv);
      }
      Defs.push_back(OperandBundleDef(tag.str(), bunds));
    } else {
      errs() << "unsupported tag " << tag << " for " << *orig << "\n";
      llvm_unreachable("unsupported tag");
    }
  }
  return Defs;
}

Value *GradientUtils::getNewIfOriginal(Value *originst) const {
  assert(originst);
  auto f = originalToNewFn.find(originst);
  if (f == originalToNewFn.end()) {
    return originst;
  }
  assert(f != originalToNewFn.end());
  if (f->second == nullptr) {
    llvm::errs() << *oldFunc << "\n";
    llvm::errs() << *newFunc << "\n";
    llvm::errs() << *originst << "\n";
  }
  assert(f->second);
  return f->second;
}

Value *GradientUtils::ompThreadId() {
  if (tid)
    return tid;
  IRBuilder<> B(inversionAllocs);

  auto FT = FunctionType::get(Type::getInt64Ty(B.getContext()),
                              ArrayRef<Type *>(), false);
  auto FN = newFunc->getParent()->getOrInsertFunction("omp_get_thread_num", FT);
  auto CI = B.CreateCall(FN);
  if (auto F = getFunctionFromCall(CI)) {
#if LLVM_VERSION_MAJOR >= 16
    F->setOnlyAccessesInaccessibleMemory();
    F->setOnlyReadsMemory();
#else
    F->addFnAttr(Attribute::InaccessibleMemOnly);
    F->addFnAttr(Attribute::ReadOnly);
#endif
  }
#if LLVM_VERSION_MAJOR >= 16
  CI->setOnlyAccessesInaccessibleMemory();
  CI->setOnlyReadsMemory();
#else
  CI->addAttribute(AttributeList::FunctionIndex,
                   Attribute::InaccessibleMemOnly);
  CI->addAttribute(AttributeList::FunctionIndex, Attribute::ReadOnly);
#endif
  return tid = CI;
}

Value *GradientUtils::ompNumThreads() {
  if (numThreads)
    return numThreads;
  IRBuilder<> B(inversionAllocs);

  auto FT = FunctionType::get(Type::getInt64Ty(B.getContext()),
                              ArrayRef<Type *>(), false);
  auto FN =
      newFunc->getParent()->getOrInsertFunction("omp_get_max_threads", FT);
  auto CI = B.CreateCall(FN);
  if (auto F = getFunctionFromCall(CI)) {
#if LLVM_VERSION_MAJOR >= 16
    F->setOnlyAccessesInaccessibleMemory();
    F->setOnlyReadsMemory();
#else
    F->addFnAttr(Attribute::InaccessibleMemOnly);
    F->addFnAttr(Attribute::ReadOnly);
#endif
  }
#if LLVM_VERSION_MAJOR >= 16
  CI->setOnlyAccessesInaccessibleMemory();
  CI->setOnlyReadsMemory();
#else
  CI->addAttribute(AttributeList::FunctionIndex,
                   Attribute::InaccessibleMemOnly);
  CI->addAttribute(AttributeList::FunctionIndex, Attribute::ReadOnly);
#endif
  return numThreads = CI;
}

Value *GradientUtils::getOrInsertTotalMultiplicativeProduct(Value *val,
                                                            LoopContext &lc) {
  // TODO optimize if val is invariant to loopContext
  assert(val->getType()->isFPOrFPVectorTy());
  for (auto &I : *lc.header) {
    if (auto PN = dyn_cast<PHINode>(&I)) {
      if (PN->getType() != val->getType())
        continue;
      if (fictiousPHIs.find(PN) != fictiousPHIs.end())
        continue;

      int Idx = PN->getBasicBlockIndex(lc.preheader);
      if (Idx < 0) {

        std::string str;
        raw_string_ostream ss(str);

        ss << " Could not find block for index, PN: " << *PN << "\n";
        ss << " preheader: " << *lc.preheader << "\n";
        ss << " header: " << *lc.header << "\n";
        ss << " fn: " << *lc.header->getParent() << "\n";

        if (CustomErrorHandler) {
          CustomErrorHandler(str.c_str(), wrap(PN), ErrorType::InternalError,
                             nullptr, nullptr, nullptr);
        } else {
          EmitFailure("GetIndexError", PN->getDebugLoc(), PN, ss.str());
        }
      }
      Value *ival = PN->getIncomingValue(Idx);
      if (auto CDV = dyn_cast<ConstantDataVector>(ival)) {
        if (CDV->isSplat())
          ival = CDV->getSplatValue();
      }
      if (auto C = dyn_cast<ConstantFP>(ival)) {
        if (!C->isExactlyValue(APFloat(C->getType()->getFltSemantics(), "1"))) {
          continue;
        }
      } else
        continue;
      for (auto IB : PN->blocks()) {
        if (IB == lc.preheader)
          continue;

        if (auto BO =
                dyn_cast<BinaryOperator>(PN->getIncomingValueForBlock(IB))) {
          if (BO->getOpcode() != BinaryOperator::FMul)
            goto continueOutermost;
          if (BO->getOperand(0) == PN && BO->getOperand(1) == val)
            return BO;
          if (BO->getOperand(1) == PN && BO->getOperand(0) == val)
            return BO;
        } else
          goto continueOutermost;
      }
    } else
      break;
  continueOutermost:;
  }

  IRBuilder<> lbuilder(lc.header, lc.header->begin());
  auto PN = lbuilder.CreatePHI(val->getType(), 2);
  Constant *One = ConstantFP::get(val->getType()->getScalarType(), "1");
  if (VectorType *VTy = dyn_cast<VectorType>(val->getType())) {
    One = ConstantVector::getSplat(VTy->getElementCount(), One);
  }
  PN->addIncoming(One, lc.preheader);
  lbuilder.SetInsertPoint(lc.header->getFirstNonPHI());
  if (auto inst = dyn_cast<Instruction>(val)) {
    if (DT.dominates(PN, inst))
      lbuilder.SetInsertPoint(inst->getNextNode());
  }
  Value *red = lbuilder.CreateFMul(PN, val);
  for (auto pred : predecessors(lc.header)) {
    if (pred == lc.preheader)
      continue;
    PN->addIncoming(red, pred);
  }
  return red;
}

Value *GradientUtils::getOrInsertConditionalIndex(Value *val, LoopContext &lc,
                                                  bool pickTrue) {
  assert(val->getType()->isIntOrIntVectorTy(1));
  // TODO optimize if val is invariant to loopContext
  for (auto &I : *lc.header) {
    if (auto PN = dyn_cast<PHINode>(&I)) {
      if (PN->getNumIncomingValues() == 0)
        continue;
      if (PN->getType() != lc.incvar->getType())
        continue;
      if (fictiousPHIs.find(PN) != fictiousPHIs.end())
        continue;
      Value *ival = PN->getIncomingValueForBlock(lc.preheader);
      if (auto C = dyn_cast<Constant>(ival)) {
        if (!C->isNullValue()) {
          continue;
        }
      } else
        continue;
      for (auto IB : PN->blocks()) {
        if (IB == lc.preheader)
          continue;

        if (auto SI = dyn_cast<SelectInst>(PN->getIncomingValueForBlock(IB))) {
          if (SI->getCondition() != val)
            goto continueOutermost;
          if (pickTrue && SI->getFalseValue() == PN) {
            // TODO handle vector of
            if (SI->getTrueValue() == lc.incvar)
              return SI;
          }
          if (!pickTrue && SI->getTrueValue() == PN) {
            // TODO handle vector of
            if (SI->getFalseValue() == lc.incvar)
              return SI;
          }
        } else
          goto continueOutermost;
      }
    } else
      break;
  continueOutermost:;
  }

  IRBuilder<> lbuilder(lc.header, lc.header->begin());
  auto PN = lbuilder.CreatePHI(lc.incvar->getType(), 2);
  Constant *Zero =
      Constant::getNullValue(lc.incvar->getType()->getScalarType());
  PN->addIncoming(Zero, lc.preheader);
  lbuilder.SetInsertPoint(lc.incvar->getNextNode());
  Value *red = lc.incvar;
  if (VectorType *VTy = dyn_cast<VectorType>(val->getType())) {
#if LLVM_VERSION_MAJOR >= 12
    red = lbuilder.CreateVectorSplat(VTy->getElementCount(), red);
#else
    red = lbuilder.CreateVectorSplat(VTy->getNumElements(), red);
#endif
  }
  if (auto inst = dyn_cast<Instruction>(val)) {
    if (DT.dominates(PN, inst))
      lbuilder.SetInsertPoint(inst->getNextNode());
  }
  assert(red->getType() == PN->getType());
  red = lbuilder.CreateSelect(val, pickTrue ? red : PN, pickTrue ? PN : red);
  for (auto pred : predecessors(lc.header)) {
    if (pred == lc.preheader)
      continue;
    PN->addIncoming(red, pred);
  }
  return red;
}

bool GradientUtils::assumeDynamicLoopOfSizeOne(Loop *L) const {
  if (!EnzymeInactiveDynamic)
    return false;
  auto OL = OrigLI->getLoopFor(isOriginal(L->getHeader()));
  assert(OL);
  for (auto OB : OL->getBlocks()) {
    for (auto &OI : *OB) {
      if (!isConstantInstruction(&OI))
        return false;
      if (auto SI = dyn_cast<StoreInst>(&OI)) {
        if (!isConstantValue(SI->getPointerOperand()))
          return false;
      }
      if (auto MTI = dyn_cast<MemTransferInst>(&OI)) {
        if (!isConstantValue(MTI->getArgOperand(0)))
          return false;
      }
    }
  }
  return true;
}

DebugLoc GradientUtils::getNewFromOriginal(const DebugLoc L) const {
  if (L.get() == nullptr)
    return nullptr;
  if (!oldFunc->getSubprogram())
    return L;
  assert(originalToNewFn.hasMD());
  auto opt = originalToNewFn.getMappedMD(L.getAsMDNode());
  if (!opt)
    return L;
  assert(opt);
#if LLVM_VERSION_MAJOR >= 23
  return DebugLoc(cast<DILocation>(*opt));
#else
  return DebugLoc(cast<MDNode>(*opt));
#endif
}

Value *GradientUtils::getNewFromOriginal(const Value *originst) const {
  assert(originst);
  if (isa<ConstantData>(originst))
    return const_cast<Value *>(originst);
  auto f = originalToNewFn.find(originst);
  if (f == originalToNewFn.end()) {
    errs() << *oldFunc << "\n";
    errs() << *newFunc << "\n";
    dumpMap(originalToNewFn, [&](const Value *const &v) -> bool {
      if (isa<Instruction>(originst))
        return isa<Instruction>(v);
      if (isa<BasicBlock>(originst))
        return isa<BasicBlock>(v);
      if (isa<Function>(originst))
        return isa<Function>(v);
      if (isa<Argument>(originst))
        return isa<Argument>(v);
      if (isa<Constant>(originst))
        return isa<Constant>(v);
      return true;
    });
    llvm::errs() << *originst << "\n";
  }
  assert(f != originalToNewFn.end());
  if (f->second == nullptr) {
    errs() << *oldFunc << "\n";
    errs() << *newFunc << "\n";
    errs() << *originst << "\n";
  }
  assert(f->second);
  return f->second;
}

Instruction *
GradientUtils::getNewFromOriginal(const Instruction *newinst) const {
  auto ninst = getNewFromOriginal((Value *)newinst);
  if (!isa<Instruction>(ninst)) {
    errs() << *oldFunc << "\n";
    errs() << *newFunc << "\n";
    errs() << *ninst << " - " << *newinst << "\n";
  }
  return cast<Instruction>(ninst);
}

BasicBlock *GradientUtils::getNewFromOriginal(const BasicBlock *newinst) const {
  return cast<BasicBlock>(getNewFromOriginal((Value *)newinst));
}

Value *GradientUtils::hasUninverted(const Value *inverted) const {
  for (auto v : invertedPointers) {
    if (v.second == inverted)
      return const_cast<Value *>(v.first);
  }
  return nullptr;
}

BasicBlock *GradientUtils::getOriginalFromNew(const BasicBlock *newinst) const {
  assert(newinst->getParent() == newFunc);
  auto found = newToOriginalFn.find(newinst);
  assert(found != newToOriginalFn.end());
  Value *res = found->second;
  return cast<BasicBlock>(res);
}

Value *GradientUtils::isOriginal(const Value *newinst) const {
  if (isa<Constant>(newinst) || isa<UndefValue>(newinst))
    return const_cast<Value *>(newinst);
#ifndef NDEBUG
  if (auto arg = dyn_cast<Argument>(newinst)) {
    assert(arg->getParent() == newFunc);
  }
  if (auto inst = dyn_cast<Instruction>(newinst)) {
    assert(inst->getParent()->getParent() == newFunc);
  }
#endif
  auto found = newToOriginalFn.find(newinst);
  if (found == newToOriginalFn.end())
    return nullptr;
  return found->second;
}

Instruction *GradientUtils::isOriginal(const Instruction *newinst) const {
  return cast_or_null<Instruction>(isOriginal((const Value *)newinst));
}

BasicBlock *GradientUtils::isOriginal(const BasicBlock *newinst) const {
  return cast_or_null<BasicBlock>(isOriginal((const Value *)newinst));
}

Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
                              const ValueToValueMapTy &available,
                              UnwrapMode unwrapMode, BasicBlock *scope,
                              bool permitCache) {
  assert(val);
  assert(val->getName() != "<badref>");
  assert(val->getType());

  for (auto pair : available) {
    assert(pair.first);
    assert(pair.first->getType());
    if (pair.second) {
      assert(pair.second->getType());
      assert(pair.first->getType() == pair.second->getType());
    }
  }

  if (isa<LoadInst>(val) &&
      cast<LoadInst>(val)->getMetadata("enzyme_mustcache")) {
    return val;
  }

  if (available.count(val)) {
    auto avail = available.lookup(val);
    assert(avail->getType());
    if (avail->getType() != val->getType()) {
      llvm::errs() << "val: " << *val << "\n";
      llvm::errs() << "available[val]: " << *available.lookup(val) << "\n";
    }
    assert(available.lookup(val)->getType() == val->getType());
    return available.lookup(val);
  }

  if (auto inst = dyn_cast<Instruction>(val)) {
    if (inversionAllocs && inst->getParent() == inversionAllocs) {
      return val;
    }
    // if (inst->getParent() == &newFunc->getEntryBlock()) {
    //  return inst;
    //}
    if (inst->getParent()->getParent() == newFunc &&
        isOriginalBlock(*BuilderM.GetInsertBlock())) {
      if (BuilderM.GetInsertBlock()->size() &&
          BuilderM.GetInsertPoint() != BuilderM.GetInsertBlock()->end()) {
        if (DT.dominates(inst, &*BuilderM.GetInsertPoint())) {
          // llvm::errs() << "allowed " << *inst << "from domination\n";
          assert(inst->getType() == val->getType());
          return inst;
        }
      } else {
        if (DT.dominates(inst, &*BuilderM.GetInsertPoint())) {
          // llvm::errs() << "allowed " << *inst << "from block domination\n";
          assert(inst->getType() == val->getType());
          return inst;
        }
      }
    }
    assert(!TapesToPreventRecomputation.count(inst));
  }

  std::pair<Value *, BasicBlock *> idx = std::make_pair(val, scope);
  // assert(!startsWith(val->getName(), "$tapeload"));
  if (permitCache) {
    auto found0 = unwrap_cache.find(BuilderM.GetInsertBlock());
    if (found0 != unwrap_cache.end()) {
      auto found1 = found0->second.find(idx.first);
      if (found1 != found0->second.end()) {
        auto found2 = found1->second.find(idx.second);
        if (found2 != found1->second.end()) {

          auto cachedValue = found2->second;
          if (cachedValue == nullptr) {
            found1->second.erase(idx.second);
            if (found1->second.size() == 0) {
              found0->second.erase(idx.first);
            }
          } else {
            if (cachedValue->getType() != val->getType()) {
              llvm::errs() << "newFunc: " << *newFunc << "\n";
              llvm::errs() << "val: " << *val << "\n";
              llvm::errs() << "unwrap_cache[cidx]: " << *cachedValue << "\n";
            }
            assert(cachedValue->getType() == val->getType());
            return cachedValue;
          }
        }
      }
    }
  }

  if (this->mode == DerivativeMode::ReverseModeGradient ||
      this->mode == DerivativeMode::ForwardModeSplit ||
      this->mode == DerivativeMode::ReverseModeCombined)
    if (auto inst = dyn_cast<Instruction>(val)) {
      if (inst->getParent()->getParent() == newFunc) {
        if (unwrapMode == UnwrapMode::LegalFullUnwrap &&
            this->mode != DerivativeMode::ReverseModeCombined) {
          // TODO this isOriginal is a bottleneck, the new mapping of
          // knownRecompute should be precomputed and maintained to lookup
          // instead
          Instruction *orig = isOriginal(inst);
          // If a given value has been chosen to be cached, do not compute the
          // operands to unwrap it, instead simply emit a placeholder to be
          // replaced by the cache load later. This placeholder should only be
          // returned when the original value would be recomputed (e.g. this
          // function would not return null). Since this case assumes everything
          // can be recomputed, simply return the placeholder.
          if (orig && knownRecomputeHeuristic.find(orig) !=
                          knownRecomputeHeuristic.end()) {
            if (!knownRecomputeHeuristic[orig]) {
              assert(inst->getParent()->getParent() == newFunc);
              auto placeholder = BuilderM.CreatePHI(
                  val->getType(), 0, val->getName() + "_krcLFUreplacement");
              unwrappedLoads[placeholder] = inst;
              SmallVector<Metadata *, 1> avail;
              for (auto pair : available)
                if (pair.second)
                  avail.push_back(MDNode::get(
                      placeholder->getContext(),
                      {ValueAsMetadata::get(const_cast<Value *>(pair.first)),
                       ValueAsMetadata::get(pair.second)}));
              placeholder->setMetadata(
                  "enzyme_available",
                  MDNode::get(placeholder->getContext(), avail));
              if (!permitCache)
                return placeholder;
              return unwrap_cache[BuilderM.GetInsertBlock()][idx.first]
                                 [idx.second] = placeholder;
            }
          }
        } else if (unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup) {
          // TODO this isOriginal is a bottleneck, the new mapping of
          // knownRecompute should be precomputed and maintained to lookup
          // instead
          Instruction *orig = isOriginal(inst);
          // If a given value has been chosen to be cached, do not compute the
          // operands to unwrap it, instead simply emit a placeholder to be
          // replaced by the cache load later. This placeholder should only be
          // returned when the original value would be recomputed (e.g. this
          // function would not return null). See note below about the condition
          // as applied to this case.
          if (orig) {
            auto found = knownRecomputeHeuristic.find(orig);
            if (found != knownRecomputeHeuristic.end()) {
              if (!found->second) {
                if (mode == DerivativeMode::ReverseModeCombined) {
                  // Don't unnecessarily cache a value if the caching
                  // heuristic says we should preserve this precise (and not
                  // an lcssa wrapped) value
                  if (!isOriginalBlock(*BuilderM.GetInsertBlock())) {
                    Value *nval = inst;
                    if (scope)
                      nval = fixLCSSA(inst, scope);
                    if (nval == inst)
                      goto endCheck;
                  }
                } else {
                  // Note that this logic (original load must dominate or
                  // alternatively be in the reverse block) is only valid iff
                  // when applicable (here if in split mode), an overwritten
                  // load cannot be hoisted outside of a loop to be used as a
                  // loop limit. This optimization is currently done in the
                  // combined mode (e.g. if a load isn't modified between a
                  // prior insertion point and the actual load, it is legal to
                  // recompute).
                  if (!isOriginalBlock(*BuilderM.GetInsertBlock()) ||
                      DT.dominates(inst, &*BuilderM.GetInsertPoint())) {
                    assert(inst->getParent()->getParent() == newFunc);
                    auto placeholder = BuilderM.CreatePHI(
                        val->getType(), 0,
                        val->getName() + "_krcAFUWLreplacement");
                    unwrappedLoads[placeholder] = inst;
                    SmallVector<Metadata *, 1> avail;
                    for (auto pair : available)
                      if (pair.second)
                        avail.push_back(
                            MDNode::get(placeholder->getContext(),
                                        {ValueAsMetadata::get(
                                             const_cast<Value *>(pair.first)),
                                         ValueAsMetadata::get(pair.second)}));
                    placeholder->setMetadata(
                        "enzyme_available",
                        MDNode::get(placeholder->getContext(), avail));
                    if (!permitCache)
                      return placeholder;
                    return unwrap_cache[BuilderM.GetInsertBlock()][idx.first]
                                       [idx.second] = placeholder;
                  }
                }
              }
            }
          }
        } else if (unwrapMode != UnwrapMode::LegalFullUnwrapNoTapeReplace &&
                   mode != DerivativeMode::ReverseModeCombined) {
          // TODO this isOriginal is a bottleneck, the new mapping of
          // knownRecompute should be precomputed and maintained to lookup
          // instead

          // If a given value has been chosen to be cached, do not compute the
          // operands to unwrap it if it is not legal to do so. This prevents
          // the creation of unused versions of the instruction's operand, which
          // may be assumed to never be used and thus cause an error when they
          // are inadvertantly cached.
          Value *orig = isOriginal(val);
          if (orig && knownRecomputeHeuristic.find(orig) !=
                          knownRecomputeHeuristic.end()) {
            if (!knownRecomputeHeuristic[orig]) {
              return nullptr;
            }
          }
        }
      }
    }

#define getOpFullest(Builder, vtmp, frominst, lookupInst, check)               \
  ({                                                                           \
    Value *v = vtmp;                                                           \
    Type *vty = v->getType();                                                  \
    BasicBlock *origParent = frominst;                                         \
    Value *___res;                                                             \
    if (unwrapMode == UnwrapMode::LegalFullUnwrap ||                           \
        unwrapMode == UnwrapMode::LegalFullUnwrapNoTapeReplace ||              \
        unwrapMode == UnwrapMode::AttemptFullUnwrap ||                         \
        unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup) {               \
      if (v == val)                                                            \
        ___res = nullptr;                                                      \
      else                                                                     \
        ___res = unwrapM(v, Builder, available, unwrapMode, origParent,        \
                         permitCache);                                         \
      if (!___res && unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup) {  \
        bool noLookup = false;                                                 \
        auto found = available.find(v);                                        \
        if (found != available.end() && !found->second)                        \
          noLookup = true;                                                     \
        if (auto opinst = dyn_cast<Instruction>(v))                            \
          if (isOriginalBlock(*Builder.GetInsertBlock())) {                    \
            if (!DT.dominates(opinst, &*Builder.GetInsertPoint()))             \
              noLookup = true;                                                 \
          }                                                                    \
        origParent = lookupInst;                                               \
        if (!noLookup)                                                         \
          ___res = lookupM(v, Builder, available, v != val, origParent);       \
      }                                                                        \
      if (___res)                                                              \
        assert(___res->getType() == vty && "uw");                              \
    } else {                                                                   \
      origParent = lookupInst;                                                 \
      assert(unwrapMode == UnwrapMode::AttemptSingleUnwrap);                   \
      auto found = available.find(v);                                          \
      if (found != available.end() && !found->second)                          \
        ___res = nullptr;                                                      \
      else {                                                                   \
        ___res = lookupM(v, Builder, available, v != val, origParent);         \
        if (___res && ___res->getType() != vty) {                              \
          llvm::errs() << *newFunc << "\n";                                    \
          llvm::errs() << " v = " << *v << " res = " << *___res << "\n";       \
        }                                                                      \
        if (___res)                                                            \
          assert(___res->getType() == vty && "lu");                            \
      }                                                                        \
    }                                                                          \
    ___res;                                                                    \
  })
#define getOpFull(Builder, vtmp, frominst)                                     \
  ({                                                                           \
    BasicBlock *parent = scope;                                                \
    if (parent == nullptr)                                                     \
      if (auto originst = dyn_cast<Instruction>(val))                          \
        parent = originst->getParent();                                        \
    getOpFullest(Builder, vtmp, frominst, parent, true);                       \
  })
#define getOpUnchecked(vtmp)                                                   \
  ({                                                                           \
    BasicBlock *parent = scope;                                                \
    getOpFullest(BuilderM, vtmp, parent, parent, false);                       \
  })
#define getOp(vtmp)                                                            \
  ({                                                                           \
    BasicBlock *parent = scope;                                                \
    if (parent == nullptr)                                                     \
      if (auto originst = dyn_cast<Instruction>(val))                          \
        parent = originst->getParent();                                        \
    getOpFullest(BuilderM, vtmp, parent, parent, true);                        \
  })

  if (isa<Argument>(val) || isa<Constant>(val)) {
    return val;
  } else if (auto op = dyn_cast<FreezeInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateFreeze(op0, op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<CastInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateCast(op->getOpcode(), op0, op->getDestTy(),
                                        op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<ExtractValueInst>(val)) {
    auto op0 = getOp(op->getAggregateOperand());
    if (op0 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateExtractValue(op0, op->getIndices(),
                                                op->getName() + "_unwrap");
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<InsertValueInst>(val)) {
    // Unwrapped Aggregate, Indices, parent
    SmallVector<std::tuple<Value *, ArrayRef<unsigned>, InsertValueInst *>, 1>
        insertElements;

    Value *agg = op;
    while (auto op1 = dyn_cast<InsertValueInst>(agg)) {
      if (Value *orig = isOriginal(op1)) {
        if (knownRecomputeHeuristic.count(orig)) {
          if (!knownRecomputeHeuristic[orig]) {
            break;
          }
        }
      }
      Value *valOp = op1->getInsertedValueOperand();
      valOp = getOp(valOp);
      if (valOp == nullptr)
        goto endCheck;
      insertElements.push_back({valOp, op1->getIndices(), op1});
      agg = op1->getAggregateOperand();
    }

    Value *toreturn = getOp(agg);
    if (toreturn == nullptr)
      goto endCheck;
    for (auto &&[valOp, idcs, parent] : reverse(insertElements)) {
      toreturn = BuilderM.CreateInsertValue(toreturn, valOp, idcs,
                                            parent->getName() + "_unwrap");

      if (permitCache)
        unwrap_cache[BuilderM.GetInsertBlock()][parent][idx.second] = toreturn;
      if (auto newi = dyn_cast<Instruction>(toreturn)) {
        newi->copyIRFlags(parent);
        unwrappedLoads[newi] = val;
        if (newi->getParent()->getParent() != parent->getParent()->getParent())
          newi->setDebugLoc(nullptr);
      }
    }

    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<ExtractElementInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    auto toreturn =
        BuilderM.CreateExtractElement(op0, op1, op->getName() + "_unwrap");
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<InsertElementInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    auto op2 = getOp(op->getOperand(2));
    if (op2 == nullptr)
      goto endCheck;
    auto toreturn =
        BuilderM.CreateInsertElement(op0, op1, op2, op->getName() + "_unwrap");
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<ShuffleVectorInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateShuffleVector(
        op0, op1, op->getShuffleMaskForBitcode(), op->getName() + "'_unwrap");
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<BinaryOperator>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    if (op0->getType() != op1->getType()) {
      llvm::errs() << " op: " << *op << " op0: " << *op0 << " op1: " << *op1
                   << " p0: " << *op->getOperand(0)
                   << "  p1: " << *op->getOperand(1) << "\n";
    }
    assert(op0->getType() == op1->getType());
    auto toreturn = BuilderM.CreateBinOp(op->getOpcode(), op0, op1,
                                         op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<ICmpInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateICmp(op->getPredicate(), op0, op1,
                                        op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<FCmpInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateFCmp(op->getPredicate(), op0, op1,
                                        op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (isa<FPMathOperator>(val) &&
             cast<FPMathOperator>(val)->getOpcode() == Instruction::FNeg) {
    auto op = cast<FPMathOperator>(val);
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto toreturn = BuilderM.CreateFNeg(op0, op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() !=
          cast<Instruction>(val)->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<SelectInst>(val)) {
    auto op0 = getOp(op->getOperand(0));
    if (op0 == nullptr)
      goto endCheck;
    auto op1 = getOp(op->getOperand(1));
    if (op1 == nullptr)
      goto endCheck;
    auto op2 = getOp(op->getOperand(2));
    if (op2 == nullptr)
      goto endCheck;
    auto toreturn =
        BuilderM.CreateSelect(op0, op1, op2, op->getName() + "_unwrap");
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(op);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != op->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto inst = dyn_cast<GetElementPtrInst>(val)) {
    auto ptr = getOp(inst->getPointerOperand());
    if (ptr == nullptr)
      goto endCheck;
    SmallVector<Value *, 4> ind;
    // llvm::errs() << "inst: " << *inst << "\n";
    for (unsigned i = 0; i < inst->getNumIndices(); ++i) {
      Value *a = inst->getOperand(1 + i);
      auto op = getOp(a);
      if (op == nullptr)
        goto endCheck;
      ind.push_back(op);
    }
    auto toreturn = BuilderM.CreateGEP(inst->getSourceElementType(), ptr, ind,
                                       inst->getName() + "_unwrap");
    if (isa<GetElementPtrInst>(toreturn))
      cast<GetElementPtrInst>(toreturn)->setIsInBounds(inst->isInBounds());
    if (auto newi = dyn_cast<Instruction>(toreturn)) {
      newi->copyIRFlags(inst);
      unwrappedLoads[newi] = val;
      if (newi->getParent()->getParent() != inst->getParent()->getParent())
        newi->setDebugLoc(nullptr);
    }
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto load = dyn_cast<LoadInst>(val)) {
    if (load->getMetadata("enzyme_noneedunwrap"))
      return load;

    bool legalMove = unwrapMode == UnwrapMode::LegalFullUnwrap ||
                     unwrapMode == UnwrapMode::LegalFullUnwrapNoTapeReplace;
    if (!legalMove) {
      BasicBlock *parent = nullptr;
      if (isOriginalBlock(*BuilderM.GetInsertBlock()))
        parent = BuilderM.GetInsertBlock();
      if (!parent ||
          LI.getLoopFor(parent) == LI.getLoopFor(load->getParent()) ||
          DT.dominates(load, parent)) {
        legalMove = legalRecompute(load, available, &BuilderM);
      } else {
        legalMove =
            legalRecompute(load, available, &BuilderM, /*reverse*/ false,
                           /*legalRecomputeCache*/ false);
      }
    }
    if (!legalMove) {
      auto &warnMap = UnwrappedWarnings[load];
      if (!warnMap.count(BuilderM.GetInsertBlock())) {
        EmitWarning("UncacheableUnwrap", *load, "Load cannot be unwrapped ",
                    *load, " in ", BuilderM.GetInsertBlock()->getName(), " - ",
                    BuilderM.GetInsertBlock()->getParent()->getName(), " mode ",
                    unwrapMode);
        warnMap.insert(BuilderM.GetInsertBlock());
      }
      goto endCheck;
    }

    Value *pidx = getOp(load->getOperand(0));

    if (pidx == nullptr) {
      goto endCheck;
    }

    if (pidx->getType() != load->getOperand(0)->getType()) {
      llvm::errs() << "load: " << *load << "\n";
      llvm::errs() << "load->getOperand(0): " << *load->getOperand(0) << "\n";
      llvm::errs() << "idx: " << *pidx << " unwrapping: " << *val
                   << " mode=" << unwrapMode << "\n";
    }
    assert(pidx->getType() == load->getOperand(0)->getType());

    auto toreturn =
        BuilderM.CreateLoad(load->getType(), pidx, load->getName() + "_unwrap");
    llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
    ToCopy2.push_back(LLVMContext::MD_noalias);
    ToCopy2.push_back(LLVMContext::MD_alias_scope);
    toreturn->copyMetadata(*load, ToCopy2);
    toreturn->copyIRFlags(load);
    if (load->getParent()->getParent() == newFunc)
      if (auto orig = isOriginal(load)) {
        SmallVector<Metadata *, 1> scopeMD = {
            getDerivativeAliasScope(orig->getOperand(0), -1)};
        if (auto prev = orig->getMetadata(LLVMContext::MD_alias_scope)) {
          for (auto &M : cast<MDNode>(prev)->operands()) {
            scopeMD.push_back(M);
          }
        }
        auto scope = MDNode::get(orig->getContext(), scopeMD);
        toreturn->setMetadata(LLVMContext::MD_alias_scope, scope);

        SmallVector<Metadata *, 1> MDs;
        for (size_t j = 0; j < getWidth(); j++) {
          MDs.push_back(getDerivativeAliasScope(orig->getOperand(0), j));
        }
        if (auto prev = orig->getMetadata(LLVMContext::MD_noalias)) {
          for (auto &M : cast<MDNode>(prev)->operands()) {
            MDs.push_back(M);
          }
        }
        auto noscope = MDNode::get(orig->getContext(), MDs);
        toreturn->setMetadata(LLVMContext::MD_noalias, noscope);
      }
    unwrappedLoads[toreturn] = load;
    if (toreturn->getParent()->getParent() != load->getParent()->getParent())
      toreturn->setDebugLoc(nullptr);
    else
      toreturn->setDebugLoc(getNewFromOriginal(load->getDebugLoc()));
    toreturn->setAlignment(load->getAlign());
    toreturn->setVolatile(load->isVolatile());
    toreturn->setOrdering(load->getOrdering());
    toreturn->setSyncScopeID(load->getSyncScopeID());
    if (toreturn->getParent()->getParent() != load->getParent()->getParent())
      toreturn->setDebugLoc(nullptr);
    else
      toreturn->setDebugLoc(getNewFromOriginal(load->getDebugLoc()));
    toreturn->setMetadata(LLVMContext::MD_tbaa,
                          load->getMetadata(LLVMContext::MD_tbaa));
    auto invar_group = load->getMetadata(LLVMContext::MD_invariant_group);
    if (!invar_group) {
      bool legal = true;
      if (load->getParent()->getParent() != newFunc)
        legal = false;
      else if (auto norig = isOriginal(load))
        for (const auto &pair : rematerializableAllocations) {
          for (auto V : pair.second.loads)
            if (V == norig) {
              legal = false;
              break;
            }
          if (!legal)
            break;
        }
      if (legal) {
        invar_group = MDNode::getDistinct(load->getContext(), {});
        load->setMetadata(LLVMContext::MD_invariant_group, invar_group);
      }
    }
    toreturn->setMetadata(LLVMContext::MD_invariant_group, invar_group);
    // TODO adding to cache only legal if no alias of any future writes
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    assert(val->getType() == toreturn->getType());
    return toreturn;
  } else if (auto op = dyn_cast<CallInst>(val)) {

    bool legalMove = unwrapMode == UnwrapMode::LegalFullUnwrap ||
                     unwrapMode == UnwrapMode::LegalFullUnwrapNoTapeReplace;
    if (!legalMove) {
      legalMove = legalRecompute(op, available, &BuilderM);
    }
    if (!legalMove)
      goto endCheck;

    SmallVector<Value *, 4> args;
#if LLVM_VERSION_MAJOR >= 14
    for (unsigned i = 0; i < op->arg_size(); ++i)
#else
    for (unsigned i = 0; i < op->getNumArgOperands(); ++i)
#endif
    {
      args.push_back(getOp(op->getArgOperand(i)));
      if (args[i] == nullptr)
        goto endCheck;
    }

    Value *fn = getOp(op->getCalledOperand());
    if (fn == nullptr)
      goto endCheck;

    auto toreturn =
        cast<CallInst>(BuilderM.CreateCall(op->getFunctionType(), fn, args));
    toreturn->copyIRFlags(op);
    toreturn->setAttributes(op->getAttributes());
    toreturn->setCallingConv(op->getCallingConv());
    toreturn->setTailCallKind(op->getTailCallKind());
    if (toreturn->getParent()->getParent() == op->getParent()->getParent())
      toreturn->setDebugLoc(getNewFromOriginal(op->getDebugLoc()));
    else
      toreturn->setDebugLoc(nullptr);
    if (permitCache)
      unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = toreturn;
    unwrappedLoads[toreturn] = val;
    return toreturn;
  } else if (auto phi = dyn_cast<PHINode>(val)) {
    if (phi->getNumIncomingValues() == 0) {
      auto uninverted = hasUninverted(phi);
      if (uninverted) {

        // This is a placeholder shadow for a load, rather than falling
        // back to the uncached variant, use the proper procedure for
        // an inverted load
        if (auto dli = dyn_cast<LoadInst>(uninverted)) {
          // Almost identical code to unwrap load (replacing use of shadow
          // where appropriate)
          if (dli->getMetadata("enzyme_noneedunwrap"))
            return dli;

          bool legalMove =
              unwrapMode == UnwrapMode::LegalFullUnwrap ||
              unwrapMode == UnwrapMode::LegalFullUnwrapNoTapeReplace;
          if (!legalMove) {
            // TODO actually consider whether this is legal to move to the new
            // location, rather than recomputable anywhere
            legalMove = legalRecompute(dli, available, &BuilderM);
          }
          if (!legalMove) {
            auto &warnMap = UnwrappedWarnings[phi];
            if (!warnMap.count(BuilderM.GetInsertBlock())) {
              EmitWarning("UncacheableUnwrap", *dli,
                          "Differential Load cannot be unwrapped ", *dli,
                          " in ", BuilderM.GetInsertBlock()->getName(),
                          " mode ", unwrapMode);
              warnMap.insert(BuilderM.GetInsertBlock());
            }
            return nullptr;
          }

          Value *pidx = nullptr;

          if (isOriginalBlock(*BuilderM.GetInsertBlock())) {
            pidx = invertPointerM(dli->getOperand(0), BuilderM);
          } else {
            pidx = lookupM(invertPointerM(dli->getOperand(0), BuilderM),
                           BuilderM, available);
          }

          if (pidx == nullptr)
            goto endCheck;

          if (pidx->getType() != getShadowType(dli->getOperand(0)->getType())) {
            llvm::errs() << "dli: " << *dli << "\n";
            llvm::errs() << "dli->getOperand(0): " << *dli->getOperand(0)
                         << "\n";
            llvm::errs() << "pidx: " << *pidx << "\n";
          }
          assert(pidx->getType() ==
                 getShadowType(dli->getOperand(0)->getType()));

          size_t s_idx = 0;
          Value *toreturn = applyChainRule(
              dli->getType(), BuilderM,
              [&](Value *pidx) {
                auto toreturn = BuilderM.CreateLoad(dli->getType(), pidx,
                                                    phi->getName() + "_unwrap");
                if (auto newi = dyn_cast<Instruction>(toreturn)) {
                  newi->copyIRFlags(dli);
                  unwrappedLoads[toreturn] = dli;
                }
                toreturn->setAlignment(dli->getAlign());
                toreturn->setVolatile(dli->isVolatile());
                toreturn->setOrdering(dli->getOrdering());
                toreturn->setSyncScopeID(dli->getSyncScopeID());
                llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
                toreturn->copyMetadata(*dli, ToCopy2);
                SmallVector<Metadata *, 1> scopeMD = {
                    getDerivativeAliasScope(dli->getOperand(0), s_idx)};
                if (auto prev = dli->getMetadata(LLVMContext::MD_alias_scope)) {
                  for (auto &M : cast<MDNode>(prev)->operands()) {
                    scopeMD.push_back(M);
                  }
                }
                auto scope = MDNode::get(dli->getContext(), scopeMD);
                toreturn->setMetadata(LLVMContext::MD_alias_scope, scope);

                SmallVector<Metadata *, 1> MDs;
                for (ssize_t j = -1; j < getWidth(); j++) {
                  if (j != (ssize_t)s_idx)
                    MDs.push_back(
                        getDerivativeAliasScope(dli->getOperand(0), j));
                }
                if (auto prev = dli->getMetadata(LLVMContext::MD_noalias)) {
                  for (auto &M : cast<MDNode>(prev)->operands()) {
                    MDs.push_back(M);
                  }
                }
                if (MDs.size()) {
                  auto noscope = MDNode::get(dli->getContext(), MDs);
                  toreturn->setMetadata(LLVMContext::MD_noalias, noscope);
                }
                toreturn->setDebugLoc(getNewFromOriginal(dli->getDebugLoc()));
                s_idx++;
                return toreturn;
              },
              pidx);

          // TODO adding to cache only legal if no alias of any future writes
          if (permitCache)
            unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] =
                toreturn;
          assert(val->getType() == toreturn->getType());
          return toreturn;
        }

        if (auto dli = dyn_cast<Instruction>(uninverted)) {
          if (hasNoCache(dli)) {
            auto found = invertedPointers.find(uninverted);
            assert(found != invertedPointers.end());
            assert(found->second == phi);
            invertedPointers.erase(found);
            auto ip = invertPointerM(dli, BuilderM);
            replaceAWithB(phi, ip);
            erase(phi);
            return unwrapM(ip, BuilderM, available, unwrapMode, scope,
                           permitCache);
          }
        }
      }

      goto endCheck;
    }
    assert(phi->getNumIncomingValues() != 0);

    // If requesting loop bound and are requesting the total size.
    // Rather than generating a new lcssa variable, use the existing loop exact
    // bound var
    BasicBlock *ivctx = scope;
    if (!ivctx)
      ivctx = BuilderM.GetInsertBlock();
    if (newFunc == ivctx->getParent() && !isOriginalBlock(*ivctx)) {
      ivctx = originalForReverseBlock(*ivctx);
    }
    if ((ivctx == phi->getParent() || DT.dominates(phi, ivctx)) &&
        (!isOriginalBlock(*BuilderM.GetInsertBlock()) ||
         DT.dominates(phi, &*BuilderM.GetInsertPoint()))) {
      LoopContext lc;
      bool loopVar = false;
      if (getContext(phi->getParent(), lc) && lc.var == phi) {
        loopVar = true;
      } else {
        Value *V = nullptr;
        bool legal = true;
        for (auto &val : phi->incoming_values()) {
          if (isa<UndefValue>(val))
            continue;
          if (V == nullptr)
            V = val;
          else if (V != val) {
            legal = false;
            break;
          }
        }
        if (legal) {
          if (auto I = dyn_cast_or_null<PHINode>(V)) {
            if (getContext(I->getParent(), lc) && lc.var == I) {
              loopVar = true;
            }
          }
        }
      }
      if (loopVar) {
        if (!lc.dynamic) {
          Value *lim = getOp(lc.trueLimit);
          if (lim) {
            unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] =
                lim;
            return lim;
          }
        } else if (unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup &&
                   reverseBlocks.size() > 0) {
          // Must be in a reverse pass fashion for a lookup to index bound to be
          // legal
          assert(/*ReverseLimit*/ reverseBlocks.size() > 0);
          LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                            lc.preheader);
          Value *lim = lookupValueFromCache(
              lc.var->getType(),
              /*forwardPass*/ false, BuilderM, lctx,
              getDynamicLoopLimit(LI.getLoopFor(lc.header)),
              /*isi1*/ false, available);
          unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] = lim;
          return lim;
        }
      }
    }

    auto parent = phi->getParent();

    // Don't attempt to unroll a loop induction variable in other
    // circumstances
    std::set<BasicBlock *> prevIteration;
    BasicBlock *origParent = isOriginal(parent);
    assert(origParent);
    if (OrigLI->isLoopHeader(origParent)) {
      if (phi->getNumIncomingValues() != 2) {
        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
        goto endCheck;
      }
      auto OrigL = OrigLI->getLoopFor(origParent);
      for (auto PH : predecessors(parent)) {
        BasicBlock *origPH = isOriginal(PH);
        assert(origPH);
        if (OrigL->contains(origPH))
          prevIteration.insert(PH);
      }
      if (prevIteration.size() && !legalRecompute(phi, available, &BuilderM)) {
        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
        goto endCheck;
      }
    }
    for (auto &val : phi->incoming_values()) {
      auto inst = dyn_cast<Instruction>(val);
      if (!inst)
        continue;
      auto origInstParent = isOriginal(inst->getParent());
      assert(origInstParent);
      const llvm::Loop *InstLoop = OrigLI->getLoopFor(origInstParent);
      if (!InstLoop) {
        continue;
      }
      bool isParentLoop = false;
      for (const llvm::Loop *L = OrigLI->getLoopFor(origParent); L;
           L = L->getParentLoop()) {
        if (L == InstLoop) {
          isParentLoop = true;
          break;
        }
      }
      if (isParentLoop)
        continue;
      if (unwrapMode == UnwrapMode::LegalFullUnwrap) {
        llvm::errs() << " module: " << *newFunc->getParent() << "\n";
        llvm::errs() << " newFunc: " << *newFunc << "\n";
        llvm::errs() << " parent: " << *parent << "\n";
        llvm::errs() << " val: " << *val << "\n";
      }
      assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
      goto endCheck;
    }

    if (phi->getNumIncomingValues() == 1) {
      assert(phi->getIncomingValue(0) != phi);
      auto toreturn = getOpUnchecked(phi->getIncomingValue(0));
      if (toreturn == nullptr || toreturn == phi) {
        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
        goto endCheck;
      }
      assert(val->getType() == toreturn->getType());
      return toreturn;
    }

    std::set<BasicBlock *> targetToPreds;
    // Map of function edges to list of values possible
    std::map<std::pair</*pred*/ BasicBlock *, /*successor*/ BasicBlock *>,
             std::set<BasicBlock *>>
        done;
    {
      std::deque<std::tuple<
          std::pair</*pred*/ BasicBlock *, /*successor*/ BasicBlock *>,
          BasicBlock *>>
          Q; // newblock, target

      for (unsigned i = 0; i < phi->getNumIncomingValues(); ++i) {
        Q.push_back(
            std::make_pair(std::make_pair(phi->getIncomingBlock(i), parent),
                           phi->getIncomingBlock(i)));
        targetToPreds.insert(phi->getIncomingBlock(i));
      }

      for (std::tuple<
               std::pair</*pred*/ BasicBlock *, /*successor*/ BasicBlock *>,
               BasicBlock *>
               trace;
           Q.size() > 0;) {
        trace = Q.front();
        Q.pop_front();
        auto edge = std::get<0>(trace);
        auto block = edge.first;
        auto target = std::get<1>(trace);

        if (done[edge].count(target))
          continue;
        done[edge].insert(target);

        if (DT.dominates(block, phi->getParent()))
          continue;

        Loop *blockLoop = LI.getLoopFor(block);

        for (BasicBlock *Pred : predecessors(block)) {
          // Don't go up the backedge as we can use the last value if desired
          // via lcssa
          if (blockLoop && blockLoop->getHeader() == block &&
              blockLoop == LI.getLoopFor(Pred))
            continue;

          Q.push_back(
              std::tuple<std::pair<BasicBlock *, BasicBlock *>, BasicBlock *>(
                  std::make_pair(Pred, block), target));
        }
      }
    }

    std::set<BasicBlock *> blocks;
    for (auto pair : done) {
      const auto &edge = pair.first;
      blocks.insert(edge.first);
    }

    BasicBlock *oldB = BuilderM.GetInsertBlock();
    if (BuilderM.GetInsertPoint() != oldB->end()) {
      assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
      goto endCheck;
    }

    BasicBlock *fwd = oldB;
    bool inReverseBlocks = false;
    if (!isOriginalBlock(*fwd)) {
      auto found = reverseBlockToPrimal.find(oldB);
      if (found == reverseBlockToPrimal.end()) {
        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
        goto endCheck;
      }
      fwd = found->second;
      inReverseBlocks =
          std::find(reverseBlocks[fwd].begin(), reverseBlocks[fwd].end(),
                    oldB) != reverseBlocks[fwd].end();
    }

    auto eraseBlocks = [&](ArrayRef<BasicBlock *> blocks, BasicBlock *bret) {
      SmallVector<BasicBlock *, 2> revtopo;
      {
        SmallPtrSet<BasicBlock *, 2> seen;
        std::function<void(BasicBlock *)> dfs = [&](BasicBlock *B) {
          if (seen.count(B))
            return;
          seen.insert(B);
          if (hasTerminator(B))
            for (auto S : successors(B))
              if (!seen.count(S))
                dfs(S);
          revtopo.push_back(B);
        };
        for (auto B : blocks)
          dfs(B);
        if (!seen.count(bret))
          revtopo.insert(revtopo.begin(), bret);
      }

      SmallVector<Instruction *, 4> toErase;
      for (auto B : revtopo) {
        if (B == bret)
          continue;
        for (auto &I : llvm::reverse(*B)) {
          toErase.push_back(&I);
        }
        unwrap_cache.erase(B);
        lookup_cache.erase(B);
        if (reverseBlocks.size() > 0) {
          auto tfwd = reverseBlockToPrimal[B];
          assert(tfwd);
          auto rfound = reverseBlocks.find(tfwd);
          assert(rfound != reverseBlocks.end());
          auto &tlst = rfound->second;
          auto found = std::find(tlst.begin(), tlst.end(), B);
          if (found != tlst.end())
            tlst.erase(found);
          reverseBlockToPrimal.erase(B);
        }
      }
      for (auto I : toErase) {
        erase(I);
      }
      for (auto B : revtopo)
        B->eraseFromParent();
    };

    if (targetToPreds.size() == 3) {
      for (auto block : blocks) {
        if (!DT.dominates(block, phi->getParent()))
          continue;
        std::set<BasicBlock *> foundtargets;
        std::set<BasicBlock *> uniqueTargets;
        for (BasicBlock *succ : successors(block)) {
          auto edge = std::make_pair(block, succ);
          for (BasicBlock *target : done[edge]) {
            if (foundtargets.find(target) != foundtargets.end()) {
              goto rnextpair;
            }
            foundtargets.insert(target);
            if (done[edge].size() == 1)
              uniqueTargets.insert(target);
          }
        }
        if (foundtargets.size() != 3)
          goto rnextpair;
        if (uniqueTargets.size() != 1)
          goto rnextpair;

        {
          BasicBlock *subblock = nullptr;
          for (auto block2 : blocks) {
            {
              // The second split block must not have a parent with an edge
              // to a block other than to itself, which can reach any of its
              // two targets.
              // TODO verify this
              for (auto P : predecessors(block2)) {
                for (auto S : successors(P)) {
                  if (S == block2)
                    continue;
                  auto edge = std::make_pair(P, S);
                  if (done.find(edge) != done.end()) {
                    for (auto target : done[edge]) {
                      if (foundtargets.find(target) != foundtargets.end() &&
                          uniqueTargets.find(target) == uniqueTargets.end())
                        goto nextblock;
                    }
                  }
                }
              }
              std::set<BasicBlock *> seen2;
              for (BasicBlock *succ : successors(block2)) {
                auto edge = std::make_pair(block2, succ);
                if (done[edge].size() != 1) {
                  // llvm::errs() << " -- failed from noonesize\n";
                  goto nextblock;
                }
                for (BasicBlock *target : done[edge]) {
                  if (seen2.find(target) != seen2.end()) {
                    // llvm::errs() << " -- failed from not uniqueTargets\n";
                    goto nextblock;
                  }
                  seen2.insert(target);
                  if (foundtargets.find(target) == foundtargets.end()) {
                    // llvm::errs() << " -- failed from not unknown target\n";
                    goto nextblock;
                  }
                  if (uniqueTargets.find(target) != uniqueTargets.end()) {
                    // llvm::errs() << " -- failed from not same target\n";
                    goto nextblock;
                  }
                }
              }
              if (seen2.size() != 2) {
                // llvm::errs() << " -- failed from not 2 seen\n";
                goto nextblock;
              }
              subblock = block2;
              break;
            }
          nextblock:;
          }

          if (subblock == nullptr)
            goto rnextpair;

          {
            auto bi1 = dyn_cast<BranchInst>(block->getTerminator());
            if (!bi1) {
              goto endCheck;
            }

            auto cond1 = getOp(bi1->getCondition());
            if (cond1 == nullptr) {
              assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
              goto endCheck;
            }
            auto bi2 = cast<BranchInst>(subblock->getTerminator());
            auto cond2 = getOp(bi2->getCondition());
            if (cond2 == nullptr) {
              assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
              goto endCheck;
            }

            SmallVector<BasicBlock *, 3> predBlocks = {bi2->getSuccessor(0),
                                                       bi2->getSuccessor(1)};
            for (int i = 0; i < 2; i++) {
              auto edge = std::make_pair(block, bi1->getSuccessor(i));
              if (done[edge].size() == 1) {
                predBlocks.push_back(bi1->getSuccessor(i));
              }
            }

            SmallVector<Value *, 2> vals;

            SmallVector<BasicBlock *, 2> blocks;
            SmallVector<BasicBlock *, 2> endingBlocks;

            BasicBlock *last = oldB;

            BasicBlock *bret = BasicBlock::Create(
                val->getContext(), oldB->getName() + "_phimerge", newFunc);

            for (size_t i = 0; i < predBlocks.size(); i++) {
              BasicBlock *valparent = (i < 2) ? subblock : block;
              assert(done.find(std::make_pair(valparent, predBlocks[i])) !=
                     done.end());
              assert(done[std::make_pair(valparent, predBlocks[i])].size() ==
                     1);
              blocks.push_back(BasicBlock::Create(
                  val->getContext(), oldB->getName() + "_phirc", newFunc));
              blocks[i]->moveAfter(last);
              last = blocks[i];
              if (inReverseBlocks)
                reverseBlocks[fwd].push_back(blocks[i]);
              reverseBlockToPrimal[blocks[i]] = fwd;
              IRBuilder<> B(blocks[i]);

              for (auto pair : unwrap_cache[oldB])
                unwrap_cache[blocks[i]].insert(pair);
              for (auto pair : lookup_cache[oldB])
                lookup_cache[blocks[i]].insert(pair);
              auto PB = *done[std::make_pair(valparent, predBlocks[i])].begin();

              if (auto inst = dyn_cast<Instruction>(
                      phi->getIncomingValueForBlock(PB))) {
                // Recompute the phi computation with the conditional if:
                // 1) the instruction may read from memory AND does not
                //    dominate the current insertion point (thereby
                //    potentially making such recomputation without the
                //    condition illegal)
                // 2) the value is a call or load and option is set to not
                //    speculatively recompute values within a phi
                //            OR
                // 3) the value comes from a previous iteration.
                BasicBlock *nextScope = PB;
                // if (inst->getParent() == nextScope) nextScope =
                // phi->getParent();
                if (prevIteration.count(PB)) {
                  assert(0 && "tri block prev iteration unhandled");
                } else if (!DT.dominates(inst->getParent(), phi->getParent()) ||
                           (!EnzymeSpeculatePHIs &&
                            (isa<CallInst>(inst) || isa<LoadInst>(inst))))
                  vals.push_back(getOpFull(B, inst, nextScope));
                else
                  vals.push_back(getOpFull(BuilderM, inst, nextScope));
              } else
                vals.push_back(
                    getOpFull(BuilderM, phi->getIncomingValueForBlock(PB), PB));

              if (!vals[i]) {
                eraseBlocks(blocks, bret);
                assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
                goto endCheck;
              }
              assert(val->getType() == vals[i]->getType());
              B.CreateBr(bret);
              endingBlocks.push_back(B.GetInsertBlock());
            }

            bret->moveAfter(last);

            BasicBlock *bsplit = BasicBlock::Create(
                val->getContext(), oldB->getName() + "_phisplt", newFunc);
            bsplit->moveAfter(oldB);
            if (inReverseBlocks)
              reverseBlocks[fwd].push_back(bsplit);
            reverseBlockToPrimal[bsplit] = fwd;
            BuilderM.CreateCondBr(
                cond1,
                (done[std::make_pair(block, bi1->getSuccessor(0))].size() == 1)
                    ? blocks[2]
                    : bsplit,
                (done[std::make_pair(block, bi1->getSuccessor(1))].size() == 1)
                    ? blocks[2]
                    : bsplit);

            BuilderM.SetInsertPoint(bsplit);
            BuilderM.CreateCondBr(cond2, blocks[0], blocks[1]);

            BuilderM.SetInsertPoint(bret);
            if (inReverseBlocks)
              reverseBlocks[fwd].push_back(bret);
            reverseBlockToPrimal[bret] = fwd;
            auto toret = BuilderM.CreatePHI(val->getType(), vals.size());
            for (size_t i = 0; i < vals.size(); i++)
              toret->addIncoming(vals[i], endingBlocks[i]);
            assert(val->getType() == toret->getType());
            if (permitCache) {
              unwrap_cache[bret][idx.first][idx.second] = toret;
            }
            unwrappedLoads[toret] = val;
            for (auto pair : unwrap_cache[oldB])
              unwrap_cache[bret].insert(pair);
            for (auto pair : lookup_cache[oldB])
              lookup_cache[bret].insert(pair);
            return toret;
          }
        }
      rnextpair:;
      }
    }

    Instruction *equivalentTerminator = nullptr;

    if (prevIteration.size() == 1) {
      if (phi->getNumIncomingValues() == 2) {

        ValueToValueMapTy prevAvailable;
        for (const auto &pair : available)
          prevAvailable.insert(pair);
        LoopContext ctx;
        getContext(parent, ctx);
        Value *prevIdx;
        if (prevAvailable.count(ctx.var))
          prevIdx = prevAvailable[ctx.var];
        else {
          if (!isOriginalBlock(*BuilderM.GetInsertBlock())) {
            // If we are using the phi in the reverse pass of a block inside the
            // loop itself the previous index variable (aka the previous inc) is
            // equivalent to the current load of antivaralloc
            if (LI.getLoopFor(ctx.header)->contains(fwd)) {
              prevIdx =
                  BuilderM.CreateLoad(ctx.var->getType(), ctx.antivaralloc);
            } else {
              // However, if we are using the phi of the reverse pass of a block
              // outside the loop we must be in the reverse pass of a block
              // after the loop. In which case, the previous index variable (aka
              // previous inc) is the total loop iteration count-1, aka the
              // trueLimit.
              Value *lim = nullptr;
              if (ctx.dynamic) {
                // Must be in a reverse pass fashion for a lookup to index bound
                // to be legal
                assert(/*ReverseLimit*/ reverseBlocks.size() > 0);
                LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                                  ctx.preheader);
                lim = lookupValueFromCache(
                    ctx.var->getType(),
                    /*forwardPass*/ false, BuilderM, lctx,
                    getDynamicLoopLimit(LI.getLoopFor(ctx.header)),
                    /*isi1*/ false, /*available*/ prevAvailable);
              } else {
                lim = lookupM(ctx.trueLimit, BuilderM, prevAvailable);
              }
              prevIdx = lim;
            }
          } else {
            prevIdx = ctx.var;
          }
        }
        // Prevent recursive unroll.
        prevAvailable[phi] = nullptr;
        SmallVector<Value *, 2> vals;

        SmallVector<BasicBlock *, 2> blocks;
        SmallVector<BasicBlock *, 2> endingBlocks;
        BasicBlock *last = oldB;

        BasicBlock *bret = BasicBlock::Create(
            val->getContext(), oldB->getName() + "_phimerge", newFunc);

        SmallVector<BasicBlock *, 2> preds(predecessors(phi->getParent()));

        for (auto tup : llvm::enumerate(preds)) {
          auto i = tup.index();
          BasicBlock *PB = tup.value();
          blocks.push_back(BasicBlock::Create(
              val->getContext(), oldB->getName() + "_phirc", newFunc));
          blocks[i]->moveAfter(last);
          last = blocks[i];
          if (reverseBlocks.size() > 0) {
            if (inReverseBlocks)
              reverseBlocks[fwd].push_back(blocks[i]);
            reverseBlockToPrimal[blocks[i]] = fwd;
          }
          IRBuilder<> B(blocks[i]);

          if (!prevIteration.count(PB)) {
            for (auto pair : unwrap_cache[oldB])
              unwrap_cache[blocks[i]].insert(pair);
            for (auto pair : lookup_cache[oldB])
              lookup_cache[blocks[i]].insert(pair);
          }

          if (auto inst =
                  dyn_cast<Instruction>(phi->getIncomingValueForBlock(PB))) {
            // Recompute the phi computation with the conditional if:
            // 1) the instruction may read from memory AND does not dominate
            //    the current insertion point (thereby potentially making such
            //    recomputation without the condition illegal)
            // 2) the value is a call or load and option is set to not
            //    speculatively recompute values within a phi
            //                OR
            // 3) the value comes from a previous iteration.
            BasicBlock *nextScope = PB;
            // if (inst->getParent() == nextScope) nextScope = phi->getParent();
            if (prevIteration.count(PB)) {
              prevAvailable[ctx.incvar] = prevIdx;
              prevAvailable[ctx.var] =
                  B.CreateSub(prevIdx, ConstantInt::get(prevIdx->getType(), 1),
                              "", /*NUW*/ true, /*NSW*/ false);
              Value *___res;
              if (unwrapMode == UnwrapMode::LegalFullUnwrap ||
                  unwrapMode == UnwrapMode::LegalFullUnwrapNoTapeReplace ||
                  unwrapMode == UnwrapMode::AttemptFullUnwrap ||
                  unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup) {
                ___res = unwrapM(inst, B, prevAvailable, unwrapMode, nextScope,
                                 /*permitCache*/ false);
                if (!___res &&
                    unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup) {
                  bool noLookup = false;
                  if (isOriginalBlock(*B.GetInsertBlock())) {
                    if (!DT.dominates(inst, &*B.GetInsertPoint()))
                      noLookup = true;
                  }
                  if (!noLookup) {
                    BasicBlock *nS2 = nextScope;
                    Value *v = inst;
                    ___res = lookupM(v, B, prevAvailable, v != val, nS2);
                  }
                }
                if (___res)
                  assert(___res->getType() == inst->getType() && "uw");
              } else {
                BasicBlock *nS2 = nextScope;
                Value *v = inst;
                ___res = lookupM(v, B, prevAvailable, v != val, nS2);
                if (___res && ___res->getType() != v->getType()) {
                  llvm::errs() << *newFunc << "\n";
                  llvm::errs() << " v = " << *v << " res = " << *___res << "\n";
                }
                if (___res)
                  assert(___res->getType() == inst->getType() && "lu");
              }
              vals.push_back(___res);
            } else if (!DT.dominates(inst->getParent(), phi->getParent()) ||
                       (!EnzymeSpeculatePHIs &&
                        (isa<CallInst>(inst) || isa<LoadInst>(inst))))
              vals.push_back(getOpFull(B, inst, nextScope));
            else
              vals.push_back(getOpFull(BuilderM, inst, nextScope));
          } else
            vals.push_back(phi->getIncomingValueForBlock(PB));

          if (!vals[i]) {
            eraseBlocks(blocks, bret);
            assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
            goto endCheck;
          }
          assert(val->getType() == vals[i]->getType());
          B.CreateBr(bret);
          endingBlocks.push_back(B.GetInsertBlock());
        }

        // Coming from a previous iteration is equivalent to the current
        // iteration at zero.
        Value *cond;
        if (prevIteration.count(preds[0]))
          cond = BuilderM.CreateICmpNE(prevIdx,
                                       ConstantInt::get(prevIdx->getType(), 0));
        else
          cond = BuilderM.CreateICmpEQ(prevIdx,
                                       ConstantInt::get(prevIdx->getType(), 0));

        if (blocks[0]->size() == 1 && blocks[1]->size() == 1) {
          if (auto B1 = dyn_cast<BranchInst>(blocks[0]->getTerminator()))
            if (auto B2 = dyn_cast<BranchInst>(blocks[1]->getTerminator()))
              if (B1->isUnconditional() && B2->isUnconditional() &&
                  B1->getSuccessor(0) == bret && B2->getSuccessor(0) == bret) {
                eraseBlocks(blocks, bret);
                Value *toret = BuilderM.CreateSelect(
                    cond, vals[0], vals[1], phi->getName() + "_unwrap");
                if (permitCache) {
                  unwrap_cache[BuilderM.GetInsertBlock()][idx.first]
                              [idx.second] = toret;
                }
                if (auto instRet = dyn_cast<Instruction>(toret)) {
                  unwrappedLoads[instRet] = val;
                }
                return toret;
              }
        }

        bret->moveAfter(last);
        BuilderM.CreateCondBr(cond, blocks[0], blocks[1]);

        BuilderM.SetInsertPoint(bret);
        if (inReverseBlocks)
          reverseBlocks[fwd].push_back(bret);
        reverseBlockToPrimal[bret] = fwd;
        auto toret = BuilderM.CreatePHI(val->getType(), vals.size());
        for (size_t i = 0; i < vals.size(); i++)
          toret->addIncoming(vals[i], endingBlocks[i]);
        assert(val->getType() == toret->getType());
        if (permitCache) {
          unwrap_cache[bret][idx.first][idx.second] = toret;
        }
        for (auto pair : unwrap_cache[oldB])
          unwrap_cache[bret].insert(pair);
        for (auto pair : lookup_cache[oldB])
          lookup_cache[bret].insert(pair);
        unwrappedLoads[toret] = val;
        return toret;
      }
    }
    if (prevIteration.size() != 0) {
      llvm::errs() << "prev iteration: " << *phi << "\n";
      assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
      goto endCheck;
    }

    for (auto block : blocks) {
      if (!DT.dominates(block, phi->getParent()))
        continue;
      std::set<BasicBlock *> foundtargets;
      for (BasicBlock *succ : successors(block)) {
        auto edge = std::make_pair(block, succ);
        if (done[edge].size() != 1) {
          goto nextpair;
        }
        BasicBlock *target = *done[edge].begin();
        if (foundtargets.find(target) != foundtargets.end()) {
          goto nextpair;
        }
        foundtargets.insert(target);
      }
      if (foundtargets.size() != targetToPreds.size()) {
        goto nextpair;
      }

      if (DT.dominates(block, parent)) {
        equivalentTerminator = block->getTerminator();
        goto fast;
      }
    nextpair:;
    }
    goto endCheck;

  fast:;
    assert(equivalentTerminator);

    if (isa<BranchInst>(equivalentTerminator) ||
        isa<SwitchInst>(equivalentTerminator)) {
      BasicBlock *oldB = BuilderM.GetInsertBlock();

      SmallVector<BasicBlock *, 2> predBlocks;
      Value *cond = nullptr;
      if (auto branch = dyn_cast<BranchInst>(equivalentTerminator)) {
        cond = branch->getCondition();
        predBlocks.push_back(branch->getSuccessor(0));
        predBlocks.push_back(branch->getSuccessor(1));
      } else {
        auto SI = cast<SwitchInst>(equivalentTerminator);
        cond = SI->getCondition();
        predBlocks.push_back(SI->getDefaultDest());
        for (auto scase : SI->cases()) {
          predBlocks.push_back(scase.getCaseSuccessor());
        }
      }
      cond = getOp(cond);
      if (!cond) {
        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
        goto endCheck;
      }

      SmallVector<Value *, 2> vals;

      SmallVector<BasicBlock *, 2> blocks;
      SmallVector<BasicBlock *, 2> endingBlocks;

      BasicBlock *last = oldB;

      assert(prevIteration.size() == 0);

      BasicBlock *bret = BasicBlock::Create(
          val->getContext(), oldB->getName() + "_phimerge", newFunc);

      for (size_t i = 0; i < predBlocks.size(); i++) {
        assert(done.find(std::make_pair(equivalentTerminator->getParent(),
                                        predBlocks[i])) != done.end());
        assert(done[std::make_pair(equivalentTerminator->getParent(),
                                   predBlocks[i])]
                   .size() == 1);
        BasicBlock *PB = *done[std::make_pair(equivalentTerminator->getParent(),
                                              predBlocks[i])]
                              .begin();
        blocks.push_back(BasicBlock::Create(
            val->getContext(), oldB->getName() + "_phirc", newFunc));
        blocks[i]->moveAfter(last);
        last = blocks[i];
        if (reverseBlocks.size() > 0) {
          if (inReverseBlocks)
            reverseBlocks[fwd].push_back(blocks[i]);
          reverseBlockToPrimal[blocks[i]] = fwd;
        }
        IRBuilder<> B(blocks[i]);

        for (auto pair : unwrap_cache[oldB])
          unwrap_cache[blocks[i]].insert(pair);
        for (auto pair : lookup_cache[oldB])
          lookup_cache[blocks[i]].insert(pair);

        if (auto inst =
                dyn_cast<Instruction>(phi->getIncomingValueForBlock(PB))) {
          // Recompute the phi computation with the conditional if:
          // 1) the instruction may reat from memory AND does not dominate
          //    the current insertion point (thereby potentially making such
          //    recomputation without the condition illegal)
          // 2) the value is a call or load and option is set to not
          //    speculatively recompute values within a phi
          //                OR
          // 3) the value comes from a previous iteration.
          BasicBlock *nextScope = PB;
          // if (inst->getParent() == nextScope) nextScope = phi->getParent();
          if (!DT.dominates(inst->getParent(), phi->getParent()) ||
              (!EnzymeSpeculatePHIs &&
               (isa<CallInst>(inst) || isa<LoadInst>(inst))))
            vals.push_back(getOpFull(B, inst, nextScope));
          else
            vals.push_back(getOpFull(BuilderM, inst, nextScope));
        } else
          vals.push_back(phi->getIncomingValueForBlock(PB));

        if (!vals[i]) {
          eraseBlocks(blocks, bret);
          assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
          goto endCheck;
        }
        assert(val->getType() == vals[i]->getType());
        B.CreateBr(bret);
        endingBlocks.push_back(B.GetInsertBlock());
      }

      // Fast path to not make a split block if no additional instructions
      // were made in the two blocks
      if (isa<BranchInst>(equivalentTerminator) && blocks[0]->size() == 1 &&
          blocks[1]->size() == 1) {
        if (auto B1 = dyn_cast<BranchInst>(blocks[0]->getTerminator()))
          if (auto B2 = dyn_cast<BranchInst>(blocks[1]->getTerminator()))
            if (B1->isUnconditional() && B2->isUnconditional() &&
                B1->getSuccessor(0) == bret && B2->getSuccessor(0) == bret) {
              eraseBlocks(blocks, bret);
              Value *toret = BuilderM.CreateSelect(cond, vals[0], vals[1],
                                                   phi->getName() + "_unwrap");
              if (permitCache) {
                unwrap_cache[BuilderM.GetInsertBlock()][idx.first][idx.second] =
                    toret;
              }
              if (auto instRet = dyn_cast<Instruction>(toret)) {
                unwrappedLoads[instRet] = val;
              }
              return toret;
            }
      }

      if (BuilderM.GetInsertPoint() != oldB->end()) {
        eraseBlocks(blocks, bret);
        assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
        goto endCheck;
      }

      bret->moveAfter(last);
      if (isa<BranchInst>(equivalentTerminator)) {
        BuilderM.CreateCondBr(cond, blocks[0], blocks[1]);
      } else {
        auto SI = cast<SwitchInst>(equivalentTerminator);
        auto NSI = BuilderM.CreateSwitch(cond, blocks[0], SI->getNumCases());
        size_t idx = 1;
        for (auto scase : SI->cases()) {
          NSI->addCase(scase.getCaseValue(), blocks[idx]);
          idx++;
        }
      }
      BuilderM.SetInsertPoint(bret);
      if (inReverseBlocks)
        reverseBlocks[fwd].push_back(bret);
      reverseBlockToPrimal[bret] = fwd;
      auto toret = BuilderM.CreatePHI(val->getType(), vals.size());
      for (size_t i = 0; i < vals.size(); i++)
        toret->addIncoming(vals[i], endingBlocks[i]);
      assert(val->getType() == toret->getType());
      if (permitCache) {
        unwrap_cache[bret][idx.first][idx.second] = toret;
      }
      for (auto pair : unwrap_cache[oldB])
        unwrap_cache[bret].insert(pair);
      for (auto pair : lookup_cache[oldB])
        lookup_cache[bret].insert(pair);
      unwrappedLoads[toret] = val;
      return toret;
    }
    assert(unwrapMode != UnwrapMode::LegalFullUnwrap);
    goto endCheck;
  }

endCheck:
  assert(val);
  if (unwrapMode == UnwrapMode::LegalFullUnwrap ||
      unwrapMode == UnwrapMode::LegalFullUnwrapNoTapeReplace ||
      unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup) {
    assert(val->getName() != "<badref>");
    Value *nval = val;
    if (auto opinst = dyn_cast<Instruction>(nval))
      if (isOriginalBlock(*BuilderM.GetInsertBlock())) {
        if (!DT.dominates(opinst, &*BuilderM.GetInsertPoint())) {
          if (unwrapMode != UnwrapMode::AttemptFullUnwrapWithLookup) {
            llvm::errs() << " oldF: " << *oldFunc << "\n";
            llvm::errs() << " opParen: " << *opinst->getParent()->getParent()
                         << "\n";
            llvm::errs() << " newF: " << *newFunc << "\n";
            llvm::errs() << " - blk: " << *BuilderM.GetInsertBlock();
            llvm::errs() << " opInst: " << *opinst << " mode=" << unwrapMode
                         << "\n";
          }
          assert(unwrapMode == UnwrapMode::AttemptFullUnwrapWithLookup);
          return nullptr;
        }
      }
    auto toreturn = lookupM(nval, BuilderM, available,
                            /*tryLegalRecomputeCheck*/ false, scope);
    assert(val->getType() == toreturn->getType());
    return toreturn;
  }

  if (auto inst = dyn_cast<Instruction>(val)) {
    if (isOriginalBlock(*BuilderM.GetInsertBlock())) {
      if (BuilderM.GetInsertBlock()->size() &&
          BuilderM.GetInsertPoint() != BuilderM.GetInsertBlock()->end()) {
        if (DT.dominates(inst, &*BuilderM.GetInsertPoint())) {
          assert(inst->getType() == val->getType());
          return inst;
        }
      } else {
        if (DT.dominates(inst, BuilderM.GetInsertBlock())) {
          assert(inst->getType() == val->getType());
          return inst;
        }
      }
    }
    assert(val->getName() != "<badref>");
    auto &warnMap = UnwrappedWarnings[inst];
    if (!warnMap.count(BuilderM.GetInsertBlock())) {
      EmitWarning("NoUnwrap", *inst, "Cannot unwrap ", *val, " in ",
                  BuilderM.GetInsertBlock()->getName());
      warnMap.insert(BuilderM.GetInsertBlock());
    }
  }
  return nullptr;
}

void GradientUtils::ensureLookupCached(Instruction *inst, bool shouldFree,
                                       BasicBlock *scope, MDNode *TBAA) {
  assert(inst);
  if (scopeMap.find(inst) != scopeMap.end())
    return;
  if (shouldFree)
    assert(reverseBlocks.size());

  if (scope == nullptr)
    scope = inst->getParent();

  LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0, scope);

  AllocaInst *cache =
      createCacheForScope(lctx, inst->getType(), inst->getName(), shouldFree);
  assert(cache);
  Value *Val = inst;
  insert_or_assign(
      scopeMap, Val,
      std::pair<AssertingVH<AllocaInst>, LimitContext>(cache, lctx));
  storeInstructionInCache(lctx, inst, cache, TBAA);
}

Value *GradientUtils::fixLCSSA(Instruction *inst, BasicBlock *forwardBlock,
                               bool legalInBlock) {
  assert(inst->getName() != "<badref>");

  if (auto lcssaPHI = dyn_cast<PHINode>(inst)) {
    auto found = lcssaPHIToOrig.find(lcssaPHI);
    if (found != lcssaPHIToOrig.end())
      inst = cast<Instruction>(found->second);
  }

  if (inst->getParent() == inversionAllocs)
    return inst;

  if (!isOriginalBlock(*forwardBlock)) {
    forwardBlock = originalForReverseBlock(*forwardBlock);
  }

  bool containsLastLoopValue = isPotentialLastLoopValue(inst, forwardBlock, LI);

  // If the instruction cannot represent a loop value, return the original
  // instruction if it either is guaranteed to be available within the block,
  // or it is not needed to guaranteed availability.
  if (!containsLastLoopValue) {
    if (!legalInBlock)
      return inst;
    if (forwardBlock == inst->getParent() || DT.dominates(inst, forwardBlock))
      return inst;
  }

  // llvm::errs() << " inst: " << *inst << "\n";
  // llvm::errs() << " seen: " << *inst->getParent() << "\n";
  assert(inst->getParent() != inversionAllocs);
  assert(isOriginalBlock(*inst->getParent()));

  if (lcssaFixes.find(inst) == lcssaFixes.end()) {
    lcssaFixes[inst][inst->getParent()] = inst;
    SmallPtrSet<BasicBlock *, 4> seen;
    std::deque<BasicBlock *> todo = {inst->getParent()};
    while (todo.size()) {
      BasicBlock *cur = todo.front();
      todo.pop_front();
      if (seen.count(cur))
        continue;
      seen.insert(cur);
      if (!hasTerminator(cur))
        continue;
      auto terminator = cur->getTerminator();
      for (auto Succ : successors(terminator)) {
        todo.push_back(Succ);
      }
    }
    for (auto &BB : *inst->getParent()->getParent()) {
      if (!seen.count(&BB) ||
          (inst->getParent() != &BB && DT.dominates(&BB, inst->getParent()))) {
        // OrigPDT.dominates(isOriginal(inst->getParent()),
        //                  isOriginal(&BB)))) {
        lcssaFixes[inst][&BB] = UndefValue::get(inst->getType());
      }
    }
  }

  if (lcssaFixes[inst].find(forwardBlock) != lcssaFixes[inst].end()) {
    return lcssaFixes[inst][forwardBlock];
  }

  // TODO replace forwardBlock with the first block dominated by inst,
  // that dominates (or is) forwardBlock to ensuring maximum reuse
  auto inspos = forwardBlock->front().getIterator();
#if LLVM_VERSION_MAJOR >= 18
#if LLVM_VERSION_MAJOR >= 21
#else
  if (forwardBlock->IsNewDbgInfoFormat)
#endif
  {
    if (!inspos.getHeadBit()) {
      auto srcmarker = forwardBlock->getMarker(inspos);
      if (srcmarker && !srcmarker->empty()) {
        inspos.setHeadBit(true);
      }
    }
  }
#endif
  IRBuilder<> lcssa(forwardBlock, inspos);
  auto lcssaPHI =
      lcssa.CreatePHI(inst->getType(), 1, inst->getName() + "!manual_lcssa");
  lcssaFixes[inst][forwardBlock] = lcssaPHI;
  lcssaPHIToOrig[lcssaPHI] = inst;
  for (auto pred : predecessors(forwardBlock)) {
    Value *val = nullptr;
    if (inst->getParent() == pred || DT.dominates(inst, pred)) {
      val = inst;
    }
    if (val == nullptr) {
      val = fixLCSSA(inst, pred, /*legalInBlock*/ true);
      assert(val->getType() == inst->getType());
    }
    assert(val->getType() == inst->getType());
    lcssaPHI->addIncoming(val, pred);
  }

  SmallPtrSet<Value *, 2> vals;
  SmallVector<Value *, 2> todo(lcssaPHI->incoming_values().begin(),
                               lcssaPHI->incoming_values().end());
  while (todo.size()) {
    Value *v = todo.back();
    todo.pop_back();
    if (v == lcssaPHI)
      continue;
    vals.insert(v);
  }
  assert(vals.size() > 0);

  if (vals.size() > 1) {
    todo.append(vals.begin(), vals.end());
    vals.clear();
    while (todo.size()) {
      Value *v = todo.back();
      todo.pop_back();

      if (auto PN = dyn_cast<PHINode>(v))
        if (lcssaPHIToOrig.find(PN) != lcssaPHIToOrig.end()) {
          v = lcssaPHIToOrig[PN];
        }
      vals.insert(v);
    }
  }
  assert(vals.size() > 0);
  Value *val = nullptr;
  if (vals.size() == 1)
    val = *vals.begin();

  if (val && (!legalInBlock || !isa<Instruction>(val) ||
              DT.dominates(cast<Instruction>(val), lcssaPHI))) {

    if (!isPotentialLastLoopValue(val, forwardBlock, LI)) {
      bool nonSelfUse = false;
      for (auto u : lcssaPHI->users()) {
        if (u != lcssaPHI) {
          nonSelfUse = true;
          break;
        }
      }
      if (!nonSelfUse) {
        lcssaFixes[inst].erase(forwardBlock);
        while (lcssaPHI->getNumOperands())
          lcssaPHI->removeIncomingValue(lcssaPHI->getNumOperands() - 1, false);
        lcssaPHIToOrig.erase(lcssaPHI);
        lcssaPHI->eraseFromParent();
      }
      return val;
    }
  }
  return lcssaPHI;
}

Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
                                      int idx, bool replace) {
  assert(malloc);
  assert(BuilderQ.GetInsertBlock()->getParent() == newFunc);
  assert(isOriginalBlock(*BuilderQ.GetInsertBlock()));
  assert(!hasNoCache(malloc));
  if (mode == DerivativeMode::ReverseModeCombined) {
    assert(!tape);
    return malloc;
  }

#ifndef NDEBUG
  if (auto CI = dyn_cast<CallInst>(malloc)) {
    if (auto F = CI->getCalledFunction()) {
      assert(F->getName() != "omp_get_thread_num");
    }
  }
#endif

  if (malloc->getType()->isTokenTy()) {
    llvm::errs() << " oldFunc: " << *oldFunc << "\n";
    llvm::errs() << " newFunc: " << *newFunc << "\n";
    llvm::errs() << " malloc: " << *malloc << "\n";
  }
  assert(!malloc->getType()->isTokenTy());
  {
    CountTrackedPointers T(malloc->getType());
    if (T.derived) {
      llvm::errs() << " oldFunc: " << *oldFunc << "\n";
      llvm::errs() << " newFunc: " << *newFunc << "\n";
      llvm::errs() << " malloc: " << *malloc << "\n";
    }
    assert(!T.derived);
  }

  if (tape) {
    if (idx == IndexMappingError) {
      assert(malloc);
      return UndefValue::get(malloc->getType());
    }
    if (idx >= 0 && !tape->getType()->isStructTy()) {
      llvm::errs() << "cacheForReverse incorrect tape type: " << *tape
                   << " idx: " << idx << "\n";
    }
    assert(idx < 0 || tape->getType()->isStructTy());
    if (idx >= 0 &&
        (unsigned)idx >= cast<StructType>(tape->getType())->getNumElements()) {
      llvm::errs() << "oldFunc: " << *oldFunc << "\n";
      llvm::errs() << "newFunc: " << *newFunc << "\n";
      if (malloc)
        llvm::errs() << "malloc: " << *malloc << "\n";
      llvm::errs() << "tape: " << *tape << "\n";
      llvm::errs() << "idx: " << idx << "\n";
    }
    assert(idx < 0 ||
           (unsigned)idx < cast<StructType>(tape->getType())->getNumElements());
    Value *ret =
        (idx < 0) ? tape : BuilderQ.CreateExtractValue(tape, {(unsigned)idx});

    if (ret->getType()->isEmptyTy()) {
      if (auto inst = dyn_cast_or_null<Instruction>(malloc)) {
        if (inst->getType() != ret->getType()) {
          llvm::errs() << "oldFunc: " << *oldFunc << "\n";
          llvm::errs() << "newFunc: " << *newFunc << "\n";
          llvm::errs() << "inst==malloc: " << *inst << "\n";
          llvm::errs() << "ret: " << *ret << "\n";
        }
        assert(inst->getType() == ret->getType());
        if (replace) {
          inst->replaceAllUsesWith(UndefValue::get(ret->getType()));
          erase(inst);
        }
      }
      Type *retType = ret->getType();
      if (replace)
        if (auto ri = dyn_cast<Instruction>(ret))
          erase(ri);
      return UndefValue::get(retType);
    }

    LimitContext ctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                     BuilderQ.GetInsertBlock());
    if (auto inst = dyn_cast<Instruction>(malloc))
      ctx = LimitContext(/*ReverseLimit*/ reverseBlocks.size() > 0,
                         inst->getParent());
    if (auto found = findInMap(scopeMap, malloc)) {
      ctx = found->second;
    }
    assert(isOriginalBlock(*ctx.Block));

    bool inLoop;
    if (ctx.ForceSingleIteration) {
      inLoop = true;
      ctx.ForceSingleIteration = false;
    } else {
      LoopContext lc;
      inLoop = getContext(ctx.Block, lc);
    }

    if (!inLoop) {
      ret->setName(malloc->getName() + "_fromtape");
      if (omp) {
        Value *tid = ompThreadId();
        Value *tPtr = BuilderQ.CreateInBoundsGEP(malloc->getType(), ret,
                                                 ArrayRef<Value *>(tid));
        ret = BuilderQ.CreateLoad(malloc->getType(), tPtr);
      }
    } else {
      if (idx >= 0)
        erase(cast<Instruction>(ret));
      IRBuilder<> entryBuilder(inversionAllocs);
      entryBuilder.setFastMathFlags(getFast());
      ret = (idx < 0) ? tape
                      : entryBuilder.CreateExtractValue(tape, {(unsigned)idx});

      assert(malloc);

      Type *innerType = nullptr;

#if LLVM_VERSION_MAJOR < 17
#if LLVM_VERSION_MAJOR >= 15
      if (ret->getContext().supportsTypedPointers()) {
#endif
        innerType = ret->getType();
        for (size_t i = 0,
                    limit = getSubLimits(
                                /*inForwardPass*/ true, nullptr,
                                LimitContext(
                                    /*ReverseLimit*/ reverseBlocks.size() > 0,
                                    BuilderQ.GetInsertBlock()))
                                .size();
             i < limit; ++i) {
          if (!isa<PointerType>(innerType)) {
            llvm::errs() << "mod: "
                         << *BuilderQ.GetInsertBlock()->getParent()->getParent()
                         << "\n";
            llvm::errs() << "fn: " << *BuilderQ.GetInsertBlock()->getParent()
                         << "\n";
            llvm::errs() << "bq insertblock: " << *BuilderQ.GetInsertBlock()
                         << "\n";
            llvm::errs() << "ret: " << *ret << " type: " << *ret->getType()
                         << "\n";
            llvm::errs() << "innerType: " << *innerType << "\n";
            if (malloc)
              llvm::errs() << " malloc: " << *malloc << " i=" << i
                           << " / lim = " << limit << "\n";
          }
          assert(isa<PointerType>(innerType));
          innerType = innerType->getPointerElementType();
        }
#if LLVM_VERSION_MAJOR >= 15
      } else {
        if (EfficientBoolCache && malloc->getType()->isIntegerTy() &&
            cast<IntegerType>(malloc->getType())->getBitWidth() == 1)
          innerType = Type::getInt8Ty(malloc->getContext());
        else
          innerType = malloc->getType();
      }
#endif
#else
      if (EfficientBoolCache && malloc->getType()->isIntegerTy() &&
          cast<IntegerType>(malloc->getType())->getBitWidth() == 1)
        innerType = Type::getInt8Ty(malloc->getContext());
      else
        innerType = malloc->getType();
#endif

      if (EfficientBoolCache && malloc->getType()->isIntegerTy() &&
          cast<IntegerType>(malloc->getType())->getBitWidth() == 1 &&
          innerType != ret->getType()) {
        assert(innerType == Type::getInt8Ty(malloc->getContext()));
      } else {
        if (innerType != malloc->getType()) {
          std::string str;
          raw_string_ostream ss(str);
          ss << "Illegal loop cache type:\n";
          ss << *oldFunc << "\n";
          ss << *newFunc << "\n";
          ss << "innerType: " << *innerType << "\n";
          ss << "malloc->getType(): " << *malloc->getType() << "\n";
          ss << "ret: " << *ret << " - " << *ret->getType() << "\n";
          ss << "malloc: " << *malloc << "\n";
          if (CustomErrorHandler) {
            CustomErrorHandler(str.c_str(), wrap(malloc),
                               ErrorType::InternalError, nullptr, nullptr,
                               nullptr);
          } else {
            DebugLoc loc;
            if (auto I = dyn_cast<Instruction>(malloc))
              EmitFailure("LoopCache", I->getDebugLoc(), I, ss.str());
            else
              EmitFailure("LoopCache", DebugLoc(), newFunc, ss.str());
          }
          return UndefValue::get(malloc->getType());
        }
      }

      LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                        BuilderQ.GetInsertBlock());
      AllocaInst *cache =
          createCacheForScope(lctx, innerType, "mdyncache_fromtape",
                              ((DiffeGradientUtils *)this)->FreeMemory, false);
      assert(malloc);
      bool isi1 = malloc->getType()->isIntegerTy() &&
                  cast<IntegerType>(malloc->getType())->getBitWidth() == 1;
      assert(isa<PointerType>(cache->getType()));
#if LLVM_VERSION_MAJOR < 17
#if LLVM_VERSION_MAJOR >= 15
      if (cache->getContext().supportsTypedPointers()) {
#endif
        assert(cache->getType()->getPointerElementType() == ret->getType());
#if LLVM_VERSION_MAJOR >= 15
      }
#endif
#endif
      entryBuilder.CreateStore(ret, cache);

      auto v =
          lookupValueFromCache(innerType, /*forwardPass*/ true, BuilderQ, lctx,
                               cache, isi1, /*available*/ ValueToValueMapTy());
      if (malloc) {
        assert(v->getType() == malloc->getType());
      }
      insert_or_assign(scopeMap, v,
                       std::make_pair(AssertingVH<AllocaInst>(cache), ctx));
      ret = cast<Instruction>(v);
    }

    if (malloc && !isa<UndefValue>(malloc)) {
      if (malloc->getType() != ret->getType()) {
        llvm::errs() << *oldFunc << "\n";
        llvm::errs() << *newFunc << "\n";
        llvm::errs() << *malloc << "\n";
        llvm::errs() << *ret << "\n";
      }
      assert(malloc->getType() == ret->getType());

      if (replace) {
        auto found = newToOriginalFn.find(malloc);
        if (found != newToOriginalFn.end()) {
          Value *orig = found->second;
          originalToNewFn[orig] = ret;
          newToOriginalFn.erase(malloc);
          newToOriginalFn[ret] = orig;
        }
      }

      if (auto found = findInMap(scopeMap, malloc)) {
        // There already exists an alloaction for this, we should fully remove
        // it
        if (!inLoop) {

          // Remove stores into
          SmallVector<Instruction *, 3> stores(
              scopeInstructions[found->first].begin(),
              scopeInstructions[found->first].end());
          scopeInstructions.erase(found->first);
          for (int i = stores.size() - 1; i >= 0; i--) {
            erase(stores[i]);
          }

          SmallVector<User *, 4> users;
          for (auto u : found->first->users()) {
            users.push_back(u);
          }
          for (auto u : users) {
            if (auto li = dyn_cast<LoadInst>(u)) {
              IRBuilder<> lb(li);
              if (replace) {

                Value *replacewith =
                    (idx < 0) ? tape
                              : lb.CreateExtractValue(tape, {(unsigned)idx});
                if (!inLoop && omp) {
                  Value *tid = ompThreadId();
                  Value *tPtr = lb.CreateInBoundsGEP(li->getType(), replacewith,
                                                     ArrayRef<Value *>(tid));
                  replacewith = lb.CreateLoad(li->getType(), tPtr);
                }
                if (li->getType() != replacewith->getType()) {
                  llvm::errs() << " oldFunc: " << *oldFunc << "\n";
                  llvm::errs() << " newFunc: " << *newFunc << "\n";
                  llvm::errs() << " malloc: " << *malloc << "\n";
                  llvm::errs() << " li: " << *li << "\n";
                  llvm::errs() << " u: " << *u << "\n";
                  llvm::errs() << " replacewith: " << *replacewith
                               << " idx=" << idx << " - tape=" << *tape << "\n";
                }
                assert(li->getType() == replacewith->getType());
                li->replaceAllUsesWith(replacewith);
              } else {
                auto phi =
                    lb.CreatePHI(li->getType(), 0, li->getName() + "_cfrphi");
                unwrappedLoads[phi] = malloc;
                li->replaceAllUsesWith(phi);
              }
              erase(li);
            } else {
              llvm::errs() << "newFunc: " << *newFunc << "\n";
              llvm::errs() << "malloc: " << *malloc << "\n";
              llvm::errs() << "scopeMap[malloc]: " << *found->first << "\n";
              llvm::errs() << "u: " << *u << "\n";
              assert(0 && "illegal use for out of loop scopeMap1");
            }
          }

          {
            AllocaInst *preerase = found->first;
            scopeMap.erase(malloc);
            erase(preerase);
          }
        } else {
          // Remove allocations for scopealloc since it is already allocated
          // by the augmented forward pass
          // Remove stores into
          SmallVector<Instruction *, 3> stores(
              scopeInstructions[found->first].begin(),
              scopeInstructions[found->first].end());
          scopeInstructions.erase(found->first);
          scopeAllocs.erase(found->first);
          for (int i = stores.size() - 1; i >= 0; i--) {
            erase(stores[i]);
          }

          // Remove frees
          SmallVector<CallInst *, 3> tofree(scopeFrees[found->first].begin(),
                                            scopeFrees[found->first].end());
          scopeFrees.erase(found->first);
          for (auto freeinst : tofree) {
            // This deque contains a list of operations
            // we can erasing upon erasing the free (and so on).
            // Since multiple operations can have the same operand,
            // this deque can contain the same value multiple times.
            // To remedy this we use a tracking value handle which will
            // be set to null when erased.
            std::deque<WeakTrackingVH> ops = {freeinst->getArgOperand(0)};
            erase(freeinst);

            while (ops.size()) {
              auto z = dyn_cast_or_null<Instruction>(ops[0]);
              ops.pop_front();
              if (z && z->getNumUses() == 0 && !z->isUsedByMetadata()) {
                for (unsigned i = 0; i < z->getNumOperands(); ++i) {
                  ops.push_back(z->getOperand(i));
                }
                erase(z);
              }
            }
          }

          // uses of the alloc
          SmallVector<User *, 4> users;
          for (auto u : found->first->users()) {
            users.push_back(u);
          }
          for (auto u : users) {
            if (auto li = dyn_cast<LoadInst>(u)) {
              // even with replace off, this can be replaced
              // as since we're in a loop this load is a load of cache
              // not of the final value (thereby overwriting the new
              // inst
              IRBuilder<> lb(li);
              auto replacewith =
                  (idx < 0) ? tape
                            : lb.CreateExtractValue(tape, {(unsigned)idx});
              li->replaceAllUsesWith(replacewith);
              erase(li);
            } else {
              llvm::errs() << "newFunc: " << *newFunc << "\n";
              llvm::errs() << "malloc: " << *malloc << "\n";
              llvm::errs() << "scopeMap[malloc]: " << *found->first << "\n";
              llvm::errs() << "u: " << *u << "\n";
              assert(0 && "illegal use for out of loop scopeMap2");
            }
          }

          AllocaInst *preerase = found->first;
          scopeMap.erase(malloc);
          if (replace)
            erase(preerase);
        }
      }
      if (replace)
        cast<Instruction>(malloc)->replaceAllUsesWith(ret);
      ret->takeName(malloc);
      if (replace) {
        auto malloci = cast<Instruction>(malloc);
        if (malloci == &*BuilderQ.GetInsertPoint()) {
          BuilderQ.SetInsertPoint(malloci->getNextNode());
        }
        erase(malloci);
      }
    }
    return ret;
  } else {
    assert(malloc);

    assert(idx >= 0 && (unsigned)idx == addedTapeVals.size());

    if (isa<UndefValue>(malloc)) {
      addedTapeVals.push_back(malloc);
      return malloc;
    }

    LimitContext ctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                     BuilderQ.GetInsertBlock());
    if (auto inst = dyn_cast<Instruction>(malloc))
      ctx = LimitContext(/*ReverseLimit*/ reverseBlocks.size() > 0,
                         inst->getParent());
    if (auto found = findInMap(scopeMap, malloc)) {
      ctx = found->second;
    }

    bool inLoop;

    if (ctx.ForceSingleIteration) {
      inLoop = true;
      ctx.ForceSingleIteration = false;
    } else {
      LoopContext lc;
      inLoop = getContext(ctx.Block, lc);
    }

    if (!inLoop) {
      Value *toStoreInTape = malloc;
      if (omp) {
        Value *numThreads = ompNumThreads();
        Value *tid = ompThreadId();
        IRBuilder<> entryBuilder(inversionAllocs);

        auto firstallocation =
            CreateAllocation(entryBuilder, malloc->getType(), numThreads,
                             malloc->getName() + "_malloccache");
        Value *tPtr = entryBuilder.CreateInBoundsGEP(
            malloc->getType(), firstallocation, ArrayRef<Value *>(tid));
        if (auto inst = dyn_cast<Instruction>(malloc)) {
          entryBuilder.SetInsertPoint(inst->getNextNode());
        }
        entryBuilder.CreateStore(malloc, tPtr);
        toStoreInTape = firstallocation;
      }
      addedTapeVals.push_back(toStoreInTape);
      return malloc;
    }

    ensureLookupCached(
        cast<Instruction>(malloc),
        /*shouldFree=*/reverseBlocks.size() > 0,
        /*scope*/ nullptr,
        cast<Instruction>(malloc)->getMetadata(LLVMContext::MD_tbaa));
    auto found2 = scopeMap.find(malloc);
    assert(found2 != scopeMap.end());
    assert(found2->second.first);

    Value *toadd;
    toadd = scopeAllocs[found2->second.first][0];
    for (auto u : toadd->users()) {
      if (auto ci = dyn_cast<CastInst>(u)) {
        toadd = ci;
        break;
      }
    }

    // llvm::errs() << " malloc: " << *malloc << "\n";
    // llvm::errs() << " toadd: " << *toadd << "\n";
#if LLVM_VERSION_MAJOR < 17
#if LLVM_VERSION_MAJOR >= 15
    if (toadd->getContext().supportsTypedPointers()) {
#endif
      Type *innerType = toadd->getType();
      for (size_t i = 0,
                  limit = getSubLimits(
                              /*inForwardPass*/ true, nullptr,
                              LimitContext(
                                  /*ReverseLimit*/ reverseBlocks.size() > 0,
                                  BuilderQ.GetInsertBlock()))
                              .size();
           i < limit; ++i) {
        innerType = innerType->getPointerElementType();
      }
      if (EfficientBoolCache && malloc->getType()->isIntegerTy() &&
          toadd->getType() != innerType &&
          cast<IntegerType>(malloc->getType())->getBitWidth() == 1) {
        assert(innerType == Type::getInt8Ty(toadd->getContext()));
      } else {
        if (innerType != malloc->getType()) {
          llvm::errs() << "oldFunc:" << *oldFunc << "\n";
          llvm::errs() << "newFunc: " << *newFunc << "\n";
          llvm::errs() << " toadd: " << *toadd << "\n";
          llvm::errs() << "innerType: " << *innerType << "\n";
          llvm::errs() << "malloc: " << *malloc << "\n";
        }
        assert(innerType == malloc->getType());
      }
#if LLVM_VERSION_MAJOR >= 15
    }
#endif
#endif
    addedTapeVals.push_back(toadd);
    return malloc;
  }
  llvm::errs()
      << "Fell through on cacheForReverse. This should never happen.\n";
  assert(false);
}

BasicBlock *GradientUtils::prepRematerializedLoopEntry(LoopContext &lc) {
  auto header = lc.header;
  SmallPtrSet<Instruction *, 1> loopRematerializations;
  SmallPtrSet<Instruction *, 1> loopReallocations;
  SmallPtrSet<Instruction *, 1> loopShadowReallocations;
  SmallSetVector<Instruction *, 1> loopShadowZeroInits;
  SmallPtrSet<Instruction *, 1> loopShadowRematerializations;
  Loop *origLI = nullptr;
  for (auto pair : rematerializableAllocations) {
    if (pair.second.LI &&
        getNewFromOriginal(pair.second.LI->getHeader()) == header) {
      if (allocationsToBeRematerialized.count(pair.first)) {
        if (auto inst = dyn_cast<Instruction>(pair.first))
          if (pair.second.LI->contains(inst->getParent())) {
            loopReallocations.insert(inst);
          }
        for (auto I : pair.second.stores)
          loopRematerializations.insert(I);
        origLI = pair.second.LI;
      }
    }
  }
  for (auto pair : backwardsOnlyShadows) {
    if (pair.second.LI &&
        getNewFromOriginal(pair.second.LI->getHeader()) == header) {
      if (auto inst = dyn_cast<Instruction>(pair.first)) {
        bool restoreStores = false;
        if (pair.second.LI->contains(inst->getParent())) {
          // TODO later make it so primalInitialize can be restored
          // rather than cached from primal
          if (!pair.second.primalInitialize) {
            loopShadowReallocations.insert(inst);
            restoreStores = true;
          }
        } else {
          // if (pair.second.primalInitialize) {
          //  loopShadowZeroInits.insert(inst);
          //}
          restoreStores = true;
        }
        if (restoreStores) {
          for (auto I : pair.second.stores) {
            loopShadowRematerializations.insert(I);
          }
        }
        origLI = pair.second.LI;
      }
    }
  }
  if (loopRematerializations.size() != 0 || loopReallocations.size() != 0 ||
      loopShadowRematerializations.size() != 0 ||
      loopShadowReallocations.size() != 0 || loopShadowZeroInits.size() != 0) {
    auto found = rematerializedLoops_cache.find(header);
    if (found != rematerializedLoops_cache.end()) {
      return found->second;
    }

    BasicBlock *enterB =
        BasicBlock::Create(header->getContext(), "remat_enter", newFunc);
    rematerializedLoops_cache[header] = enterB;
    std::map<BasicBlock *, BasicBlock *> origToNewForward;
    for (auto B : origLI->getBlocks()) {
      BasicBlock *newB = BasicBlock::Create(
          B->getContext(), "remat_" + header->getName() + "_" + B->getName(),
          newFunc);
      origToNewForward[B] = newB;
      reverseBlockToPrimal[newB] = getNewFromOriginal(B);
      if (B == origLI->getHeader()) {
        IRBuilder<> NB(newB);
        for (auto inst : loopShadowZeroInits) {
          auto anti = lookupM(invertPointerM(inst, NB), NB);
          StringRef funcName;
          SmallVector<Value *, 8> args;
          if (auto orig = dyn_cast<CallInst>(inst)) {
#if LLVM_VERSION_MAJOR >= 14
            for (auto &arg : orig->args())
#else
            for (auto &arg : orig->arg_operands())
#endif
            {
              args.push_back(lookupM(getNewFromOriginal(arg), NB));
            }
            funcName = getFuncNameFromCall(orig);
          } else if (auto AI = dyn_cast<AllocaInst>(inst)) {
            funcName = "malloc";
            Value *sz = lookupM(getNewFromOriginal(AI->getArraySize()), NB);

            auto ci = ConstantInt::get(
                sz->getType(),
                B->getParent()
                        ->getParent()
                        ->getDataLayout()
                        .getTypeAllocSizeInBits(AI->getAllocatedType()) /
                    8);
            sz = NB.CreateMul(sz, ci);
            args.push_back(sz);
          }
          assert(funcName.size());

          applyChainRule(
              NB,
              [&](Value *anti) {
                zeroKnownAllocation(NB, anti, args, funcName, TLI,
                                    dyn_cast<CallInst>(inst));
              },
              anti);
        }
      }
    }

    ValueToValueMapTy available;

    {
      IRBuilder<> NB(enterB);
      NB.CreateBr(origToNewForward[origLI->getHeader()]);
    }

    std::function<void(Loop *, bool)> handleLoop = [&](Loop *OL, bool subLoop) {
      if (subLoop) {
        auto Header = OL->getHeader();
        IRBuilder<> NB(origToNewForward[Header]);
        LoopContext flc;
        getContext(getNewFromOriginal(Header), flc);

        auto iv = NB.CreatePHI(flc.var->getType(), 2, "fiv");
        auto inc = NB.CreateAdd(iv, ConstantInt::get(iv->getType(), 1));

        for (auto PH : predecessors(Header)) {
          if (notForAnalysis.count(PH))
            continue;

          if (OL->contains(PH))
            iv->addIncoming(inc, origToNewForward[PH]);
          else
            iv->addIncoming(ConstantInt::get(iv->getType(), 0),
                            origToNewForward[PH]);
        }
        available[flc.var] = iv;
        available[flc.incvar] = inc;
      }
      for (auto SL : OL->getSubLoops())
        handleLoop(SL, /*subLoop*/ true);
    };
    handleLoop(origLI, /*subLoop*/ false);

    for (auto B : origLI->getBlocks()) {
      auto newB = origToNewForward[B];
      IRBuilder<> NB(newB);

      // TODO fill available with relevant IV's surrounding and
      // IV's of inner loop phi's

      for (auto &I : *B) {
        // Only handle store, memset, and julia.write_barrier
        if (loopRematerializations.count(&I)) {
          if (auto SI = dyn_cast<StoreInst>(&I)) {
            auto ts = NB.CreateStore(
                lookupM(getNewFromOriginal(SI->getValueOperand()), NB,
                        available),
                lookupM(getNewFromOriginal(SI->getPointerOperand()), NB,
                        available));
            llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
            ToCopy2.push_back(LLVMContext::MD_noalias);
            ToCopy2.push_back(LLVMContext::MD_alias_scope);
            ts->copyMetadata(*SI, ToCopy2);
            ts->setAlignment(SI->getAlign());
            ts->setVolatile(SI->isVolatile());
            ts->setOrdering(SI->getOrdering());
            ts->setSyncScopeID(SI->getSyncScopeID());
            ts->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
          } else if (auto CI = dyn_cast<CallInst>(&I)) {
            StringRef funcName = getFuncNameFromCall(CI);
            if (funcName == "enzyme_zerotype")
              continue;
            if (funcName == "julia.write_barrier" ||
                funcName == "julia.write_barrier_binding" ||
                isa<MemSetInst>(&I) || isa<MemTransferInst>(&I)) {

              // TODO
              SmallVector<Value *, 2> args;
#if LLVM_VERSION_MAJOR >= 14
              for (auto &arg : CI->args())
#else
              for (auto &arg : CI->arg_operands())
#endif
                args.push_back(lookupM(getNewFromOriginal(arg), NB, available));

              SmallVector<ValueType, 2> BundleTypes(args.size(),
                                                    ValueType::Primal);

              auto Defs = getInvertedBundles(CI, BundleTypes, NB,
                                             /*lookup*/ true, available);
              auto cal = NB.CreateCall(CI->getFunctionType(),
                                       CI->getCalledOperand(), args, Defs);
              cal->setAttributes(CI->getAttributes());
              cal->setCallingConv(CI->getCallingConv());
              cal->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
            } else {
              assert(isDeallocationFunction(funcName, TLI));
              continue;
            }
          } else {
            assert(0 && "unhandlable loop rematerialization instruction");
          }
        } else if (loopReallocations.count(&I)) {
          LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                            &newFunc->getEntryBlock());

          auto inst = getNewFromOriginal((Value *)&I);

          auto found = scopeMap.find(inst);
          if (found == scopeMap.end()) {
            AllocaInst *cache = createCacheForScope(
                lctx, inst->getType(), inst->getName(), /*shouldFree*/ true);
            assert(cache);
            found = insert_or_assign(
                scopeMap, inst,
                std::pair<AssertingVH<AllocaInst>, LimitContext>(cache, lctx));
          }
          auto cache = found->second.first;
          if (auto MD = hasMetadata(&I, "enzyme_fromstack")) {
            auto replacement = NB.CreateAlloca(
                Type::getInt8Ty(I.getContext()),
                lookupM(getNewFromOriginal(I.getOperand(0)), NB, available));
            for (auto MD : {"enzyme_active", "enzyme_inactive", "enzyme_type",
                            "enzymejl_allocart", "enzymejl_allocart_name",
                            "enzymejl_gc_alloc_rt"})
              if (auto M = I.getMetadata(MD))
                replacement->setMetadata(MD, M);
            auto Alignment =
                cast<ConstantInt>(
                    cast<ConstantAsMetadata>(MD->getOperand(0))->getValue())
                    ->getLimitedValue();
            replacement->setAlignment(Align(Alignment));
            replacement->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
            storeInstructionInCache(lctx, NB, replacement, cache);
          } else if (auto CI = dyn_cast<CallInst>(&I)) {
            SmallVector<Value *, 2> args;
#if LLVM_VERSION_MAJOR >= 14
            for (auto &arg : CI->args())
#else
            for (auto &arg : CI->arg_operands())
#endif
              args.push_back(lookupM(getNewFromOriginal(arg), NB, available));

            SmallVector<ValueType, 2> BundleTypes(args.size(),
                                                  ValueType::Primal);

            auto Defs = getInvertedBundles(CI, BundleTypes, NB,
                                           /*lookup*/ true, available);
            auto cal = NB.CreateCall(CI->getCalledFunction(), args, Defs);
            llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
            ToCopy2.push_back(LLVMContext::MD_noalias);
            ToCopy2.push_back(LLVMContext::MD_alias_scope);
            cal->copyMetadata(*CI, ToCopy2);
            cal->setName("remat_" + CI->getName());
            cal->setAttributes(CI->getAttributes());
            cal->setCallingConv(CI->getCallingConv());
            cal->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
            storeInstructionInCache(lctx, NB, cal, cache);
          } else {
            llvm::errs() << " realloc: " << I << "\n";
            llvm_unreachable("Unknown loop reallocation");
          }
        }
        if (loopShadowRematerializations.count(&I)) {
          if (auto SI = dyn_cast<StoreInst>(&I)) {
            Value *orig_ptr = SI->getPointerOperand();
            Value *orig_val = SI->getValueOperand();
            Type *valType = orig_val->getType();
            assert(!isConstantValue(orig_ptr));

            auto &DL = newFunc->getParent()->getDataLayout();

            bool constantval = isConstantValue(orig_val) ||
                               parseTBAA(I, DL, nullptr)[{-1}].isIntegral();

            // TODO allow recognition of other types that could contain
            // pointers [e.g. {void*, void*} or <2 x i64> ]
            auto storeSize = DL.getTypeSizeInBits(valType) / 8;

            //! Storing a floating point value
            Type *FT = nullptr;
            if (valType->isFPOrFPVectorTy()) {
              FT = valType->getScalarType();
            } else if (!valType->isPointerTy()) {
              if (looseTypeAnalysis) {
                auto fp = TR.firstPointer(storeSize, orig_ptr, &I, this,
                                          /*errifnotfound*/ nullptr,
                                          /*pointerIntSame*/ true);
                if (fp.isKnown()) {
                  FT = fp.isFloat();
                  llvm::errs() << "assuming type as " << fp.str()
                               << " for store: " << I << "\n";
                } else if (isa<ConstantInt>(orig_val) ||
                           valType->isIntOrIntVectorTy()) {
                  llvm::errs()
                      << "assuming type as integral for store: " << I << "\n";
                  FT = nullptr;
                } else {
                  TR.firstPointer(storeSize, orig_ptr, &I, this,
                                  /*errifnotfound*/ &NB,
                                  /*pointerIntSame*/ true);
                }
              } else {
                FT = TR.firstPointer(storeSize, orig_ptr, &I, this,
                                     /*errifnotfound*/ &NB,
                                     /*pointerIntSame*/ true)
                         .isFloat();
              }
            }
            if (!FT) {
              Value *valueop = nullptr;
              if (constantval) {
                Value *val =
                    lookupM(getNewFromOriginal(orig_val), NB, available);
                valueop = val;
                if (getWidth() > 1) {
                  Value *array = UndefValue::get(getShadowType(val->getType()));
                  for (unsigned i = 0; i < getWidth(); ++i) {
                    array = NB.CreateInsertValue(array, val, {i});
                  }
                  valueop = array;
                }
              } else {
                valueop = lookupM(invertPointerM(orig_val, NB), NB, available);
              }
              SmallVector<Metadata *, 1> prevScopes;
              if (auto prev = SI->getMetadata(LLVMContext::MD_alias_scope)) {
                for (auto &M : cast<MDNode>(prev)->operands()) {
                  prevScopes.push_back(M);
                }
              }
              SmallVector<Metadata *, 1> prevNoAlias;
              if (auto prev = SI->getMetadata(LLVMContext::MD_noalias)) {
                for (auto &M : cast<MDNode>(prev)->operands()) {
                  prevNoAlias.push_back(M);
                }
              }
              auto align = SI->getAlign();
              setPtrDiffe(SI, orig_ptr, valueop, NB, align, 0, storeSize,
                          SI->isVolatile(), SI->getOrdering(),
                          SI->getSyncScopeID(),
                          /*mask*/ nullptr, prevNoAlias, prevScopes);
            }
            // TODO shadow memtransfer
          } else if (auto MS = dyn_cast<MemSetInst>(&I)) {
            if (!isConstantValue(MS->getArgOperand(0))) {
              Value *args[4] = {
                  lookupM(invertPointerM(MS->getArgOperand(0), NB), NB,
                          available),
                  lookupM(getNewFromOriginal(MS->getArgOperand(1)), NB,
                          available),
                  lookupM(getNewFromOriginal(MS->getArgOperand(2)), NB,
                          available),
                  lookupM(getNewFromOriginal(MS->getArgOperand(3)), NB,
                          available)};

              ValueType BundleTypes[4] = {ValueType::Shadow, ValueType::Primal,
                                          ValueType::Primal, ValueType::Primal};
              auto Defs = getInvertedBundles(MS, BundleTypes, NB,
                                             /*lookup*/ true, available);
              auto cal = NB.CreateCall(MS->getCalledFunction(), args, Defs);
              llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
              ToCopy2.push_back(LLVMContext::MD_noalias);
              ToCopy2.push_back(LLVMContext::MD_alias_scope);
              cal->copyMetadata(*MS, ToCopy2);
              cal->setAttributes(MS->getAttributes());
              cal->setCallingConv(MS->getCallingConv());
              cal->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
            }
          } else if (auto CI = dyn_cast<CallInst>(&I)) {
            StringRef funcName = getFuncNameFromCall(CI);
            if (funcName == "julia.write_barrier" ||
                funcName == "julia.write_barrier_binding") {

              // TODO
              SmallVector<Value *, 2> args;
#if LLVM_VERSION_MAJOR >= 14
              for (auto &arg : CI->args())
#else
              for (auto &arg : CI->arg_operands())
#endif
                if (!isConstantValue(arg))
                  args.push_back(
                      lookupM(invertPointerM(arg, NB), NB, available));

              if (args.size()) {
                SmallVector<ValueType, 2> BundleTypes(args.size(),
                                                      ValueType::Primal);

                auto Defs = getInvertedBundles(CI, BundleTypes, NB,
                                               /*lookup*/ true, available);
                auto cal = NB.CreateCall(CI->getFunctionType(),
                                         CI->getCalledOperand(), args, Defs);
                cal->setAttributes(CI->getAttributes());
                cal->setCallingConv(CI->getCallingConv());
                cal->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
              }
            } else {
              assert(isDeallocationFunction(funcName, TLI));
              continue;
            }
          } else {
            assert(0 &&
                   "unhandlable loop shadow rematerialization instruction");
          }
        } else if (loopShadowReallocations.count(&I)) {

          LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                            &newFunc->getEntryBlock());
          auto ipfound = invertedPointers.find(&I);
          PHINode *placeholder = cast<PHINode>(&*ipfound->second);

          auto found = scopeMap.find(placeholder);
          if (found == scopeMap.end()) {
            AllocaInst *cache = createCacheForScope(
                lctx, placeholder->getType(), placeholder->getName(),
                /*shouldFree*/ true);
            assert(cache);
            Value *placeholder_tmp = placeholder;
            found = insert_or_assign(
                scopeMap, placeholder_tmp,
                std::pair<AssertingVH<AllocaInst>, LimitContext>(cache, lctx));
          }
          auto cache = found->second.first;
          Value *anti = nullptr;

          if (auto orig = dyn_cast<CallInst>(&I)) {
            StringRef funcName = getFuncNameFromCall(orig);
            assert(funcName.size());

            SmallVector<Value *, 8> args;
#if LLVM_VERSION_MAJOR >= 14
            for (auto &arg : orig->args())
#else
            for (auto &arg : orig->arg_operands())
#endif
            {
              args.push_back(lookupM(getNewFromOriginal(arg), NB));
            }

            placeholder->setName("");
            if (shadowHandlers.find(funcName) != shadowHandlers.end()) {

              anti = shadowHandlers[funcName](NB, orig, args, this);
            } else {
              auto rule = [&]() {
                Value *anti = NB.CreateCall(orig->getFunctionType(),
                                            orig->getCalledOperand(), args,
                                            orig->getName() + "'mi");
                cast<CallInst>(anti)->setAttributes(orig->getAttributes());
                cast<CallInst>(anti)->setCallingConv(orig->getCallingConv());
                cast<CallInst>(anti)->setDebugLoc(
                    getNewFromOriginal(I.getDebugLoc()));

                cast<CallInst>(anti)->addAttribute(AttributeList::ReturnIndex,
                                                   Attribute::NoAlias);
                cast<CallInst>(anti)->addAttribute(AttributeList::ReturnIndex,
                                                   Attribute::NonNull);
                return anti;
              };

              anti = applyChainRule(orig->getType(), NB, rule);

              if (auto MD = hasMetadata(orig, "enzyme_fromstack")) {
                auto rule = [&](Value *anti) {
                  AllocaInst *replacement = NB.CreateAlloca(
                      Type::getInt8Ty(orig->getContext()), args[0]);
                  for (auto MD :
                       {"enzyme_active", "enzyme_inactive", "enzyme_type",
                        "enzymejl_allocart", "enzymejl_allocart_name"})
                    if (auto M = I.getMetadata(MD))
                      replacement->setMetadata(MD, M);
                  replacement->takeName(anti);
                  auto Alignment = cast<ConstantInt>(cast<ConstantAsMetadata>(
                                                         MD->getOperand(0))
                                                         ->getValue())
                                       ->getLimitedValue();
                  replacement->setAlignment(Align(Alignment));
                  replacement->setDebugLoc(getNewFromOriginal(I.getDebugLoc()));
                  return replacement;
                };

                Value *replacement = applyChainRule(
                    Type::getInt8Ty(orig->getContext()), NB, rule, anti);

                replaceAWithB(cast<Instruction>(anti), replacement);
                erase(cast<Instruction>(anti));
                anti = replacement;
              }

              applyChainRule(
                  NB,
                  [&](Value *anti) {
                    zeroKnownAllocation(NB, anti, args, funcName, TLI, orig);
                  },
                  anti);
            }
          } else {
            llvm_unreachable("Unknown shadow rematerialization value");
          }
          assert(anti);
          storeInstructionInCache(lctx, NB, anti, cache);
        }
      }

      llvm::SmallPtrSet<llvm::BasicBlock *, 8> origExitBlocks;
      getExitBlocks(origLI, origExitBlocks);
      // Remap a branch to the header to enter the incremented
      // reverse of that block.
      auto remap = [&](BasicBlock *rB) {
        // Remap of an exit branch is to go to the reverse
        // exiting block.
        if (origExitBlocks.count(rB)) {
          return reverseBlocks[getNewFromOriginal(B)].front();
        }
        // Reverse of an incrementing branch is go to the
        // reverse of the branching block.
        if (rB == origLI->getHeader())
          return reverseBlocks[getNewFromOriginal(B)].front();
        auto found = origToNewForward.find(rB);
        if (found == origToNewForward.end()) {
          llvm::errs() << *newFunc << "\n";
          llvm::errs() << *origLI << "\n";
          llvm::errs() << *rB << "\n";
        }
        assert(found != origToNewForward.end());
        return found->second;
      };

      // TODO clone terminator
      auto TI = B->getTerminator();
      assert(TI);
      if (notForAnalysis.count(B)) {
        NB.CreateUnreachable();
      } else if (auto BI = dyn_cast<BranchInst>(TI)) {
        if (BI->isUnconditional()) {
          if (notForAnalysis.count(BI->getSuccessor(0)))
            NB.CreateUnreachable();
          else
            NB.CreateBr(remap(BI->getSuccessor(0)));
        } else {
          if (notForAnalysis.count(BI->getSuccessor(0))) {
            if (notForAnalysis.count(BI->getSuccessor(1))) {
              NB.CreateUnreachable();
            } else {
              NB.CreateBr(remap(BI->getSuccessor(1)));
            }
          } else if (notForAnalysis.count(BI->getSuccessor(1))) {
            NB.CreateBr(remap(BI->getSuccessor(0)));
          } else {
            NB.CreateCondBr(
                lookupM(getNewFromOriginal(BI->getCondition()), NB, available),
                remap(BI->getSuccessor(0)), remap(BI->getSuccessor(1)));
          }
        }
      } else if (auto SI = dyn_cast<SwitchInst>(TI)) {
        BasicBlock *newDest = nullptr;
        if (!notForAnalysis.count(SI->getDefaultDest()))
          newDest = remap(SI->getDefaultDest());
        else {
          for (auto cas : SI->cases()) {
            if (!notForAnalysis.count(cas.getCaseSuccessor()))
              newDest = remap(cas.getCaseSuccessor());
            break;
          }
        }
        if (!newDest) {
          NB.CreateUnreachable();
        } else {
          auto NSI = NB.CreateSwitch(
              lookupM(getNewFromOriginal(SI->getCondition()), NB, available),
              newDest);
          for (auto cas : SI->cases()) {
            if (!notForAnalysis.count(cas.getCaseSuccessor()))
              NSI->addCase(cas.getCaseValue(), remap(cas.getCaseSuccessor()));
          }
        }
      } else {
        assert(isa<UnreachableInst>(TI));
        NB.CreateUnreachable();
      }
      // Fixup phi nodes that may have their predecessors now changed by
      // the phi unwrapping
      if (!notForAnalysis.count(B) &&
          NB.GetInsertBlock() != origToNewForward[B]) {
        for (auto S0 : successors(B)) {
          if (!origToNewForward.count(S0))
            continue;
          auto S = origToNewForward[S0];
          assert(S);
          for (auto I = S->begin(), E = S->end(); I != E; ++I) {
            PHINode *orig = dyn_cast<PHINode>(&*I);
            if (orig == nullptr)
              break;
            for (unsigned Op = 0, NumOps = orig->getNumOperands(); Op != NumOps;
                 ++Op)
              if (orig->getIncomingBlock(Op) == origToNewForward[B])
                orig->setIncomingBlock(Op, NB.GetInsertBlock());
          }
        }
      }
    }
    return enterB;
  }
  return nullptr;
}

/// Given an edge from BB to branchingBlock get the corresponding block to
/// branch to in the reverse pass
BasicBlock *GradientUtils::getReverseOrLatchMerge(BasicBlock *BB,
                                                  BasicBlock *branchingBlock) {
  assert(BB);
  // BB should be a forward pass block, assert that
  if (reverseBlocks.find(BB) == reverseBlocks.end()) {
    llvm::errs() << *oldFunc << "\n";
    llvm::errs() << *newFunc << "\n";
    llvm::errs() << "BB: " << *BB << "\n";
    llvm::errs() << "branchingBlock: " << *branchingBlock << "\n";
  }
  assert(reverseBlocks.find(BB) != reverseBlocks.end());
  assert(reverseBlocks.find(branchingBlock) != reverseBlocks.end());
  LoopContext lc;
  bool inLoop = getContext(BB, lc);

  LoopContext branchingContext;
  bool inLoopContext = getContext(branchingBlock, branchingContext);

  if (!inLoop)
    return reverseBlocks[BB].front();

  auto tup = std::make_tuple(BB, branchingBlock);
  if (newBlocksForLoop_cache.find(tup) != newBlocksForLoop_cache.end())
    return newBlocksForLoop_cache[tup];

  // If we're reversing a latch edge.
  bool incEntering = inLoopContext && branchingBlock == lc.header &&
                     lc.header == branchingContext.header;

  auto L = LI.getLoopFor(BB);
  auto latches = getLatches(L, lc.exitBlocks);
  // If we're reverseing a loop exit.
  bool exitEntering =
      std::find(latches.begin(), latches.end(), BB) != latches.end() &&
      std::find(lc.exitBlocks.begin(), lc.exitBlocks.end(), branchingBlock) !=
          lc.exitBlocks.end();

  // It is illegal to be both an increment into a loop, and exiting the loop.
  assert(!(incEntering && exitEntering));

  // If we're re-entering a loop, prepare a loop-level forward pass to
  // rematerialize any loop-scope rematerialization.

  if (incEntering) {
    BasicBlock *resumeblock = reverseBlocks[BB].front();
    auto tmp_resumeblock = prepRematerializedLoopEntry(lc);
    if (tmp_resumeblock)
      resumeblock = tmp_resumeblock;
    BasicBlock *incB = BasicBlock::Create(
        BB->getContext(), "inc" + reverseBlocks[lc.header].front()->getName(),
        BB->getParent());
    incB->moveAfter(reverseBlocks[lc.header].back());

    IRBuilder<> tbuild(incB);

    Value *av = tbuild.CreateLoad(lc.var->getType(), lc.antivaralloc);
    Value *sub = tbuild.CreateAdd(av, ConstantInt::get(av->getType(), -1), "",
                                  /*NUW*/ false, /*NSW*/ true);
    tbuild.CreateStore(sub, lc.antivaralloc);
    tbuild.CreateBr(resumeblock);
    return newBlocksForLoop_cache[tup] = incB;
  }

  if (exitEntering) {
    SmallVector<LoopContext, 1> exitingContexts = {lc};

    auto L2 = L;
    while ((L2 = L2->getParentLoop())) {
      LoopContext lc2;
      bool inLoop = getContext(L2->getHeader(), lc2);
      if (!inLoop)
        break;

      auto latches2 = getLatches(L2, lc2.exitBlocks);

      // If we're reverseing a loop exit.
      bool exitEntering2 =
          std::find(latches2.begin(), latches2.end(), BB) != latches2.end() &&
          std::find(lc2.exitBlocks.begin(), lc2.exitBlocks.end(),
                    branchingBlock) != lc2.exitBlocks.end();
      if (exitEntering2) {
        exitingContexts.push_back(lc2);
      } else
        break;
    }

    BasicBlock *resumeblock = reverseBlocks[BB].front();
    BasicBlock *prevBlock = reverseBlocks[branchingBlock].back();

    BasicBlock *outerMerge = nullptr;

    BasicBlock *incB = BasicBlock::Create(
        BB->getContext(),
        "merge" + reverseBlocks[lc.header].front()->getName() + "_" +
            branchingBlock->getName(),
        BB->getParent());
    if (!outerMerge)
      outerMerge = incB;
    incB->moveAfter(prevBlock);

    IRBuilder<> tbuild(prevBlock);

    SmallVector<std::pair<Value *, Value *>, 1> lims;
    ValueToValueMapTy available;
    for (auto I = exitingContexts.rbegin(), E = exitingContexts.rend(); I != E;
         I++) {
      auto &lc = *I;
      auto L = LI.getLoopFor(lc.header);
      Value *lim = nullptr;
      if (lc.dynamic && assumeDynamicLoopOfSizeOne(L)) {
        lim = ConstantInt::get(lc.var->getType(), 0);
      } else if (lc.dynamic) {
        // Must be in a reverse pass fashion for a lookup to index bound to be
        // legal
        assert(/*ReverseLimit*/ reverseBlocks.size() > 0);
        LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                          lc.preheader);
        lim = lookupValueFromCache(lc.var->getType(),
                                   /*forwardPass*/ false, tbuild, lctx,
                                   getDynamicLoopLimit(L),
                                   /*isi1*/ false, available);
      } else {
        lim = lookupM(lc.trueLimit, tbuild, available);
      }
      available[lc.var] = lim;
      lims.push_back(std::make_pair(lim, (Value *)lc.antivaralloc));
    }

    tbuild.SetInsertPoint(incB);
    for (auto &pair : lims) {
      tbuild.CreateStore(pair.first, pair.second);
    }

    auto tmp_resumeblock = prepRematerializedLoopEntry(exitingContexts.back());
    if (tmp_resumeblock)
      resumeblock = tmp_resumeblock;

    tbuild.CreateBr(resumeblock);

    return newBlocksForLoop_cache[tup] = incB;
  }

  return newBlocksForLoop_cache[tup] = reverseBlocks[BB].front();
}

void GradientUtils::forceContexts() {
  for (auto BB : originalBlocks) {
    LoopContext lc;
    getContext(BB, lc);
  }
}

bool GradientUtils::legalRecompute(const Value *val,
                                   const ValueToValueMapTy &available,
                                   IRBuilder<> *BuilderM, bool reverse,
                                   bool legalRecomputeCache) const {
  {
    auto found = available.find(val);
    if (found != available.end()) {
      if (found->second)
        return true;
      else {
        return false;
      }
    }
  }

  if (isa<AtomicRMWInst>(val))
    return false;

  if (auto phi = dyn_cast<PHINode>(val)) {
    if (auto uiv = hasUninverted(val)) {
      if (auto dli = dyn_cast_or_null<LoadInst>(uiv)) {
        return legalRecompute(
            dli, available, BuilderM,
            reverse); // TODO ADD && !TR.intType(getOriginal(dli),
                      // /*mustfind*/false).isPossibleFloat();
      }
      if (auto ci = dyn_cast<CallInst>(uiv)) {
        auto called = ci->getCalledFunction();
        if (ci->hasFnAttr("enzyme_shouldrecompute") ||
            (called && called->hasFnAttribute("enzyme_shouldrecompute")))
          return true;
      }
      if (phi->getNumIncomingValues() == 0) {
        return false;
      }
    }

    auto found = fictiousPHIs.find(const_cast<llvm::PHINode *>(phi));
    if (found != fictiousPHIs.end()) {
      auto orig = found->second;
      if (isa<AtomicRMWInst>(orig))
        return false;
    }

    if (phi->getNumIncomingValues() == 0) {
      std::string str;
      raw_string_ostream ss(str);
      ss << "oldFunc: " << *oldFunc << "\n";
      ss << "newFunc: " << *newFunc << "\n";
      ss << "phi: " << *phi << "\n";
      ss << "Invalid legalRecompute query on ficticious phi\n";
      if (CustomErrorHandler) {
        CustomErrorHandler(str.c_str(), wrap(phi), ErrorType::InternalError,
                           nullptr, nullptr, nullptr);
      } else {
        EmitFailure("InvalidLegalRecompute", phi->getDebugLoc(), phi, ss.str());
      }
    }
    auto parent = phi->getParent();
    struct {
      Function *func;
      const LoopInfo &FLI;
    } options[2] = {{newFunc, LI}, {oldFunc, *OrigLI}};
    for (const auto &tup : options) {
      if (parent->getParent() == tup.func) {
        for (auto &val : phi->incoming_values()) {
          if (isPotentialLastLoopValue(val, parent, tup.FLI)) {
            return false;
          }
        }
        if (tup.FLI.isLoopHeader(parent)) {
          // Currently can only recompute header
          // with two incoming values
          if (phi->getNumIncomingValues() != 2)
            return false;
          auto L = tup.FLI.getLoopFor(parent);

          // Only recomputable if non recursive.
          SmallPtrSet<Instruction *, 2> seen;
          SmallVector<Instruction *, 1> todo;
          for (auto PH : predecessors(parent)) {
            // Prior iterations must be recomputable without
            // this value.
            if (L->contains(PH)) {
              if (auto I =
                      dyn_cast<Instruction>(phi->getIncomingValueForBlock(PH)))
                if (L->contains(I->getParent()))
                  todo.push_back(I);
            }
          }

          while (todo.size()) {
            auto cur = todo.back();
            todo.pop_back();
            if (seen.count(cur))
              continue;
            seen.insert(cur);
            if (cur == phi)
              return false;
            for (auto &op : cur->operands()) {
              if (auto I = dyn_cast<Instruction>(op)) {
                if (L->contains(I->getParent()))
                  todo.push_back(I);
              }
            }
          }
        }
        return true;
      }
    }
    return false;
  }

  if (isa<Instruction>(val) &&
      cast<Instruction>(val)->getMetadata("enzyme_mustcache")) {
    return false;
  }

  // If this is a load from cache already, dont force a cache of this
  if (legalRecomputeCache && isa<LoadInst>(val) &&
      CacheLookups.count(cast<LoadInst>(val))) {
    return true;
  }

  // TODO consider callinst here

  if (auto li = dyn_cast<Instruction>(val)) {

    const IntrinsicInst *II;
    if (isa<LoadInst>(li) || isNVLoad(li) ||
        ((II = dyn_cast<IntrinsicInst>(li)) &&
         (II->getIntrinsicID() == Intrinsic::masked_load))) {
      // If this is an already unwrapped value, legal to recompute again.
      if (unwrappedLoads.find(li) != unwrappedLoads.end())
        return legalRecompute(unwrappedLoads.find(li)->second, available,
                              BuilderM, reverse);

      const Instruction *orig = nullptr;
      if (li->getParent()->getParent() == oldFunc) {
        orig = li;
      } else if (li->getParent()->getParent() == newFunc) {
        orig = isOriginal(li);
        // todo consider when we pass non original queries
        if (orig && !isa<LoadInst>(orig)) {
          return legalRecompute(orig, available, BuilderM, reverse,
                                legalRecomputeCache);
        }
      } else {
        llvm::errs() << " newFunc: " << *newFunc << "\n";
        llvm::errs() << " parent: " << *li->getParent()->getParent() << "\n";
        llvm::errs() << " li: " << *li << "\n";
        assert(0 && "illegal load legalRecopmute query");
      }

      if (orig) {
        assert(can_modref_map);
        auto found = can_modref_map->find(const_cast<Instruction *>(orig));
        if (found == can_modref_map->end()) {
          llvm::errs() << *newFunc << "\n";
          llvm::errs() << *oldFunc << "\n";
          llvm::errs() << "can_modref_map:\n";
          for (auto &pair : *can_modref_map) {
            llvm::errs() << " + " << *pair.first << ": " << pair.second
                         << " of func "
                         << pair.first->getParent()->getParent()->getName()
                         << "\n";
          }
          llvm::errs() << "couldn't find in can_modref_map: " << *li << " - "
                       << *orig << " in fn: "
                       << orig->getParent()->getParent()->getName();
        }
        assert(found != can_modref_map->end());
        if (!found->second)
          return true;
        // if insertion block of this function:
        BasicBlock *fwdBlockIfReverse = nullptr;
        if (BuilderM) {
          fwdBlockIfReverse = BuilderM->GetInsertBlock();
          if (!reverse) {
            auto found = reverseBlockToPrimal.find(BuilderM->GetInsertBlock());
            if (found != reverseBlockToPrimal.end()) {
              fwdBlockIfReverse = found->second;
              reverse = true;
            }
          }
          if (fwdBlockIfReverse->getParent() != oldFunc)
            fwdBlockIfReverse =
                cast_or_null<BasicBlock>(isOriginal(fwdBlockIfReverse));
        }
        if (mode == DerivativeMode::ReverseModeCombined && fwdBlockIfReverse) {
          if (reverse) {
            bool failed = false;
            allFollowersOf(
                const_cast<Instruction *>(orig), [&](Instruction *I) -> bool {
                  if (I->mayWriteToMemory() &&
                      writesToMemoryReadBy(
                          &TR, *OrigAA, TLI,
                          /*maybeReader*/ const_cast<Instruction *>(orig),
                          /*maybeWriter*/ I)) {
                    failed = true;
                    EmitWarning(
                        "UncacheableLoad", *orig, "Load must be recomputed ",
                        *orig, " in reverse_",
                        BuilderM->GetInsertBlock()->getName(), " due to ", *I);
                    return /*earlyBreak*/ true;
                  }
                  return /*earlyBreak*/ false;
                });
            if (!failed)
              return true;
          } else {
            Instruction *origStart = &*BuilderM->GetInsertPoint();
            do {
              if (Instruction *og = isOriginal(origStart)) {
                origStart = og;
                break;
              }
              origStart = origStart->getNextNode();
            } while (true);
            if (OrigDT->dominates(origStart, const_cast<Instruction *>(orig))) {
              bool failed = false;

              allInstructionsBetween(
                  const_cast<GradientUtils *>(this)->LI, origStart,
                  const_cast<Instruction *>(orig), [&](Instruction *I) -> bool {
                    if (I->mayWriteToMemory() &&
                        writesToMemoryReadBy(
                            &TR, *OrigAA, TLI,
                            /*maybeReader*/ const_cast<Instruction *>(orig),
                            /*maybeWriter*/ I)) {
                      failed = true;
                      EmitWarning("UncacheableLoad", *orig,
                                  "Load must be recomputed ", *orig, " in ",
                                  BuilderM->GetInsertBlock()->getName(),
                                  " due to ", *I);
                      return /*earlyBreak*/ true;
                    }
                    return /*earlyBreak*/ false;
                  });
              if (!failed)
                return true;
            }
          }
        }
        return false;
      } else {
        if (auto dli = dyn_cast_or_null<LoadInst>(hasUninverted(li))) {
          return legalRecompute(dli, available, BuilderM, reverse);
        }

        // TODO mark all the explicitly legal nodes (caches, etc)
        return true;
        llvm::errs() << *li << " orig: " << orig
                     << " parent: " << li->getParent()->getParent()->getName()
                     << "\n";
        llvm_unreachable("unknown load to redo!");
      }
    }
  }

  if (auto ci = dyn_cast<CallInst>(val)) {
    auto n = getFuncNameFromCall(const_cast<CallInst *>(ci));
    auto called = ci->getCalledFunction();
    Intrinsic::ID ID = Intrinsic::not_intrinsic;

    if (ci->hasFnAttr("enzyme_shouldrecompute") ||
        (called && called->hasFnAttribute("enzyme_shouldrecompute")) ||
        isMemFreeLibMFunction(n, &ID) || n == "lgamma_r" || n == "lgammaf_r" ||
        n == "lgammal_r" || n == "__lgamma_r_finite" ||
        n == "__lgammaf_r_finite" || n == "__lgammal_r_finite" || n == "tanh" ||
        n == "tanhf" || n == "__pow_finite" ||
        n == "julia.pointer_from_objref" || startsWith(n, "enzyme_wrapmpi$$") ||
        n == "omp_get_thread_num" || n == "omp_get_max_threads" ||
        n.contains("__enzyme_ignore_derivatives")) {
      return true;
    }
#if LLVM_VERSION_MAJOR >= 14
    if (ci->doesNotAccessMemory())
#else
    if (ci->hasFnAttr(Attribute::ReadNone) ||
        (called && called->hasFnAttribute(Attribute::ReadNone)))
#endif
      return true;
    if (isPointerArithmeticInst(ci))
      return true;
  }

  if (auto inst = dyn_cast<Instruction>(val)) {
    if (inst->mayReadOrWriteMemory()) {
      return false;
    }
  }

  return true;
}

//! Given the option to recompute a value or re-use an old one, return true if
//! it is faster to recompute this value from scratch
bool GradientUtils::shouldRecompute(const Value *val,
                                    const ValueToValueMapTy &available,
                                    IRBuilder<> *BuilderM) {
  if (available.count(val))
    return true;
  // TODO: remake such that this returns whether a load to a cache is more
  // expensive than redoing the computation.

  // If this is a load from cache already, just reload this
  if (isa<LoadInst>(val) &&
      cast<LoadInst>(val)->getMetadata("enzyme_fromcache"))
    return true;

  if (!isa<Instruction>(val))
    return true;

  const Instruction *inst = cast<Instruction>(val);

  if (TapesToPreventRecomputation.count(inst))
    return false;

  if (knownRecomputeHeuristic.find(inst) != knownRecomputeHeuristic.end()) {
    return knownRecomputeHeuristic[inst];
  }
  if (auto OrigInst = isOriginal(inst)) {
    if (knownRecomputeHeuristic.find(OrigInst) !=
        knownRecomputeHeuristic.end()) {
      return knownRecomputeHeuristic[OrigInst];
    }
  }

  if (isa<CastInst>(val) || isa<GetElementPtrInst>(val))
    return true;

  if (EnzymeNewCache && !EnzymeMinCutCache) {
    // if this has operands that need to be loaded and haven't already been
    // loaded
    // TODO, just cache this
    for (auto &op : inst->operands()) {
      if (!legalRecompute(op, available, BuilderM)) {

        // If this is a load from cache already, dont force a cache of this
        if (isa<LoadInst>(op) && CacheLookups.count(cast<LoadInst>(op)))
          continue;

        // If a previously cached this operand, don't let it trigger the
        // heuristic for caching this value instead.
        if (scopeMap.find(op) != scopeMap.end())
          continue;

        // If the actually overwritten operand is in a different loop scope
        // don't cache this value instead as it may require more memory
        LoopContext lc1;
        LoopContext lc2;
        bool inLoop1 =
            getContext(const_cast<Instruction *>(inst)->getParent(), lc1);
        bool inLoop2 = getContext(cast<Instruction>(op)->getParent(), lc2);
        if (inLoop1 != inLoop2 || (inLoop1 && (lc1.header != lc2.header))) {
          continue;
        }

        // If a placeholder phi for inversion (and we know from above not
        // recomputable)
        if (!isa<PHINode>(op) &&
            dyn_cast_or_null<LoadInst>(hasUninverted(op))) {
          goto forceCache;
        }

        // Even if cannot recompute (say a phi node), don't force a reload if it
        // is possible to just use this instruction from forward pass without
        // issue
        if (auto i2 = dyn_cast<Instruction>(op)) {
          if (!i2->mayReadOrWriteMemory()) {
            LoopContext lc;
            bool inLoop = const_cast<GradientUtils *>(this)->getContext(
                i2->getParent(), lc);
            if (!inLoop) {
              // TODO upgrade this to be all returns that this could enter from
              BasicBlock *orig = isOriginal(i2->getParent());
              assert(orig);
              bool legal = BlocksDominatingAllReturns.count(orig);
              if (legal) {
                continue;
              }
            }
          }
        }
      forceCache:;
        EmitWarning("ChosenCache", *inst, "Choosing to cache use ", *inst,
                    " due to ", *op);
        return false;
      }
    }
  }

  if (auto op = dyn_cast<IntrinsicInst>(val)) {
    if (!op->mayReadOrWriteMemory() || isReadNone(op) || isNVLoad(op))
      return true;
    switch (op->getIntrinsicID()) {
    case Intrinsic::sin:
    case Intrinsic::cos:
    case Intrinsic::exp:
#if LLVM_VERSION_MAJOR >= 19
    case Intrinsic::tanh:
    case Intrinsic::cosh:
    case Intrinsic::sinh:
#endif
    case Intrinsic::log:
      return true;
    default:
      return false;
    }
  }

  if (auto ci = dyn_cast<CallInst>(val)) {
    auto called = ci->getCalledFunction();
    auto n = getFuncNameFromCall(const_cast<CallInst *>(ci));
    Intrinsic::ID ID = Intrinsic::not_intrinsic;
    if ((called && called->hasFnAttribute("enzyme_shouldrecompute")) ||
        isMemFreeLibMFunction(n, &ID) || n == "lgamma_r" || n == "lgammaf_r" ||
        n == "lgammal_r" || n == "__lgamma_r_finite" ||
        n == "__lgammaf_r_finite" || n == "__lgammal_r_finite" || n == "tanh" ||
        n == "tanhf" || n == "__pow_finite" ||
        n == "julia.pointer_from_objref" || startsWith(n, "enzyme_wrapmpi$$") ||
        n == "omp_get_thread_num" || n == "omp_get_max_threads" ||
        startsWith(n, "_ZN4libm4math3log") ||
        n.contains("__enzyme_ignore_derivatives")) {
      return true;
    }
    if (isPointerArithmeticInst(ci))
      return true;
  }

  // cache a call, assuming its longer to run that
  if (isa<CallInst>(val)) {
    llvm::errs() << " caching call: " << *val << "\n";
    // cast<CallInst>(val)->getCalledFunction()->dump();
    return false;
  }

  return true;
}

MDNode *GradientUtils::getDerivativeAliasScope(const Value *origptr,
                                               ssize_t newptr) {
  origptr = getBaseObject(origptr);

  auto found = differentialAliasScopeDomains.find(origptr);
  if (found == differentialAliasScopeDomains.end()) {
    MDBuilder MDB(oldFunc->getContext());
    MDNode *scope = MDB.createAnonymousAliasScopeDomain(
        (" diff: %" + origptr->getName()).str());
    // vec.first = scope;
    // found = differentialAliasScope.find(origptr);
    found = differentialAliasScopeDomains.insert(std::make_pair(origptr, scope))
                .first;
  }
  auto &mp = differentialAliasScope[origptr];
  auto found2 = mp.find(newptr);
  if (found2 == mp.end()) {
    MDBuilder MDB(oldFunc->getContext());
    std::string name;
    if (newptr == -1)
      name = "primal";
    else
      name = "shadow_" + std::to_string(newptr);
    found2 = mp.insert(std::make_pair(newptr, MDB.createAnonymousAliasScope(
                                                  found->second, name)))
                 .first;
  }
  return found2->second;
}

GradientUtils *GradientUtils::CreateFromClone(
    EnzymeLogic &Logic, bool runtimeActivity, bool strongZero, unsigned width,
    Function *todiff, TargetLibraryInfo &TLI, TypeAnalysis &TA,
    FnTypeInfo &oldTypeInfo, DIFFE_TYPE retType,
    ArrayRef<DIFFE_TYPE> constant_args, bool returnUsed, bool shadowReturnUsed,
    std::map<AugmentedStruct, int> &returnMapping, bool omp) {
  Function *oldFunc = todiff;

  // Since this is forward pass this should always return the tape (at index 0)
  returnMapping[AugmentedStruct::Tape] = 0;

  int returnCount = 0;

  if (returnUsed) {
    assert(!todiff->getReturnType()->isEmptyTy());
    assert(!todiff->getReturnType()->isVoidTy());
    returnMapping[AugmentedStruct::Return] = returnCount + 1;
    ++returnCount;
  }

  // We don't need to differentially return something that we know is not a
  // pointer (or somehow needed for shadow analysis)
  if (shadowReturnUsed) {
    assert(!todiff->getReturnType()->isEmptyTy());
    assert(!todiff->getReturnType()->isVoidTy());
    returnMapping[AugmentedStruct::DifferentialReturn] = returnCount + 1;
    ++returnCount;
  }

  ValueToValueMapTy invertedPointers;
  SmallPtrSet<Instruction *, 4> constants;
  SmallPtrSet<Instruction *, 20> nonconstant;
  SmallPtrSet<Value *, 2> returnvals;
  llvm::ValueMap<const llvm::Value *, AssertingReplacingVH> originalToNew;

  SmallPtrSet<Value *, 4> constant_values;
  SmallPtrSet<Value *, 4> nonconstant_values;

  std::string prefix = "fakeaugmented";
  if (width > 1)
    prefix += std::to_string(width);
  prefix += "_";
  prefix += todiff->getName().str();

  auto newFunc = Logic.PPC.CloneFunctionWithReturns(
      DerivativeMode::ReverseModePrimal, width, oldFunc, invertedPointers,
      constant_args, constant_values, nonconstant_values, returnvals,
      /*returnTape*/ true, /*returnPrimal*/ returnUsed,
      /*returnShadow*/ shadowReturnUsed, prefix, &originalToNew,
      /*diffeReturnArg*/ false, /*additionalArg*/ nullptr);

  // Convert overwritten args from the input function to the preprocessed
  // function

  FnTypeInfo typeInfo(oldFunc);
  {
    auto toarg = todiff->arg_begin();
    auto olarg = oldFunc->arg_begin();
    for (; toarg != todiff->arg_end(); ++toarg, ++olarg) {

      {
        auto fd = oldTypeInfo.Arguments.find(toarg);
        assert(fd != oldTypeInfo.Arguments.end());
        typeInfo.Arguments.insert(
            std::pair<Argument *, TypeTree>(olarg, fd->second));
      }

      {
        auto cfd = oldTypeInfo.KnownValues.find(toarg);
        assert(cfd != oldTypeInfo.KnownValues.end());
        typeInfo.KnownValues.insert(
            std::pair<Argument *, std::set<int64_t>>(olarg, cfd->second));
      }
    }
    typeInfo.Return = oldTypeInfo.Return;
  }

  TypeResults TR = TA.analyzeFunction(typeInfo);
  if (!oldFunc->empty())
    assert(TR.getFunction() == oldFunc);

  auto res = new GradientUtils(
      Logic, newFunc, oldFunc, TLI, TA, TR, invertedPointers, constant_values,
      nonconstant_values, retType, shadowReturnUsed, constant_args,
      originalToNew, DerivativeMode::ReverseModePrimal, runtimeActivity,
      strongZero, width, omp);
  return res;
}

DIFFE_TYPE GradientUtils::getReturnDiffeType(llvm::Value *orig,
                                             bool *primalReturnUsedP,
                                             bool *shadowReturnUsedP) const {
  return getReturnDiffeType(orig, primalReturnUsedP, shadowReturnUsedP, mode);
}

DIFFE_TYPE GradientUtils::getReturnDiffeType(llvm::Value *orig,
                                             bool *primalReturnUsedP,
                                             bool *shadowReturnUsedP,
                                             DerivativeMode cmode) const {
  bool shadowReturnUsed = false;

  DIFFE_TYPE subretType;
  if (isConstantValue(orig)) {
    subretType = DIFFE_TYPE::CONSTANT;
  } else {
    if (cmode == DerivativeMode::ForwardMode ||
        cmode == DerivativeMode::ForwardModeError ||
        cmode == DerivativeMode::ForwardModeSplit) {
      if (DifferentialUseAnalysis::is_value_needed_in_reverse<
              QueryType::Shadow>(this, orig, cmode, notForAnalysis)) {
        subretType = DIFFE_TYPE::DUP_ARG;
        shadowReturnUsed = true;
      } else {
        subretType = DIFFE_TYPE::CONSTANT;
      }
    } else {
      if (!orig->getType()->isFPOrFPVectorTy() && TR.anyPointer(orig)) {
        if (DifferentialUseAnalysis::is_value_needed_in_reverse<
                QueryType::Shadow>(this, orig, cmode, notForAnalysis)) {
          subretType = DIFFE_TYPE::DUP_ARG;
          shadowReturnUsed = true;
        } else
          subretType = DIFFE_TYPE::CONSTANT;
      } else {
        subretType = DIFFE_TYPE::OUT_DIFF;
      }
    }
  }

  if (primalReturnUsedP) {
    bool subretused = !unnecessaryValuesP || unnecessaryValuesP->find(orig) ==
                                                 unnecessaryValuesP->end();
    auto found = knownRecomputeHeuristic.find(orig);
    if (found != knownRecomputeHeuristic.end()) {
      if (!found->second) {
        subretused = true;
      }
    }
    *primalReturnUsedP = subretused;
  }

  if (shadowReturnUsedP)
    *shadowReturnUsedP = shadowReturnUsed;
  return subretType;
}

DIFFE_TYPE GradientUtils::getDiffeType(Value *v, bool foreignFunction) const {
  if (isConstantValue(v) && !foreignFunction) {
    return DIFFE_TYPE::CONSTANT;
  }

  auto argType = v->getType();

  if (!argType->isFPOrFPVectorTy() && (TR.anyPointer(v) || foreignFunction)) {
    if (argType->isPointerTy()) {
      auto at = getBaseObject(v);
      if (auto arg = dyn_cast<Argument>(at)) {
        if (ArgDiffeTypes[arg->getArgNo()] == DIFFE_TYPE::DUP_NONEED) {
          return DIFFE_TYPE::DUP_NONEED;
        }
      } else if ((isa<AllocaInst>(at) || isAllocationCall(at, TLI)) &&
                 unnecessaryValuesP) {
        if (unnecessaryValuesP->count(at)) {
          // Just because we chose to cahce the variable (and thus the value is
          // unnecessary) for saving, does not mean we are no need.
          auto found = knownRecomputeHeuristic.find(at);
          if (found == knownRecomputeHeuristic.end() || found->second)
            return DIFFE_TYPE::DUP_NONEED;
        }
      }
    }
    return DIFFE_TYPE::DUP_ARG;
  } else {
    if (foreignFunction)
      assert(!argType->isIntOrIntVectorTy());
    if (mode == DerivativeMode::ForwardMode ||
        mode == DerivativeMode::ForwardModeError ||
        mode == DerivativeMode::ForwardModeSplit)
      return DIFFE_TYPE::DUP_ARG;
    else
      return DIFFE_TYPE::OUT_DIFF;
  }
}

Constant *GradientUtils::GetOrCreateShadowConstant(
    RequestContext context, EnzymeLogic &Logic, TargetLibraryInfo &TLI,
    TypeAnalysis &TA, Constant *oval, DerivativeMode mode, bool runtimeActivity,
    bool strongZero, unsigned width, bool AtomicAdd) {
  if (isa<ConstantPointerNull>(oval)) {
    return oval;
  } else if (isa<UndefValue>(oval)) {
    return oval;
  } else if (isa<ConstantInt>(oval)) {
    return oval;
  } else if (auto CD = dyn_cast<ConstantDataArray>(oval)) {
    SmallVector<Constant *, 1> Vals;
    for (size_t i = 0, len = CD->getNumElements(); i < len; i++) {
      Vals.push_back(GetOrCreateShadowConstant(
          context, Logic, TLI, TA, CD->getElementAsConstant(i), mode,
          runtimeActivity, strongZero, width, AtomicAdd));
    }
    return ConstantArray::get(CD->getType(), Vals);
  } else if (auto CD = dyn_cast<ConstantArray>(oval)) {
    SmallVector<Constant *, 1> Vals;
    for (size_t i = 0, len = CD->getNumOperands(); i < len; i++) {
      Vals.push_back(GetOrCreateShadowConstant(
          context, Logic, TLI, TA, CD->getOperand(i), mode, runtimeActivity,
          strongZero, width, AtomicAdd));
    }
    return ConstantArray::get(CD->getType(), Vals);
  } else if (auto CD = dyn_cast<ConstantStruct>(oval)) {
    SmallVector<Constant *, 1> Vals;
    for (size_t i = 0, len = CD->getNumOperands(); i < len; i++) {
      Vals.push_back(GetOrCreateShadowConstant(
          context, Logic, TLI, TA, CD->getOperand(i), mode, runtimeActivity,
          strongZero, width, AtomicAdd));
    }
    return ConstantStruct::get(CD->getType(), Vals);
  } else if (auto CD = dyn_cast<ConstantVector>(oval)) {
    SmallVector<Constant *, 1> Vals;
    for (size_t i = 0, len = CD->getNumOperands(); i < len; i++) {
      Vals.push_back(GetOrCreateShadowConstant(
          context, Logic, TLI, TA, CD->getOperand(i), mode, runtimeActivity,
          strongZero, width, AtomicAdd));
    }
    return ConstantVector::get(Vals);
  } else if (auto F = dyn_cast<Function>(oval)) {
    return GetOrCreateShadowFunction(context, Logic, TLI, TA, F, mode,
                                     runtimeActivity, strongZero, width,
                                     AtomicAdd);
  } else if (auto arg = dyn_cast<ConstantExpr>(oval)) {
    auto C = GetOrCreateShadowConstant(
        context, Logic, TLI, TA, arg->getOperand(0), mode, runtimeActivity,
        strongZero, width, AtomicAdd);
    if (arg->isCast() || arg->getOpcode() == Instruction::GetElementPtr ||
        arg->getOpcode() == Instruction::Add) {
      SmallVector<Constant *, 8> NewOps;
      for (unsigned i = 0, e = arg->getNumOperands(); i != e; ++i)
        NewOps.push_back(i == 0 ? C : arg->getOperand(i));
      return arg->getWithOperands(NewOps);
    }
  } else if (auto arg = dyn_cast<GlobalAlias>(oval)) {
    return GetOrCreateShadowConstant(context, Logic, TLI, TA, arg->getAliasee(),
                                     mode, runtimeActivity, strongZero, width,
                                     AtomicAdd);
  } else if (auto arg = dyn_cast<GlobalVariable>(oval)) {
    if (arg->getName() == "_ZTVN10__cxxabiv120__si_class_type_infoE" ||
        arg->getName() == "_ZTVN10__cxxabiv117__class_type_infoE" ||
        arg->getName() == "_ZTVN10__cxxabiv121__vmi_class_type_infoE" ||
        startsWith(arg->getName(), "??_R")) // any of the MS RTTI manglings
      return arg;

    if (hasMetadata(arg, "enzyme_shadow")) {
      auto md = arg->getMetadata("enzyme_shadow");
      if (!isa<MDTuple>(md)) {
        llvm::errs() << *arg << "\n";
        llvm::errs() << *md << "\n";
        assert(0 && "cannot compute with global variable that doesn't have "
                    "marked shadow global");
        report_fatal_error(
            "cannot compute with global variable that doesn't "
            "have marked shadow global (metadata incorrect type)");
      }
      auto md2 = cast<MDTuple>(md);
      assert(md2->getNumOperands() == 1);
      auto gvemd = cast<ConstantAsMetadata>(md2->getOperand(0));
      return gvemd->getValue();
    }

    auto TT = llvm::Triple(arg->getParent()->getTargetTriple());
    int SharedAddrSpace = getGPUSharedAddrSpace(TT);
    int AddrSpace = cast<PointerType>(arg->getType())->getAddressSpace();
    if (isGPUArch(TT) && AddrSpace == SharedAddrSpace) {
      assert(0 && "shared memory not handled in meta global");
    }

    // Create global variable locally if not externally visible
    if (arg->isConstant() || arg->hasInternalLinkage() ||
        arg->hasPrivateLinkage() ||
        (arg->hasExternalLinkage() && arg->hasInitializer())) {
      Type *type = arg->getValueType();
      auto shadow = new GlobalVariable(
          *arg->getParent(), type, arg->isConstant(), arg->getLinkage(),
          Constant::getNullValue(type), arg->getName() + "_shadow", arg,
          arg->getThreadLocalMode(), arg->getType()->getAddressSpace(),
          arg->isExternallyInitialized());
      arg->setMetadata("enzyme_shadow",
                       MDTuple::get(shadow->getContext(),
                                    {ConstantAsMetadata::get(shadow)}));
      shadow->setAlignment(arg->getAlign());
      shadow->setUnnamedAddr(arg->getUnnamedAddr());
      if (arg->hasInitializer())
        shadow->setInitializer(GetOrCreateShadowConstant(
            context, Logic, TLI, TA, cast<Constant>(arg->getOperand(0)), mode,
            runtimeActivity, strongZero, width, AtomicAdd));
      return shadow;
    }
  }
  llvm::errs() << " unknown constant to create shadow of: " << *oval << "\n";
  llvm_unreachable("unknown constant to create shadow of");
}

Constant *GradientUtils::GetOrCreateShadowFunction(
    RequestContext context, EnzymeLogic &Logic, TargetLibraryInfo &TLI,
    TypeAnalysis &TA, Function *fn, DerivativeMode mode, bool runtimeActivity,
    bool strongZero, unsigned width, bool AtomicAdd) {
  //! Todo allow tape propagation
  //  Note that specifically this should _not_ be called with topLevel=true
  //  (since it may not be valid to always assume we can recompute the
  //  augmented primal) However, in the absence of a way to pass tape data
  //  from an indirect augmented (and also since we dont presently allow
  //  indirect augmented calls), topLevel MUST be true otherwise subcalls will
  //  not be able to lookup the augmenteddata/subdata (triggering an assertion
  //  failure, among much worse)
  bool isRealloc = false;
  if (fn->empty()) {
    if (hasMetadata(fn, "enzyme_callwrapper")) {
      auto md = fn->getMetadata("enzyme_callwrapper");
      if (!isa<MDTuple>(md)) {
        llvm::errs() << *fn << "\n";
        llvm::errs() << *md << "\n";
        assert(0 && "callwrapper of incorrect type");
        report_fatal_error("callwrapper of incorrect type");
      }
      auto md2 = cast<MDTuple>(md);
      assert(md2->getNumOperands() == 1);
      auto gvemd = cast<ConstantAsMetadata>(md2->getOperand(0));
      fn = cast<Function>(gvemd->getValue());
    } else {
      auto oldfn = fn;
      fn = Function::Create(oldfn->getFunctionType(), Function::InternalLinkage,
                            "callwrap_" + oldfn->getName(), oldfn->getParent());
      BasicBlock *entry = BasicBlock::Create(fn->getContext(), "entry", fn);
      IRBuilder<> B(entry);
      SmallVector<Value *, 4> args;
      for (auto &a : fn->args())
        args.push_back(&a);
      auto res = B.CreateCall(oldfn, args);
      if (fn->getReturnType()->isVoidTy())
        B.CreateRetVoid();
      else
        B.CreateRet(res);
      oldfn->setMetadata(
          "enzyme_callwrapper",
          MDTuple::get(oldfn->getContext(), {ConstantAsMetadata::get(fn)}));
      if (oldfn->getName() == "realloc")
        isRealloc = true;
    }
  }

  bool subsequent_calls_may_write = mode != DerivativeMode::ForwardMode &&
                                    mode != DerivativeMode::ForwardModeError;
  std::vector<bool> overwritten_args;
  FnTypeInfo type_args(fn);
  if (isRealloc) {
    llvm::errs() << "warning: assuming realloc only creates pointers\n";
    type_args.Return.insert({-1, -1}, BaseType::Pointer);
  }

  // conservatively assume that we can only cache existing floating types
  // (i.e. that all args are overwritten)
  std::vector<DIFFE_TYPE> types;
  for (auto &a : fn->args()) {
    overwritten_args.push_back(!a.getType()->isFPOrFPVectorTy());
    TypeTree TT;
    if (a.getType()->isFPOrFPVectorTy())
      TT.insert({-1}, ConcreteType(a.getType()->getScalarType()));
    type_args.Arguments.insert(std::pair<Argument *, TypeTree>(&a, TT));
    type_args.KnownValues.insert(
        std::pair<Argument *, std::set<int64_t>>(&a, {}));
    DIFFE_TYPE typ;
    if (a.getType()->isFPOrFPVectorTy()) {
      typ = (mode == DerivativeMode::ForwardMode ||
             mode == DerivativeMode::ForwardModeError)
                ? DIFFE_TYPE::DUP_ARG
                : DIFFE_TYPE::OUT_DIFF;
    } else if (a.getType()->isIntegerTy() &&
               cast<IntegerType>(a.getType())->getBitWidth() < 16) {
      typ = DIFFE_TYPE::CONSTANT;
    } else if (a.getType()->isVoidTy() || a.getType()->isEmptyTy()) {
      typ = DIFFE_TYPE::CONSTANT;
    } else {
      typ = DIFFE_TYPE::DUP_ARG;
    }
    types.push_back(typ);
  }

  DIFFE_TYPE retType = fn->getReturnType()->isFPOrFPVectorTy() &&
                               mode != DerivativeMode::ForwardMode &&
                               mode != DerivativeMode::ForwardModeError
                           ? DIFFE_TYPE::OUT_DIFF
                           : DIFFE_TYPE::DUP_ARG;

  if (fn->getReturnType()->isVoidTy() || fn->getReturnType()->isEmptyTy() ||
      (fn->getReturnType()->isIntegerTy() &&
       cast<IntegerType>(fn->getReturnType())->getBitWidth() < 16))
    retType = DIFFE_TYPE::CONSTANT;

  if (mode != DerivativeMode::ForwardMode &&
      mode != DerivativeMode::ForwardModeError &&
      retType == DIFFE_TYPE::DUP_ARG) {
    if (auto ST = dyn_cast<StructType>(fn->getReturnType())) {
      size_t numflt = 0;

      for (unsigned i = 0; i < ST->getNumElements(); ++i) {
        auto midTy = ST->getElementType(i);
        if (midTy->isFPOrFPVectorTy())
          numflt++;
      }
      if (numflt == ST->getNumElements())
        retType = DIFFE_TYPE::OUT_DIFF;
    }
  }

  std::vector<bool> nowrite_shadows(fn->arg_size(), false);

  switch (mode) {
  case DerivativeMode::ForwardModeError:
  case DerivativeMode::ForwardMode: {
    Constant *newf = Logic.CreateForwardDiff(
        context, fn, retType, types, TA, false, mode, /*freeMemory*/ true,
        runtimeActivity, strongZero, width, nullptr, type_args,
        subsequent_calls_may_write, overwritten_args,
        /*augmented*/ nullptr);

    assert(newf);

    std::string prefix = (mode == DerivativeMode::ForwardMode)
                             ? "_enzyme_forward"
                             : "_enzyme_forwarderror";

    if (width > 1) {
      prefix += std::to_string(width);
    }

    std::string globalname = (prefix + "_" + fn->getName() + "'").str();
    auto GV = fn->getParent()->getNamedValue(globalname);

    if (GV == nullptr) {
      GV = new GlobalVariable(*fn->getParent(), newf->getType(), true,
                              GlobalValue::LinkageTypes::InternalLinkage, newf,
                              globalname);
    }

    return ConstantExpr::getPointerCast(GV, fn->getType());
  }
  case DerivativeMode::ForwardModeSplit: {
    auto &augdata = Logic.CreateAugmentedPrimal(
        context, fn, retType, /*constant_args*/ types, TA,
        /*returnUsed*/ !fn->getReturnType()->isEmptyTy() &&
            !fn->getReturnType()->isVoidTy(),
        /*shadowReturnUsed*/ false, type_args, subsequent_calls_may_write,
        overwritten_args, nowrite_shadows,
        /*forceAnonymousTape*/ true, runtimeActivity, strongZero, width,
        AtomicAdd);
    Constant *newf = Logic.CreateForwardDiff(
        context, fn, retType, types, TA, false, mode, /*freeMemory*/ true,
        runtimeActivity, strongZero, width, nullptr, type_args,
        subsequent_calls_may_write, overwritten_args,
        /*augmented*/ &augdata);

    assert(newf);

    std::string prefix = "_enzyme_forwardsplit";

    if (width > 1) {
      prefix += std::to_string(width);
    }

    auto cdata = ConstantStruct::get(
        StructType::get(newf->getContext(),
                        {augdata.fn->getType(), newf->getType()}),
        {augdata.fn, newf});

    std::string globalname = (prefix + "_" + fn->getName() + "'").str();
    auto GV = fn->getParent()->getNamedValue(globalname);

    if (GV == nullptr) {
      GV = new GlobalVariable(*fn->getParent(), cdata->getType(), true,
                              GlobalValue::LinkageTypes::InternalLinkage, cdata,
                              globalname);
    }

    return ConstantExpr::getPointerCast(GV, fn->getType());
  }
  case DerivativeMode::ReverseModeCombined:
  case DerivativeMode::ReverseModeGradient:
  case DerivativeMode::ReverseModePrimal: {
    // TODO re atomic add consider forcing it to be atomic always as fallback if
    // used in a parallel context
    bool returnUsed =
        !fn->getReturnType()->isEmptyTy() && !fn->getReturnType()->isVoidTy();
    bool shadowReturnUsed = returnUsed && (retType == DIFFE_TYPE::DUP_ARG ||
                                           retType == DIFFE_TYPE::DUP_NONEED);
    auto &augdata = Logic.CreateAugmentedPrimal(
        context, fn, retType, /*constant_args*/ types, TA, returnUsed,
        shadowReturnUsed, type_args, subsequent_calls_may_write,
        overwritten_args, nowrite_shadows,
        /*forceAnonymousTape*/ true, runtimeActivity, strongZero, width,
        AtomicAdd);
    Constant *newf = Logic.CreatePrimalAndGradient(
        context,
        (ReverseCacheKey){.todiff = fn,
                          .retType = retType,
                          .constant_args = types,
                          .subsequent_calls_may_write =
                              subsequent_calls_may_write,
                          .overwritten_args = overwritten_args,
                          .returnUsed = false,
                          .shadowReturnUsed = false,
                          .mode = DerivativeMode::ReverseModeGradient,
                          .width = width,
                          .freeMemory = true,
                          .AtomicAdd = AtomicAdd,
                          .additionalType = getInt8PtrTy(fn->getContext()),
                          .forceAnonymousTape = true,
                          .typeInfo = type_args,
                          .runtimeActivity = runtimeActivity,
                          .strongZero = strongZero},
        TA,
        /*map*/ &augdata);
    assert(newf);
    auto cdata = ConstantStruct::get(
        StructType::get(newf->getContext(),
                        {augdata.fn->getType(), newf->getType()}),
        {augdata.fn, newf});
    std::string globalname = ("_enzyme_reverse_" + fn->getName() + "'").str();
    auto GV = fn->getParent()->getNamedValue(globalname);

    if (GV == nullptr) {
      GV = new GlobalVariable(*fn->getParent(), cdata->getType(), true,
                              GlobalValue::LinkageTypes::InternalLinkage, cdata,
                              globalname);
    }
    return ConstantExpr::getPointerCast(GV, fn->getType());
  }
  }
  llvm_unreachable("Illegal state: unknown mode for GetOrCreateShadowFunction");
}

void GradientUtils::getReverseBuilder(IRBuilder<> &Builder2, bool original) {
  assert(reverseBlocks.size());
  BasicBlock *BB = Builder2.GetInsertBlock();
  if (original)
    BB = getNewFromOriginal(BB);
  assert(reverseBlocks.find(BB) != reverseBlocks.end());
  BasicBlock *BB2 = reverseBlocks[BB].back();
  if (!BB2) {
    llvm::errs() << "oldFunc: " << oldFunc << "\n";
    llvm::errs() << "newFunc: " << newFunc << "\n";
    llvm::errs() << "could not invert " << *BB;
  }
  assert(BB2);

  if (hasTerminator(BB2))
    Builder2.SetInsertPoint(BB2->getTerminator());
  else
    Builder2.SetInsertPoint(BB2);
  Builder2.SetCurrentDebugLocation(
      getNewFromOriginal(Builder2.getCurrentDebugLocation()));
  Builder2.setFastMathFlags(getFast());
}

void GradientUtils::getForwardBuilder(IRBuilder<> &Builder2) {
  Instruction *insert = &*Builder2.GetInsertPoint();
  Instruction *nInsert = getNewFromOriginal(insert);

  assert(nInsert);

  Builder2.SetInsertPoint(getNextNonDebugInstruction(nInsert));
  Builder2.SetCurrentDebugLocation(
      getNewFromOriginal(Builder2.getCurrentDebugLocation()));
  Builder2.setFastMathFlags(getFast());
}

void GradientUtils::setPtrDiffe(Instruction *orig, Value *ptr, Value *newval,
                                IRBuilder<> &BuilderM, MaybeAlign align,
                                unsigned start, unsigned size, bool isVolatile,
                                AtomicOrdering ordering,
                                SyncScope::ID syncScope, Value *mask,
                                ArrayRef<Metadata *> noAlias,
                                ArrayRef<Metadata *> scopes,
                                bool needs_post_cache) {
#ifndef NDEBUG
  if (auto inst = dyn_cast<Instruction>(ptr)) {
    assert(inst->getParent()->getParent() == oldFunc);
  }
  if (auto arg = dyn_cast<Argument>(ptr)) {
    assert(arg->getParent() == oldFunc);
  }
#endif

  Value *origptr = ptr;

  ptr = invertPointerM(ptr, BuilderM);
  if (!isOriginalBlock(*BuilderM.GetInsertBlock()) &&
      mode != DerivativeMode::ForwardMode &&
      mode != DerivativeMode::ForwardModeError)
    ptr = lookupM(ptr, BuilderM);

  if (mask && !isOriginalBlock(*BuilderM.GetInsertBlock()) &&
      mode != DerivativeMode::ForwardMode &&
      mode != DerivativeMode::ForwardModeError)
    mask = lookupM(mask, BuilderM);

  size_t idx = 0;

  auto &DL = oldFunc->getParent()->getDataLayout();

  Value *invertedBarrier = nullptr;
  if (needs_post_cache && EnzymeJuliaAddrLoad &&
      anyJuliaObjects(newval->getType())) {
    auto obj = origptr;
    while (true) {
      if (auto CI = dyn_cast<CastInst>(obj)) {
        obj = CI->getOperand(0);
        continue;
      }
      if (auto GO = dyn_cast<GetElementPtrInst>(obj)) {
        obj = GO->getOperand(0);
        continue;
      }
      if (auto CI = dyn_cast<CallInst>(obj)) {
        if (getFuncNameFromCall(CI) == "julia.gc_loaded") {
          obj = CI->getArgOperand(0);
          continue;
        }
      }
      if (auto PT = dyn_cast<PointerType>(obj->getType())) {
        if (PT->getAddressSpace() == 13) {
          if (auto LI = dyn_cast<LoadInst>(obj)) {
            obj = LI->getOperand(0);
            continue;
          }
        }
      }
      break;
    }
    auto PT = cast<PointerType>(obj->getType());
    assert(PT->getAddressSpace() != 11 && PT->getAddressSpace() != 13);
    if (PT->getAddressSpace() == 10) {
      obj = invertPointerM(obj, BuilderM);

      if (!isOriginalBlock(*BuilderM.GetInsertBlock()) &&
          mode != DerivativeMode::ForwardMode &&
          mode != DerivativeMode::ForwardModeError)
        obj = lookupM(obj, BuilderM);

      invertedBarrier = obj;
    }
  }

  auto rule = [&](Value *ptr, Value *newval, Value *invertedBarrier) {
    auto storeSize = (DL.getTypeSizeInBits(newval->getType()) + 7) / 8;
    if (!mask) {

      if (size != storeSize) {
        IRBuilder<> A(inversionAllocs);
        Value *valptr = A.CreateAlloca(newval->getType());
        BuilderM.CreateStore(newval, valptr);

        auto i8 = Type::getInt8Ty(ptr->getContext());

        if (start != 0) {
          ptr = BuilderM.CreatePointerCast(
              ptr,
              getPointerType(
                  i8, cast<PointerType>(ptr->getType())->getAddressSpace()));
          auto off =
              ConstantInt::get(Type::getInt64Ty(ptr->getContext()), start);
          ptr = BuilderM.CreateInBoundsGEP(i8, ptr, off);

          valptr = BuilderM.CreatePointerCast(
              valptr,
              getPointerType(
                  i8, cast<PointerType>(valptr->getType())->getAddressSpace()));
          valptr = BuilderM.CreateInBoundsGEP(i8, valptr, off);
        }

        Type *ty = nullptr;

        if (size == 8)
          ty = BuilderM.getInt64Ty();
        else if (size % 8 == 0)
          ty = ArrayType::get(BuilderM.getInt64Ty(), size / 8);
        else if (size == 4)
          ty = BuilderM.getInt32Ty();
        else if (size % 4 == 0)
          ty = ArrayType::get(BuilderM.getInt32Ty(), size / 4);
        else
          ty = ArrayType::get(i8, size);

        ptr = BuilderM.CreatePointerCast(
            ptr, getPointerType(
                     ty, cast<PointerType>(ptr->getType())->getAddressSpace()));
        valptr = BuilderM.CreatePointerCast(
            valptr,
            getPointerType(
                ty, cast<PointerType>(valptr->getType())->getAddressSpace()));
        newval = BuilderM.CreateLoad(ty, valptr);
      }

      auto ts = BuilderM.CreateStore(newval, ptr);
      if (align)
        ts->setAlignment(*align);

      if (invertedBarrier) {
        auto T_jlvalue = StructType::get(ptr->getContext(), {});
        auto T_prjlvalue = getPointerType(T_jlvalue, 10);

        if (invertedBarrier->getType() != T_prjlvalue) {
          invertedBarrier =
              BuilderM.CreateBitCast(invertedBarrier, T_prjlvalue);
        }

        auto FT = FunctionType::get(Type::getVoidTy(newval->getContext()),
                                    {T_prjlvalue}, true);
        auto wb = BuilderM.GetInsertBlock()
                      ->getParent()
                      ->getParent()
                      ->getOrInsertFunction("julia.write_barrier", FT);

        auto subvals = getJuliaObjects(newval, BuilderM);

        subvals.insert(subvals.begin(), invertedBarrier);
        BuilderM.CreateCall(wb, subvals);
      }

      ts->setVolatile(isVolatile);
      ts->setOrdering(ordering);
      ts->setSyncScopeID(syncScope);
      SmallVector<Metadata *, 1> scopeMD = {
          getDerivativeAliasScope(origptr, idx)};
      for (auto M : scopes)
        scopeMD.push_back(M);
      auto scope = MDNode::get(ts->getContext(), scopeMD);
      ts->setMetadata(LLVMContext::MD_alias_scope, scope);

      if (start == 0 && size == storeSize) {
        ts->setMetadata(LLVMContext::MD_tbaa,
                        orig->getMetadata(LLVMContext::MD_tbaa));
        ts->setMetadata(LLVMContext::MD_tbaa_struct,
                        orig->getMetadata(LLVMContext::MD_tbaa_struct));
      }
      ts->setDebugLoc(getNewFromOriginal(orig->getDebugLoc()));

      SmallVector<Metadata *, 1> MDs;
      for (ssize_t j = -1; j < getWidth(); j++) {
        if (j != (ssize_t)idx)
          MDs.push_back(getDerivativeAliasScope(origptr, j));
      }
      for (auto M : noAlias)
        MDs.push_back(M);
      if (MDs.size()) {
        auto noscope = MDNode::get(ptr->getContext(), MDs);
        ts->setMetadata(LLVMContext::MD_noalias, noscope);
      }
    } else {
      assert(start == 0 && size == storeSize);
      Type *tys[] = {newval->getType(), ptr->getType()};
      auto F = getIntrinsicDeclaration(oldFunc->getParent(),
                                       Intrinsic::masked_store, tys);
      assert(align);
      assert(!needs_post_cache);
      Value *alignv =
          ConstantInt::get(Type::getInt32Ty(ptr->getContext()), align->value());
      Value *args[] = {newval, ptr, alignv, mask};
      auto ts = BuilderM.CreateCall(F, args);
      ts->setCallingConv(F->getCallingConv());
      ts->setMetadata(LLVMContext::MD_tbaa,
                      orig->getMetadata(LLVMContext::MD_tbaa));
      ts->setMetadata(LLVMContext::MD_tbaa_struct,
                      orig->getMetadata(LLVMContext::MD_tbaa_struct));
      ts->setDebugLoc(getNewFromOriginal(orig->getDebugLoc()));
    }
    idx++;
  };

  applyChainRule(BuilderM, rule, ptr, newval, invertedBarrier);
}

Type *GradientUtils::getShadowType(Type *ty, unsigned width) {
  if (width > 1) {
    if (ty->isVoidTy())
      return ty;
    return ArrayType::get(ty, width);
  } else {
    return ty;
  }
}

Type *GradientUtils::getShadowType(Type *ty) {
  return getShadowType(ty, width);
}

Type *GradientUtils::extractMeta(Type *T, ArrayRef<unsigned> off) {
  for (auto idx : off) {
    if (auto AT = dyn_cast<ArrayType>(T)) {
      T = AT->getElementType();
      continue;
    }
    if (auto ST = dyn_cast<StructType>(T)) {
      T = ST->getElementType(idx);
      continue;
    }
    assert(false && "could not sub index into type");
  }
  return T;
}

Value *GradientUtils::extractMeta(IRBuilder<> &Builder, Value *Agg,
                                  unsigned off, const Twine &name) {
  return extractMeta(Builder, Agg, ArrayRef<unsigned>({off}), name);
}

Value *GradientUtils::extractMeta(IRBuilder<> &Builder, Value *Agg,
                                  ArrayRef<unsigned> off_init,
                                  const Twine &name, bool fallback) {
  std::vector<unsigned> off(off_init.begin(), off_init.end());
  while (off.size() != 0) {
    if (auto Ins = dyn_cast<InsertValueInst>(Agg)) {
      size_t until = Ins->getNumIndices();
      if (off.size() < until)
        until = off.size();
      bool subset = true;
      for (size_t i = 0; i < until; i++) {
        if (Ins->getIndices()[i] != off[i]) {
          subset = false;
          break;
        }
      }
      if (!subset) {
        Agg = Ins->getAggregateOperand();
        continue;
      } else if (until < Ins->getNumIndices()) {
        break;
      } else {
        off.erase(off.begin(), off.begin() + until);
        Agg = Ins->getInsertedValueOperand();
        continue;
      }
    }
    if (auto ext = dyn_cast<ExtractValueInst>(Agg)) {
      off.insert(off.begin(), ext->getIndices().begin(),
                 ext->getIndices().end());
      Agg = ext->getAggregateOperand();
      continue;
    }
    if (auto CA = dyn_cast<ConstantAggregateZero>(Agg)) {
      Agg = CA->getElementValue(off[0]);
      off.erase(off.begin(), off.begin() + 1);
    }
    break;
  }
  if (off.size() == 0)
    return Agg;

  if (!fallback)
    return nullptr;

  if (Agg->getType()->isVectorTy() && off.size() == 1)
    return Builder.CreateExtractElement(Agg, off[0], name);

  return Builder.CreateExtractValue(Agg, off, name);
}

llvm::Value *GradientUtils::recursiveFAdd(llvm::IRBuilder<> &B,
                                          llvm::Value *lhs, llvm::Value *rhs,
                                          llvm::ArrayRef<unsigned> lhs_off,
                                          llvm::ArrayRef<unsigned> rhs_off,
                                          llvm::Value *prev, bool vectorLayer) {
  llvm::Type *lhs_ty = lhs->getType();
  if (!vectorLayer) {
    for (auto idx : lhs_off)
      lhs_ty = getSubType(lhs_ty, idx);
    llvm::Type *rhs_ty = rhs->getType();
    for (auto idx : rhs_off)
      rhs_ty = getSubType(rhs_ty, idx);
    assert(lhs_ty == rhs_ty);
  }
  if (lhs_ty->isFPOrFPVectorTy()) {
    if (lhs_off.size())
      lhs = extractMeta(B, lhs, lhs_off);
    if (rhs_off.size())
      rhs = extractMeta(B, rhs, rhs_off);
    llvm::Value *res = nullptr;
    if (auto fp = llvm::dyn_cast<llvm::ConstantFP>(lhs)) {
      if (fp->isZero())
        res = rhs;
    }
    if (auto fp = llvm::dyn_cast<llvm::ConstantFP>(rhs)) {
      if (fp->isZero())
        res = lhs;
    }
    if (!res) {
      if (auto *FPMO = dyn_cast<FPMathOperator>(rhs))
        if (FPMO->getOpcode() == Instruction::FNeg) {
          res = B.CreateFSub(lhs, FPMO->getOperand(0));
        }
    }
    if (!res) {
      if (auto *S = dyn_cast<BinaryOperator>(rhs)) {
        if (S->getOpcode() == Instruction::FSub) {
          if (auto C = dyn_cast<ConstantFP>(S->getOperand(0)))
            if (C->isZero())
              res = B.CreateFSub(lhs, S->getOperand(1));
        }
      }
    }
    if (!res) {
      res = B.CreateFAdd(lhs, rhs);
    }
    if (lhs_off.size()) {
      assert(prev);
      res = B.CreateInsertValue(prev, res, lhs_off);
    }
    return res;
  } else if (isa<ArrayType>(lhs_ty) || isa<StructType>(lhs_ty)) {
    if (prev == nullptr)
      prev = llvm::UndefValue::get(lhs_ty);

    size_t size;
    if (auto AT = dyn_cast<ArrayType>(lhs_ty))
      size = AT->getNumElements();
    else
      size = cast<StructType>(lhs_ty)->getNumElements();

    for (size_t i = 0; i < size; ++i) {
      llvm::SmallVector<unsigned, 1> nlhs_off(lhs_off.begin(), lhs_off.end());
      if (vectorLayer)
        nlhs_off.insert(nlhs_off.begin(), i);
      else
        nlhs_off.push_back(i);
      llvm::SmallVector<unsigned, 1> nrhs_off(rhs_off.begin(), rhs_off.end());
      if (vectorLayer)
        nrhs_off.insert(nrhs_off.begin(), i);
      else
        nrhs_off.push_back(i);
      prev = recursiveFAdd(B, lhs, rhs, nlhs_off, nrhs_off, prev);
    }
    return prev;
  }
  llvm_unreachable("Unknown type to recursively accumulate");
}

static bool allNullOrUndef(Value *C, const DataLayout &dl, TypeTree TT) {
  if (!TT.anyPointer(C, dl)) {
    return true;
  }
  if (isa<UndefValue>(C) || isa<ConstantPointerNull>(C) ||
      isa<ConstantAggregateZero>(C)) {
    return true;
  }
  if (auto CF = dyn_cast<ConstantFP>(C)) {
    if (CF->isZero())
      return true;
  }
  if (auto CInt = dyn_cast<ConstantInt>(C)) {
    if (CInt->isZero())
      return true;
  }
  if (auto CS = dyn_cast<ConstantStruct>(C)) {
    const StructLayout *Layout = dl.getStructLayout(CS->getType());
    for (unsigned i = 0; i < CS->getNumOperands(); ++i) {
      auto el = CS->getOperand(i);
      auto Off = Layout->getElementOffset(i);
      auto ObjSize = (dl.getTypeSizeInBits(el->getType()) + 7) / 8;
      TypeTree subTT = TT.ShiftIndices(dl, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, dl);
      if (!allNullOrUndef(el, dl, subTT)) {
        return false;
      }
    }
    return true;
  }
  if (auto CA = dyn_cast<ConstantArray>(C)) {
    auto ElTy = CA->getType()->getElementType();
    auto ObjSize = (dl.getTypeSizeInBits(ElTy) + 7) / 8;
    for (unsigned i = 0; i < CA->getNumOperands(); ++i) {
      auto el = CA->getOperand(i);
      auto Off = i * ObjSize;
      TypeTree subTT = TT.ShiftIndices(dl, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, dl);
      if (!allNullOrUndef(el, dl, subTT)) {
        return false;
      }
    }
    return true;
  }
  if (auto CV = dyn_cast<ConstantVector>(C)) {
    auto ElTy = CV->getType()->getElementType();
    auto ObjSize = (dl.getTypeSizeInBits(ElTy) + 7) / 8;
    for (unsigned i = 0; i < CV->getNumOperands(); ++i) {
      auto el = CV->getOperand(i);
      auto Off = i * ObjSize;
      TypeTree subTT = TT.ShiftIndices(dl, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, dl);
      if (!allNullOrUndef(el, dl, subTT)) {
        return false;
      }
    }
    return true;
  }
  if (auto CDS = dyn_cast<ConstantDataSequential>(C)) {
    auto ElTy = CDS->getElementType();
    auto ObjSize = (dl.getTypeSizeInBits(ElTy) + 7) / 8;
    for (unsigned i = 0; i < CDS->getNumElements(); ++i) {
      auto el = CDS->getElementAsConstant(i);
      auto Off = i * ObjSize;
      TypeTree subTT = TT.ShiftIndices(dl, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, dl);
      if (!allNullOrUndef(el, dl, subTT)) {
        return false;
      }
    }
    return true;
  }
  return false;
}

Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM) {
  return invertPointerM(oval, BuilderM, TR.query(oval));
}

bool GradientUtils::isAtomic(Value *origptr) const {
  return ::isAtomic(origptr, AtomicAdd, newFunc);
}

Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM,
                                     TypeTree TT) {
  assert(oval);
#ifndef NDEBUG
  if (auto inst = dyn_cast<Instruction>(oval)) {
    assert(inst->getParent()->getParent() == oldFunc);
  }
  if (auto arg = dyn_cast<Argument>(oval)) {
    assert(arg->getParent() == oldFunc);
  }
#endif

  auto &DL = oldFunc->getParent()->getDataLayout();
  if (isa<ConstantPointerNull>(oval) || isa<UndefValue>(oval) ||
      isa<ConstantInt>(oval) || isa<ConstantAggregateZero>(oval) ||
      isa<PoisonValue>(oval)) {
    if (isa<ConstantPointerNull>(oval) || isa<UndefValue>(oval) ||
        isa<PoisonValue>(oval) || isa<ConstantAggregateZero>(oval) ||
        (isa<ConstantInt>(oval) && cast<ConstantInt>(oval)->isZero()) ||
        TT.allFloat(oval, DL))
      return Constant::getNullValue(getShadowType(oval->getType()));
    else if (!TT.anyFloat(oval, DL))
      return applyChainRule(oval->getType(), BuilderM, [&]() { return oval; });
  }
  if (auto CD = dyn_cast<ConstantDataArray>(oval)) {
    SmallVector<Constant *, 1> Vals;
    auto ElTy = CD->getType()->getElementType();
    auto ObjSize = (DL.getTypeSizeInBits(ElTy) + 7) / 8;
    for (size_t i = 0, len = CD->getNumElements(); i < len; i++) {
      auto el = CD->getElementAsConstant(i);
      auto Off = i * ObjSize;
      TypeTree subTT = TT.ShiftIndices(DL, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, DL);
      Value *val = invertPointerM(el, BuilderM, subTT);
      Vals.push_back(cast<Constant>(val));
    }
    auto rule = [&CD](ArrayRef<Constant *> Vals) {
      return ConstantArray::get(CD->getType(), Vals);
    };
    return applyChainRule(CD->getType(), Vals, BuilderM, rule);
  } else if (auto CD = dyn_cast<ConstantArray>(oval)) {
    SmallVector<Constant *, 1> Vals;
    auto ElTy = CD->getType()->getElementType();
    auto ObjSize = (DL.getTypeSizeInBits(ElTy) + 7) / 8;
    for (size_t i = 0, len = CD->getNumOperands(); i < len; i++) {
      auto el = CD->getOperand(i);
      auto Off = i * ObjSize;
      TypeTree subTT = TT.ShiftIndices(DL, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, DL);
      Value *val = invertPointerM(el, BuilderM, subTT);
      Vals.push_back(cast<Constant>(val));
    }

    auto rule = [&CD](ArrayRef<Constant *> Vals) {
      return ConstantArray::get(CD->getType(), Vals);
    };

    return applyChainRule(CD->getType(), Vals, BuilderM, rule);
  } else if (auto CD = dyn_cast<ConstantStruct>(oval)) {
    SmallVector<Constant *, 1> Vals;
    auto StructTy = CD->getType();
    const StructLayout *Layout = DL.getStructLayout(StructTy);
    for (size_t i = 0, len = CD->getNumOperands(); i < len; i++) {
      auto el = CD->getOperand(i);
      auto Off = Layout->getElementOffset(i);
      auto ObjSize = (DL.getTypeSizeInBits(el->getType()) + 7) / 8;
      TypeTree subTT = TT.ShiftIndices(DL, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, DL);
      Vals.push_back(cast<Constant>(invertPointerM(el, BuilderM, subTT)));
    }

    auto rule = [&CD](ArrayRef<Constant *> Vals) {
      return ConstantStruct::get(CD->getType(), Vals);
    };
    return applyChainRule(CD->getType(), Vals, BuilderM, rule);
  } else if (auto CD = dyn_cast<ConstantVector>(oval)) {
    SmallVector<Constant *, 1> Vals;
    auto ElTy = CD->getType()->getElementType();
    auto ObjSize = (DL.getTypeSizeInBits(ElTy) + 7) / 8;
    for (size_t i = 0, len = CD->getNumOperands(); i < len; i++) {
      auto el = CD->getOperand(i);
      auto Off = i * ObjSize;
      TypeTree subTT = TT.ShiftIndices(DL, Off, ObjSize, 0);
      subTT.CanonicalizeInPlace(ObjSize, DL);
      Vals.push_back(cast<Constant>(invertPointerM(el, BuilderM, subTT)));
    }

    auto rule = [](ArrayRef<Constant *> Vals) {
      return ConstantVector::get(Vals);
    };

    return applyChainRule(CD->getType(), Vals, BuilderM, rule);
  } else if (isa<ConstantData>(oval) &&
             TT.allFloat(oval, DL, /*anythingIsFloat*/ true)) {
    return Constant::getNullValue(getShadowType(oval->getType()));
  }

  bool shouldNullShadow = isConstantValue(oval);
  if (shouldNullShadow) {
    if (isa<InsertValueInst>(oval) || isa<ExtractValueInst>(oval) ||
        isa<InsertElementInst>(oval) || isa<ExtractElementInst>(oval)) {
      shouldNullShadow = false;
      auto orig = cast<Instruction>(oval);
      if (knownRecomputeHeuristic.count(orig)) {
        if (!knownRecomputeHeuristic[orig]) {
          shouldNullShadow = true;
        }
      }
    }
  }

  // NOTE, this is legal and the correct resolution, however, our activity
  // analysis honeypot no longer exists

  // Nulling the shadow for a constant is only necessary if any of the data
  // could contain a float (e.g. should not be applied to pointers).
  if (shouldNullShadow) {
    size_t size = (DL.getTypeSizeInBits(oval->getType()) + 7) / 8;
    if (TT.anyFloat(oval, DL)) {
      if (TT.allFloat(oval, DL, /*anythingIsFloat*/ true))
        return Constant::getNullValue(getShadowType(oval->getType()));
      else {
        IRBuilder<> bb(inversionAllocs);
        if (auto arg = dyn_cast<Instruction>(oval)) {
          arg = getNewFromOriginal(arg);
          // Go one after since otherwise we won't be able
          // to use in the store.
          arg = arg->getNextNode();
          while (auto PN = dyn_cast<PHINode>(arg)) {
            if (PN->getNumIncomingValues() == 0)
              break;
            arg = PN->getNextNode();
          }
          bb.SetInsertPoint(arg);
        }
        auto alloc = bb.CreateAlloca(oval->getType());
        auto AT = ArrayType::get(bb.getInt8Ty(), size);
        bb.CreateStore(getNewFromOriginal(oval), alloc);
        Value *cur = bb.CreatePointerCast(alloc, getUnqual(AT));
        size_t i = 0;
        assert(size > 0);
        for (; i < size;) {
          auto CT2 = TT[{(int)i}];
          if (CT2 == BaseType::Pointer) {
            i += DL.getPointerSize(0);
            continue;
          } else if (auto flt = CT2.isFloat()) {
            auto ptr = bb.CreateConstInBoundsGEP2_32(AT, cur, 0, i);
            ptr = bb.CreatePointerCast(ptr, getUnqual(flt));
            bb.CreateStore(Constant::getNullValue(flt), ptr);
            size_t chunk = DL.getTypeSizeInBits(flt) / 8;
            i += chunk;
          } else if (CT2 != BaseType::Integer) {
            auto ptr = bb.CreateConstInBoundsGEP2_32(AT, cur, 0, i);
            bb.CreateStore(Constant::getNullValue(bb.getInt8Ty()), ptr);
            i++;
          } else {
            i++;
          }
        }
        auto res = bb.CreateLoad(oval->getType(), alloc);
        auto rule = [&res]() { return res; };
        auto res2 = applyChainRule(oval->getType(), BuilderM, rule);
        invertedPointers.insert(
            std::make_pair((const Value *)oval, InvertedPointerVH(this, res2)));
        return res2;
      }
    }

    if (isa<ConstantExpr>(oval) || isa<GlobalValue>(oval)) {
      auto rule = [&oval]() { return oval; };
      return applyChainRule(oval->getType(), BuilderM, rule);
    }

    Value *newval = getNewFromOriginal(oval);

    auto rule = [&]() { return newval; };

    return applyChainRule(oval->getType(), BuilderM, rule);
  }

  auto M = oldFunc->getParent();
  assert(oval);

  {
    auto ifound = invertedPointers.find(oval);
    if (ifound != invertedPointers.end()) {
      return &*ifound->second;
    }
  }

  if (mode != DerivativeMode::ForwardMode &&
      mode != DerivativeMode::ForwardModeError &&
      mode != DerivativeMode::ForwardModeSplit) {
    if (TT.allFloat(oval, DL)) {
      return Constant::getNullValue(getShadowType(oval->getType()));
    }
  }

  if (isa<Argument>(oval) && !TT.anyPointer(oval, DL)) {
    return Constant::getNullValue(getShadowType(oval->getType()));
  } else if (isa<Argument>(oval) && cast<Argument>(oval)->hasByValAttr()) {
    IRBuilder<> bb(inversionAllocs);

    Type *subType = nullptr;
    auto attr = cast<Argument>(oval)->getAttribute(Attribute::ByVal);
    subType = attr.getValueAsType();

    auto rule1 = [&]() {
      AllocaInst *antialloca = bb.CreateAlloca(
          subType, cast<PointerType>(oval->getType())->getPointerAddressSpace(),
          nullptr, oval->getName() + "'ipa");

      auto dst_arg =
          bb.CreateBitCast(antialloca, getInt8PtrTy(oval->getContext()));
      auto val_arg = ConstantInt::get(Type::getInt8Ty(oval->getContext()), 0);
      auto len_arg = ConstantInt::get(
          Type::getInt64Ty(oval->getContext()),
          M->getDataLayout().getTypeAllocSizeInBits(subType) / 8);
      auto volatile_arg = ConstantInt::getFalse(oval->getContext());

      Value *args[] = {dst_arg, val_arg, len_arg, volatile_arg};
      Type *tys[] = {dst_arg->getType(), len_arg->getType()};
      bb.CreateCall(getIntrinsicDeclaration(M, Intrinsic::memset, tys), args);

      return antialloca;
    };

    Value *antialloca = applyChainRule(oval->getType(), bb, rule1);

    invertedPointers.insert(std::make_pair(
        (const Value *)oval, InvertedPointerVH(this, antialloca)));

    return antialloca;
  } else if (auto arg = dyn_cast<GlobalAlias>(oval)) {
    Value *aliasTarget = arg->getAliasee();
    return invertPointerM(aliasTarget, BuilderM, TT);
  } else if (auto arg = dyn_cast<GlobalVariable>(oval)) {
    if (!hasMetadata(arg, "enzyme_shadow")) {

      if ((mode == DerivativeMode::ReverseModeCombined ||
           mode == DerivativeMode::ForwardMode ||
           mode == DerivativeMode::ForwardModeError) &&
          arg->getType()->getPointerAddressSpace() == 0) {
        auto CT = TR.query(arg)[{-1, -1}];
        // Can only localy replace a global variable if it is
        // known not to contain a pointer, which may be initialized
        // outside of this function to contain other memory which
        // will not have a shadow within the current function.
        if (CT.isKnown() && CT != BaseType::Pointer) {
          bool seen = false;
          MemoryLocation
#if LLVM_VERSION_MAJOR >= 12
              Loc = MemoryLocation(oval, LocationSize::beforeOrAfterPointer());
#else
              Loc = MemoryLocation(oval, LocationSize::unknown());
#endif
          for (CallInst *CI : originalCalls) {
            if (isa<IntrinsicInst>(CI))
              continue;
            if (!isConstantInstruction(CI)) {
              auto F = getFunctionFromCall(CI);
              if (F && isMemFreeLibMFunction(F->getName())) {
                continue;
              }
              if (llvm::isModOrRefSet(OrigAA->getModRefInfo(CI, Loc))) {
                seen = true;
                llvm::errs() << " cannot shadow-inline global " << *oval
                             << " due to " << *CI << "\n";
                goto endCheck;
              }
            }
          }
        endCheck:;
          if (!seen) {
            IRBuilder<> bb(inversionAllocs);
            Type *allocaTy = arg->getValueType();

            auto rule1 = [&]() {
              AllocaInst *antialloca = bb.CreateAlloca(
                  allocaTy, arg->getType()->getPointerAddressSpace(), nullptr,
                  arg->getName() + "'ipa");
              if (arg->getAlignment()) {
                antialloca->setAlignment(Align(arg->getAlignment()));
              }
              return antialloca;
            };

            Value *antialloca = applyChainRule(arg->getType(), bb, rule1);

            invertedPointers.insert(std::make_pair(
                (const Value *)oval, InvertedPointerVH(this, antialloca)));

            auto rule2 = [&](Value *antialloca) {
              auto dst_arg =
                  bb.CreateBitCast(antialloca, getInt8PtrTy(arg->getContext()));
              auto val_arg =
                  ConstantInt::get(Type::getInt8Ty(arg->getContext()), 0);
              auto len_arg =
                  ConstantInt::get(Type::getInt64Ty(arg->getContext()),
                                   M->getDataLayout().getTypeAllocSizeInBits(
                                       arg->getValueType()) /
                                       8);
              auto volatile_arg = ConstantInt::getFalse(oval->getContext());

              Value *args[] = {dst_arg, val_arg, len_arg, volatile_arg};
              Type *tys[] = {dst_arg->getType(), len_arg->getType()};
              auto memset = cast<CallInst>(bb.CreateCall(
                  getIntrinsicDeclaration(M, Intrinsic::memset, tys), args));
              if (arg->getAlignment()) {
                memset->addParamAttr(
                    0, Attribute::getWithAlignment(arg->getContext(),
                                                   Align(arg->getAlignment())));
              }
              memset->addParamAttr(0, Attribute::NonNull);
              assert((width > 1 && antialloca->getType() ==
                                       ArrayType::get(arg->getType(), width)) ||
                     antialloca->getType() == arg->getType());
              return antialloca;
            };

            return applyChainRule(arg->getType(), bb, rule2, antialloca);
          }
        }
      }

      auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple());
      int SharedAddrSpace = getGPUSharedAddrSpace(TT);
      int AddrSpace = cast<PointerType>(arg->getType())->getAddressSpace();
      if (isGPUArch(TT) && AddrSpace == SharedAddrSpace) {
        llvm::errs() << "warning found shared memory\n";
        Type *type = arg->getValueType();
        // TODO this needs initialization by entry
        auto shadow = new GlobalVariable(
            *arg->getParent(), type, arg->isConstant(), arg->getLinkage(),
            UndefValue::get(type), arg->getName() + "_shadow", arg,
            arg->getThreadLocalMode(), arg->getType()->getAddressSpace(),
            arg->isExternallyInitialized());
        arg->setMetadata("enzyme_shadow",
                         MDTuple::get(shadow->getContext(),
                                      {ConstantAsMetadata::get(shadow)}));
        shadow->setMetadata("enzyme_internalshadowglobal",
                            MDTuple::get(shadow->getContext(), {}));
        shadow->setAlignment(arg->getAlign());
        shadow->setUnnamedAddr(arg->getUnnamedAddr());
        return shadow;
      }

      // Create global variable locally if not externally visible
      //  If a variable is constant, for forward mode it will also
      //  only be read, so invert initializing is fine.
      //  For reverse mode, any floats will be +='d into, but never
      //  read, and any pointers will be used as expected. The never
      //  read means even if two globals for floats, that's fine.
      //  As long as the pointers point to equivalent places (which
      //  they should from the same initialization), it is also ok.
      if (arg->hasInternalLinkage() || arg->hasPrivateLinkage() ||
          (arg->hasExternalLinkage() && arg->hasInitializer()) ||
          arg->isConstant()) {
        Type *elemTy = arg->getValueType();
        IRBuilder<> B(inversionAllocs);

        auto rule = [&]() {
          auto shadow = new GlobalVariable(
              *arg->getParent(), elemTy, arg->isConstant(), arg->getLinkage(),
              Constant::getNullValue(elemTy), arg->getName() + "_shadow", arg,
              arg->getThreadLocalMode(), arg->getType()->getAddressSpace(),
              arg->isExternallyInitialized());
          shadow->setAlignment(arg->getAlign());
          shadow->setUnnamedAddr(arg->getUnnamedAddr());

          return shadow;
        };

        Value *shadow = applyChainRule(oval->getType(), BuilderM, rule);
        arg->setMetadata(
            "enzyme_shadow",
            MDTuple::get(shadow->getContext(),
                         {ConstantAsMetadata::get(cast<Constant>(shadow))}));
        if (getWidth() != 1) {
          BuilderM.Insert(InsertValueInst::Create(shadow, arg, {0}), "tmp");
        }

        if (arg->hasInitializer()) {
          size_t tsize =
              (DL.getTypeSizeInBits(arg->getInitializer()->getType()) + 7) / 8;
          applyChainRule(
              BuilderM,
              [&](Value *shadow, Value *ip) {
                cast<GlobalVariable>(shadow)->setInitializer(
                    cast<Constant>(ip));
              },
              shadow,
              invertPointerM(arg->getInitializer(), B,
                             TR.query(oval).Lookup(tsize, DL)));
        }

        invertedPointers.insert(std::make_pair(
            (const Value *)oval, InvertedPointerVH(this, shadow)));
        return shadow;
      }

      std::string s;
      llvm::raw_string_ostream ss(s);
      ss << "cannot compute with global variable that doesn't have marked "
            "shadow global\n";
      ss << *arg << "\n";
      if (CustomErrorHandler) {
        return unwrap(CustomErrorHandler(ss.str().c_str(), wrap(arg),
                                         ErrorType::NoShadow, this, nullptr,
                                         wrap(&BuilderM)));
      } else {
        EmitFailure("InvertGlobal", BuilderM.getCurrentDebugLocation(), oldFunc,
                    ss.str());
      }
      return UndefValue::get(getShadowType(arg->getType()));
    }
    auto md = arg->getMetadata("enzyme_shadow");
    if (!isa<MDTuple>(md)) {
      llvm::errs() << *arg << "\n";
      llvm::errs() << *md << "\n";
      std::string s;
      llvm::raw_string_ostream ss(s);
      ss << "cannot compute with global variable that doesn't have marked "
            "shadow global as mdtuple\n";
      ss << *arg << "\n";
      ss << " md: " << *md << "\n";
      if (CustomErrorHandler) {
        return unwrap(CustomErrorHandler(ss.str().c_str(), wrap(arg),
                                         ErrorType::NoShadow, this, nullptr,
                                         wrap(&BuilderM)));
      } else {
        EmitFailure("InvertGlobal", BuilderM.getCurrentDebugLocation(), oldFunc,
                    ss.str());
      }
      return UndefValue::get(getShadowType(arg->getType()));
    }
    auto md2 = cast<MDTuple>(md);
    assert(md2->getNumOperands() == 1);
    auto gvemd = cast<ConstantAsMetadata>(md2->getOperand(0));
    auto cs = cast<Constant>(gvemd->getValue());

    if (width > 1) {
      SmallVector<Constant *, 2> Vals;
      for (unsigned i = 0; i < width; ++i) {

        Constant *idxs[] = {
            ConstantInt::get(Type::getInt32Ty(cs->getContext()), 0),
            ConstantInt::get(Type::getInt32Ty(cs->getContext()), i)};
        Constant *elem = ConstantExpr::getInBoundsGetElementPtr(
            getShadowType(arg->getValueType()), cs, idxs);
        Vals.push_back(elem);
      }

      auto agg = ConstantArray::get(
          cast<ArrayType>(getShadowType(arg->getType())), Vals);

      invertedPointers.insert(
          std::make_pair((const Value *)oval, InvertedPointerVH(this, agg)));
      return agg;
    } else {
      invertedPointers.insert(
          std::make_pair((const Value *)oval, InvertedPointerVH(this, cs)));
      return cs;
    }
  } else if (auto fn = dyn_cast<Function>(oval)) {
    Constant *shadow = GetOrCreateShadowFunction(
        RequestContext(nullptr, &BuilderM), Logic, TLI, TA, fn, mode,
        runtimeActivity, strongZero, width, AtomicAdd);
    if (width > 1) {
      SmallVector<Constant *, 3> arr;
      for (unsigned i = 0; i < width; ++i) {
        arr.push_back(shadow);
      }
      ArrayType *arrTy = ArrayType::get(shadow->getType(), width);
      shadow = ConstantArray::get(arrTy, arr);
    }
    return shadow;
  } else if (auto arg = dyn_cast<CastInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *invertOp = invertPointerM(arg->getOperand(0), bb, TT);
    Type *shadowTy = arg->getDestTy();

    auto rule = [&](Value *invertOp) {
      return bb.CreateCast(arg->getOpcode(), invertOp, shadowTy,
                           arg->getName() + "'ipc");
    };

    Value *shadow = applyChainRule(shadowTy, bb, rule, invertOp);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<FreezeInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *invertOp = invertPointerM(arg->getOperand(0), bb, TT);
    Type *shadowTy = arg->getType();

    if (mode == DerivativeMode::ReverseModeCombined ||
        mode == DerivativeMode::ReverseModePrimal ||
        mode == DerivativeMode::ReverseModeGradient) {
      if (TT.allFloat(oval, DL)) {
        return Constant::getNullValue(getShadowType(oval->getType()));
      }
    }
    assert(!arg->getType()->isDoubleTy());

    auto rule = [&](Value *invertOp) {
      return bb.CreateFreeze(invertOp, arg->getName() + "'ipf");
    };

    Value *shadow = applyChainRule(shadowTy, bb, rule, invertOp);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<ConstantExpr>(oval)) {
    IRBuilder<> bb(inversionAllocs);
    if (arg->getOpcode() == Instruction::Add) {
      if (isa<ConstantInt>(arg->getOperand(0))) {
        auto rule = [&arg](Value *ip) {
          Constant *invops[2] = {arg->getOperand(0), cast<Constant>(ip)};
          return arg->getWithOperands(invops);
        };
        auto ip = invertPointerM(arg->getOperand(1), bb, TT);
        return applyChainRule(arg->getType(), bb, rule, ip);
      }
      if (isa<ConstantInt>(arg->getOperand(1))) {
        auto rule = [&arg](Value *ip) {
          Constant *invops[2] = {cast<Constant>(ip), arg->getOperand(1)};
          return arg->getWithOperands(invops);
        };
        auto ip = invertPointerM(arg->getOperand(0), bb, TT);
        return applyChainRule(arg->getType(), bb, rule, ip);
      }
    }
    auto ip = invertPointerM(arg->getOperand(0), bb, TT);

    if (arg->isCast()) {
#if LLVM_VERSION_MAJOR < 17
      if (auto PT = dyn_cast<PointerType>(arg->getType())) {
        if (isConstantValue(arg->getOperand(0)) &&
            PT->getPointerElementType()->isFunctionTy()) {
          goto end;
        }
      }
#endif
      if (isa<Constant>(ip)) {
        auto rule = [&arg](Value *ip) {
          return ConstantExpr::getCast(arg->getOpcode(), cast<Constant>(ip),
                                       arg->getType());
        };

        return applyChainRule(arg->getType(), bb, rule, ip);

      } else {
        auto rule = [&](Value *ip) {
          return bb.CreateCast((Instruction::CastOps)arg->getOpcode(), ip,
                               arg->getType(), arg->getName() + "'ipc");
        };

        Value *shadow = applyChainRule(arg->getType(), bb, rule, ip);

        invertedPointers.insert(std::make_pair(
            (const Value *)oval, InvertedPointerVH(this, shadow)));

        return shadow;
      }
    } else if (arg->getOpcode() == Instruction::GetElementPtr) {
      if (auto C = dyn_cast<Constant>(ip)) {
        auto rule = [&arg, &C]() {
          SmallVector<Constant *, 8> NewOps;
          for (unsigned i = 0, e = arg->getNumOperands(); i != e; ++i)
            NewOps.push_back(i == 0 ? C : arg->getOperand(i));
          return cast<Value>(arg->getWithOperands(NewOps));
        };

        return applyChainRule(arg->getType(), bb, rule);
      } else {
        SmallVector<Value *, 4> invertargs;
        for (unsigned i = 0; i < arg->getNumOperands() - 1; ++i) {
          Value *b = getNewFromOriginal(arg->getOperand(1 + i));
          invertargs.push_back(b);
        }

        auto rule = [&bb, &arg, &invertargs](Value *ip) {
          // TODO mark this the same inbounds as the original
          return bb.CreateGEP(cast<GEPOperator>(ip)->getSourceElementType(), ip,
                              invertargs, arg->getName() + "'ipg");
        };

        Value *shadow = applyChainRule(arg->getType(), bb, rule, ip);

        invertedPointers.insert(std::make_pair(
            (const Value *)oval, InvertedPointerVH(this, shadow)));
        return shadow;
      }
    } else {
      llvm::errs() << *arg << "\n";
      assert(0 && "unhandled");
    }
    goto end;
  } else if (auto arg = dyn_cast<ExtractValueInst>(oval)) {
    auto newi = getNewFromOriginal(arg);
    IRBuilder<> bb(newi->getNextNode());

    auto AggTy = arg->getAggregateOperand()->getType();
    SmallVector<Value *, 4> vec;
    vec.push_back(ConstantInt::get(Type::getInt64Ty(arg->getContext()), 0));
    for (auto ind : arg->getIndices()) {
      vec.push_back(ConstantInt::get(Type::getInt32Ty(arg->getContext()), ind));
    }
    auto ud = UndefValue::get(getUnqual(AggTy));
    auto g2 = GetElementPtrInst::Create(AggTy, ud, vec);
    APInt ai(DL.getIndexSizeInBits(g2->getPointerAddressSpace()), 0);
    g2->accumulateConstantOffset(DL, ai);
    delete g2;

    unsigned Off = (unsigned)ai.getLimitedValue();
    auto ObjSize = (DL.getTypeSizeInBits(arg->getType()) + 7) / 8;
    auto AggSize = (DL.getTypeSizeInBits(AggTy) + 7) / 8;

    TypeTree agg_look = TR.query(arg->getAggregateOperand());
    if (TT.isKnown()) {
      agg_look = agg_look.Clear(Off, Off + ObjSize, AggSize);
    }
    agg_look |= TT.ShiftIndices(DL, 0, ObjSize, Off);
    agg_look.CanonicalizeInPlace(AggSize, DL);

    auto ip = invertPointerM(arg->getOperand(0), bb, agg_look);

    auto rule = [&bb, &arg, &newi, this](Value *ip) -> llvm::Value * {
      if (ip == getNewFromOriginal(arg->getOperand(0)))
        return newi;
      return bb.CreateExtractValue(ip, arg->getIndices(),
                                   arg->getName() + "'ipev");
    };

    Value *shadow = applyChainRule(arg->getType(), bb, rule, ip);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<InsertValueInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *ivops[2] = {nullptr, nullptr};

    auto AggTy = arg->getAggregateOperand()->getType();
    auto InsertedTy = arg->getInsertedValueOperand()->getType();
    SmallVector<Value *, 4> vec;
    vec.push_back(ConstantInt::get(Type::getInt64Ty(arg->getContext()), 0));
    for (auto ind : arg->getIndices()) {
      vec.push_back(ConstantInt::get(Type::getInt32Ty(arg->getContext()), ind));
    }
    auto ud = UndefValue::get(getUnqual(AggTy));
    auto g2 = GetElementPtrInst::Create(AggTy, ud, vec);
    APInt ai(DL.getIndexSizeInBits(g2->getPointerAddressSpace()), 0);
    g2->accumulateConstantOffset(DL, ai);
    delete g2;

    unsigned Off = (unsigned)ai.getLimitedValue();
    auto ObjSize = (DL.getTypeSizeInBits(InsertedTy) + 7) / 8;

    for (int i = 0; i < 2; i++) {
      auto op = arg->getOperand(i);
      TypeTree subTT;
      if (i == 0) {
        subTT = TT;
      } else {
        subTT = TT.ShiftIndices(DL, Off, ObjSize, 0);
        if (auto MD = hasMetadata(arg, "enzyme_truetype")) {
          for (size_t i = 0; i < MD->getNumOperands(); i += 2) {
            ConcreteType base(
                llvm::cast<llvm::MDString>(MD->getOperand(i))->getString(),
                MD->getContext());
            auto offset =
                llvm::cast<llvm::ConstantInt>(
                    llvm::cast<llvm::ConstantAsMetadata>(MD->getOperand(i + 1))
                        ->getValue())
                    ->getSExtValue();
            if (offset < Off || offset >= Off + ObjSize) {
              continue;
            }
            subTT.insert({(int)(offset - Off)}, base);
          }
        }
        subTT.CanonicalizeInPlace(ObjSize, DL);
      }

      if (!runtimeActivity && !isa<InsertValueInst>(op)) {
        if (isConstantValue(op)) {
          if (subTT.anyPointer(op, DL) &&
              subTT[{-1, -1}] != BaseType::Integer) {
            if (!allNullOrUndef(op, DL, subTT)) {
              std::string str;
              raw_string_ostream ss(str);
              ss << "Mismatched activity for: " << *arg
                 << " const val: " << *op;
              if (CustomErrorHandler)
                ivops[i] = unwrap(CustomErrorHandler(
                    str.c_str(), wrap(arg), ErrorType::MixedActivityError, this,
                    wrap(op), wrap(&bb)));
              else
                EmitWarning("MixedActivityError", *arg, ss.str());
            }
          }
        }
      }
      if (!ivops[i]) {
        ivops[i] = invertPointerM(op, bb, subTT);
      }
    }

    auto rule = [&bb, &arg](Value *ip0, Value *ip1) {
      return bb.CreateInsertValue(ip0, ip1, arg->getIndices(),
                                  arg->getName() + "'ipiv");
    };

    Value *shadow =
        applyChainRule(arg->getType(), bb, rule, ivops[0], ivops[1]);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<ExtractElementInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    auto VecTy = arg->getVectorOperand()->getType();
    auto ElTy = cast<VectorType>(VecTy)->getElementType();
    auto ObjSize = (DL.getTypeSizeInBits(ElTy) + 7) / 8;
    auto VecSize = (DL.getTypeSizeInBits(VecTy) + 7) / 8;

    TypeTree vec_look = TR.query(arg->getVectorOperand());
    if (auto CI = dyn_cast<ConstantInt>(arg->getIndexOperand())) {
      unsigned Off = (CI->getZExtValue() * DL.getTypeSizeInBits(ElTy)) / 8;
      if (TT.isKnown()) {
        vec_look = vec_look.Clear(Off, Off + ObjSize, VecSize);
        vec_look |= TT.ShiftIndices(DL, 0, ObjSize, Off);
        vec_look.CanonicalizeInPlace(VecSize, DL);
      }
    }

    auto ip = invertPointerM(arg->getVectorOperand(), bb, vec_look);

    auto rule = [&](Value *ip) {
      return bb.CreateExtractElement(ip,
                                     getNewFromOriginal(arg->getIndexOperand()),
                                     arg->getName() + "'ipee");
      ;
    };

    Value *shadow = applyChainRule(arg->getType(), bb, rule, ip);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<InsertElementInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *op0 = arg->getOperand(0);
    Value *op1 = arg->getOperand(1);
    Value *op2 = arg->getOperand(2);
    auto VecTy = op0->getType();
    auto ElTy = cast<VectorType>(VecTy)->getElementType();
    auto ObjSize = (DL.getTypeSizeInBits(ElTy) + 7) / 8;
    TypeTree subTT1 = TT;
    if (auto CI = dyn_cast<ConstantInt>(op2)) {
      unsigned Off = (CI->getZExtValue() * DL.getTypeSizeInBits(ElTy)) / 8;
      subTT1 = TT.ShiftIndices(DL, Off, ObjSize, 0);
      subTT1.CanonicalizeInPlace(ObjSize, DL);
    }
    auto ip0 = invertPointerM(op0, bb, TT);
    auto ip1 = invertPointerM(op1, bb, subTT1);

    auto rule = [&](Value *ip0, Value *ip1) {
      return bb.CreateInsertElement(ip0, ip1, getNewFromOriginal(op2),
                                    arg->getName() + "'ipie");
    };

    Value *shadow = applyChainRule(arg->getType(), bb, rule, ip0, ip1);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<ShuffleVectorInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *op0 = arg->getOperand(0);
    Value *op1 = arg->getOperand(1);
    auto ip0 = invertPointerM(op0, bb, TT);
    auto ip1 = invertPointerM(op1, bb, TT);

    auto rule = [&bb, &arg](Value *ip0, Value *ip1) {
      return bb.CreateShuffleVector(ip0, ip1, arg->getShuffleMaskForBitcode(),
                                    arg->getName() + "'ipsv");
    };

    Value *shadow = applyChainRule(arg->getType(), bb, rule, ip0, ip1);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<SelectInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    bb.setFastMathFlags(getFast());

    Value *itval = nullptr;
    {
      auto tval = arg->getTrueValue();
      if (!runtimeActivity && TT.anyPointer(tval, DL) &&
          !isa<UndefValue>(tval) && !isa<ConstantPointerNull>(tval) &&
          !isa<ConstantAggregateZero>(tval) && isConstantValue(tval)) {
        std::string str;
        raw_string_ostream ss(str);
        ss << "Mismatched activity for: " << *arg << " const val: " << *tval;
        if (CustomErrorHandler)
          itval = unwrap(CustomErrorHandler(str.c_str(), wrap(arg),
                                            ErrorType::MixedActivityError, this,
                                            wrap(tval), wrap(&bb)));
        else
          EmitWarning("MixedActivityError", *arg, ss.str());
      }
      if (!itval) {
        itval = invertPointerM(tval, bb, TT);
      }
    }
    Value *ifval = nullptr;
    {
      auto fval = arg->getFalseValue();
      if (!runtimeActivity && TT[{-1}].isPossiblePointer() &&
          !isa<UndefValue>(fval) && !isa<ConstantPointerNull>(fval) &&
          !isa<ConstantAggregateZero>(fval) && isConstantValue(fval)) {
        std::string str;
        raw_string_ostream ss(str);
        ss << "Mismatched activity for: " << *arg << " const val: " << *fval;
        if (CustomErrorHandler)
          ifval = unwrap(CustomErrorHandler(str.c_str(), wrap(arg),
                                            ErrorType::MixedActivityError, this,
                                            wrap(fval), wrap(&bb)));
        else
          EmitWarning("MixedActivityError", *arg, ss.str());
      }
      if (!ifval) {
        ifval = invertPointerM(fval, bb, TT);
      }
    }

    Value *shadow = applyChainRule(
        arg->getType(), bb,
        [&](Value *tv, Value *fv) {
          return bb.CreateSelect(getNewFromOriginal(arg->getCondition()), tv,
                                 fv, arg->getName() + "'ipse");
        },
        itval, ifval);
    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto arg = dyn_cast<LoadInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *op0 = arg->getOperand(0);
    Value *ip = invertPointerM(op0, bb);

    SmallVector<Metadata *, 1> prevScopes;
    if (auto prev = arg->getMetadata(LLVMContext::MD_alias_scope)) {
      for (auto &M : cast<MDNode>(prev)->operands()) {
        prevScopes.push_back(M);
      }
    }
    SmallVector<Metadata *, 1> prevNoAlias;
    if (auto prev = arg->getMetadata(LLVMContext::MD_noalias)) {
      for (auto &M : cast<MDNode>(prev)->operands()) {
        prevNoAlias.push_back(M);
      }
    }
    size_t idx = 0;
    auto rule = [&](Value *ip) {
      auto li = bb.CreateLoad(arg->getType(), ip, arg->getName() + "'ipl");
      llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
      li->copyMetadata(*arg, ToCopy2);
      li->copyIRFlags(arg);

      SmallVector<Metadata *, 1> scopeMD = {getDerivativeAliasScope(op0, idx)};
      for (auto M : prevScopes)
        scopeMD.push_back(M);
      auto scope = MDNode::get(li->getContext(), scopeMD);
      li->setMetadata(LLVMContext::MD_alias_scope, scope);

      SmallVector<Metadata *, 1> MDs;
      for (ssize_t j = -1; j < getWidth(); j++) {
        if (j != (ssize_t)idx)
          MDs.push_back(getDerivativeAliasScope(op0, j));
      }
      for (auto M : prevNoAlias)
        MDs.push_back(M);
      if (MDs.size()) {
        auto noscope = MDNode::get(li->getContext(), MDs);
        li->setMetadata(LLVMContext::MD_noalias, noscope);
      }

      li->setAlignment(arg->getAlign());
      li->setDebugLoc(getNewFromOriginal(arg->getDebugLoc()));
      li->setVolatile(arg->isVolatile());
      li->setOrdering(arg->getOrdering());
      li->setSyncScopeID(arg->getSyncScopeID());
      idx++;
      return li;
    };

    Value *li = applyChainRule(arg->getType(), bb, rule, ip);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, li)));
    return li;

  } else if (auto arg = dyn_cast<BinaryOperator>(oval)) {
    if (mode == DerivativeMode::ReverseModePrimal ||
        mode == DerivativeMode::ReverseModeCombined ||
        mode == DerivativeMode::ReverseModeGradient) {
      if (TT.allFloat(arg, DL)) {
        return Constant::getNullValue(getShadowType(arg->getType()));
      }
    }

    if (!arg->getType()->isIntOrIntVectorTy()) {
      llvm::errs() << *oval << "\n";
    }
    assert(arg->getType()->isIntOrIntVectorTy());
    IRBuilder<> bb(getNewFromOriginal(arg));
    Value *val0 = nullptr;
    Value *val1 = nullptr;

    val0 = invertPointerM(arg->getOperand(0), bb, TT);
    val1 = invertPointerM(arg->getOperand(1), bb, TT);
    assert(val0->getType() == val1->getType());

    auto rule = [&bb, &arg](Value *val0, Value *val1) {
      auto li = bb.CreateBinOp(arg->getOpcode(), val0, val1, arg->getName());
      if (auto BI = dyn_cast<BinaryOperator>(li))
        BI->copyIRFlags(arg);
      return li;
    };

    Value *li = applyChainRule(arg->getType(), bb, rule, val0, val1);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, li)));
    return li;
  } else if (auto arg = dyn_cast<GetElementPtrInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(arg));
    SmallVector<Value *, 4> invertargs;
    for (unsigned i = 0; i < arg->getNumIndices(); ++i) {
      Value *b = getNewFromOriginal(arg->getOperand(1 + i));
      invertargs.push_back(b);
    }
    Value *ip = invertPointerM(arg->getPointerOperand(), bb);

    auto rule = [&](Value *ip) {
      auto shadow = bb.CreateGEP(arg->getSourceElementType(), ip, invertargs,
                                 arg->getName() + "'ipg");

      if (auto gep = dyn_cast<GetElementPtrInst>(shadow))
        gep->setIsInBounds(arg->isInBounds());

      return shadow;
    };

    Value *shadow = applyChainRule(arg->getType(), bb, rule, ip);

    invertedPointers.insert(
        std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
    return shadow;
  } else if (auto inst = dyn_cast<AllocaInst>(oval)) {
    IRBuilder<> bb(getNewFromOriginal(inst));
    Value *asize = getNewFromOriginal(inst->getArraySize());

    auto rule1 = [&]() {
      AllocaInst *antialloca = bb.CreateAlloca(
          inst->getAllocatedType(), inst->getType()->getPointerAddressSpace(),
          asize, inst->getName() + "'ipa");
      antialloca->setAlignment(inst->getAlign());
      return antialloca;
    };

    Value *antialloca = applyChainRule(oval->getType(), bb, rule1);

    invertedPointers.insert(std::make_pair(
        (const Value *)oval, InvertedPointerVH(this, antialloca)));

    if (auto ci = dyn_cast<ConstantInt>(asize)) {
      if (ci->isOne()) {

        auto rule = [&](Value *antialloca) {
          StoreInst *st = bb.CreateStore(
              Constant::getNullValue(inst->getAllocatedType()), antialloca);
          cast<StoreInst>(st)->setAlignment(inst->getAlign());
        };

        applyChainRule(bb, rule, antialloca);

        return antialloca;
      } else {
        // TODO handle alloca of size > 1
      }
    }

    auto rule2 = [&](Value *antialloca) {
      auto dst_arg =
          bb.CreateBitCast(antialloca, getInt8PtrTy(oval->getContext()));
      auto val_arg = ConstantInt::get(Type::getInt8Ty(oval->getContext()), 0);
      auto len_arg = bb.CreateMul(
          bb.CreateZExtOrTrunc(asize, Type::getInt64Ty(oval->getContext())),
          ConstantInt::get(Type::getInt64Ty(oval->getContext()),
                           M->getDataLayout().getTypeAllocSizeInBits(
                               inst->getAllocatedType()) /
                               8),
          "", true, true);
      auto volatile_arg = ConstantInt::getFalse(oval->getContext());

      Value *args[] = {dst_arg, val_arg, len_arg, volatile_arg};
      Type *tys[] = {dst_arg->getType(), len_arg->getType()};
      auto memset = cast<CallInst>(bb.CreateCall(
          getIntrinsicDeclaration(M, Intrinsic::memset, tys), args));
      memset->addParamAttr(
          0, Attribute::getWithAlignment(inst->getContext(), inst->getAlign()));
      memset->addParamAttr(0, Attribute::NonNull);
    };

    applyChainRule(bb, rule2, antialloca);

    return antialloca;
  } else if (auto II = dyn_cast<IntrinsicInst>(oval)) {
    if (isIntelSubscriptIntrinsic(*II)) {
      IRBuilder<> bb(getNewFromOriginal(II));

      const std::array<size_t, 4> idxArgsIndices{{0, 1, 2, 4}};
      const size_t ptrArgIndex = 3;

      SmallVector<Value *, 5> invertArgs(5);
      for (auto i : idxArgsIndices) {
        Value *idx = getNewFromOriginal(II->getOperand(i));
        invertArgs[i] = idx;
      }
      Value *invertPtrArg = invertPointerM(II->getOperand(ptrArgIndex), bb);
      invertArgs[ptrArgIndex] = invertPtrArg;

      auto rule = [&](Value *ip) {
        auto shadow = bb.CreateCall(II->getCalledFunction(), invertArgs);
        assert(isa<CallInst>(shadow));
#if LLVM_VERSION_MAJOR >= 13
        auto CI = cast<CallInst>(shadow);
        // Must copy the elementtype attribute as it is needed by the intrinsic
        CI->addParamAttr(
            ptrArgIndex,
            II->getParamAttr(ptrArgIndex, Attribute::AttrKind::ElementType));
#endif
        return shadow;
      };

      Value *shadow = applyChainRule(II->getType(), bb, rule, invertPtrArg);

      invertedPointers.insert(
          std::make_pair((const Value *)oval, InvertedPointerVH(this, shadow)));
      return shadow;
    }

    IRBuilder<> bb(getNewFromOriginal(II));
    bb.setFastMathFlags(getFast());
    switch (II->getIntrinsicID()) {
    default:
      goto end;
#if LLVM_VERSION_MAJOR < 20
    case Intrinsic::nvvm_ldg_global_i:
    case Intrinsic::nvvm_ldg_global_p:
    case Intrinsic::nvvm_ldg_global_f:
#endif
    case Intrinsic::nvvm_ldu_global_i:
    case Intrinsic::nvvm_ldu_global_p:
    case Intrinsic::nvvm_ldu_global_f: {
      return applyChainRule(
          II->getType(), bb,
          [&](Value *ptr) {
            Value *args[] = {ptr, getNewFromOriginal(II->getArgOperand(1))};
            auto li = bb.CreateCall(II->getCalledFunction(), args);
            llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
            ToCopy2.push_back(LLVMContext::MD_noalias);
            li->copyMetadata(*II, ToCopy2);
            li->setDebugLoc(getNewFromOriginal(II->getDebugLoc()));
            return li;
          },
          invertPointerM(II->getArgOperand(0), bb));
    case Intrinsic::masked_load: {
      auto invDefault = invertPointerM(II->getArgOperand(3), bb, TT);
      auto invPtr = invertPointerM(II->getArgOperand(0), bb);
      return applyChainRule(
          II->getType(), bb,
          [&](Value *ptr, Value *defaultV) {
            Value *args[] = {ptr, getNewFromOriginal(II->getArgOperand(1)),
                             getNewFromOriginal(II->getArgOperand(2)),
                             defaultV};
            llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
            ToCopy2.push_back(LLVMContext::MD_noalias);
            auto li = bb.CreateCall(II->getCalledFunction(), args);
            li->copyMetadata(*II, ToCopy2);
            li->setDebugLoc(getNewFromOriginal(II->getDebugLoc()));
            return li;
          },
          invPtr, invDefault);
    }
    }
    }
  } else if (auto phi = dyn_cast<PHINode>(oval)) {

    if (phi->getNumIncomingValues() == 0) {
      dumpMap(invertedPointers);
      assert(0 && "illegal iv of phi");
    }
    std::map<Value *, std::set<BasicBlock *>> mapped;
    for (unsigned int i = 0; i < phi->getNumIncomingValues(); ++i) {
      mapped[phi->getIncomingValue(i)].insert(phi->getIncomingBlock(i));
    }

    if (false && mapped.size() == 1) {
      return invertPointerM(phi->getIncomingValue(0), BuilderM, TT);
    }
#if 0
     else if (false && mapped.size() == 2) {
         IRBuilder <> bb(phi);
         auto which = bb.CreatePHI(Type::getInt1Ty(phi->getContext()), phi->getNumIncomingValues());
         //TODO this is not recursive

         int cnt = 0;
         Value* vals[2];
         for(auto v : mapped) {
            assert( cnt <= 1 );
            vals[cnt] = v.first;
            for (auto b : v.second) {
                which->addIncoming(ConstantInt::get(which->getType(), cnt), b);
            }
            ++cnt;
         }
         auto inv0 = invertPointerM(vals[0], BuilderM, TT);
         auto inv1 = invertPointerM(vals[1], BuilderM, TT);
         auto result = BuilderM.CreateSelect(which, inv1, inv0);
         return result;
     }
#endif

    else {
      auto NewV = getNewFromOriginal(phi);
      IRBuilder<> bb(NewV);
      bb.setFastMathFlags(getFast());
      // Note if the original phi node get's scev'd in NewF, it may
      // no longer be a phi and we need a new place to insert this phi
      // Note that if scev'd this can still be a phi with 0 incoming indicating
      // an unnecessary value to be replaced
      // TODO consider allowing the inverted pointer to become a scev
      if (!isa<PHINode>(NewV) ||
          cast<PHINode>(NewV)->getNumIncomingValues() == 0) {
        bb.SetInsertPoint(bb.GetInsertBlock(), bb.GetInsertBlock()->begin());
      }

      if (EnzymeVectorSplitPhi && width > 1) {
        IRBuilder<> postPhi(NewV->getParent()->getFirstNonPHI());
        Type *shadowTy = getShadowType(phi->getType());
        PHINode *tmp = bb.CreatePHI(shadowTy, phi->getNumIncomingValues());

        invertedPointers.insert(
            std::make_pair((const Value *)oval, InvertedPointerVH(this, tmp)));

        Type *wrappedType = ArrayType::get(phi->getType(), width);
        Value *res = UndefValue::get(wrappedType);

        SmallVector<Value *, 1> invertedVals;
        for (unsigned int j = 0; j < phi->getNumIncomingValues(); ++j) {
          IRBuilder<> pre(
              cast<BasicBlock>(getNewFromOriginal(phi->getIncomingBlock(j)))
                  ->getTerminator());
          Value *preval = phi->getIncomingValue(j);

          Value *val = nullptr;
          if (!runtimeActivity && TT[{-1}].isPossiblePointer() &&
              !isa<UndefValue>(preval) && !isa<ConstantPointerNull>(preval) &&
              isConstantValue(preval)) {
            std::string str;
            raw_string_ostream ss(str);
            ss << "Mismatched activity for: " << *phi
               << " const val: " << *preval;
            if (CustomErrorHandler)
              val = unwrap(CustomErrorHandler(str.c_str(), wrap(phi),
                                              ErrorType::MixedActivityError,
                                              this, wrap(preval), wrap(&pre)));
            else
              EmitWarning("MixedActivityError", *phi, ss.str());
          }
          if (!val) {
            val = invertPointerM(preval, pre, TT);
          }
          invertedVals.push_back(val);
        }

        for (unsigned int i = 0; i < getWidth(); ++i) {
          PHINode *which =
              bb.CreatePHI(phi->getType(), phi->getNumIncomingValues());
          which->setDebugLoc(getNewFromOriginal(phi->getDebugLoc()));

          // Avoid re-extracting from the same value, since multiple
          // entries to the same phi from the same block must have the
          // same value;
          DenseMap<BasicBlock *, Value *> samePHI;
          for (unsigned int j = 0; j < phi->getNumIncomingValues(); ++j) {
            IRBuilder<> pre(
                cast<BasicBlock>(getNewFromOriginal(phi->getIncomingBlock(j)))
                    ->getTerminator());
            Value *val = invertedVals[j];
            Value *extracted_diff;
            auto found = samePHI.find(phi->getIncomingBlock(j));
            if (found == samePHI.end()) {
              extracted_diff = extractMeta(pre, val, i);
              samePHI[phi->getIncomingBlock(j)] = extracted_diff;
            } else {
              extracted_diff = found->second;
            }
            which->addIncoming(
                extracted_diff,
                cast<BasicBlock>(getNewFromOriginal(phi->getIncomingBlock(j))));
          }

          res = postPhi.CreateInsertValue(res, which, {i});
        }
        invertedPointers.erase((const Value *)oval);
        replaceAWithB(tmp, res);
        erase(tmp);

        invertedPointers.insert(
            std::make_pair((const Value *)oval, InvertedPointerVH(this, res)));

        return res;
      } else {
        Type *shadowTy = getShadowType(phi->getType());
        PHINode *which = bb.CreatePHI(shadowTy, phi->getNumIncomingValues());
        which->setDebugLoc(getNewFromOriginal(phi->getDebugLoc()));

        invertedPointers.insert(std::make_pair((const Value *)oval,
                                               InvertedPointerVH(this, which)));

        for (unsigned int i = 0; i < phi->getNumIncomingValues(); ++i) {
          IRBuilder<> pre(
              cast<BasicBlock>(getNewFromOriginal(phi->getIncomingBlock(i)))
                  ->getTerminator());

          Value *preval = phi->getIncomingValue(i);

          Value *val = nullptr;
          if (!runtimeActivity && TT[{-1}].isPossiblePointer() &&
              !isa<UndefValue>(preval) && !isa<ConstantPointerNull>(preval) &&
              isConstantValue(preval)) {
            std::string str;
            raw_string_ostream ss(str);
            ss << "Mismatched activity for: " << *phi
               << " const val: " << *preval;
            if (CustomErrorHandler)
              val = unwrap(CustomErrorHandler(str.c_str(), wrap(phi),
                                              ErrorType::MixedActivityError,
                                              this, wrap(preval), wrap(&pre)));
            else
              EmitWarning("MixedActivityError", *phi, ss.str());
          }
          if (!val) {
            val = invertPointerM(preval, pre, TT);
          }

          which->addIncoming(val, cast<BasicBlock>(getNewFromOriginal(
                                      phi->getIncomingBlock(i))));
        }
        return which;
      }
    }
  } else if (auto CB = dyn_cast<CallBase>(oval)) {
    auto funcName = getFuncNameFromCall(CB);
    if (funcName == "julia.gc_loaded") {
      IRBuilder<> bb(getNewFromOriginal(CB));
      bb.setFastMathFlags(getFast());
      Value *ptr0shadow = invertPointerM(CB->getArgOperand(0), bb);
      Value *ptr1shadow = invertPointerM(CB->getArgOperand(1), bb);
      Value *res = applyChainRule(
          CB->getType(), bb,
          [&](Value *v1, Value *v2) -> Value * {
            Value *args[2] = {v1, v2};
            return bb.CreateCall(CB->getCalledFunction(), args);
          },
          ptr0shadow, ptr1shadow);

      assert(invertedPointers.find(oval) == invertedPointers.end());

      invertedPointers.insert(
          std::make_pair((const Value *)oval, InvertedPointerVH(this, res)));

      return res;
    }
  } else if (auto FPMO = dyn_cast<FPMathOperator>(oval)) {
    if (FPMO->getOpcode() == Instruction::FNeg) {
      if (mode == DerivativeMode::ReverseModePrimal ||
          mode == DerivativeMode::ReverseModeCombined ||
          mode == DerivativeMode::ReverseModeGradient) {
        if (TT.allFloat(FPMO, DL)) {
          return Constant::getNullValue(getShadowType(FPMO->getType()));
        }
      }
    }
  }

end:;
  assert(BuilderM.GetInsertBlock());
  assert(BuilderM.GetInsertBlock()->getParent());
  assert(oval);

  if (isa<CallBase>(oval) && TT.allFloat(oval, DL)) {
    return Constant::getNullValue(getShadowType(oval->getType()));
  }

  if (CustomErrorHandler) {
    std::string str;
    raw_string_ostream ss(str);
    ss << "cannot find shadow for " << *oval;
    auto iv =
        unwrap(CustomErrorHandler(str.c_str(), wrap(oval), ErrorType::NoShadow,
                                  this, nullptr, wrap(&BuilderM)));
    if (iv) {
      invertedPointers.insert(
          std::make_pair((const Value *)oval, InvertedPointerVH(this, iv)));
      return iv;
    }
  }

  llvm::errs() << *newFunc->getParent() << "\n";
  llvm::errs() << "fn:" << *newFunc << "\noval=" << *oval
               << " icv=" << isConstantValue(oval) << "\n";
  for (auto z : invertedPointers) {
    llvm::errs() << "available inversion for " << *z.first << " of "
                 << *z.second << "\n";
  }
  assert(0 && "cannot find deal with ptr that isnt arg");
  report_fatal_error("cannot find deal with ptr that isnt arg");
}

Value *GradientUtils::lookupM(Value *val, IRBuilder<> &BuilderM,
                              const ValueToValueMapTy &incoming_available,
                              bool tryLegalRecomputeCheck, BasicBlock *scope) {

  assert(mode == DerivativeMode::ReverseModePrimal ||
         mode == DerivativeMode::ReverseModeGradient ||
         mode == DerivativeMode::ReverseModeCombined);

  assert(val->getName() != "<badref>");
  {
    auto found = incoming_available.find(val);
    if (found != incoming_available.end())
      return found->second;
  }
  if (isa<Constant>(val)) {
    return val;
  }
  if (isa<BasicBlock>(val)) {
    return val;
  }
  if (isa<Function>(val)) {
    return val;
  }
  if (isa<UndefValue>(val)) {
    return val;
  }
  if (isa<Argument>(val)) {
    return val;
  }
  if (isa<MetadataAsValue>(val)) {
    return val;
  }
  if (isa<InlineAsm>(val)) {
    return val;
  }

  if (!isa<Instruction>(val)) {
    llvm::errs() << *val << "\n";
  }

  auto inst = cast<Instruction>(val);
  if (inversionAllocs && inst->getParent() == inversionAllocs) {
    return val;
  }
  assert(inst->getParent()->getParent() == newFunc);
  assert(BuilderM.GetInsertBlock()->getParent() == newFunc);
  if (scope == nullptr)
    scope = BuilderM.GetInsertBlock();
  assert(scope->getParent() == newFunc);

  bool reduceRegister = false;

  if (EnzymeRegisterReduce) {
    if (isNVLoad(inst)) {
      reduceRegister = true;
    }
    if (auto LI = dyn_cast<LoadInst>(inst)) {
      auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple());
      unsigned int SharedAddrSpace = getGPUSharedAddrSpace(TT);
      if (cast<PointerType>(LI->getPointerOperand()->getType())
              ->getAddressSpace() == SharedAddrSpace) {
        reduceRegister |= tryLegalRecomputeCheck &&
                          legalRecompute(LI, incoming_available, &BuilderM) &&
                          shouldRecompute(LI, incoming_available, &BuilderM);
      }
    }
    if (!inst->mayReadOrWriteMemory()) {
      reduceRegister |= tryLegalRecomputeCheck &&
                        legalRecompute(inst, incoming_available, &BuilderM) &&
                        shouldRecompute(inst, incoming_available, &BuilderM);
    }
    if (this->isOriginalBlock(*BuilderM.GetInsertBlock()))
      reduceRegister = false;
  }

  if (!reduceRegister) {
    if (isOriginalBlock(*BuilderM.GetInsertBlock())) {
      if (BuilderM.GetInsertBlock()->size() &&
          BuilderM.GetInsertPoint() != BuilderM.GetInsertBlock()->end()) {
        Instruction *use = &*BuilderM.GetInsertPoint();
        while (isa<PHINode>(use))
          use = use->getNextNode();
        if (DT.dominates(inst, use)) {
          return inst;
        } else {
          llvm::errs() << *BuilderM.GetInsertBlock()->getParent() << "\n";
          llvm::errs() << "didn't dominate inst: " << *inst
                       << "  point: " << *BuilderM.GetInsertPoint()
                       << "\nbb: " << *BuilderM.GetInsertBlock() << "\n";
        }
      } else {
        if (inst->getParent() == BuilderM.GetInsertBlock() ||
            DT.dominates(inst, BuilderM.GetInsertBlock())) {
          // allowed from block domination
          return inst;
        } else {
          llvm::errs() << *BuilderM.GetInsertBlock()->getParent() << "\n";
          llvm::errs() << "didn't dominate inst: " << *inst
                       << "\nbb: " << *BuilderM.GetInsertBlock() << "\n";
        }
      }
      // This is a reverse block
    } else if (BuilderM.GetInsertBlock() != inversionAllocs) {
      // Something in the entry (or anything that dominates all returns, doesn't
      // need caching)
      BasicBlock *orig = isOriginal(inst->getParent());
      if (!orig) {
        llvm::errs() << "oldFunc: " << *oldFunc << "\n";
        llvm::errs() << "newFunc: " << *newFunc << "\n";
        llvm::errs() << "insertBlock: " << *BuilderM.GetInsertBlock() << "\n";
        llvm::errs() << "instP: " << *inst->getParent() << "\n";
        llvm::errs() << "inst: " << *inst << "\n";
      }
      assert(orig);

      // TODO upgrade this to be all returns that this could enter from
      bool legal = BlocksDominatingAllReturns.count(orig);
      if (legal) {

        BasicBlock *forwardBlock =
            isOriginal(originalForReverseBlock(*BuilderM.GetInsertBlock()));
        assert(forwardBlock);

        // Don't allow this if we're not definitely using the last iteration of
        // this value
        //   + either because the value isn't in a loop
        //   + or because the forward of the block usage location isn't in a
        //   loop (thus last iteration)
        //   + or because the loop nests share no ancestry

        bool loopLegal = true;
        for (Loop *idx = OrigLI->getLoopFor(orig); idx != nullptr;
             idx = idx->getParentLoop()) {
          for (Loop *fdx = OrigLI->getLoopFor(forwardBlock); fdx != nullptr;
               fdx = fdx->getParentLoop()) {
            if (idx == fdx) {
              loopLegal = false;
              break;
            }
          }
        }

        if (loopLegal) {
          return inst;
        }
      }
    }
  }

  if (lookup_cache[BuilderM.GetInsertBlock()].find(val) !=
      lookup_cache[BuilderM.GetInsertBlock()].end()) {
    auto result = lookup_cache[BuilderM.GetInsertBlock()][val];
    if (result == nullptr) {
      lookup_cache[BuilderM.GetInsertBlock()].erase(val);
    } else {
      assert(result);
      assert(result->getType());
      result = BuilderM.CreateBitCast(result, val->getType());
      assert(result->getType() == inst->getType());
      return result;
    }
  }

  ValueToValueMapTy available;
  for (auto pair : incoming_available) {
    if (pair.second)
      assert(pair.first->getType() == pair.second->getType());
    available[pair.first] = pair.second;
  }

  {
    BasicBlock *forwardPass = BuilderM.GetInsertBlock();
    if (forwardPass != inversionAllocs && !isOriginalBlock(*forwardPass)) {
      forwardPass = originalForReverseBlock(*forwardPass);
    }
    LoopContext lc;
    bool inLoop = getContext(forwardPass, lc);

    if (inLoop) {
      bool first = true;
      for (LoopContext idx = lc;; getContext(idx.parent->getHeader(), idx)) {
        if (available.count(idx.var) == 0) {
          if (!isOriginalBlock(*BuilderM.GetInsertBlock())) {
            available[idx.var] =
                BuilderM.CreateLoad(idx.var->getType(), idx.antivaralloc);
          } else {
            available[idx.var] = idx.var;
          }
        }
        if (!first && idx.var == inst)
          return available[idx.var];
        if (first) {
          first = false;
        }
        if (idx.parent == nullptr)
          break;
      }
    }
  }

  if (available.count(inst)) {
    assert(available[inst]->getType() == inst->getType());
    return available[inst];
  }

  // If requesting loop bound and not available from index per above
  // we must be requesting the total size. Rather than generating
  // a new lcssa variable, use the existing loop exact bound var
  {
    LoopContext lc;
    bool loopVar = false;
    if (getContext(inst->getParent(), lc) && lc.var == inst) {
      loopVar = true;
    } else if (auto phi = dyn_cast<PHINode>(inst)) {
      Value *V = nullptr;
      bool legal = true;
      for (auto &val : phi->incoming_values()) {
        if (isa<UndefValue>(val))
          continue;
        if (V == nullptr)
          V = val;
        else if (V != val) {
          legal = false;
          break;
        }
      }
      if (legal) {
        if (auto I = dyn_cast_or_null<PHINode>(V)) {
          if (getContext(I->getParent(), lc) && lc.var == I) {
            loopVar = true;
          }
        }
      }
    }
    if (loopVar) {
      Value *lim = nullptr;
      if (lc.dynamic) {
        // Must be in a reverse pass fashion for a lookup to index bound to be
        // legal
        assert(/*ReverseLimit*/ reverseBlocks.size() > 0);
        LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                          lc.preheader);
        lim = lookupValueFromCache(
            lc.var->getType(), /*forwardPass*/ false, BuilderM, lctx,
            getDynamicLoopLimit(LI.getLoopFor(lc.header)),
            /*isi1*/ false, available);
      } else {
        lim = lookupM(lc.trueLimit, BuilderM);
      }
      lookup_cache[BuilderM.GetInsertBlock()][val] = lim;
      return lim;
    }
  }

  Instruction *prelcssaInst = inst;

  assert(inst->getName() != "<badref>");
  val = fixLCSSA(inst, scope);
  if (isa<UndefValue>(val)) {
    llvm::errs() << *oldFunc << "\n";
    llvm::errs() << *newFunc << "\n";
    llvm::errs() << *BuilderM.GetInsertBlock() << "\n";
    llvm::errs() << *scope << "\n";
    llvm::errs() << *val << " inst " << *inst << "\n";
    assert(0 && "undef value upon lcssa");
  }
  inst = cast<Instruction>(val);
  assert(prelcssaInst->getType() == inst->getType());
  assert(!this->isOriginalBlock(*BuilderM.GetInsertBlock()));

  // Update index and caching per lcssa
  if (lookup_cache[BuilderM.GetInsertBlock()].find(val) !=
      lookup_cache[BuilderM.GetInsertBlock()].end()) {
    auto result = lookup_cache[BuilderM.GetInsertBlock()][val];
    if (result == nullptr) {
      lookup_cache[BuilderM.GetInsertBlock()].erase(val);
    } else {
      assert(result);
      assert(result->getType());
      result = BuilderM.CreateBitCast(result, val->getType());
      assert(result->getType() == inst->getType());
      return result;
    }
  }

  // TODO consider call as part of
  bool lrc = false, src = false;
  if (tryLegalRecomputeCheck &&
      (lrc = legalRecompute(prelcssaInst, available, &BuilderM))) {
    if ((src = shouldRecompute(prelcssaInst, available, &BuilderM))) {
      auto ity = inst->getType();
      // unwrap of a placeholder shadow may replace the placeholder with a
      // real value, in which case the original value may have deleted.
      // Store the information to correctly store the cache.
      Value *uninverted = nullptr;
      if (auto phi = dyn_cast<PHINode>(val)) {
        if (phi->getNumIncomingValues() == 0) {
          uninverted = hasUninverted(phi);
        }
      }
      auto op = unwrapM(prelcssaInst, BuilderM, available,
                        UnwrapMode::AttemptSingleUnwrap, scope);
      if (op) {
        assert(op);
        assert(op->getType());
        if (op->getType() != ity) {
          llvm::errs() << " op: " << *op << " inst: " << *inst << "\n";
        }
        assert(op->getType() == ity);
        if (!reduceRegister) {
          auto cache_key = val;
          if (uninverted) {
            auto found = invertedPointers.find(uninverted);
            if (found != invertedPointers.end()) {
              cache_key = found->second;
            }
          }
          lookup_cache[BuilderM.GetInsertBlock()][cache_key] = op;
        }
        return op;
      }
    } else {
      if (isa<LoadInst>(prelcssaInst)) {
      }
    }
  }

  if (auto li = dyn_cast<LoadInst>(inst))
    if (auto origInst = dyn_cast_or_null<LoadInst>(isOriginal(inst))) {
      auto liobj = getBaseObject(li->getPointerOperand());

      auto orig_liobj = getBaseObject(origInst->getPointerOperand());

      if (scopeMap.find(inst) == scopeMap.end()) {
        for (auto pair : scopeMap) {
          if (auto li2 = dyn_cast<LoadInst>(const_cast<Value *>(pair.first))) {

            auto li2obj = getBaseObject(li2->getPointerOperand());

            if (liobj == li2obj && DT.dominates(li2, li)) {
              auto orig2 = dyn_cast_or_null<LoadInst>(isOriginal(li2));
              if (!orig2)
                continue;

              bool failed = false;

              // llvm::errs() << "found potential candidate loads: oli:"
              //             << *origInst << " oli2: " << *orig2 << "\n";

              auto scev1 = OrigSE->getSCEV(origInst->getPointerOperand());
              auto scev2 = OrigSE->getSCEV(orig2->getPointerOperand());
              // llvm::errs() << " scev1: " << *scev1 << " scev2: " << *scev2
              //             << "\n";

              allInstructionsBetween(
                  *OrigLI, orig2, origInst, [&](Instruction *I) -> bool {
                    if (I->mayWriteToMemory() &&
                        writesToMemoryReadBy(&TR, *OrigAA, TLI,
                                             /*maybeReader*/ origInst,
                                             /*maybeWriter*/ I)) {
                      failed = true;
                      // llvm::errs() << "FAILED: " << *I << "\n";
                      return /*earlyBreak*/ true;
                    }
                    return /*earlyBreak*/ false;
                  });
              if (failed)
                continue;

              if (auto ar1 = dyn_cast<SCEVAddRecExpr>(scev1)) {
                if (auto ar2 = dyn_cast<SCEVAddRecExpr>(scev2)) {
                  if (ar1->getStart() != OrigSE->getCouldNotCompute() &&
                      ar1->getStart() == ar2->getStart() &&
                      ar1->getStepRecurrence(*OrigSE) !=
                          OrigSE->getCouldNotCompute() &&
                      ar1->getStepRecurrence(*OrigSE) ==
                          ar2->getStepRecurrence(*OrigSE)) {

                    LoopContext l1;
                    getContext(getNewFromOriginal(ar1->getLoop()->getHeader()),
                               l1);
                    LoopContext l2;
                    getContext(getNewFromOriginal(ar2->getLoop()->getHeader()),
                               l2);
                    if (l1.dynamic || l2.dynamic)
                      continue;

                    // TODO IF len(ar2) >= len(ar1) then we can replace li with
                    // li2
                    if (SE.getSCEV(l1.trueLimit) != SE.getCouldNotCompute() &&
                        SE.getSCEV(l1.trueLimit) == SE.getSCEV(l2.trueLimit)) {
                      // llvm::errs()
                      //    << " step1: " << *ar1->getStepRecurrence(SE)
                      //    << " step2: " << *ar2->getStepRecurrence(SE) <<
                      //    "\n";

                      inst = li2;
                      break;
                    }
                  }
                }
              }
            }
          }
        }

        auto scev1 = OrigSE->getSCEV(origInst->getPointerOperand());

        auto TT = llvm::Triple(newFunc->getParent()->getTargetTriple());
        unsigned int SharedAddrSpace = getGPUSharedAddrSpace(TT);
        if (EnzymeSharedForward && scev1 != OrigSE->getCouldNotCompute() &&
            cast<PointerType>(orig_liobj->getType())->getAddressSpace() ==
                SharedAddrSpace) {
          Value *resultValue = nullptr;
          ValueToValueMapTy newavail;
          for (const auto &pair : available) {
            assert(pair.first->getType() == pair.second->getType());
            newavail[pair.first] = pair.second;
          }
          allDomPredecessorsOf(origInst, *OrigDT, [&](Instruction *pred) {
            if (auto SI = dyn_cast<StoreInst>(pred)) {
              // auto NewSI = cast<StoreInst>(getNewFromOriginal(SI));
              auto si2obj = getBaseObject(SI->getPointerOperand());

              if (si2obj != orig_liobj)
                return false;

              bool lastStore = true;
              bool interveningSync = false;
              allInstructionsBetween(
                  *OrigLI, SI, origInst, [&](Instruction *potentialAlias) {
                    if (!potentialAlias->mayWriteToMemory())
                      return false;
                    if (!writesToMemoryReadBy(&TR, *OrigAA, TLI, origInst,
                                              potentialAlias))
                      return false;

                    if (auto II = dyn_cast<IntrinsicInst>(potentialAlias)) {
#if LLVM_VERSION_MAJOR > 20
                      if (II->getIntrinsicID() ==
                              Intrinsic::nvvm_barrier_cta_sync_aligned_all ||
                          II->getIntrinsicID() == Intrinsic::amdgcn_s_barrier) {
#else
                      if (II->getIntrinsicID() == Intrinsic::nvvm_barrier0 ||
                          II->getIntrinsicID() == Intrinsic::amdgcn_s_barrier) {
#endif
                        interveningSync = OrigDT->dominates(SI, II) &&
                                          OrigDT->dominates(II, origInst);
                        allUnsyncdPredecessorsOf(
                            II,
                            [&](Instruction *mid) {
                              if (!mid->mayWriteToMemory())
                                return false;

                              if (mid == SI)
                                return false;

                              if (!writesToMemoryReadBy(&TR, *OrigAA, TLI,
                                                        origInst, mid)) {
                                return false;
                              }
                              lastStore = false;
                              return true;
                            },
                            [&]() {
                              // if gone past entry
                              if (mode != DerivativeMode::ReverseModeCombined) {
                                lastStore = false;
                              }
                            });
                        if (!lastStore)
                          return true;
                        else
                          return false;
                      }
                    }

                    lastStore = false;
                    return true;
                  });

              if (!lastStore)
                return false;

              auto scev2 = OrigSE->getSCEV(SI->getPointerOperand());
              bool legal = scev1 == scev2;
              if (auto ar2 = dyn_cast<SCEVAddRecExpr>(scev2)) {
                if (auto ar1 = dyn_cast<SCEVAddRecExpr>(scev1)) {
                  if (ar2->getStart() != OrigSE->getCouldNotCompute() &&
                      ar1->getStart() == ar2->getStart() &&
                      ar2->getStepRecurrence(*OrigSE) !=
                          OrigSE->getCouldNotCompute() &&
                      ar1->getStepRecurrence(*OrigSE) ==
                          ar2->getStepRecurrence(*OrigSE)) {

                    LoopContext l1;
                    getContext(getNewFromOriginal(ar1->getLoop()->getHeader()),
                               l1);
                    LoopContext l2;
                    getContext(getNewFromOriginal(ar2->getLoop()->getHeader()),
                               l2);
                    if (!l1.dynamic && !l2.dynamic) {
                      // TODO IF len(ar2) >= len(ar1) then we can replace li
                      // with li2
                      if (l1.trueLimit == l2.trueLimit) {
                        const Loop *L1 = ar1->getLoop();
                        while (L1) {
                          if (L1 == ar2->getLoop())
                            return false;
                          L1 = L1->getParentLoop();
                        }
                        newavail[l2.var] = available[l1.var];
                        legal = true;
                      }
                    }
                  }
                }
              }
              if (!legal) {
                Value *sval = SI->getPointerOperand();
                Value *lval = origInst->getPointerOperand();
                while (auto CI = dyn_cast<CastInst>(sval))
                  sval = CI->getOperand(0);
                while (auto CI = dyn_cast<CastInst>(lval))
                  lval = CI->getOperand(0);
                if (auto sgep = dyn_cast<GetElementPtrInst>(sval)) {
                  if (auto lgep = dyn_cast<GetElementPtrInst>(lval)) {
                    if (sgep->getPointerOperand() ==
                        lgep->getPointerOperand()) {
                      SmallVector<Value *, 3> svals;
                      for (auto &v : sgep->indices()) {
                        Value *q = v;
                        while (auto CI = dyn_cast<CastInst>(q))
                          q = CI->getOperand(0);
                        svals.push_back(q);
                      }
                      SmallVector<Value *, 3> lvals;
                      for (auto &v : lgep->indices()) {
                        Value *q = v;
                        while (auto CI = dyn_cast<CastInst>(q))
                          q = CI->getOperand(0);
                        lvals.push_back(q);
                      }
                      ValueToValueMapTy ThreadLookup;
                      bool legal = true;
                      for (size_t i = 0; i < svals.size(); i++) {
                        auto ss = OrigSE->getSCEV(svals[i]);
                        auto ls = OrigSE->getSCEV(lvals[i]);
                        if (cast<IntegerType>(ss->getType())->getBitWidth() >
                            cast<IntegerType>(ls->getType())->getBitWidth()) {
                          ls = OrigSE->getZeroExtendExpr(ls, ss->getType());
                        }
                        if (cast<IntegerType>(ss->getType())->getBitWidth() <
                            cast<IntegerType>(ls->getType())->getBitWidth()) {
                          ls = OrigSE->getTruncateExpr(ls, ss->getType());
                        }
                        if (ls != ss) {
                          if (auto II = dyn_cast<IntrinsicInst>(svals[i])) {
                            switch (II->getIntrinsicID()) {
                            case Intrinsic::nvvm_read_ptx_sreg_tid_x:
                            case Intrinsic::nvvm_read_ptx_sreg_tid_y:
                            case Intrinsic::nvvm_read_ptx_sreg_tid_z:
                            case Intrinsic::amdgcn_workitem_id_x:
                            case Intrinsic::amdgcn_workitem_id_y:
                            case Intrinsic::amdgcn_workitem_id_z:
                              ThreadLookup[getNewFromOriginal(II)] =
                                  BuilderM.CreateZExtOrTrunc(
                                      lookupM(getNewFromOriginal(lvals[i]),
                                              BuilderM, available),
                                      II->getType());
                              break;
                            default:
                              legal = false;
                              break;
                            }
                          } else {
                            legal = false;
                            break;
                          }
                        }
                      }
                      if (legal) {
                        for (auto pair : newavail) {
                          assert(pair.first->getType() ==
                                 pair.second->getType());
                          ThreadLookup[pair.first] = pair.second;
                        }
                        Value *recomp = unwrapM(
                            getNewFromOriginal(SI->getValueOperand()), BuilderM,
                            ThreadLookup, UnwrapMode::AttemptFullUnwrap, scope,
                            /*permitCache*/ false);
                        if (recomp) {
                          resultValue = recomp;
                          return true;
                          ;
                        }
                      }
                    }
                  }
                }
              }
              if (!legal)
                return false;
              return true;
            }
            return false;
          });

          if (resultValue) {
            if (resultValue->getType() != val->getType())
              resultValue = BuilderM.CreateBitCast(resultValue, val->getType());
            return resultValue;
          }
        }
      }

      auto loadSize = (li->getParent()
                           ->getParent()
                           ->getParent()
                           ->getDataLayout()
                           .getTypeAllocSizeInBits(li->getType()) +
                       7) /
                      8;

      // this is guarded because havent told cacheForReverse how to move
      if (mode == DerivativeMode::ReverseModeCombined)
        if (!li->isVolatile() && EnzymeLoopInvariantCache) {
          if (auto AI = dyn_cast<AllocaInst>(liobj)) {
            assert(isa<AllocaInst>(orig_liobj));
            if (auto AT = dyn_cast<ArrayType>(AI->getAllocatedType()))
              if (auto GEP =
                      dyn_cast<GetElementPtrInst>(li->getPointerOperand())) {
                if (GEP->getPointerOperand() == AI) {
                  LoopContext l1;
                  if (!getContext(li->getParent(), l1))
                    goto noSpeedCache;

                  BasicBlock *ctx = l1.preheader;

                  auto origPH = cast_or_null<BasicBlock>(isOriginal(ctx));
                  assert(origPH);
                  if (OrigPDT->dominates(origPH, origInst->getParent())) {
                    goto noSpeedCache;
                  }

                  Instruction *origTerm = origPH->getTerminator();
                  if (!origTerm)
                    llvm::errs() << *origPH << "\n";
                  assert(origTerm);
                  IRBuilder<> OB(origTerm);
                  LoadInst *tmpload = OB.CreateLoad(AT, orig_liobj, "'tmpload");

                  bool failed = false;
                  allInstructionsBetween(
                      *OrigLI, &*origTerm, origInst,
                      [&](Instruction *I) -> bool {
                        if (I->mayWriteToMemory() &&
                            writesToMemoryReadBy(&TR, *OrigAA, TLI,
                                                 /*maybeReader*/ tmpload,
                                                 /*maybeWriter*/ I)) {
                          failed = true;
                          return /*earlyBreak*/ true;
                        }
                        return /*earlyBreak*/ false;
                      });
                  if (failed) {
                    tmpload->eraseFromParent();
                    goto noSpeedCache;
                  }
                  while (Loop *L = LI.getLoopFor(ctx)) {
                    BasicBlock *nctx = L->getLoopPreheader();
                    assert(nctx);
                    bool failed = false;
                    auto origPH = cast_or_null<BasicBlock>(isOriginal(nctx));
                    assert(origPH);
                    if (OrigPDT->dominates(origPH, origInst->getParent())) {
                      break;
                    }
                    Instruction *origTerm = origPH->getTerminator();
                    allInstructionsBetween(
                        *OrigLI, &*origTerm, origInst,
                        [&](Instruction *I) -> bool {
                          if (I->mayWriteToMemory() &&
                              writesToMemoryReadBy(&TR, *OrigAA, TLI,
                                                   /*maybeReader*/ tmpload,
                                                   /*maybeWriter*/ I)) {
                            failed = true;
                            return /*earlyBreak*/ true;
                          }
                          return /*earlyBreak*/ false;
                        });
                    if (failed)
                      break;
                    ctx = nctx;
                  }

                  tmpload->eraseFromParent();

                  IRBuilder<> v(ctx->getTerminator());

                  AllocaInst *cache = nullptr;

                  LoopContext tmp;
                  bool forceSingleIter = false;
                  if (!getContext(ctx, tmp)) {
                    forceSingleIter = true;
                  }
                  LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                                    ctx, forceSingleIter);

                  if (auto found = findInMap(scopeMap, (Value *)liobj)) {
                    cache = found->first;
                  } else {
                    // if freeing reverseblocks must exist
                    assert(reverseBlocks.size());
                    cache = createCacheForScope(lctx, AT, li->getName(),
                                                /*shouldFree*/ true,
                                                /*allocate*/ true);
                    assert(cache);
                    scopeMap.insert(
                        std::make_pair(AI, std::make_pair(cache, lctx)));

                    v.setFastMathFlags(getFast());
                    assert(isOriginalBlock(*v.GetInsertBlock()));
                    Value *outer =
                        getCachePointer(AT,
                                        /*inForwardPass*/ true, v, lctx, cache,
                                        /*storeinstorecache*/ true,
                                        /*available*/ ValueToValueMapTy(),
                                        /*extraSize*/ nullptr);

                    auto ld = v.CreateLoad(AT, AI);
                    ld->setAlignment(AI->getAlign());
                    scopeInstructions[cache].push_back(ld);
                    auto st = v.CreateStore(ld, outer);
                    auto bsize = newFunc->getParent()
                                     ->getDataLayout()
                                     .getTypeAllocSizeInBits(AT) /
                                 8;
                    if ((bsize & (bsize - 1)) == 0) {
                      st->setAlignment(Align(bsize));
                    }
                    scopeInstructions[cache].push_back(st);
                    for (auto post : PostCacheStore(st, v)) {
                      scopeInstructions[cache].push_back(post);
                    }
                  }

                  assert(!isOriginalBlock(*BuilderM.GetInsertBlock()));
                  Value *outer = getCachePointer(
                      AT,
                      /*inForwardPass*/ false, BuilderM, lctx, cache,
                      /*storeinstorecache*/ true, available,
                      /*extraSize*/ nullptr);
                  SmallVector<Value *, 2> idxs;
                  for (auto &idx : GEP->indices()) {
                    idxs.push_back(lookupM(idx, BuilderM, available,
                                           tryLegalRecomputeCheck));
                  }

                  auto cptr = BuilderM.CreateGEP(GEP->getSourceElementType(),
                                                 outer, idxs);
                  cast<GetElementPtrInst>(cptr)->setIsInBounds(true);

                  // Retrieve the actual result
                  auto result = loadFromCachePointer(val->getType(), BuilderM,
                                                     cptr, cache);

                  assert(result->getType() == inst->getType());
                  lookup_cache[BuilderM.GetInsertBlock()][val] = result;
                  return result;
                }
              }
          }

          auto scev1 = OrigSE->getSCEV(origInst->getPointerOperand());
          // Store in memcpy opt
          Value *lim = nullptr;
          BasicBlock *ctx = nullptr;
          Value *start = nullptr;
          Value *offset = nullptr;
          if (auto ar1 = dyn_cast<SCEVAddRecExpr>(scev1)) {
            if (auto step =
                    dyn_cast<SCEVConstant>(ar1->getStepRecurrence(*OrigSE))) {
              if (step->getAPInt() != loadSize)
                goto noSpeedCache;

              LoopContext l1;
              getContext(getNewFromOriginal(ar1->getLoop()->getHeader()), l1);

              if (l1.dynamic)
                goto noSpeedCache;

              offset = available[l1.var];
              ctx = l1.preheader;

              IRBuilder<> v(ctx->getTerminator());

              auto origPH = cast_or_null<BasicBlock>(isOriginal(ctx));
              assert(origPH);
              if (OrigPDT->dominates(origPH, origInst->getParent())) {
                goto noSpeedCache;
              }

              lim = unwrapM(l1.trueLimit, v,
                            /*available*/ ValueToValueMapTy(),
                            UnwrapMode::AttemptFullUnwrapWithLookup);
              if (!lim) {
                goto noSpeedCache;
              }
              lim = v.CreateAdd(lim, ConstantInt::get(lim->getType(), 1), "",
                                true, true);

              {
#if LLVM_VERSION_MAJOR >= 12
                Value *start0;
                SmallVector<Instruction *, 32> InsertedInstructions;
                {
#if LLVM_VERSION_MAJOR >= 22
                  SCEVExpander OrigExp(*OrigSE, "enzyme",
                                       /*PreserveLCSSA = */ false);
#else
                  SCEVExpander OrigExp(
                      *OrigSE, ctx->getParent()->getParent()->getDataLayout(),
                      "enzyme", /*PreserveLCSSA = */ false);
#endif

                  OrigExp.setInsertPoint(
                      isOriginal(l1.header)->getTerminator());

                  start0 = OrigExp.expandCodeFor(
                      ar1->getStart(), li->getPointerOperand()->getType());
                  InsertedInstructions = OrigExp.getAllInsertedInstructions();
                }

                ValueToValueMapTy available;
                for (const auto &pair : originalToNewFn) {
                  if (pair.first->getType() == pair.second->getType())
                    available[pair.first] = pair.second;
                }

                // Sort so that later instructions do not dominate earlier
                // instructions.
                llvm::stable_sort(InsertedInstructions,
                                  [this](Instruction *A, Instruction *B) {
                                    return OrigDT->dominates(A, B);
                                  });
                for (auto a : InsertedInstructions) {
                  if (isa<PHINode>(a)) {
                    std::string str;
                    raw_string_ostream ss(str);
                    ss << "oldFunc: " << *oldFunc << "\n";
                    ss << "newFunc: " << *newFunc << "\n";
                    ss << "li: " << *li << "\n";
                    ss << "start0: " << *start0 << "\n";
                    ss << "Inserted a phi node (" << *a
                       << ") during unwrap of SCEV: " << *ar1->getStart()
                       << "\n";
                    if (CustomErrorHandler) {
                      CustomErrorHandler(str.c_str(), wrap(li),
                                         ErrorType::InternalError, nullptr,
                                         nullptr, nullptr);
                    } else {
                      EmitFailure("InsertedPHISCEV", li->getDebugLoc(), li,
                                  ss.str());
                    }
                  }
                  auto uwV =
                      unwrapM(a, v, available, UnwrapMode::AttemptSingleUnwrap,
                              /*scope*/ nullptr, /*cache*/ false);
                  auto uw = dyn_cast<Instruction>(uwV);
                  assert(uwV->getType() == a->getType());
#ifndef NDEBUG
                  if (uw) {
                    for (size_t i = 0; i < uw->getNumOperands(); i++) {
                      auto op = uw->getOperand(i);
                      if (auto arg = dyn_cast<Argument>(op))
                        assert(arg->getParent() == newFunc);
                      else if (auto inst = dyn_cast<Instruction>(op))
                        assert(inst->getParent()->getParent() == newFunc);
                    }
                    assert(uw->getParent()->getParent() == newFunc);
                  }
#endif
                  available[a] = uwV;
                  if (uw)
                    unwrappedLoads.erase(uw);
                }

                start =
                    isa<Constant>(start0) ? start0 : (Value *)available[start0];
                if (!start) {
                  llvm::errs() << "old: " << *oldFunc << "\n";
                  llvm::errs() << "new: " << *newFunc << "\n";
                  llvm::errs() << "start0: " << *start0 << "\n";
                }
                assert(start);

                available.clear();
                for (auto I : llvm::reverse(InsertedInstructions)) {
                  assert(I->getNumUses() == 0);
                  OrigSE->forgetValue(I);
                  I->eraseFromParent();
                }
#endif
              }

              if (!start)
                goto noSpeedCache;

              Instruction *origTerm = origPH->getTerminator();

              bool failed = false;
              allInstructionsBetween(
                  *OrigLI, &*origTerm, origInst, [&](Instruction *I) -> bool {
                    if (I->mayWriteToMemory() &&
                        writesToMemoryReadBy(&TR, *OrigAA, TLI,
                                             /*maybeReader*/ origInst,
                                             /*maybeWriter*/ I)) {
                      failed = true;
                      return /*earlyBreak*/ true;
                    }
                    return /*earlyBreak*/ false;
                  });
              if (failed)
                goto noSpeedCache;
            }
          }

          if (ctx && lim && start && offset) {
            Value *firstLim = lim;
            Value *firstStart = start;
            while (Loop *L = LI.getLoopFor(ctx)) {
              BasicBlock *nctx = L->getLoopPreheader();
              assert(nctx);
              bool failed = false;
              auto origPH = cast_or_null<BasicBlock>(isOriginal(nctx));
              assert(origPH);
              if (OrigPDT->dominates(origPH, origInst->getParent())) {
                break;
              }
              Instruction *origTerm = origPH->getTerminator();
              allInstructionsBetween(
                  *OrigLI, &*origTerm, origInst, [&](Instruction *I) -> bool {
                    if (I->mayWriteToMemory() &&
                        writesToMemoryReadBy(&TR, *OrigAA, TLI,
                                             /*maybeReader*/ origInst,
                                             /*maybeWriter*/ I)) {
                      failed = true;
                      return /*earlyBreak*/ true;
                    }
                    return /*earlyBreak*/ false;
                  });
              if (failed)
                break;
              IRBuilder<> nv(nctx->getTerminator());
              Value *nlim = unwrapM(firstLim, nv,
                                    /*available*/ ValueToValueMapTy(),
                                    UnwrapMode::AttemptFullUnwrapWithLookup);
              if (!nlim)
                break;
              Value *nstart = unwrapM(firstStart, nv,
                                      /*available*/ ValueToValueMapTy(),
                                      UnwrapMode::AttemptFullUnwrapWithLookup);
              if (!nstart)
                break;
              lim = nlim;
              start = nstart;
              ctx = nctx;
            }
            IRBuilder<> v(ctx->getTerminator());
            bool isi1 = val->getType()->isIntegerTy() &&
                        cast<IntegerType>(li->getType())->getBitWidth() == 1;

            AllocaInst *cache = nullptr;

            LoopContext tmp;
            bool forceSingleIter = false;
            if (!getContext(ctx, tmp)) {
              forceSingleIter = true;
            } else if (auto inst = dyn_cast<Instruction>(lim)) {
              if (inst->getParent() == ctx ||
                  !DT.dominates(inst->getParent(), ctx)) {
                forceSingleIter = true;
              }
            }
            LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0, ctx,
                              forceSingleIter);

            if (auto found = findInMap(scopeMap, (Value *)inst)) {
              cache = found->first;
            } else {
              // if freeing reverseblocks must exist
              assert(reverseBlocks.size());
              cache = createCacheForScope(lctx, li->getType(), li->getName(),
                                          /*shouldFree*/ true,
                                          /*allocate*/ true, /*extraSize*/ lim);
              assert(cache);
              scopeMap.insert(
                  std::make_pair(inst, std::make_pair(cache, lctx)));

              v.setFastMathFlags(getFast());
              assert(isOriginalBlock(*v.GetInsertBlock()));
              Value *outer =
                  getCachePointer(li->getType(),
                                  /*inForwardPass*/ true, v, lctx, cache,
                                  /*storeinstorecache*/ true,
                                  /*available*/ ValueToValueMapTy(),
                                  /*extraSize*/ lim);

              auto dst_arg = v.CreateBitCast(
                  outer,
                  getInt8PtrTy(
                      inst->getContext(),
                      cast<PointerType>(outer->getType())->getAddressSpace()));
              scopeInstructions[cache].push_back(cast<Instruction>(dst_arg));
              auto src_arg = v.CreateBitCast(
                  start,
                  getInt8PtrTy(
                      inst->getContext(),
                      cast<PointerType>(start->getType())->getAddressSpace()));
              auto len_arg =
                  v.CreateMul(ConstantInt::get(lim->getType(), loadSize), lim,
                              "", true, true);
              if (Instruction *I = dyn_cast<Instruction>(len_arg))
                scopeInstructions[cache].push_back(I);
              auto volatile_arg = ConstantInt::getFalse(inst->getContext());

              Value *nargs[] = {dst_arg, src_arg, len_arg, volatile_arg};

              Type *tys[] = {dst_arg->getType(), src_arg->getType(),
                             len_arg->getType()};

              auto memcpyF = getIntrinsicDeclaration(newFunc->getParent(),
                                                     Intrinsic::memcpy, tys);
              auto mem = cast<CallInst>(v.CreateCall(memcpyF, nargs));

              mem->addParamAttr(0, Attribute::NonNull);
              mem->addParamAttr(1, Attribute::NonNull);

              auto bsize =
                  newFunc->getParent()->getDataLayout().getTypeAllocSizeInBits(
                      li->getType()) /
                  8;
              if ((bsize & (bsize - 1)) == 0) {
                mem->addParamAttr(0, Attribute::getWithAlignment(
                                         memcpyF->getContext(), Align(bsize)));
              }

              mem->addParamAttr(1, Attribute::getWithAlignment(
                                       memcpyF->getContext(), li->getAlign()));
              scopeInstructions[cache].push_back(mem);
            }

            assert(!isOriginalBlock(*BuilderM.GetInsertBlock()));
            Value *result = lookupValueFromCache(
                inst->getType(),
                /*isForwardPass*/ false, BuilderM, lctx, cache, isi1, available,
                /*extraSize*/ lim, offset);
            assert(result->getType() == inst->getType());
            lookup_cache[BuilderM.GetInsertBlock()][val] = result;

            EmitWarning("Uncacheable", *inst, "Caching instruction ", *inst,
                        " legalRecompute: ", lrc, " shouldRecompute: ", src,
                        " tryLegalRecomputeCheck: ", tryLegalRecomputeCheck);
            return result;
          }
        }
    noSpeedCache:;
    }

  if (scopeMap.find(inst) == scopeMap.end()) {
    EmitWarning("Uncacheable", *inst, "Caching instruction ", *inst,
                " legalRecompute: ", lrc, " shouldRecompute: ", src,
                " tryLegalRecomputeCheck: ", tryLegalRecomputeCheck);
  }

  BasicBlock *scopeI = inst->getParent();
  if (auto origInst = isOriginal(inst)) {
    auto found = rematerializableAllocations.find(origInst);
    if (found != rematerializableAllocations.end())
      if (found->second.LI && found->second.LI->contains(origInst)) {
        // If not caching whole allocation and rematerializing the allocation
        // within the loop, force an entry-level scope so there is no need
        // to cache.
        if (!needsCacheWholeAllocation(origInst))
          scopeI = &newFunc->getEntryBlock();
      }
  } else {
    for (auto pair : backwardsOnlyShadows) {
      if (auto pinst = dyn_cast<Instruction>(pair.first))
        if (!pair.second.primalInitialize && pair.second.LI &&
            pair.second.LI->contains(pinst->getParent())) {
          auto found = invertedPointers.find(pair.first);
          if (found != invertedPointers.end() && found->second == inst) {
            scopeI = &newFunc->getEntryBlock();

            // Prevent the phi node from being stored into the cache by creating
            // it before the ensureLookupCached.
            if (scopeMap.find(inst) == scopeMap.end()) {
              LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0,
                                scopeI);

              AllocaInst *cache = createCacheForScope(
                  lctx, inst->getType(), inst->getName(), /*shouldFree*/ true);
              assert(cache);
              Value *inst_tmp = inst;
              insert_or_assign(scopeMap, inst_tmp,
                               std::pair<AssertingVH<AllocaInst>, LimitContext>(
                                   cache, lctx));
            }
            break;
          }
        }
    }
  }

  ensureLookupCached(inst, /*shouldFree*/ true, scopeI,
                     inst->getMetadata(LLVMContext::MD_tbaa));
  bool isi1 = inst->getType()->isIntegerTy() &&
              cast<IntegerType>(inst->getType())->getBitWidth() == 1;
  assert(!isOriginalBlock(*BuilderM.GetInsertBlock()));
  auto found = findInMap(scopeMap, (Value *)inst);
  Value *result =
      lookupValueFromCache(inst->getType(), /*isForwardPass*/ false, BuilderM,
                           found->second, found->first, isi1, available);
  if (auto LI2 = dyn_cast<LoadInst>(result))
    if (auto LI1 = dyn_cast<LoadInst>(inst)) {
      llvm::SmallVector<unsigned int, 9> ToCopy2(MD_ToCopy);
      ToCopy2.push_back(LLVMContext::MD_noalias);
      ToCopy2.push_back(LLVMContext::MD_alias_scope);
      LI2->copyMetadata(*LI1, ToCopy2);
    }
  if (result->getType() != inst->getType()) {
    llvm::errs() << "newFunc: " << *newFunc << "\n";
    llvm::errs() << "result: " << *result << "\n";
    llvm::errs() << "inst: " << *inst << "\n";
    llvm::errs() << "val: " << *val << "\n";
  }
  assert(result->getType() == inst->getType());
  lookup_cache[BuilderM.GetInsertBlock()][val] = result;
  assert(result);
  if (result->getType() != val->getType()) {
    result = BuilderM.CreateBitCast(result, val->getType());
  }
  assert(result->getType() == val->getType());
  assert(result->getType());
  return result;
}

BasicBlock *GradientUtils::originalForReverseBlock(BasicBlock &BB2) const {
  auto found = reverseBlockToPrimal.find(&BB2);
  if (found == reverseBlockToPrimal.end()) {
    errs() << "newFunc: " << *newFunc << "\n";
    errs() << BB2 << "\n";
  }
  assert(found != reverseBlockToPrimal.end());
  return found->second;
}

//! Given a map of edges we could have taken to desired target, compute a value
//! that determines which target should be branched to
//  This function attempts to determine an equivalent condition from earlier in
//  the code and use that if possible, falling back to creating a phi node of
//  which edge was taken if necessary This function can be used in two ways:
//   * If replacePHIs is null (usual case), this function does the branch
//   * If replacePHIs isn't null, do not perform the branch and instead replace
//   the PHI's with the derived condition as to whether we should branch to a
//   particular target
void GradientUtils::branchToCorrespondingTarget(
    BasicBlock *ctx, IRBuilder<> &BuilderM,
    const std::map<BasicBlock *,
                   std::vector<std::pair</*pred*/ BasicBlock *,
                                         /*successor*/ BasicBlock *>>>
        &targetToPreds,
    const std::map<BasicBlock *, PHINode *> *replacePHIs) {
  assert(targetToPreds.size() > 0);
  if (replacePHIs) {
    if (replacePHIs->size() == 0)
      return;

#ifndef NDEBUG
    for (auto x : *replacePHIs) {
      assert(targetToPreds.find(x.first) != targetToPreds.end());
    }
#endif
  }

  if (targetToPreds.size() == 1) {
    if (replacePHIs == nullptr) {
      if (!(BuilderM.GetInsertBlock()->size() == 0 ||
            !isa<BranchInst>(BuilderM.GetInsertBlock()->back()))) {
        llvm::errs() << *oldFunc << "\n";
        llvm::errs() << *newFunc << "\n";
        llvm::errs() << *BuilderM.GetInsertBlock() << "\n";
      }
      assert(BuilderM.GetInsertBlock()->size() == 0 ||
             !isa<BranchInst>(BuilderM.GetInsertBlock()->back()));
      BuilderM.CreateBr(targetToPreds.begin()->first);
    } else {
      for (auto pair : *replacePHIs) {
        pair.second->replaceAllUsesWith(
            ConstantInt::getTrue(pair.second->getContext()));
        pair.second->eraseFromParent();
      }
    }
    return;
  }

  // Map of function edges to list of targets this can branch to we have
  std::map<std::pair</*pred*/ BasicBlock *, /*successor*/ BasicBlock *>,
           std::set<BasicBlock *>>
      done;
  {
    std::deque<
        std::tuple<std::pair</*pred*/ BasicBlock *, /*successor*/ BasicBlock *>,
                   BasicBlock *>>
        Q; // newblock, target

    for (auto pair : targetToPreds) {
      for (auto pred_edge : pair.second) {
        Q.push_back(std::make_pair(pred_edge, pair.first));
      }
    }

    for (std::tuple<
             std::pair</*pred*/ BasicBlock *, /*successor*/ BasicBlock *>,
             BasicBlock *>
             trace;
         Q.size() > 0;) {
      trace = Q.front();
      Q.pop_front();
      auto edge = std::get<0>(trace);
      auto block = edge.first;
      auto target = std::get<1>(trace);

      if (done[edge].count(target))
        continue;
      done[edge].insert(target);

      // If this block dominates the context, don't go back up as any
      // predecessors won't contain the conditions.
      if (DT.dominates(block, ctx))
        continue;

      Loop *blockLoop = LI.getLoopFor(block);

      for (BasicBlock *Pred : predecessors(block)) {
        // Don't go up the backedge as we can use the last value if desired via
        // lcssa
        if (blockLoop && blockLoop->getHeader() == block &&
            blockLoop == LI.getLoopFor(Pred))
          continue;

        Q.push_back(
            std::tuple<std::pair<BasicBlock *, BasicBlock *>, BasicBlock *>(
                std::make_pair(Pred, block), target));
      }
    }
  }

  IntegerType *T;
  if (targetToPreds.size() == 2)
    T = Type::getInt1Ty(BuilderM.getContext());
  else if (targetToPreds.size() < 256)
    T = Type::getInt8Ty(BuilderM.getContext());
  else
    T = Type::getInt32Ty(BuilderM.getContext());

  Instruction *equivalentTerminator = nullptr;

  std::set<BasicBlock *> blocks;

  // llvm::errs() << "\n\n<DONE = " << ctx->getName() << ">\n";
  for (auto pair : done) {
    const auto &edge = pair.first;
    blocks.insert(edge.first);
    // llvm::errs() << " edge  (" << edge.first->getName() << ", "
    //             << edge.second->getName() << ") : [";
    // for (auto s : pair.second)
    //  llvm::errs() << s->getName() << ",";
    // llvm::errs() << "]\n";
  }
  // llvm::errs() << "</DONE>\n";

  if (targetToPreds.size() == 3) {
    // Try `block` as a potential first split point.
    for (auto block : blocks) {
      {
        // The original split block must not have a parent with an edge
        // to a block other than to itself, which can reach any targets.
        if (!DT.dominates(block, ctx))
          continue;

        // For all successors and thus edges (block, succ):
        // 1) Ensure that no successors have overlapping potential
        // destinations (a list of destinations previously seen is in
        // foundtargets).
        // 2) The block branches to all 3 destinations (foundTargets==3)
        std::set<BasicBlock *> foundtargets;
        // 3) The unique target split off from the others is stored in
        //   uniqueTarget.
        std::set<BasicBlock *> uniqueTargets;
        for (BasicBlock *succ : successors(block)) {
          auto edge = std::make_pair(block, succ);
          for (BasicBlock *target : done[edge]) {
            if (foundtargets.find(target) != foundtargets.end()) {
              goto rnextpair;
            }
            foundtargets.insert(target);
            if (done[edge].size() == 1)
              uniqueTargets.insert(target);
          }
        }
        if (foundtargets.size() != 3)
          goto rnextpair;
        if (uniqueTargets.size() != 1)
          goto rnextpair;

        // Only handle cases where the split was due to a conditional
        // branch. This branch, `bi`, splits off uniqueTargets[0] from
        // the remainder of foundTargets.
        auto bi1 = dyn_cast<BranchInst>(block->getTerminator());
        if (!bi1)
          goto rnextpair;

        {
          // Find a second block `subblock` which splits the two merged
          // targets from each other.
          BasicBlock *subblock = nullptr;
          for (auto block2 : blocks) {
            {
              // The second split block must not have a parent with an edge
              // to a block other than to itself, which can reach any of its two
              // targets.
              // TODO verify this
              for (auto P : predecessors(block2)) {
                for (auto S : successors(P)) {
                  if (S == block2)
                    continue;
                  auto edge = std::make_pair(P, S);
                  if (done.find(edge) != done.end()) {
                    for (auto target : done[edge]) {
                      if (foundtargets.find(target) != foundtargets.end() &&
                          uniqueTargets.find(target) == uniqueTargets.end()) {
                        goto nextblock;
                      }
                    }
                  }
                }
              }

              // Again, a successful split must have unique targets.
              std::set<BasicBlock *> seen2;
              for (BasicBlock *succ : successors(block2)) {
                auto edge = std::make_pair(block2, succ);
                // Since there are only two targets, a successful split
                // condition has only 1 target per successor of block2.
                if (done[edge].size() != 1) {
                  goto nextblock;
                }
                for (BasicBlock *target : done[edge]) {
                  // block2 has non-unique targets.
                  if (seen2.find(target) != seen2.end()) {
                    goto nextblock;
                  }
                  seen2.insert(target);
                  // block2 has a target which is not part of the two needing
                  // to be split. The two needing to be split is equal to
                  //    foundtargets-uniqueTargets.
                  if (foundtargets.find(target) == foundtargets.end()) {
                    goto nextblock;
                  }
                  if (uniqueTargets.find(target) != uniqueTargets.end()) {
                    goto nextblock;
                  }
                }
              }
              // If we didn't find two valid successors, continue.
              if (seen2.size() != 2) {
                // llvm::errs() << " -- failed from not 2 seen\n";
                goto nextblock;
              }
              subblock = block2;
              break;
            }
          nextblock:;
          }

          // If no split block was found, try again.
          if (subblock == nullptr)
            goto rnextpair;

          // This branch, `bi2`, splits off the two blocks in
          // (foundTargets-uniqueTargets) from each other.
          auto bi2 = dyn_cast<BranchInst>(subblock->getTerminator());
          if (!bi2)
            goto rnextpair;

          // Condition cond1 splits off uniqueTargets[0] from
          // the remainder of foundTargets.
          auto cond1 = lookupM(bi1->getCondition(), BuilderM);

          // Condition cond2 splits off the two blocks in
          // (foundTargets-uniqueTargets) from each other.
          auto cond2 = lookupM(bi2->getCondition(), BuilderM);

          if (replacePHIs == nullptr) {
            BasicBlock *staging =
                BasicBlock::Create(oldFunc->getContext(), "staging", newFunc);
            auto stagingIfNeeded = [&](BasicBlock *B) {
              auto edge = std::make_pair(block, B);
              if (done[edge].size() == 1) {
                return *done[edge].begin();
              } else {
                assert(done[edge].size() == 2);
                return staging;
              }
            };
            BuilderM.CreateCondBr(cond1, stagingIfNeeded(bi1->getSuccessor(0)),
                                  stagingIfNeeded(bi1->getSuccessor(1)));
            BuilderM.SetInsertPoint(staging);
            BuilderM.CreateCondBr(
                cond2,
                *done[std::make_pair(subblock, bi2->getSuccessor(0))].begin(),
                *done[std::make_pair(subblock, bi2->getSuccessor(1))].begin());
          } else {
            Value *otherBranch = nullptr;
            for (unsigned i = 0; i < 2; ++i) {
              Value *val = cond1;
              if (i == 1)
                val = BuilderM.CreateNot(val, "anot1_");
              auto edge = std::make_pair(block, bi1->getSuccessor(i));
              if (done[edge].size() == 1) {
                auto found = replacePHIs->find(*done[edge].begin());
                if (found == replacePHIs->end())
                  continue;
                if (&*BuilderM.GetInsertPoint() == found->second) {
                  if (found->second->getNextNode())
                    BuilderM.SetInsertPoint(found->second->getNextNode());
                  else
                    BuilderM.SetInsertPoint(found->second->getParent());
                }
                found->second->replaceAllUsesWith(val);
                found->second->eraseFromParent();
              } else {
                otherBranch = val;
              }
            }

            for (unsigned i = 0; i < 2; ++i) {
              auto edge = std::make_pair(subblock, bi2->getSuccessor(i));
              auto found = replacePHIs->find(*done[edge].begin());
              if (found == replacePHIs->end())
                continue;

              Value *val = cond2;
              if (i == 1)
                val = BuilderM.CreateNot(val, "bnot1_");
              val = BuilderM.CreateAnd(val, otherBranch, "andVal" + Twine(i));
              if (&*BuilderM.GetInsertPoint() == found->second) {
                if (found->second->getNextNode())
                  BuilderM.SetInsertPoint(found->second->getNextNode());
                else
                  BuilderM.SetInsertPoint(found->second->getParent());
              }
              found->second->replaceAllUsesWith(val);
              found->second->eraseFromParent();
            }
          }

          return;
        }
      }
    rnextpair:;
    }
  }

  BasicBlock *forwardBlock = BuilderM.GetInsertBlock();

  if (!isOriginalBlock(*forwardBlock)) {
    forwardBlock = originalForReverseBlock(*forwardBlock);
  }

  for (auto block : blocks) {
    {
      // The original split block must not have a parent with an edge
      // to a block other than to itself, which can reach any targets.
      if (!DT.dominates(block, ctx))
        for (auto P : predecessors(block)) {
          for (auto S : successors(P)) {
            if (S == block)
              continue;
            auto edge = std::make_pair(P, S);
            if (done.find(edge) != done.end() && done[edge].size())
              goto nextpair;
          }
        }

      std::set<BasicBlock *> foundtargets;
      for (BasicBlock *succ : successors(block)) {
        auto edge = std::make_pair(block, succ);
        if (done[edge].size() != 1) {
          goto nextpair;
        }
        BasicBlock *target = *done[edge].begin();
        if (foundtargets.find(target) != foundtargets.end()) {
          goto nextpair;
        }
        foundtargets.insert(target);
      }
      if (foundtargets.size() != targetToPreds.size()) {
        goto nextpair;
      }

      if (forwardBlock == block || DT.dominates(block, forwardBlock)) {
        equivalentTerminator = block->getTerminator();
        goto fast;
      }
    }
  nextpair:;
  }
  goto nofast;

fast:;
  assert(equivalentTerminator);

  if (auto branch = dyn_cast<BranchInst>(equivalentTerminator)) {
    BasicBlock *block = equivalentTerminator->getParent();
    assert(branch->getCondition());

    assert(branch->getCondition()->getType() == T);

    if (replacePHIs == nullptr) {
      if (!(BuilderM.GetInsertBlock()->size() == 0 ||
            !isa<BranchInst>(BuilderM.GetInsertBlock()->back()))) {
        llvm::errs() << "newFunc : " << *newFunc << "\n";
        llvm::errs() << "blk : " << *BuilderM.GetInsertBlock() << "\n";
      }
      assert(BuilderM.GetInsertBlock()->size() == 0 ||
             !isa<BranchInst>(BuilderM.GetInsertBlock()->back()));
      BuilderM.CreateCondBr(
          lookupM(branch->getCondition(), BuilderM),
          *done[std::make_pair(block, branch->getSuccessor(0))].begin(),
          *done[std::make_pair(block, branch->getSuccessor(1))].begin());
    } else {
      for (auto pair : *replacePHIs) {
        Value *phi = lookupM(branch->getCondition(), BuilderM);
        Value *val = nullptr;
        if (pair.first ==
            *done[std::make_pair(block, branch->getSuccessor(0))].begin()) {
          val = phi;
        } else if (pair.first ==
                   *done[std::make_pair(block, branch->getSuccessor(1))]
                        .begin()) {
          val = BuilderM.CreateNot(phi);
        } else {
          llvm::errs() << *pair.first->getParent() << "\n";
          llvm::errs() << *pair.first << "\n";
          llvm::errs() << *branch << "\n";
          llvm_unreachable("unknown successor for replacephi");
        }
        if (&*BuilderM.GetInsertPoint() == pair.second) {
          if (pair.second->getNextNode())
            BuilderM.SetInsertPoint(pair.second->getNextNode());
          else
            BuilderM.SetInsertPoint(pair.second->getParent());
        }
        pair.second->replaceAllUsesWith(val);
        pair.second->eraseFromParent();
      }
    }
  } else if (auto si = dyn_cast<SwitchInst>(equivalentTerminator)) {
    BasicBlock *block = equivalentTerminator->getParent();

    IRBuilder<> pbuilder(equivalentTerminator);
    pbuilder.setFastMathFlags(getFast());

    if (replacePHIs == nullptr) {
      SwitchInst *swtch = BuilderM.CreateSwitch(
          lookupM(si->getCondition(), BuilderM),
          *done[std::make_pair(block, si->getDefaultDest())].begin());
      for (auto switchcase : si->cases()) {
        swtch->addCase(
            switchcase.getCaseValue(),
            *done[std::make_pair(block, switchcase.getCaseSuccessor())]
                 .begin());
      }
    } else {
      for (auto pair : *replacePHIs) {
        Value *cas = nullptr;
        for (auto c : si->cases()) {
          if (pair.first ==
              *done[std::make_pair(block, c.getCaseSuccessor())].begin()) {
            cas = c.getCaseValue();
            break;
          }
        }
        if (cas == nullptr) {
          assert(pair.first ==
                 *done[std::make_pair(block, si->getDefaultDest())].begin());
        }
        Value *val = nullptr;
        Value *phi = lookupM(si->getCondition(), BuilderM);

        if (cas) {
          val = BuilderM.CreateICmpEQ(cas, phi);
        } else {
          // default case
          val = ConstantInt::getFalse(pair.second->getContext());
          for (auto switchcase : si->cases()) {
            val = BuilderM.CreateOr(
                val, BuilderM.CreateICmpEQ(switchcase.getCaseValue(), phi));
          }
          val = BuilderM.CreateNot(val);
        }
        if (&*BuilderM.GetInsertPoint() == pair.second) {
          if (pair.second->getNextNode())
            BuilderM.SetInsertPoint(pair.second->getNextNode());
          else
            BuilderM.SetInsertPoint(pair.second->getParent());
        }
        pair.second->replaceAllUsesWith(val);
        pair.second->eraseFromParent();
      }
    }
  } else {
    llvm::errs() << "unknown equivalent terminator\n";
    llvm::errs() << *equivalentTerminator << "\n";
    llvm_unreachable("unknown equivalent terminator");
  }
  return;

nofast:;

  // if freeing reverseblocks must exist
  assert(reverseBlocks.size());
  LimitContext lctx(/*ReverseLimit*/ reverseBlocks.size() > 0, ctx);
  AllocaInst *cache = createCacheForScope(lctx, T, "", /*shouldFree*/ true);
  SmallVector<BasicBlock *, 4> targets;
  {
    size_t idx = 0;
    std::map<BasicBlock * /*storingblock*/,
             std::map<ConstantInt * /*target*/,
                      std::vector<BasicBlock *> /*predecessors*/>>
        storing;
    for (const auto &pair : targetToPreds) {
      for (auto pred : pair.second) {
        storing[pred.first][ConstantInt::get(T, idx)].push_back(pred.second);
      }
      targets.push_back(pair.first);
      ++idx;
    }
    assert(targets.size() > 0);

    for (const auto &pair : storing) {
      IRBuilder<> pbuilder(pair.first);

      if (hasTerminator(pair.first))
        pbuilder.SetInsertPoint(pair.first->getTerminator());

      pbuilder.setFastMathFlags(getFast());

      Value *tostore = ConstantInt::get(T, 0);

      if (pair.second.size() == 1) {
        tostore = pair.second.begin()->first;
      } else {
        assert(0 && "multi exit edges not supported");
        exit(1);
        // for(auto targpair : pair.second) {
        //     tostore = pbuilder.CreateOr(tostore, pred);
        //}
      }
      storeInstructionInCache(lctx, pbuilder, tostore, cache);
    }
  }

  bool isi1 = T->isIntegerTy() && cast<IntegerType>(T)->getBitWidth() == 1;
  Value *which = lookupValueFromCache(
      T,
      /*forwardPass*/ isOriginalBlock(*BuilderM.GetInsertBlock()), BuilderM,
      LimitContext(/*reversePass*/ reverseBlocks.size() > 0, ctx), cache, isi1,
      /*available*/ ValueToValueMapTy());
  assert(which);
  assert(which->getType() == T);

  if (replacePHIs == nullptr) {
    if (targetToPreds.size() == 2) {
      assert(BuilderM.GetInsertBlock()->size() == 0 ||
             !isa<BranchInst>(BuilderM.GetInsertBlock()->back()));
      BuilderM.CreateCondBr(which, /*true*/ targets[1], /*false*/ targets[0]);
    } else {
      assert(targets.size() > 0);
      auto swit =
          BuilderM.CreateSwitch(which, targets.back(), targets.size() - 1);
      for (unsigned i = 0; i < targets.size() - 1; ++i) {
        swit->addCase(ConstantInt::get(T, i), targets[i]);
      }
    }
  } else {
    for (unsigned i = 0; i < targets.size(); ++i) {
      auto found = replacePHIs->find(targets[i]);
      if (found == replacePHIs->end())
        continue;

      Value *val = nullptr;
      if (targets.size() == 2 && i == 0) {
        val = BuilderM.CreateNot(which);
      } else if (targets.size() == 2 && i == 1) {
        val = which;
      } else {
        val = BuilderM.CreateICmpEQ(ConstantInt::get(T, i), which);
      }
      if (&*BuilderM.GetInsertPoint() == found->second) {
        if (found->second->getNextNode())
          BuilderM.SetInsertPoint(found->second->getNextNode());
        else
          BuilderM.SetInsertPoint(found->second->getParent());
      }
      found->second->replaceAllUsesWith(val);
      found->second->eraseFromParent();
    }
  }
  return;
}

void GradientUtils::computeMinCache() {
  if (EnzymeMinCutCache) {
    SetVector<Value *> Recomputes;

    std::map<UsageKey, bool> FullSeen;
    std::map<UsageKey, bool> OneLevelSeen;

    ValueToValueMapTy Available;

    std::map<Loop *, std::set<Instruction *>> LoopAvail;

    for (BasicBlock &BB : *oldFunc) {
      if (notForAnalysis.count(&BB))
        continue;
      auto L = OrigLI->getLoopFor(&BB);

      auto invariant = [&](Value *V) {
        if (isa<Constant>(V))
          return true;
        if (isa<Argument>(V))
          return true;
        if (auto I = dyn_cast<Instruction>(V)) {
          if (!L->contains(OrigLI->getLoopFor(I->getParent())))
            return true;
        }
        return false;
      };
      for (Instruction &I : BB) {
        if (auto PN = dyn_cast<PHINode>(&I)) {
          if (!OrigLI->isLoopHeader(&BB))
            continue;
          if (PN->getType()->isIntegerTy()) {
            bool legal = true;
            SmallPtrSet<Instruction *, 4> Increment;
            for (auto B : PN->blocks()) {
              if (OrigLI->getLoopFor(B) == L) {
                if (auto BO = dyn_cast<BinaryOperator>(
                        PN->getIncomingValueForBlock(B))) {
                  if (BO->getOpcode() == BinaryOperator::Add) {
                    if ((BO->getOperand(0) == PN &&
                         invariant(BO->getOperand(1))) ||
                        (BO->getOperand(1) == PN &&
                         invariant(BO->getOperand(0)))) {
                      Increment.insert(BO);
                    } else {
                      legal = false;
                    }
                  } else if (BO->getOpcode() == BinaryOperator::Sub) {
                    if (BO->getOperand(0) == PN &&
                        invariant(BO->getOperand(1))) {
                      Increment.insert(BO);
                    } else {
                      legal = false;
                    }
                  } else {
                    legal = false;
                  }
                } else {
                  legal = false;
                }
              }
            }
            if (legal) {
              LoopAvail[L].insert(PN);
              for (auto I : Increment)
                LoopAvail[L].insert(I);
            }
          }
        } else if (auto CI = dyn_cast<CallInst>(&I)) {
          StringRef funcName = getFuncNameFromCall(CI);
          if (isAllocationFunction(funcName, TLI)) {
            bool legal = true;
            auto found = rematerializableAllocations.find(CI);
            if (found != rematerializableAllocations.end()) {
              if (found->second.nonRepeatableWritingCall)
                legal = false;
            }
            if (legal)
              Available[CI] = CI;
          }
        }
      }
    }

    SmallPtrSet<Instruction *, 3> NewLoopBoundReq;
    {
      std::deque<Instruction *> LoopBoundRequirements;

      for (auto &context : loopContexts) {
        for (auto val : {context.second.maxLimit, context.second.trueLimit}) {
          if (val)
            if (auto inst = dyn_cast<Instruction>(&*val)) {
              LoopBoundRequirements.push_back(inst);
            }
        }
      }
      SmallPtrSet<Instruction *, 3> Seen;
      while (LoopBoundRequirements.size()) {
        Instruction *val = LoopBoundRequirements.front();
        LoopBoundRequirements.pop_front();
        if (NewLoopBoundReq.count(val))
          continue;
        if (Seen.count(val))
          continue;
        Seen.insert(val);
        if (auto orig = isOriginal(val)) {
          NewLoopBoundReq.insert(orig);
        } else {
          for (auto &op : val->operands()) {
            if (auto inst = dyn_cast<Instruction>(op)) {
              LoopBoundRequirements.push_back(inst);
            }
          }
        }
      }
      for (auto inst : NewLoopBoundReq) {
        OneLevelSeen[UsageKey(inst, QueryType::Primal)] = true;
        FullSeen[UsageKey(inst, QueryType::Primal)] = true;
      }
    }

    auto minCutMode = (mode == DerivativeMode::ReverseModePrimal ||
                       mode == DerivativeMode::ReverseModeCombined)
                          ? DerivativeMode::ReverseModeGradient
                          : mode;

    for (BasicBlock &BB : *oldFunc) {
      if (notForAnalysis.count(&BB))
        continue;
      ValueToValueMapTy Available2;
      for (auto a : Available)
        Available2[a.first] = a.second;
      for (Loop *L = OrigLI->getLoopFor(&BB); L != nullptr;
           L = L->getParentLoop()) {
        for (auto v : LoopAvail[L]) {
          Available2[v] = v;
        }
      }
      for (Instruction &I : BB) {
        if (!legalRecompute(&I, Available2, nullptr)) {
          if (DifferentialUseAnalysis::is_value_needed_in_reverse<
                  QueryType::Primal>(this, &I, minCutMode, FullSeen,
                                     notForAnalysis)) {
            bool oneneed = DifferentialUseAnalysis::is_value_needed_in_reverse<
                QueryType::Primal,
                /*OneLevel*/ true>(this, &I, minCutMode, OneLevelSeen,
                                   notForAnalysis);

            bool shadowOneNeed = false;
            // even if the primal is not needed directly by its users, if the
            // primal is constant and used to create a shadow insertvalue which
            // is used, we need to save the shadow since shadow cache and primal
            // cache are the same, we force a save of cache here.
            // TODO(wsmoses): extend this to separate caching decisions for
            // primal and shadow
            if (!oneneed && isConstantValue(&I) && !TR.allFloat(&I)) {
              SmallVector<Instruction *, 1> todo;
              todo.push_back(&I);
              while (todo.size()) {
                auto cur = todo.pop_back_val();
                for (auto u : cur->users()) {
                  if (isa<InsertValueInst>(u) || isa<InsertElementInst>(u) ||
                      isa<ExtractValueInst>(u) || isa<ExtractElementInst>(u)) {
                    auto I2 = cast<Instruction>(u);
                    if (!isConstantValue(I2)) {
                      if (DifferentialUseAnalysis::is_value_needed_in_reverse<
                              QueryType::Shadow>(this, I2, minCutMode, FullSeen,
                                                 notForAnalysis)) {
                        shadowOneNeed = true;
                        goto endOneNeed;
                      }
                    } else {
                      todo.push_back(I2);
                    }
                  }
                }
              }
            endOneNeed:;
            }

            if (oneneed || shadowOneNeed) {
              knownRecomputeHeuristic[&I] = false;

              CountTrackedPointers T(I.getType());
              assert(!T.derived);
            } else {
              Recomputes.insert(&I);
            }
          }
        }
      }
    }

    SetVector<Value *> Intermediates;
    SetVector<Value *> Required;
    std::deque<Value *> todo(Recomputes.begin(), Recomputes.end());

    while (todo.size()) {
      Value *V = todo.front();
      todo.pop_front();
      if (Intermediates.count(V))
        continue;
      bool multiLevel = DifferentialUseAnalysis::is_value_needed_in_reverse<
          QueryType::Primal>(this, V, minCutMode, FullSeen, notForAnalysis);
      if (!multiLevel) {
        continue;
      }
      if (!Recomputes.count(V)) {
        ValueToValueMapTy Available2;
        for (auto a : Available)
          Available2[a.first] = a.second;
        for (Loop *L = OrigLI->getLoopFor(cast<Instruction>(V)->getParent());
             L != nullptr; L = L->getParentLoop()) {
          for (auto v : LoopAvail[L]) {
            Available2[v] = v;
          }
        }
        if (!legalRecompute(V, Available2, nullptr)) {
          // if not legal to recompute, we would've already explicitly marked
          // this for caching if it was needed in reverse pass
          continue;
        }
      }
      Intermediates.insert(V);
      bool singleLevel = DifferentialUseAnalysis::is_value_needed_in_reverse<
          QueryType::Primal, /*OneLevel*/ true>(this, V, minCutMode,
                                                OneLevelSeen, notForAnalysis);
      if (singleLevel) {
        Required.insert(V);
      } else {
        DifferentialUseAnalysis::forEachDifferentialUser(
            [&](Value *V2) { todo.push_back(V2); }, this, V);
      }
    }

    SetVector<Value *> MinReq;
    DifferentialUseAnalysis::minCut(oldFunc->getParent()->getDataLayout(),
                                    *OrigLI, Recomputes, Intermediates,
                                    Required, MinReq, this, TLI);
    SmallPtrSet<Value *, 5> NeedGraph;

    // The start value of a loopy reduction is needed in the reverse pass, but
    // is only discovered now that the min cut is known. Since the cut has
    // already been computed, such a value can no longer be selected for
    // caching -- everything placed into Intermediates from here on will be
    // marked as recomputed below. Only add it if that is legal; if it is not,
    // it has already been given an explicit cache decision by the scan above,
    // which must not be overwritten. Either way push it onto the worklist so
    // the need graph stays complete.
    auto pushLoopyPHIPreheader = [&](Value *V) {
      SmallVector<Value *, 2> preheaderVals;
      DifferentialUseAnalysis::pushLoopyPHIPreheader(this, V, preheaderVals);
      for (Value *PV : preheaderVals) {
        ValueToValueMapTy Available2;
        for (auto a : Available)
          Available2[a.first] = a.second;
        for (Loop *L = OrigLI->getLoopFor(cast<Instruction>(PV)->getParent());
             L != nullptr; L = L->getParentLoop()) {
          for (auto v : LoopAvail[L]) {
            Available2[v] = v;
          }
        }
        if (legalRecompute(PV, Available2, nullptr))
          Intermediates.insert(PV);
        todo.push_back(PV);
      }
    };

    for (Value *V : MinReq) {
      NeedGraph.insert(V);
      pushLoopyPHIPreheader(V);
    }
    for (Value *V : Required) {
      todo.push_back(V);
    }
    while (todo.size()) {
      Value *V = todo.front();
      todo.pop_front();
      if (NeedGraph.count(V))
        continue;
      NeedGraph.insert(V);
      pushLoopyPHIPreheader(V);
      auto I = dyn_cast<Instruction>(V);
      if (!I)
        continue;
      for (auto &V2 : I->operands()) {
        if (Intermediates.count(V2)) {
          todo.push_back(V2);
        }
      }
      auto found = rematerializableAllocations.find(I);
      if (found != rematerializableAllocations.end()) {
        for (auto store : found->second.stores) {
          for (auto &operand : store->operands()) {
            if (Intermediates.count(operand)) {
              todo.push_back(operand);
            }
          }
        }
      }
    }

    for (auto V : Intermediates) {
      knownRecomputeHeuristic[V] = !MinReq.count(V);
      if (!MinReq.count(V) && NeedGraph.count(V)) {
        if (auto CI = dyn_cast<CallInst>(V))
          if (getFuncNameFromCall(CI) == "julia.call")
            assert(0);

        ValueToValueMapTy Available2;
        for (auto a : Available)
          Available2[a.first] = a.second;
        for (Loop *L = OrigLI->getLoopFor(cast<Instruction>(V)->getParent());
             L != nullptr; L = L->getParentLoop()) {
          for (auto v : LoopAvail[L]) {
            Available2[v] = v;
          }
        }
        assert(legalRecompute(V, Available2, nullptr));
      }
      if (!NeedGraph.count(V)) {
        assert(!MinReq.count(V));
        unnecessaryIntermediates.insert(cast<Instruction>(V));
      }

      if (NeedGraph.count(V) && MinReq.count(V)) {
        CountTrackedPointers T(V->getType());
        if (T.derived) {
          std::string str;
          raw_string_ostream ss(str);
          ss << "Illegal cached pointer: " << *V << "\n";
          if (CustomErrorHandler) {
            CustomErrorHandler(str.c_str(), wrap((Value *)V),
                               ErrorType::InternalError, nullptr, nullptr,
                               nullptr);
          } else {
            EmitFailure(
                "CachedPointerError", cast<Instruction>(V)->getDebugLoc(),
                cast<Instruction>(V)->getParent()->getParent(), ss.str());
          }
        }
      }
    }
  }
  if (rematerializableAllocations.size()) {

    // We iterate through the instructions here to ensure a consistent order for
    // the analysis results.
    for (auto &BB : *oldFunc) {
      for (auto &I : BB) {
        auto found = rematerializableAllocations.find(&I);
        if (found == rematerializableAllocations.end())
          continue;
        std::map<UsageKey, bool> Seen = populateSeenFromKnownRecompute();
        bool primalNeededInReverse =
            DifferentialUseAnalysis::is_value_needed_in_reverse<
                QueryType::Primal>(this, &I, mode, Seen, notForAnalysis);

        {
          auto found = knownRecomputeHeuristic.find(&I);
          if (found != knownRecomputeHeuristic.end() && !found->second) {
            primalNeededInReverse = true;
          }
        }

        if (primalNeededInReverse && !needsCacheWholeAllocation(&I)) {
          allocationsToBeRematerialized.insert(&I);
        }
      }
    }
  }
}

bool GradientUtils::isOriginalBlock(const BasicBlock &BB) const {
  for (auto A : originalBlocks) {
    if (A == &BB)
      return true;
  }
  return false;
}

void GradientUtils::eraseFictiousPHIs() {
  {
    for (auto P : rematerializedPrimalOrShadowAllocations) {
      Value *replacement =
          getUndefinedValueForType(*oldFunc->getParent(), P->getType());
      P->replaceAllUsesWith(replacement);
      erase(P);
    }
  }
  SmallVector<std::pair<PHINode *, Value *>, 4> phis;
  for (auto pair : fictiousPHIs)
    phis.emplace_back(pair.first, pair.second);
  fictiousPHIs.clear();

  for (auto pair : phis) {
    auto pp = pair.first;
    if (pp->getNumUses() != 0) {
      bool skip = false;
      assert(isa<Instruction>(pair.second));
      auto *I = dyn_cast<Instruction>(pair.second);
      if (!OrigDT->isReachableFromEntry(I->getParent())) {
        skip = true;
      }

      if (!skip) {
        std::string str;
        raw_string_ostream ss(str);
        ss << "Illegal replace ficticious phi for: " << *pp << " of "
           << *pair.second << "\n";
        for (auto U : pp->users()) {
          ss << "  user: " << *U << "\n";
        }
        if (CustomErrorHandler) {
          CustomErrorHandler(str.c_str(), wrap(pp), ErrorType::InternalError,
                             nullptr, nullptr, nullptr);
        } else {
          ss << " newFunc:\n" << *newFunc << "\n";
          EmitFailure("IllegalReplacePHI", I->getDebugLoc(), I, str);
        }
      }
      Value *replacement =
          getUndefinedValueForType(*oldFunc->getParent(), pp->getType());
      pp->replaceAllUsesWith(replacement);
    }
    erase(pp);
  }
}

void GradientUtils::forceActiveDetection() {

  TimeTraceScope timeScope("Activity Analysis", oldFunc->getName());

  for (auto &Arg : oldFunc->args()) {
    ATA->isConstantValue(TR, &Arg);
  }

  for (BasicBlock &BB : *oldFunc) {
    for (Instruction &I : BB) {
      bool const_inst = ATA->isConstantInstruction(TR, &I);
      bool const_value = ATA->isConstantValue(TR, &I);
      if (EnzymePrintActivity)
        llvm::errs() << I << " cv=" << const_value << " ci=" << const_inst
                     << "\n";
    }
  }
}

bool GradientUtils::isConstantValue(Value *val) const {
  if (auto inst = dyn_cast<Instruction>(val)) {
    (void)inst;
    assert(inst->getParent()->getParent() == oldFunc);
    return ATA->isConstantValue(TR, val);
  }

  if (auto arg = dyn_cast<Argument>(val)) {
    (void)arg;
    assert(arg->getParent() == oldFunc);
    return ATA->isConstantValue(TR, val);
  }

  //! Functions must be false so we can replace function with augmentation,
  //! fallback to analysis
  if (isa<Function>(val) || isa<InlineAsm>(val) || isa<Constant>(val) ||
      isa<UndefValue>(val) || isa<MetadataAsValue>(val)) {
    // llvm::errs() << "calling icv on: " << *val << "\n";
    return ATA->isConstantValue(TR, val);
  }

  if (auto gv = dyn_cast<GlobalVariable>(val)) {
    if (hasMetadata(gv, "enzyme_shadow"))
      return false;
    if (auto md = gv->getMetadata("enzyme_activity_value")) {
      auto res = cast<MDString>(md->getOperand(0))->getString();
      if (res == "const")
        return true;
      if (res == "active")
        return false;
    }
    if (EnzymeNonmarkedGlobalsInactive)
      return true;
    goto err;
  }
  if (isa<GlobalValue>(val)) {
    if (EnzymeNonmarkedGlobalsInactive)
      return true;
    goto err;
  }

err:;
  llvm::errs() << *oldFunc << "\n";
  llvm::errs() << *newFunc << "\n";
  llvm::errs() << *val << "\n";
  llvm::errs() << "  unknown did status attribute\n";
  assert(0 && "bad");
  exit(1);
}

bool GradientUtils::isConstantInstruction(const Instruction *inst) const {
  assert(inst->getParent()->getParent() == oldFunc);
  return ATA->isConstantInstruction(TR, const_cast<Instruction *>(inst));
}

bool GradientUtils::getContext(llvm::BasicBlock *BB, LoopContext &lc) {
  return CacheUtility::getContext(BB, lc,
                                  /*ReverseLimit*/ reverseBlocks.size() > 0);
}

void GradientUtils::forceAugmentedReturns() {
  assert(TR.getFunction() == oldFunc);

  // Pass 1: create BB-level contexts for the whole loop/function
  for (BasicBlock &oBB : *oldFunc) {
    if (notForAnalysis.find(&oBB) != notForAnalysis.end())
      continue;
    LoopContext LC;
    getContext(cast<BasicBlock>(getNewFromOriginal(&oBB)), LC);
  }

  // Pass 2: instruction processing
  for (BasicBlock &oBB : *oldFunc) {
    if (notForAnalysis.find(&oBB) != notForAnalysis.end())
      continue;

    for (Instruction &I : oBB) {
      Instruction *inst = &I;

      if (inst->getType()->isEmptyTy() || inst->getType()->isVoidTy())
        continue;

      if (mode == DerivativeMode::ForwardMode ||
          mode == DerivativeMode::ForwardModeSplit ||
          mode == DerivativeMode::ForwardModeError) {
        if (!isConstantValue(inst)) {
          IRBuilder<> BuilderZ(inst);
          getForwardBuilder(BuilderZ);
#if LLVM_VERSION_MAJOR >= 18
          auto It = BuilderZ.GetInsertPoint();
          It.setHeadBit(true);
          BuilderZ.SetInsertPoint(It);
#endif
          Type *antiTy = getShadowType(inst->getType());
          PHINode *anti =
              BuilderZ.CreatePHI(antiTy, 1, inst->getName() + "'dual_phi");
          invertedPointers.insert(std::make_pair(
              (const Value *)inst, InvertedPointerVH(this, anti)));
        }
        continue;
      }

      if (inst->getType()->isFPOrFPVectorTy())
        continue; //! op->getType()->isPointerTy() &&
                  //! !op->getType()->isIntegerTy()) {

      if (!TR.query(inst)[{-1}].isPossiblePointer())
        continue;

      if (isa<LoadInst>(inst)) {
        IRBuilder<> BuilderZ(inst);
        getForwardBuilder(BuilderZ);
#if LLVM_VERSION_MAJOR >= 18
        auto It = BuilderZ.GetInsertPoint();
        It.setHeadBit(true);
        BuilderZ.SetInsertPoint(It);
#endif
        Type *antiTy = getShadowType(inst->getType());
        PHINode *anti =
            BuilderZ.CreatePHI(antiTy, 1, inst->getName() + "'il_phi");
        invertedPointers.insert(
            std::make_pair((const Value *)inst, InvertedPointerVH(this, anti)));
        continue;
      }

      if (!isa<CallInst>(inst)) {
        continue;
      }

      CallInst *op = cast<CallInst>(inst);
      Function *called = op->getCalledFunction();

      if ((mode == DerivativeMode::ReverseModeGradient ||
           mode == DerivativeMode::ReverseModeCombined) &&
          called && called->getName() == "llvm.julia.gc_preserve_begin") {
        IRBuilder<> BuilderZ(inst);
        getForwardBuilder(BuilderZ);
        auto anti = BuilderZ.CreateCall(called, ArrayRef<Value *>(),
                                        op->getName() + "'ip");
        anti->setDebugLoc(getNewFromOriginal(op->getDebugLoc()));
        invertedPointers.insert(
            std::make_pair((const Value *)inst, InvertedPointerVH(this, anti)));
        continue;
      }

      if (isa<IntrinsicInst>(inst)) {
        continue;
      }

      if (isConstantValue(inst)) {
        continue;
      }

      IRBuilder<> BuilderZ(inst);
      getForwardBuilder(BuilderZ);
#if LLVM_VERSION_MAJOR >= 18
      auto It = BuilderZ.GetInsertPoint();
      It.setHeadBit(true);
      BuilderZ.SetInsertPoint(It);
#endif

      // Shadow allocations must strictly preceede the primal, lest Julia have
      // GC issues. Consider the following: %r = gc_alloc() init %r
      // ...
      // if the shadow did not preceed
      // %r = gc_alloc()
      // %dr = gc_alloc()
      // zero %dr
      // init %r, %dr
      // ...
      // After %r, before %dr the %r memory would be uninit, so the allocator
      // inside %dr would hit garbage and segfault. However, by having the %dr
      // first, then it will be zero'd before the %r allocation, preventing the
      // issue.
      if (isAllocationCall(inst, TLI)) {
        BuilderZ.SetInsertPoint(getNewFromOriginal(inst));
#if LLVM_VERSION_MAJOR >= 18
        auto It = BuilderZ.GetInsertPoint();
        It.setHeadBit(true);
        BuilderZ.SetInsertPoint(It);
#endif
      }
      Type *antiTy = getShadowType(inst->getType());

      PHINode *anti = BuilderZ.CreatePHI(antiTy, 1, op->getName() + "'ip_phi");
      if (hasNoCache(inst)) {
        anti->setMetadata("enzyme_nocache",
                          MDNode::get(inst->getContext(), {}));
      }
      anti->setDebugLoc(getNewFromOriginal(op->getDebugLoc()));
      invertedPointers.insert(
          std::make_pair((const Value *)inst, InvertedPointerVH(this, anti)));

      if (isAllocationCall(inst, TLI)) {
        anti->setName(op->getName() + "'mi");
      }
    }
  }
}

void InvertedPointerVH::deleted() {
  llvm::errs() << *gutils->oldFunc << "\n";
  llvm::errs() << *gutils->newFunc << "\n";
  gutils->dumpPointers();
  llvm::errs() << **this << "\n";
  assert(0 && "erasing something in invertedPointers map");
}

void SubTransferHelper(GradientUtils *gutils, DerivativeMode mode,
                       Type *secretty, Intrinsic::ID intrinsic,
                       unsigned dstalign, unsigned srcalign, unsigned offset,
                       bool dstConstant, Value *shadow_dst, Value *primal_dst,
                       bool srcConstant, Value *shadow_src, Value *primal_src,
                       Value *length, Value *isVolatile, llvm::CallInst *MTI,
                       bool allowForward, bool shadowsLookedUp,
                       bool backwardsShadow) {
  // TODO offset
  if (secretty) {
    // no change to forward pass if represents floats
    if (mode == DerivativeMode::ReverseModeGradient ||
        mode == DerivativeMode::ReverseModeCombined ||
        mode == DerivativeMode::ForwardModeSplit) {
      IRBuilder<> Builder2(MTI);
      if (mode == DerivativeMode::ForwardModeSplit)
        gutils->getForwardBuilder(Builder2);
      else
        gutils->getReverseBuilder(Builder2);

      // If the src is constant simply zero d_dst and don't propagate to d_src
      // (which thus == src and may be illegal)
      if (srcConstant) {
        // Don't zero in forward mode.
        if (mode != DerivativeMode::ForwardModeSplit) {

          Value *args[] = {
              shadowsLookedUp ? shadow_dst
                              : gutils->lookupM(shadow_dst, Builder2),
              ConstantInt::get(Type::getInt8Ty(MTI->getContext()), 0),
              gutils->lookupM(length, Builder2),
              ConstantInt::getFalse(MTI->getContext())};

          if (args[0]->getType()->isIntegerTy())
            args[0] = Builder2.CreateIntToPtr(args[0],
                                              getInt8PtrTy(MTI->getContext()));

          Type *tys[] = {args[0]->getType(), args[2]->getType()};
          auto memsetIntr = getIntrinsicDeclaration(
              MTI->getParent()->getParent()->getParent(), Intrinsic::memset,
              tys);
          auto cal = Builder2.CreateCall(memsetIntr, args);
          cal->setCallingConv(memsetIntr->getCallingConv());
          if (dstalign != 0) {
            cal->addParamAttr(0, Attribute::getWithAlignment(MTI->getContext(),
                                                             Align(dstalign)));
          }
        }

      } else {
        auto dsto =
            (shadowsLookedUp || mode == DerivativeMode::ForwardModeSplit)
                ? shadow_dst
                : gutils->lookupM(shadow_dst, Builder2);
        Value *dst_inactive = nullptr;
        if (gutils->runtimeActivity) {
          if (primal_dst) {
            Value *primal_dsto =
                (shadowsLookedUp || mode == DerivativeMode::ForwardModeSplit)
                    ? primal_dst
                    : gutils->lookupM(primal_dst, Builder2);
            dst_inactive = Builder2.CreateICmpEQ(dsto, primal_dsto);
          } else {
            dst_inactive = ConstantInt::getFalse(Builder2.getContext());
          }
        }
        if (dsto->getType()->isIntegerTy())
          dsto =
              Builder2.CreateIntToPtr(dsto, getInt8PtrTy(dsto->getContext()));
        unsigned dstaddr =
            cast<PointerType>(dsto->getType())->getAddressSpace();
        if (offset != 0) {
          dsto = Builder2.CreateConstInBoundsGEP1_64(
              Type::getInt8Ty(dsto->getContext()), dsto, offset);
        }
        auto srco =
            (shadowsLookedUp || mode == DerivativeMode::ForwardModeSplit)
                ? shadow_src
                : gutils->lookupM(shadow_src, Builder2);
        Value *src_inactive = nullptr;
        if (gutils->runtimeActivity) {
          if (srcConstant) {
            src_inactive = ConstantInt::getTrue(Builder2.getContext());
          } else {
            Value *primal_srco =
                (shadowsLookedUp || mode == DerivativeMode::ForwardModeSplit)
                    ? primal_src
                    : gutils->lookupM(primal_src, Builder2);
            src_inactive = Builder2.CreateICmpEQ(srco, primal_srco);
          }
        }
        if (mode != DerivativeMode::ForwardModeSplit)
          dsto = Builder2.CreatePointerCast(dsto,
                                            getPointerType(secretty, dstaddr));
        if (srco->getType()->isIntegerTy())
          srco =
              Builder2.CreateIntToPtr(srco, getInt8PtrTy(srco->getContext()));
        unsigned srcaddr =
            cast<PointerType>(srco->getType())->getAddressSpace();
        if (offset != 0) {
          srco = Builder2.CreateConstInBoundsGEP1_64(
              Type::getInt8Ty(srco->getContext()), srco, offset);
        }
        if (mode != DerivativeMode::ForwardModeSplit)
          srco = Builder2.CreatePointerCast(srco,
                                            getPointerType(secretty, srcaddr));

        if (mode == DerivativeMode::ForwardModeSplit) {
          MaybeAlign dalign;
          if (dstalign)
            dalign = MaybeAlign(dstalign);
          MaybeAlign salign;
          if (srcalign)
            salign = MaybeAlign(srcalign);

          if (intrinsic == Intrinsic::memmove) {
            Builder2.CreateMemMove(dsto, dalign, srco, salign, length);
          } else {
            Builder2.CreateMemCpy(dsto, dalign, srco, salign, length);
          }
        } else {
          SmallVector<Value *, 5> args = {
              Builder2.CreatePointerCast(dsto,
                                         getPointerType(secretty, dstaddr)),
              Builder2.CreatePointerCast(srco,
                                         getPointerType(secretty, srcaddr)),
              Builder2.CreateUDiv(
                  gutils->lookupM(length, Builder2),
                  ConstantInt::get(length->getType(),
                                   Builder2.GetInsertBlock()
                                           ->getParent()
                                           ->getParent()
                                           ->getDataLayout()
                                           .getTypeAllocSizeInBits(secretty) /
                                       8))};

          if (gutils->runtimeActivity) {
            args.push_back(dst_inactive);
            args.push_back(src_inactive);
          }

          auto dmemcpy = ((intrinsic == Intrinsic::memcpy)
                              ? getOrInsertDifferentialFloatMemcpy
                              : getOrInsertDifferentialFloatMemmove)(
              *MTI->getParent()->getParent()->getParent(), secretty, dstalign,
              srcalign, dstaddr, srcaddr,
              cast<IntegerType>(length->getType())->getBitWidth(),
              gutils->runtimeActivity, gutils->isAtomic(primal_src));
          Builder2.CreateCall(dmemcpy, args);
        }
      }
    }
  } else {

    // if represents pointer or integer type then only need to modify forward
    // pass with the copy
    if ((allowForward && (mode == DerivativeMode::ReverseModePrimal ||
                          mode == DerivativeMode::ReverseModeCombined)) ||
        (backwardsShadow && (mode == DerivativeMode::ReverseModeGradient ||
                             mode == DerivativeMode::ForwardModeSplit))) {
      assert(!shadowsLookedUp);

      // It is questionable how the following case would even occur, but if
      // the dst is constant, we shouldn't do anything extra
      if (dstConstant) {
        return;
      }

      IRBuilder<> BuilderZ(gutils->getNewFromOriginal(MTI));

      // If src is inactive, then we should copy from the regular pointer
      // (i.e. suppose we are copying constant memory representing dimensions
      // into a tensor)
      //  to ensure that the differential tensor is well formed for use
      //  OUTSIDE the derivative generation (as enzyme doesn't need this), we
      //  should also perform the copy onto the differential. Future
      //  Optimization (not implemented): If dst can never escape Enzyme code,
      //  we may omit this copy.
      // no need to update pointers, even if dst is active
      auto dsto = shadow_dst;
      if (dsto->getType()->isIntegerTy())
        dsto = BuilderZ.CreateIntToPtr(dsto, getInt8PtrTy(MTI->getContext()));
      if (offset != 0) {
        dsto = BuilderZ.CreateConstInBoundsGEP1_64(
            Type::getInt8Ty(dsto->getContext()), dsto, offset);
      }
      auto srco = shadow_src;
      if (srco->getType()->isIntegerTy())
        srco = BuilderZ.CreateIntToPtr(srco, getInt8PtrTy(MTI->getContext()));
      if (offset != 0) {
        srco = BuilderZ.CreateConstInBoundsGEP1_64(
            Type::getInt8Ty(srco->getContext()), srco, offset);
      }
      Value *args[] = {dsto, srco, length, isVolatile};

      Type *tys[] = {args[0]->getType(), args[1]->getType(),
                     args[2]->getType()};

      auto memtransIntr =
          getIntrinsicDeclaration(gutils->newFunc->getParent(), intrinsic, tys);
      auto cal = BuilderZ.CreateCall(memtransIntr, args);
      cal->setAttributes(MTI->getAttributes());
      cal->setCallingConv(memtransIntr->getCallingConv());
      cal->setTailCallKind(MTI->getTailCallKind());

      if (dstalign != 0) {
        cal->addParamAttr(
            0, Attribute::getWithAlignment(MTI->getContext(), Align(dstalign)));
      }
      if (srcalign != 0) {
        cal->addParamAttr(
            1, Attribute::getWithAlignment(MTI->getContext(), Align(srcalign)));
      }
    }
  }
}

void GradientUtils::computeForwardingProperties(Instruction *V) {
  if (!EnzymeRematerialize)
    return;

  // For the piece of memory V allocated within this scope, it will be
  // initialized in some way by the (augmented) forward pass. Loads and other
  // load-like operations will either require the allocation V itself to be
  // preserved for the reverse pass, or alternatively the tape for those
  // operations.
  //
  // Instead, we ask here whether or not we can restore the memory state of V in
  // the reverse pass by recreating all of the stores and store-like operations
  // into the V prior to their load-like uses.
  //
  // Notably, we only need to preserve the ability to reload any values actually
  // used in the reverse pass.

  std::map<UsageKey, bool> Seen;
  bool primalNeededInReverse =
      DifferentialUseAnalysis::is_value_needed_in_reverse<QueryType::Primal>(
          this, V, DerivativeMode::ReverseModeGradient, Seen, notForAnalysis);

  SmallVector<LoadInst *, 1> loads;
  SmallVector<LoadLikeCall, 1> loadLikeCalls;
  SmallPtrSet<Instruction *, 1> stores;
  SmallPtrSet<Instruction *, 1> storingOps;
  SmallPtrSet<Instruction *, 1> frees;
  SmallPtrSet<IntrinsicInst *, 1> LifetimeStarts;
  bool promotable = true;
  bool shadowpromotable = true;

  CallInst *nonRepeatableWritingCall = nullptr;
  SmallVector<Instruction *, 1> shadowPointerLoads;

  std::set<std::pair<Instruction *, Value *>> seen;
  SmallVector<std::pair<Instruction *, Value *>, 1> todo;
  for (auto U : V->users())
    if (auto I = dyn_cast<Instruction>(U))
      todo.push_back(std::make_pair(I, V));
  while (todo.size()) {
    auto tup = todo.back();
    Instruction *cur = tup.first;
    Value *prev = tup.second;
    todo.pop_back();
    if (seen.count(tup))
      continue;
    seen.insert(tup);
    if (notForAnalysis.count(cur->getParent()))
      continue;
    if (isPointerArithmeticInst(cur)) {
      for (auto u : cur->users()) {
        if (auto I = dyn_cast<Instruction>(u))
          todo.push_back(std::make_pair(I, (Value *)cur));
      }
    } else if (auto load = dyn_cast<LoadInst>(cur)) {

      // If loaded value is an int or pointer, may need
      // to preserve initialization within the primal.
      auto TT = TR.query(load)[{-1}];
      if (!TT.isFloat()) {
        shadowPointerLoads.push_back(cur);
      }
      loads.push_back(load);
    } else if (auto store = dyn_cast<StoreInst>(cur)) {
      // TODO only add store to shadow iff non float type
      if (store->getValueOperand() == prev) {
        EmitWarning("NotPromotable", *cur, " Could not promote allocation ", *V,
                    " due to capturing store ", *cur);
        promotable = false;
        shadowpromotable = false;
        break;
      } else {
        stores.insert(store);
        storingOps.insert(store);
      }
    } else if (auto II = dyn_cast<IntrinsicInst>(cur)) {
      if (II->getCalledFunction()->getName() == "llvm.enzyme.lifetime_start") {
        LifetimeStarts.insert(II);
      } else if (II->getCalledFunction()->getName() ==
                 "llvm.enzyme.lifetime_end") {
      } else {
        switch (II->getIntrinsicID()) {
        case Intrinsic::lifetime_start:
          LifetimeStarts.insert(II);
          break;
        case Intrinsic::dbg_declare:
        case Intrinsic::dbg_value:
        case Intrinsic::dbg_label:
#if LLVM_VERSION_MAJOR <= 16
        case llvm::Intrinsic::dbg_addr:
#endif
        case Intrinsic::lifetime_end:
          break;
        case Intrinsic::memset: {
          stores.insert(II);
          storingOps.insert(II);
          break;
        }
        // TODO memtransfer(cpy/move)
        case Intrinsic::memcpy:
        case Intrinsic::memmove:
        default:
          promotable = false;
          shadowpromotable = false;
          EmitWarning("NotPromotable", *cur, " Could not promote allocation ",
                      *V, " due to unknown intrinsic ", *cur);
          break;
        }
      }
    } else if (auto CI = dyn_cast<CallInst>(cur)) {
      StringRef funcName = getFuncNameFromCall(CI);
      if (isDeallocationFunction(funcName, TLI)) {
        frees.insert(CI);
        continue;
      }
      // The allocation arg is the first arg of the write barrier.
      //  The capturing store in subsequent args should be handled by forbidding
      //  capturing stores
      if (funcName == "julia.write_barrier" ||
          funcName == "julia.write_barrier_binding") {
        if (CI->getArgOperand(0) == prev) {
          stores.insert(CI);
        }
        continue;
      }
      if (funcName == "enzyme_zerotype") {
        stores.insert(CI);
        continue;
      }

      size_t idx = 0;
      bool seenLoadLikeCall = false;
#if LLVM_VERSION_MAJOR >= 14
      for (auto &arg : CI->args())
#else
      for (auto &arg : CI->arg_operands())
#endif
      {
        if (arg != prev) {
          idx++;
          continue;
        }
        auto TT = TR.query(prev)[{-1, -1}];

        bool NoCapture = isNoCapture(CI, idx);

        bool ReadOnly = isReadOnly(CI, idx);

        bool WriteOnly = isWriteOnly(CI, idx);

        // If the pointer is captured, conservatively assume it is used in
        // nontrivial ways that make both the primal and shadow not promotable.
        if (!NoCapture) {
          shadowpromotable = false;
          promotable = false;
          EmitWarning("NotPromotable", *cur, " Could not promote allocation ",
                      *V, " due to unknown capturing call ", *cur,
                      " at idx=", idx, " prev=", *prev);
          idx++;
          continue;
        }

        // From here on out we can assume the pointer is not captured, and only
        // written to or read from.

        // If we may read from the memory, consider this a load-like call
        // that must have all writes done in preparation for any reverse-pass
        // users.
        if (!WriteOnly) {
          if (!seenLoadLikeCall) {
            loadLikeCalls.push_back(LoadLikeCall(CI, prev));
            seenLoadLikeCall = true;
          }
        }

        // If we may write to memory, we cannot promote if any values
        // need the allocation or any descendants for the reverse pass.
        if (!ReadOnly) {
          // There is an exception for Julia returnRoots which will be
          // separately handled in a GC postprocessing pass. Moreover these
          // values are never `needed` in the reverse pass (just we need to mark
          // those values as being GC'd by the function).
          bool returnRoots = false;
          //    CI->getAttributes().hasParamAttr(idx, "enzymejl_returnRoots") ||
          //    CI->getAttributes().hasParamAttr(idx, "enzymejl_returnRoots_v");
          if (primalNeededInReverse && !returnRoots) {
            promotable = false;
            EmitWarning("NotPromotable", *cur, " Could not promote allocation ",
                        *V, " due to unknown writing call ", *cur);
          }
          if (!nonRepeatableWritingCall)
            nonRepeatableWritingCall = CI;
          storingOps.insert(cur);
        }

        // Consider shadow memory now.
        //
        // If the memory is all floats, there's no issue, since besides zero
        // initialization nothing should occur for them in the forward pass
        if (TT.isFloat()) {
        } else if (WriteOnly) {
          // Don't need in the case of int/pointer stores, (should be done by
          // fwd pass), and as isFloat above described does not prevent the
          // shadow
        } else {
          shadowPointerLoads.push_back(cur);
        }

        idx++;
      }

    } else {
      promotable = false;
      shadowpromotable = false;
      EmitWarning("NotPromotable", *cur, " Could not promote allocation ", *V,
                  " due to unknown instruction ", *cur);
    }
  }

  // Find the outermost loop of all stores, and the allocation/lifetime
  Loop *outer = OrigLI->getLoopFor(V->getParent());
  if (LifetimeStarts.size() == 1) {
    outer = OrigLI->getLoopFor((*LifetimeStarts.begin())->getParent());
  }

  for (auto S : stores) {
    outer = getAncestor(outer, OrigLI->getLoopFor(S->getParent()));
  }

  // May now read pointers for storing into other pointers. Therefore we
  // need to pre initialize the shadow.
  bool primalInitializationOfShadow = shadowPointerLoads.size() > 0;

  if (shadowpromotable && !isConstantValue(V)) {
    for (auto LI : shadowPointerLoads) {
      // Is there a store which could occur after the load.
      // This subsequent store would invalidate any loads being re-performed.
      SmallVector<Instruction *, 2> results;
      mayExecuteAfter(results, LI, storingOps, outer);
      for (auto res : results) {
        if (overwritesToMemoryReadBy(&TR, *OrigAA, TLI, *OrigSE, *OrigLI,
                                     *OrigDT, LI, res, outer)) {
          EmitWarning("NotPromotable", *LI,
                      " Could not promote shadow allocation ", *V,
                      " due to pointer load ", *LI,
                      " which does not postdominates store ", *res);
          shadowpromotable = false;
          goto exitL;
        }
      }
    }
    // If there is a store not reproduced in the reverse pass (e.g. as part
    // of a write in a call), and this store is necessary to a pointer load of
    // the shadow, this is not materializable since the load will not return
    // the same value.
    {
      SmallVector<Instruction *, 2> nonReproducedStores;
      for (auto S : storingOps)
        if (!stores.count(S)) {
          SmallVector<Instruction *, 2> results;
          SmallPtrSet<Instruction *, 2> shadowPtrLoadSet(
              shadowPointerLoads.begin(), shadowPointerLoads.end());
          mayExecuteAfter(results, S, shadowPtrLoadSet, outer);
          if (results.size()) {
            EmitWarning("NotPromotable", *results[0],
                        " Could not promote shadow allocation ", *V,
                        " due to non-reproduced store ", *S,
                        " which may impact pointer load ", *results[0]);
            shadowpromotable = false;
            goto exitL;
          }
        }
    }
  exitL:;
    if (shadowpromotable) {
      backwardsOnlyShadows[V] = ShadowRematerializer(
          stores, frees, primalInitializationOfShadow, outer);
    }
  }

  if (!promotable)
    return;

  SmallPtrSet<LoadInst *, 1> rematerializable;

  // We currently require a rematerializable allocation to have
  // all of its loads be able to be performed again. Thus if
  // there is an overwriting store after a load in context,
  // it may no longer be rematerializable.
  for (auto LI : loads) {
    // Is there a store which could occur after the load.
    // In other words
    SmallVector<Instruction *, 2> results;
    mayExecuteAfter(results, LI, storingOps, outer);
    for (auto res : results) {
      if (overwritesToMemoryReadBy(&TR, *OrigAA, TLI, *OrigSE, *OrigLI, *OrigDT,
                                   LI, res, outer)) {
        EmitWarning("NotPromotable", *LI, " Could not promote allocation ", *V,
                    " due to load ", *LI,
                    " which does not postdominates store ", *res);
        return;
      }
    }
    rematerializable.insert(LI);
  }
  for (auto LI : loadLikeCalls) {
    // Is there a store which could occur after the load.
    // In other words
    SmallVector<Instruction *, 2> results;
    mayExecuteAfter(results, LI.loadCall, storingOps, outer);
    for (auto res : results) {
      if (overwritesToMemoryReadBy(&TR, *OrigAA, TLI, *OrigSE, *OrigLI, *OrigDT,
                                   LI.loadCall, res, outer)) {
        EmitWarning("NotPromotable", *LI.loadCall,
                    " Could not promote allocation ", *V,
                    " due to load-like call ", *LI.loadCall,
                    " which does not postdominates store ", *res);
        return;
      }
    }
  }
  rematerializableAllocations[V] = Rematerializer(
      loads, loadLikeCalls, stores, frees, outer, nonRepeatableWritingCall);
}

BasicBlock *GradientUtils::addReverseBlock(BasicBlock *currentBlock,
                                           Twine const &name, bool forkCache,
                                           bool push) {
  assert(reverseBlocks.size());
  auto found = reverseBlockToPrimal.find(currentBlock);
  assert(found != reverseBlockToPrimal.end());

  SmallVector<BasicBlock *, 4> &vec = reverseBlocks[found->second];
  assert(vec.size());
  if (vec.back() != currentBlock) {
    std::string str;
    raw_string_ostream ss(str);
    ss << "Error adding reverse block:\n";
    ss << "fwdBlock: " << *found->second << "\n";
    ss << "currentBlock: " << *currentBlock << "\n";
    ss << "vec.back(): " << *vec.back() << "\n";
    if (CustomErrorHandler) {
      CustomErrorHandler(str.c_str(), wrap((Value *)currentBlock),
                         ErrorType::InternalError, nullptr, nullptr, nullptr);
    } else {
      DebugLoc loc;
      if (hasTerminator(found->second)) {
        loc = found->second->getTerminator()->getDebugLoc();
      }
      EmitFailure("AddReverseBlockError", loc, found->second->getParent(),
                  ss.str());
    }
  }
  assert(vec.back() == currentBlock);

  BasicBlock *rev =
      BasicBlock::Create(currentBlock->getContext(), name, newFunc);
  rev->moveAfter(currentBlock);
  if (push)
    vec.push_back(rev);
  reverseBlockToPrimal[rev] = found->second;
  if (forkCache) {
    for (auto pair : unwrap_cache[currentBlock])
      unwrap_cache[rev].insert(pair);
    for (auto pair : lookup_cache[currentBlock])
      lookup_cache[rev].insert(pair);
  }
  return rev;
}

void GradientUtils::replaceAWithB(Value *A, Value *B, bool storeInCache) {
  if (A == B)
    return;
  assert(A->getType() == B->getType());

  if (auto iA = dyn_cast<Instruction>(A)) {
    if (unwrappedLoads.find(iA) != unwrappedLoads.end()) {
      auto iB = cast<Instruction>(B);
      unwrappedLoads[iB] = unwrappedLoads[iA];
      unwrappedLoads.erase(iA);
    }
  }

  // Check that the replacement doesn't already exist in the mapping
  // thereby resulting in a conflict.
#ifndef NDEBUG
  if (!isa<UndefValue>(B)) {
    auto found = newToOriginalFn.find(A);
    if (found != newToOriginalFn.end()) {
      auto foundB = newToOriginalFn.find(B);
      assert(foundB == newToOriginalFn.end());
    }
  }
#endif

  CacheUtility::replaceAWithB(A, B, storeInCache);
}

void GradientUtils::erase(Instruction *I) {
  assert(I);
  if (I->getParent()->getParent() != newFunc) {
    llvm::errs() << "newFunc: " << *newFunc << "\n";
    llvm::errs() << "paren: " << *I->getParent()->getParent() << "\n";
    llvm::errs() << "I: " << *I << "\n";
  }
  assert(I->getParent()->getParent() == newFunc);

  // not original, should not contain
  assert(!invertedPointers.count(I));
  // not original, should not contain
  assert(!originalToNewFn.count(I));

  originalToNewFn.erase(I);
  {
    auto found = newToOriginalFn.find(I);
    if (found != newToOriginalFn.end()) {
      Value *orig = found->second;
      newToOriginalFn.erase(found);
      originalToNewFn.erase(orig);
    }
  }
  {
    auto found = UnwrappedWarnings.find(I);
    if (found != UnwrappedWarnings.end()) {
      UnwrappedWarnings.erase(found);
    }
  }
  unwrappedLoads.erase(I);

  for (auto &pair : unwrap_cache) {
    if (pair.second.find(I) != pair.second.end())
      pair.second.erase(I);
  }

  for (auto &pair : lookup_cache) {
    if (pair.second.find(I) != pair.second.end())
      pair.second.erase(I);
  }
  CacheUtility::erase(I);
}

void GradientUtils::eraseWithPlaceholder(Instruction *I, Instruction *orig,
                                         const Twine &suffix, bool erase) {
  if (I->getType()->isTokenTy()) {
    replaceAWithB(I, UndefValue::get(I->getType()));
  } else if (!I->getType()->isVoidTy() && !I->getType()->isTokenTy()) {
    auto inspos = I->getIterator();
#if LLVM_VERSION_MAJOR >= 18
#if LLVM_VERSION_MAJOR >= 21
#else
    if (I->getParent()->IsNewDbgInfoFormat)
#endif
    {
      if (!inspos.getHeadBit()) {
        auto srcmarker = I->getParent()->getMarker(inspos);
        if (srcmarker && !srcmarker->empty()) {
          inspos.setHeadBit(true);
        }
      }
    }
#endif
    IRBuilder<> BuilderZ(I->getParent(), inspos);
    auto pn = BuilderZ.CreatePHI(I->getType(), 1, I->getName() + suffix);
    fictiousPHIs[pn] = orig;
    replaceAWithB(I, pn);
  }

  if (erase) {
    this->erase(I);
  }
}

void GradientUtils::setTape(Value *newtape) {
  assert(tape == nullptr);
  assert(newtape != nullptr);
  assert(tapeidx == 0);
  assert(addedTapeVals.size() == 0);
  tape = newtape;
}

void GradientUtils::dumpPointers() {
  errs() << "invertedPointers:\n";
  for (auto a : invertedPointers) {
    errs() << "   invertedPointers[" << *a.first << "] = " << *a.second << "\n";
  }
  errs() << "end invertedPointers\n";
}

int GradientUtils::getIndex(
    std::pair<Instruction *, CacheType> idx,
    const std::map<std::pair<Instruction *, CacheType>, int> &mapping,
    IRBuilder<> &B) {
  assert(tape);
  auto found = mapping.find(idx);
  if (found == mapping.end()) {
    std::string s;
    llvm::raw_string_ostream ss(s);
    ss << *oldFunc << "\n";
    ss << *newFunc << "\n";
    ss << " <mapping>\n";
    for (auto &p : mapping) {
      ss << "   idx: " << *p.first.first << ", " << p.first.second
         << " pos=" << p.second << "\n";
    }
    ss << " </mapping>\n";
    ss << "idx: " << *idx.first << ", " << idx.second << "\n";
    ss << " could not find index in mapping\n";
    if (CustomErrorHandler) {
      CustomErrorHandler(ss.str().c_str(), wrap(idx.first),
                         ErrorType::GetIndexError, this, nullptr, wrap(&B));
    } else {
      EmitFailure("GetIndexError", idx.first->getDebugLoc(), idx.first,
                  ss.str());
    }
    return IndexMappingError;
  }
  return found->second;
}

int GradientUtils::getIndex(
    std::pair<Instruction *, CacheType> idx,
    std::map<std::pair<Instruction *, CacheType>, int> &mapping,
    IRBuilder<> &B) {
  if (tape) {
    return getIndex(
        idx,
        (const std::map<std::pair<Instruction *, CacheType>, int> &)mapping, B);
  } else {
    if (mapping.find(idx) != mapping.end()) {
      return mapping[idx];
    }
    mapping[idx] = tapeidx;
    ++tapeidx;
    return mapping[idx];
  }
}

void GradientUtils::computeGuaranteedFrees() {
  SmallPtrSet<CallInst *, 2> allocsToPromote;

  DenseMap<Metadata *, SmallVector<CallInst *>> cache_frees;
  for (auto &BB : *oldFunc) {
    for (auto &I : BB) {
      auto CI = dyn_cast<CallInst>(&I);
      if (!CI)
        continue;
      if (auto MD = CI->getMetadata("enzyme_cache_free")) {
        if (MD->getNumOperands() > 0) {
          Metadata *id = MD->getOperand(0);
          cache_frees[id].push_back(CI);
        }
      }
    }
  }

  for (auto &BB : *oldFunc) {
    if (notForAnalysis.count(&BB))
      continue;
    for (auto &I : BB) {
      if (auto AI = dyn_cast<AllocaInst>(&I))
        computeForwardingProperties(AI);

      auto CI = dyn_cast<CallInst>(&I);
      if (!CI)
        continue;

      StringRef funcName = getFuncNameFromCall(CI);

      if (isDeallocationFunction(funcName, TLI)) {
        llvm::Value *val = getBaseObject(CI->getArgOperand(0));

        if (auto dc = dyn_cast<CallInst>(val)) {
          StringRef sfuncName = getFuncNameFromCall(dc);
          if (isAllocationFunction(sfuncName, TLI)) {

            bool hasPDFree = false;
            if (dc->getParent() == CI->getParent() ||
                OrigPDT->dominates(CI->getParent(), dc->getParent())) {
              hasPDFree = true;
            }

            if (hasPDFree) {
              allocationsWithGuaranteedFree[dc].insert(CI);
            }
          }
        }
      }
      if (auto MD = CI->getMetadata("enzyme_cache_alloc")) {
        if (MD->getNumOperands() > 0) {
          Metadata *id = MD->getOperand(0);
          if (cast<ConstantInt>(
                  cast<ConstantAsMetadata>(cast<MDNode>(id)->getOperand(0))
                      ->getValue())
                  ->isOne()) {
            for (auto otherCI : cache_frees[id]) {
              allocationsWithGuaranteedFree[CI].insert(otherCI);
            }
          }
        }
      }
      if (isAllocationFunction(funcName, TLI)) {
        allocsToPromote.insert(CI);
        if (hasMetadata(CI, "enzyme_fromstack")) {
          allocationsWithGuaranteedFree[CI].insert(CI);
        }
        // TODO: special case object managed by the GC as it is automatically
        // freed.
        if (EnzymeJuliaAddrLoad && isa<PointerType>(CI->getType()) &&
            cast<PointerType>(CI->getType())->getAddressSpace() == 10) {
        }
      }
    }
  }
  for (CallInst *V : allocsToPromote) {
    // TODO compute if an only load/store (non capture)
    // allocaion by traversing its users. If so, mark
    // all of its load/stores, as now the loads can
    // potentially be rematerialized without a cache
    // of the allocation, but the operands of all stores.
    // This info needs to be provided to minCutCache
    // the derivative of store needs to redo the store,
    // isValueNeededInReverse needs to know to preserve the
    // store operands in this case, etc
    computeForwardingProperties(V);
  }
}

/// Perform the corresponding deallocation of tofree, given it was allocated by
/// allocationfn
// For updating below one should read MemoryBuiltins.cpp, TargetLibraryInfo.cpp
llvm::CallInst *freeKnownAllocation(llvm::IRBuilder<> &builder,
                                    llvm::Value *tofree,
                                    llvm::StringRef allocationfn,
                                    const llvm::DebugLoc &debuglocation,
                                    const llvm::TargetLibraryInfo &TLI,
                                    llvm::CallInst *orig,
                                    GradientUtils *gutils) {
  assert(isAllocationFunction(allocationfn, TLI));

  if (allocationfn == "__rust_alloc" || allocationfn == "__rust_alloc_zeroed") {
    Type *VoidTy = Type::getVoidTy(tofree->getContext());
    Type *IntPtrTy = orig->getType();
    Type *RustSz = orig->getArgOperand(0)->getType();
    Type *inTys[3] = {IntPtrTy, RustSz, RustSz};

    auto FT = FunctionType::get(VoidTy, inTys, false);
    Value *freevalue = builder.GetInsertBlock()
                           ->getParent()
                           ->getParent()
                           ->getOrInsertFunction("__rust_dealloc", FT)
                           .getCallee();
    Value *vals[3];
    vals[0] = builder.CreatePointerCast(tofree, IntPtrTy);
    // size
    vals[1] = gutils->lookupM(
        gutils->getNewFromOriginal(orig->getArgOperand(0)), builder);
    // alignment
    vals[2] = gutils->lookupM(
        gutils->getNewFromOriginal(orig->getArgOperand(1)), builder);
    CallInst *freecall = cast<CallInst>(
        CallInst::Create(FT, freevalue, vals, "", builder.GetInsertBlock()));
    freecall->setDebugLoc(debuglocation);
    if (isa<CallInst>(tofree) &&
        cast<CallInst>(tofree)->getAttributes().hasAttribute(
            AttributeList::ReturnIndex, Attribute::NonNull)) {
      freecall->addAttribute(AttributeList::FirstArgIndex, Attribute::NonNull);
    }
    if (Function *F = dyn_cast<Function>(freevalue))
      freecall->setCallingConv(F->getCallingConv());
    if (freecall->getParent() == nullptr)
      builder.Insert(freecall);
    return freecall;
  }
  if (allocationfn == "julia.gc_alloc_obj" ||
      allocationfn == "jl_gc_alloc_typed" ||
      allocationfn == "ijl_gc_alloc_typed" ||
      allocationfn == "jl_alloc_array_1d" ||
      allocationfn == "ijl_alloc_array_1d" ||
      allocationfn == "jl_alloc_array_2d" ||
      allocationfn == "ijl_alloc_array_2d" ||
      allocationfn == "jl_alloc_array_3d" ||
      allocationfn == "ijl_alloc_array_3d" || allocationfn == "jl_new_array" ||
      allocationfn == "ijl_new_array" ||
      allocationfn == "jl_alloc_genericmemory" ||
      allocationfn == "ijl_alloc_genericmemory" ||
      allocationfn == "jl_alloc_genericmemory_unchecked" ||
      allocationfn == "ijl_alloc_genericmemory_unchecked")
    return nullptr;

  if (allocationfn == "enzyme_allocator") {
    auto inds = getDeallocationIndicesFromCall(orig);
    SmallVector<Value *, 2> vals;
    for (auto ind : inds) {
      if (ind == -1)
        vals.push_back(tofree);
      else
        vals.push_back(gutils->lookupM(
            gutils->getNewFromOriginal(orig->getArgOperand(ind)), builder));
    }
    auto tocall = getDeallocatorFnFromCall(orig);
    auto freecall = builder.CreateCall(tocall, vals);
    freecall->setDebugLoc(debuglocation);
    return freecall;
  }

  if (allocationfn == "swift_allocObject") {
    Type *VoidTy = Type::getVoidTy(tofree->getContext());
    Type *IntPtrTy = getInt8PtrTy(tofree->getContext());

    auto FT = FunctionType::get(VoidTy, ArrayRef<Type *>(IntPtrTy), false);
    Value *freevalue = builder.GetInsertBlock()
                           ->getParent()
                           ->getParent()
                           ->getOrInsertFunction("swift_release", FT)
                           .getCallee();
    CallInst *freecall = cast<CallInst>(CallInst::Create(
        FT, freevalue,
        ArrayRef<Value *>(builder.CreatePointerCast(tofree, IntPtrTy)), "",
        builder.GetInsertBlock()));
    freecall->setDebugLoc(debuglocation);
    if (isa<CallInst>(tofree) &&
        cast<CallInst>(tofree)->getAttributes().hasAttribute(
            AttributeList::ReturnIndex, Attribute::NonNull)) {
      freecall->addAttribute(AttributeList::FirstArgIndex, Attribute::NonNull);
    }
    if (Function *F = dyn_cast<Function>(freevalue))
      freecall->setCallingConv(F->getCallingConv());
    if (freecall->getParent() == nullptr)
      builder.Insert(freecall);
    return freecall;
  }

  if (shadowErasers.find(allocationfn) != shadowErasers.end()) {
    return shadowErasers[allocationfn](builder, tofree);
  }

  if (allocationfn == "__size_returning_new_experiment") {
    allocationfn = "malloc";
    tofree = builder.CreateExtractValue(tofree, 0);
  }

  if (tofree->getType()->isIntegerTy())
    tofree = builder.CreateIntToPtr(tofree, getInt8PtrTy(tofree->getContext()));

  llvm::LibFunc libfunc;
  if (allocationfn == "calloc" || allocationfn == "malloc" ||
      allocationfn == "_mlir_memref_to_llvm_alloc") {
    libfunc = LibFunc_malloc;
  } else {
    bool res = TLI.getLibFunc(allocationfn, libfunc);
    (void)res;
    assert(res && "ought find known allocation fn");
  }

  llvm::LibFunc freefunc;

  switch (libfunc) {
  case LibFunc_malloc: // malloc(unsigned int);
  case LibFunc_valloc: // valloc(unsigned int);
    freefunc = LibFunc_free;
    break;

  case LibFunc_Znwj:                // new(unsigned int);
  case LibFunc_ZnwjRKSt9nothrow_t:  // new(unsigned int, nothrow);
  case LibFunc_ZnwjSt11align_val_t: // new(unsigned int, align_val_t)
  case LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t: // new(unsigned int,
                                                  // align_val_t, nothrow)

  case LibFunc_Znwm:                // new(unsigned long);
  case LibFunc_ZnwmRKSt9nothrow_t:  // new(unsigned long, nothrow);
  case LibFunc_ZnwmSt11align_val_t: // new(unsigned long, align_val_t)
  case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t: // new(unsigned long,
                                                  // align_val_t, nothrow)
    freefunc = LibFunc_ZdlPv;
    break;

  case LibFunc_Znaj:                // new[](unsigned int);
  case LibFunc_ZnajRKSt9nothrow_t:  // new[](unsigned int, nothrow);
  case LibFunc_ZnajSt11align_val_t: // new[](unsigned int, align_val_t)
  case LibFunc_ZnajSt11align_val_tRKSt9nothrow_t: // new[](unsigned int,
                                                  // align_val_t, nothrow

  case LibFunc_Znam:                // new[](unsigned long);
  case LibFunc_ZnamRKSt9nothrow_t:  // new[](unsigned long, nothrow);
  case LibFunc_ZnamSt11align_val_t: // new[](unsigned long, align_val_t)
  case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t: // new[](unsigned long,
                                                  // align_val_t, nothrow)
    freefunc = LibFunc_ZdaPv;
    break;

  case LibFunc_msvc_new_longlong:         // new(unsigned long long);
  case LibFunc_msvc_new_longlong_nothrow: // new(unsigned long long, nothrow);
    freefunc = LibFunc_msvc_delete_ptr64_longlong;
    break;

  case LibFunc_msvc_new_array_longlong:         // new[](unsigned long long);
  case LibFunc_msvc_new_array_longlong_nothrow: // new[](unsigned long long,
                                                // nothrow);
    freefunc = LibFunc_msvc_delete_array_ptr64_longlong;
    break;

  case LibFunc_msvc_new_int:               // new(unsigned int);
  case LibFunc_msvc_new_int_nothrow:       // new(unsigned int, nothrow);
  case LibFunc_msvc_new_array_int:         // new[](unsigned int);
  case LibFunc_msvc_new_array_int_nothrow: // new[](unsigned int, nothrow);
    llvm_unreachable("msvc deletion not handled");

  default:
    llvm_unreachable("unknown allocation function");
  }
  llvm::StringRef freename = TLI.getName(freefunc);
  if (freefunc == LibFunc_free) {
    freename = "free";
    assert(freename == "free");
    if (freename != "free")
      llvm_unreachable("illegal free");
  }
  if (allocationfn == "_mlir_memref_to_llvm_alloc")
    freename = "_mlir_memref_to_llvm_free";

  Type *VoidTy = Type::getVoidTy(tofree->getContext());
  Type *IntPtrTy = getInt8PtrTy(tofree->getContext());

  auto FT = FunctionType::get(VoidTy, {IntPtrTy}, false);
  Value *freevalue = builder.GetInsertBlock()
                         ->getParent()
                         ->getParent()
                         ->getOrInsertFunction(freename, FT)
                         .getCallee();
  CallInst *freecall = cast<CallInst>(CallInst::Create(
      FT, freevalue, {builder.CreatePointerCast(tofree, IntPtrTy)}, "",
      builder.GetInsertBlock()));
  freecall->setDebugLoc(debuglocation);
  if (isa<CallInst>(tofree) &&
      cast<CallInst>(tofree)->getAttributes().hasAttribute(
          AttributeList::ReturnIndex, Attribute::NonNull)) {
    freecall->addAttribute(AttributeList::FirstArgIndex, Attribute::NonNull);
  }
  if (Function *F = dyn_cast<Function>(freevalue))
    freecall->setCallingConv(F->getCallingConv());
  if (freecall->getParent() == nullptr)
    builder.Insert(freecall);
  return freecall;
}

bool GradientUtils::needsCacheWholeAllocation(
    const llvm::Value *origInst) const {
  auto found = knownRecomputeHeuristic.find(origInst);
  if (found == knownRecomputeHeuristic.end())
    return false;
  if (!found->second)
    return true;
  // User, operand of input, whehter the input is the original allocation
  SmallVector<std::tuple<const Instruction *, size_t, bool>, 1> todo;
  for (auto &use : origInst->uses())
    todo.push_back(std::make_tuple(cast<Instruction>(use.getUser()),
                                   use.getOperandNo(), true));
  SmallSet<std::tuple<const Instruction *, size_t, bool>, 1> seen;
  while (todo.size()) {
    auto pair = todo.back();
    auto [cur, idx, orig] = pair;
    todo.pop_back();
    if (seen.count(pair))
      continue;
    seen.insert(pair);
    // Loads are always fine
    if (isa<LoadInst>(cur) || isNVLoad(cur))
      continue;

    if (auto II = dyn_cast<IntrinsicInst>(cur))
      if (II->getIntrinsicID() == Intrinsic::masked_load)
        continue;

    bool returnedSameValue = false;

    if (auto CI = dyn_cast<CallInst>(cur)) {
#if LLVM_VERSION_MAJOR >= 14
      if (idx < CI->arg_size())
#else
      if (idx < CI->getNumArgOperands())
#endif
      {

        // Calling a non-empty function with a julia base object, this is fine.
        // as GC will deal with any issues with.
        if (auto PT = dyn_cast<PointerType>(CI->getArgOperand(idx)->getType()))
          if (PT->getAddressSpace() == 10)
            if (EnzymeJuliaAddrLoad)
              if (auto F = getFunctionFromCall(CI))
                if (!F->empty())
                  continue;

        if (isNoCapture(CI, idx))
          continue;

        if (auto F = CI->getCalledFunction())
          if (F->getCallingConv() == CI->getCallingConv() && !F->empty()) {
            bool onlyReturnUses = true;
            bool hasReturnUse = true;

            if (CI->getFunctionType() != F->getFunctionType() ||
                idx >= F->getFunctionType()->getNumParams()) {
              onlyReturnUses = false;
            } else {
              for (auto u : F->getArg(idx)->users()) {
                if (isa<ReturnInst>(u)) {
                  hasReturnUse = true;
                  continue;
                }
                onlyReturnUses = false;
                continue;
              }
            }
            // The arg itself has no use in the function
            if (onlyReturnUses && !hasReturnUse)
              continue;

            // If this is the original allocation, we return it guaranteed, and
            // cache the return, that's still fine
            if (onlyReturnUses && orig) {
              found = knownRecomputeHeuristic.find(cur);
              if (found == knownRecomputeHeuristic.end())
                continue;

              if (!found->second)
                continue;
              returnedSameValue = true;
            }
          }
      } else {
        // Either this is the called function, or this is a jlrooted, either way
        // continue.
        continue;
      }
    }

    found = knownRecomputeHeuristic.find(cur);
    if (found == knownRecomputeHeuristic.end())
      continue;

    // If caching a julia base object, this is fine as
    // GC will deal with any issues with.
    if (auto PT = dyn_cast<PointerType>(cur->getType()))
      if (PT->getAddressSpace() == 10)
        if (EnzymeJuliaAddrLoad)
          continue;

    // If caching this user, it cannot be a gep/cast of original
    if (!found->second) {
      llvm::errs() << " mod: " << *oldFunc->getParent() << "\n";
      llvm::errs() << " oldFunc: " << *oldFunc << "\n";
      for (auto &pair : knownRecomputeHeuristic)
        llvm::errs() << " krc[" << *pair.first << "] = " << pair.second << "\n";
      llvm::errs() << " cur: " << *cur << "\n";
      llvm::errs() << " origInst: " << *origInst << "\n";
      assert(false && "caching potentially capturing/offset of allocation");
    } else {
      // if not caching this user, it is legal to recompute, consider its users
      for (auto &use : cur->uses()) {
        todo.push_back(std::make_tuple(cast<Instruction>(use.getUser()),
                                       use.getOperandNo(),
                                       returnedSameValue && orig));
      }
    }
  }
  return false;
}

void GradientUtils::replaceAndRemoveUnwrapCacheFor(llvm::Value *A,
                                                   llvm::Value *B) {
  SmallVector<Instruction *, 1> toErase;
  for (auto &pair : unwrap_cache) {
    auto found = pair.second.find(A);
    if (found != pair.second.end()) {
      for (auto &p : found->second) {
        Value *pre = p.second;
        replaceAWithB(pre, B);
        if (auto I = dyn_cast<Instruction>(pre)) {
          toErase.push_back(I);
        }
      }
      pair.second.erase(A);
    }
  }
  for (auto I : toErase) {
    erase(I);
  }
}