GradientUtils.cpp File Reference

#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"
#include "Interfaces/GradientUtils.h"
#include "Dialect/Ops.h"
#include "Interfaces/AutoDiffOpInterface.h"
#include "Interfaces/AutoDiffTypeInterface.h"
#include "Interfaces/CloneFunction.h"
#include "Interfaces/Utils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Dominance.h"
#include "llvm/ADT/BreadthFirstIterator.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Interfaces/ViewLikeInterface.h"

Include dependency graph for GradientUtils.cpp:

Go to the source code of this file.

Macros
#define	getOpFullest(Builder, vtmp, frominst, lookupInst, check)
#define	getOpFull(Builder, vtmp, frominst)
#define	getOpUnchecked(vtmp)
#define	getOp(vtmp)

Functions
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"))
static bool	isPotentialLastLoopValue (llvm::Value *val, const llvm::BasicBlock *loc, const llvm::LoopInfo &LI)
static bool	allNullOrUndef (Value *C, const DataLayout &dl, TypeTree TT)
static bool	addRecsCorrespond (const SCEVAddRecExpr *ar1, const SCEVAddRecExpr *ar2)
	Two add recurrences with the same start, step and trip count take the same sequence of values, so an access indexed by one can stand in for an access indexed by the other at the same iteration – provided the iterations are matched up.
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)
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)
	Perform the corresponding deallocation of tofree, given it was allocated by allocationfn.

Variables
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
SmallVector< unsigned int, 9 >	MD_ToCopy

Macro Definition Documentation

getOp

#define getOp

(

vtmp

)

Value:

  ({                                                                           \
    BasicBlock *parent = scope;                                                \
    if (parent == nullptr)                                                     \
      if (auto originst = dyn_cast<Instruction>(val))                          \
        parent = originst->getParent();                                        \
    getOpFullest(BuilderM, vtmp, parent, parent, true);                        \
  })

Referenced by mlir::enzyme::removalBlockExplore(), and mlir::enzyme::ForLikeEnzymeOpsRemover< FinalClass, OpName >::removeEnzymeOps().

getOpFull

#define getOpFull	(		Builder,
			vtmp,
			frominst )

Value:

  ({                                                                           \
    BasicBlock *parent = scope;                                                \
    if (parent == nullptr)                                                     \
      if (auto originst = dyn_cast<Instruction>(val))                          \
        parent = originst->getParent();                                        \
    getOpFullest(Builder, vtmp, frominst, parent, true);                       \
  })

getOpFullest

#define getOpFullest	(		Builder,
			vtmp,
			frominst,
			lookupInst,
			check )

getOpUnchecked

#define getOpUnchecked

(

vtmp

)

Value:

  ({                                                                           \
    BasicBlock *parent = scope;                                                \
    getOpFullest(BuilderM, vtmp, parent, parent, false);                       \
  })

Function Documentation

addRecsCorrespond()

static bool addRecsCorrespond ( const SCEVAddRecExpr * ar1, const SCEVAddRecExpr * ar2 )

static

Two add recurrences with the same start, step and trip count take the same sequence of values, so an access indexed by one can stand in for an access indexed by the other at the same iteration – provided the iterations are matched up.

That holds for the same loop or for loops that do not contain one another. It does not for nested loops: in the inner loop, the outer recurrence stays fixed while the inner one moves (e.g. x[i] and x[j] in for i, for j).

Definition at line 6826 of file GradientUtils.cpp.

allNullOrUndef()

static bool allNullOrUndef ( Value * C, const DataLayout & dl, TypeTree TT )

static

Definition at line 5408 of file GradientUtils.cpp.

EnzymeFreeInternalAllocations()

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

EnzymeInactiveDynamic()

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

EnzymeLoopInvariantCache()

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

EnzymeMinCutCache()

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

EnzymeNewCache()

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

EnzymePrintDiffUse()

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

EnzymeRegisterReduce()

llvm::cl::opt< bool > EnzymeRegisterReduce	(	"enzyme-register-reduce"	,
		cl::init(false)	,
		cl::Hidden	,
		cl::desc("Reduce the amount of register reduce")	)

EnzymeRematerialize()

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

EnzymeSharedForward()

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

EnzymeSpeculatePHIs()

llvm::cl::opt< bool > EnzymeSpeculatePHIs	(	"enzyme-speculate-phis"	,
		cl::init(false)	,
		cl::Hidden	,
		cl::desc("Speculatively execute phi computations")	)

EnzymeVectorSplitPhi()

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

freeKnownAllocation()

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 )

Perform the corresponding deallocation of tofree, given it was allocated by allocationfn.

Definition at line 10102 of file GradientUtils.cpp.

Referenced by AdjointGenerator::handleKnownCallDerivatives().

isPotentialLastLoopValue()

static bool isPotentialLastLoopValue ( llvm::Value * val, const llvm::BasicBlock * loc, const llvm::LoopInfo & LI )

static

Definition at line 140 of file GradientUtils.cpp.

References getShadowType().

SubTransferHelper()

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 )

Definition at line 9241 of file GradientUtils.cpp.

Referenced by EnzymeGradientUtilsSubTransferHelper(), and AdjointGenerator::visitMemTransferCommon().

Variable Documentation

customCallHandlers

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

Definition at line 79 of file GradientUtils.cpp.

Referenced by EnzymeRegisterCallHandler(), and AdjointGenerator::visitCallInst().

customFwdCallHandlers

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

Definition at line 83 of file GradientUtils.cpp.

Referenced by EnzymeRegisterFwdCallHandler(), and AdjointGenerator::visitCallInst().

MD_ToCopy

SmallVector<unsigned int, 9> MD_ToCopy

Initial value:

= {
    LLVMContext::MD_dbg,
    LLVMContext::MD_tbaa,
    LLVMContext::MD_tbaa_struct,
    LLVMContext::MD_range,
    LLVMContext::MD_nonnull,
    LLVMContext::MD_dereferenceable,
    LLVMContext::MD_dereferenceable_or_null}

Definition at line 131 of file GradientUtils.cpp.

Referenced by AdjointGenerator::visitMemSetCommon().

shadowErasers

StringMap<std::function<CallInst *(IRBuilder<> &, Value *)> > shadowErasers

Definition at line 72 of file GradientUtils.cpp.

Referenced by EnzymeRegisterAllocationHandler().

shadowHandlers

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

Definition at line 71 of file GradientUtils.cpp.

Referenced by EnzymeRegisterAllocationHandler(), AdjointGenerator::handleKnownCallDerivatives(), isAllocationFunction(), and TypeAnalyzer::visitCallBase().