FunctionUtils.h

Go to the documentation of this file.

//===- FunctionUtils.h - Declaration of function utilities ---------------===//
//
//                             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 declares utilities on LLVM Functions that are used as part of the
// AD process.
//
//===----------------------------------------------------------------------===//
#ifndef ENZYME_FUNCTION_UTILS_H
#define ENZYME_FUNCTION_UTILS_H
 
#include <deque>
#include <set>
 
#include <llvm/Config/llvm-config.h>
 
#if LLVM_VERSION_MAJOR >= 16
#define private public
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
#undef private
#else
#include "SCEV/ScalarEvolution.h"
#include "SCEV/ScalarEvolutionExpander.h"
#endif
 
#include "Utils.h"
 
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/LoopAnalysisManager.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
 
#include "llvm/IR/Function.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
 
#include "llvm/IR/Instructions.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
 
#include "llvm/ADT/STLExtras.h"
 
//;
 
extern "C" {
extern llvm::cl::opt<bool> EnzymeAlwaysInlineDiff;
}
 
// Perform an analysis to detect functions which only write to visible memory
// outside the function if an error is not throw. Such a function can touch
// inaccessible memory [e.g. the insides of malloc/etc], and the only violation
// is whether existing memory before the call is written to.
// If non-local, returning memory written to is a violation (since it writes to
// externally visible memory).
// If local, returning memory written to is fine (since existing memory before
// the call remains unchanged).
// In other words, malloc [local and non-local], calloc [local and non-local],
// copy_array [local only], and friends, are all considered
// readonly_or_throw, as they only either read externally visible state, throw
// an error, or write to inaccesible memory.
bool DetectReadonlyOrThrow(llvm::Module &M);
 
class PreProcessCache {
public:
  PreProcessCache();
  PreProcessCache(PreProcessCache &) = delete;
  // Using the default move constructor will botch the FAM/MAM proxy passes
  // since now the new location of FAM/MAM will not be used. Therefore, use a
  // custom move constructor and default initialize these, and move the
  // cache/origin maps.
  PreProcessCache(PreProcessCache &&prev) : PreProcessCache() {
    cache = std::move(prev.cache);
    CloneOrigin = std::move(prev.CloneOrigin);
  };
 
  llvm::LoopAnalysisManager LAM;
  llvm::FunctionAnalysisManager FAM;
  llvm::ModuleAnalysisManager MAM;
 
  std::map<std::pair<llvm::Function *, DerivativeMode>, llvm::Function *> cache;
  std::map<llvm::Function *, llvm::Function *> CloneOrigin;
 
  llvm::Function *preprocessForClone(llvm::Function *F, DerivativeMode mode);
 
  llvm::AAResults &getAAResultsFromFunction(llvm::Function *NewF);
 
  llvm::Function *CloneFunctionWithReturns(
      DerivativeMode mode, unsigned width, llvm::Function *&F,
      llvm::ValueToValueMapTy &ptrInputs,
      llvm::ArrayRef<DIFFE_TYPE> constant_args,
      llvm::SmallPtrSetImpl<llvm::Value *> &constants,
      llvm::SmallPtrSetImpl<llvm::Value *> &nonconstant,
      llvm::SmallPtrSetImpl<llvm::Value *> &returnvals, bool returnTape,
      bool returnPrimal, bool returnShadow, const llvm::Twine &name,
      llvm::ValueMap<const llvm::Value *, AssertingReplacingVH> *VMapO,
      bool diffeReturnArg, llvm::Type *additionalArg = nullptr);
 
  void ReplaceReallocs(llvm::Function *NewF, bool mem2reg = false);
  void LowerAllocAddr(llvm::Function *NewF);
  void AlwaysInline(llvm::Function *NewF);
  void optimizeIntermediate(llvm::Function *F);
 
  void clear();
};
 
class GradientUtils;
 
static inline void
getExitBlocks(const llvm::Loop *L,
              llvm::SmallPtrSetImpl<llvm::BasicBlock *> &ExitBlocks) {
  llvm::SmallVector<llvm::BasicBlock *, 8> PotentialExitBlocks;
  L->getExitBlocks(PotentialExitBlocks);
  for (auto a : PotentialExitBlocks) {
 
    llvm::SmallVector<llvm::BasicBlock *, 4> tocheck;
    llvm::SmallPtrSet<llvm::BasicBlock *, 4> checked;
    tocheck.push_back(a);
 
    bool isExit = false;
 
    while (tocheck.size()) {
      auto foo = tocheck.back();
      tocheck.pop_back();
      if (checked.count(foo)) {
        isExit = true;
        goto exitblockcheck;
      }
      checked.insert(foo);
      if (isAnyBranch(foo->getTerminator())) {
        auto *bTerm = foo->getTerminator();
        // Instruction::successors() postdates some supported majors; index.
        for (unsigned i = 0, e = bTerm->getNumSuccessors(); i < e; ++i) {
          auto *nb = bTerm->getSuccessor(i);
          if (L->contains(nb))
            continue;
          tocheck.push_back(nb);
        }
      } else if (llvm::isa<llvm::UnreachableInst>(foo->getTerminator())) {
        continue;
      } else {
        isExit = true;
        goto exitblockcheck;
      }
    }
 
  exitblockcheck:
    if (isExit) {
      ExitBlocks.insert(a);
    }
  }
}
 
static inline llvm::SmallVector<llvm::BasicBlock *, 3>
getLatches(const llvm::Loop *L,
           const llvm::SmallPtrSetImpl<llvm::BasicBlock *> &ExitBlocks) {
  llvm::BasicBlock *Preheader = L->getLoopPreheader();
  if (!Preheader) {
    llvm::errs() << *L->getHeader()->getParent() << "\n";
    llvm::errs() << *L->getHeader() << "\n";
    llvm::errs() << *L << "\n";
  }
  assert(Preheader && "requires preheader");
 
  // Find latch, defined as a (perhaps unique) block in loop that branches to
  // exit block
  llvm::SmallVector<llvm::BasicBlock *, 3> Latches;
  for (llvm::BasicBlock *ExitBlock : ExitBlocks) {
    for (llvm::BasicBlock *pred : llvm::predecessors(ExitBlock)) {
      if (L->contains(pred)) {
        if (std::find(Latches.begin(), Latches.end(), pred) != Latches.end())
          continue;
        Latches.push_back(pred);
      }
    }
  }
  return Latches;
}
 
// TODO note this doesn't go through [loop, unreachable], and we could get more
// performance by doing this can consider doing some domtree magic potentially
static inline llvm::SmallPtrSet<llvm::BasicBlock *, 4>
getGuaranteedUnreachable(llvm::Function *F) {
  llvm::SmallPtrSet<llvm::BasicBlock *, 4> knownUnreachables;
  if (F->empty())
    return knownUnreachables;
  std::deque<llvm::BasicBlock *> todo;
  for (auto &BB : *F) {
    todo.push_back(&BB);
  }
 
  while (!todo.empty()) {
    llvm::BasicBlock *next = todo.front();
    todo.pop_front();
 
    if (knownUnreachables.find(next) != knownUnreachables.end())
      continue;
 
    if (llvm::isa<llvm::ReturnInst>(next->getTerminator()))
      continue;
 
    if (llvm::isa<llvm::UnreachableInst>(next->getTerminator())) {
      knownUnreachables.insert(next);
      for (llvm::BasicBlock *Pred : predecessors(next)) {
        todo.push_back(Pred);
      }
      continue;
    }
 
    // Assume resumes don't happen
    // TODO consider EH
    if (llvm::isa<llvm::ResumeInst>(next->getTerminator())) {
      knownUnreachables.insert(next);
      for (llvm::BasicBlock *Pred : predecessors(next)) {
        todo.push_back(Pred);
      }
      continue;
    }
 
    bool unreachable = true;
    for (llvm::BasicBlock *Succ : llvm::successors(next)) {
      if (knownUnreachables.find(Succ) == knownUnreachables.end()) {
        unreachable = false;
        break;
      }
    }
 
    if (!unreachable)
      continue;
    knownUnreachables.insert(next);
    for (llvm::BasicBlock *Pred : llvm::predecessors(next)) {
      todo.push_back(Pred);
    }
    continue;
  }
 
  return knownUnreachables;
}
 
enum class UseReq {
  Need,
  Recur,
  Cached,
};
static inline void calculateUnusedValues(
    const llvm::Function &oldFunc,
    llvm::SmallPtrSetImpl<const llvm::Value *> &unnecessaryValues,
    llvm::SmallPtrSetImpl<const llvm::Instruction *> &unnecessaryInstructions,
    bool returnValue, llvm::function_ref<bool(const llvm::Value *)> valneeded,
    llvm::function_ref<UseReq(const llvm::Instruction *)> instneeded,
    llvm::function_ref<bool(const llvm::Instruction *, const llvm::Value *)>
        useneeded) {
 
  std::deque<const llvm::Instruction *> todo;
 
  for (const llvm::BasicBlock &BB : oldFunc) {
    if (auto ri = llvm::dyn_cast<llvm::ReturnInst>(BB.getTerminator())) {
      if (!returnValue) {
        unnecessaryInstructions.insert(ri);
      }
      unnecessaryValues.insert(ri);
    }
    for (auto &inst : BB) {
      if (&inst == BB.getTerminator())
        continue;
      todo.push_back(&inst);
    }
  }
 
  while (!todo.empty()) {
    auto inst = todo.front();
    todo.pop_front();
 
    if (unnecessaryInstructions.count(inst)) {
      assert(unnecessaryValues.count(inst));
      continue;
    }
 
    if (!unnecessaryValues.count(inst)) {
 
      if (valneeded(inst)) {
        continue;
      }
 
      bool necessaryUse = false;
 
      llvm::SmallPtrSet<const llvm::Instruction *, 4> seen;
      std::deque<const llvm::Instruction *> users;
 
      for (auto user_dtx : inst->users()) {
        if (auto cst = llvm::dyn_cast<llvm::Instruction>(user_dtx)) {
          if (useneeded(cst, inst))
            users.push_back(cst);
        }
      }
 
      while (users.size()) {
        auto val = users.front();
        users.pop_front();
 
        if (seen.count(val))
          continue;
        seen.insert(val);
 
        if (unnecessaryInstructions.count(val))
          continue;
 
        switch (instneeded(val)) {
        case UseReq::Need:
          necessaryUse = true;
          break;
        case UseReq::Recur:
          for (auto user_dtx : val->users()) {
            if (auto cst = llvm::dyn_cast<llvm::Instruction>(user_dtx)) {
              if (useneeded(cst, val))
                users.push_back(cst);
            }
          }
          break;
        case UseReq::Cached:
          break;
        }
        if (necessaryUse)
          break;
      }
 
      if (necessaryUse)
        continue;
 
      unnecessaryValues.insert(inst);
 
      for (auto user : inst->users()) {
        if (auto usedinst = llvm::dyn_cast<llvm::Instruction>(user))
          todo.push_back(usedinst);
      }
    }
 
    if (instneeded(inst) == UseReq::Need)
      continue;
 
    unnecessaryInstructions.insert(inst);
 
    for (auto &operand : inst->operands()) {
      if (auto usedinst = llvm::dyn_cast<llvm::Instruction>(operand.get())) {
        todo.push_back(usedinst);
      }
    }
  }
 
  if (false && endsWith(oldFunc.getName(), "subfn")) {
    llvm::errs() << "Prepping values for: " << oldFunc.getName()
                 << " returnValue: " << returnValue << "\n";
    for (auto v : unnecessaryInstructions) {
      llvm::errs() << "+ unnecessaryInstructions: " << *v << "\n";
    }
    for (auto v : unnecessaryValues) {
      llvm::errs() << "+ unnecessaryValues: " << *v << "\n";
    }
    llvm::errs() << "</end>\n";
  }
}
 
static inline void calculateUnusedStores(
    const llvm::Function &oldFunc,
    llvm::SmallPtrSetImpl<const llvm::Instruction *> &unnecessaryStores,
    llvm::function_ref<bool(const llvm::Instruction *)> needStore) {
 
  std::deque<const llvm::Instruction *> todo;
 
  for (const llvm::BasicBlock &BB : oldFunc) {
    for (auto &inst : BB) {
      if (&inst == BB.getTerminator())
        continue;
      todo.push_back(&inst);
    }
  }
 
  while (!todo.empty()) {
    auto inst = todo.front();
    todo.pop_front();
 
    if (unnecessaryStores.count(inst)) {
      continue;
    }
 
    if (needStore(inst))
      continue;
 
    unnecessaryStores.insert(inst);
  }
}
 
void RecursivelyReplaceAddressSpace(llvm::Value *AI, llvm::Value *rep,
                                    bool legal);
 
void ReplaceFunctionImplementation(llvm::Module &M);
 
/// Is the use of value val as an argument of call CI potentially captured
bool couldFunctionArgumentCapture(llvm::CallInst *CI, llvm::Value *val);
 
llvm::FunctionType *getFunctionTypeForClone(
    llvm::FunctionType *FTy, DerivativeMode mode, unsigned width,
    llvm::Type *additionalArg, llvm::ArrayRef<DIFFE_TYPE> constant_args,
    bool diffeReturnArg, bool returnTape, bool returnPrimal, bool returnShadow);
 
/// Lower __enzyme_todense, returning if changed.
bool LowerSparsification(llvm::Function *F, bool replaceAll = true);
 
#endif