#include "MLIR/Implementations/LoopCheckpointing.h"

Classes
struct	BinomialCacheLayout

struct	PeriodicSchedule

struct	ScaffoldLoop

class	SegmentCacheCreatorGuard

Public Types
using	CacheCreator = std::function<std::pair<Value, Value>(Type)>

Static Public Member Functions
static StringRef	enableCheckpointingAttrName ()

static StringRef	binomialCheckpointingAttrName ()

static StringRef	checkpointPeriodAttrName ()

static void	preserveAttributesButCheckpointing (Operation newOp, Operation oldOp)

static bool	hasBinomialAttr (OpName forOp)

static bool	checkpointingEnabled (OpName forOp)

static bool	supportsDynamicPeriodic ()

static bool	needsCheckpointing (OpName forOp)

static bool	needsBinomialCheckpointing (OpName forOp)

static std::optional< int64_t >	getCheckpointBudget (OpName forOp)

static ValueRange	getInits (OpName forOp)

static Value	getInductionVar (OpName forOp)

static Block *	getBodyBlock (OpName forOp)

static void	setInsertionPointToBodyEnd (OpBuilder &builder, Block *block)

static size_t	getNumRegionIterArgs (OpName forOp)

static OperandRange	getCarriedTerminatorOperands (Operation *term)

static ResultRange	getCarriedResults (Operation *loop)

static Type	getIndexLikeType (OpBuilder &builder)

static Value	emitConst (OpBuilder &b, Location loc, int64_t v)

static Value	emitAdd (OpBuilder &b, Location loc, Value l, Value r)

static Value	emitSub (OpBuilder &b, Location loc, Value l, Value r)

static Value	emitMul (OpBuilder &b, Location loc, Value l, Value r)

static Value	emitDivU (OpBuilder &b, Location loc, Value l, Value r)

static Value	emitMin (OpBuilder &b, Location loc, Value l, Value r)

static Value	getNumIterationsValue (OpBuilder &builder, Location loc, OpName forOp, MGradientUtilsReverse *gutils)

static Value	castToType (OpBuilder &builder, Location loc, Value v, Type targetType)

static MemRefType	checkpointBufferType (int64_t budget, Type t)

static Value	checkpointRow (OpBuilder &b, Location loc, Value buf, Value slot, MemRefType rowTy)

static void	storeCheckpoint (OpBuilder &b, Location loc, Value buf, Value slot, Value val)

static Value	loadCheckpoint (OpBuilder &b, Location loc, Value buf, Value slot, Type valTy)

static bool	publishesCopiedSeeds ()

static bool	isDefinedInBlock (Block *b, Value v)

static void	copyBlockWithoutTerminator (OpBuilder &builder, Block b, MGradientUtilsReverse gutils, IRMapping &mapping)

static void	splitOutsideRefs (OpName forOp, SmallVectorImpl< Value > &mutableRefs, SmallVectorImpl< Value > &immutableRefs)

static bool	storesAreLoopCarried ()

static Value	createStore (OpBuilder &b, Location loc, int64_t budget, Type valTy)

static Value	storeSlot (OpBuilder &b, Location loc, Value store, Value slot, Value val)

static Value	loadSlot (OpBuilder &b, Location loc, Value store, Value slot, Type valTy)

static void	destroyStore (OpBuilder &b, Location loc, Value store)

static ScaffoldLoop	createScaffoldForLoop (OpBuilder &b, Location loc, Value lb, Value ub, Value step, ValueRange inits)

static ScaffoldLoop	createScaffoldWhileLoop (OpBuilder &b, Location loc, ValueRange inits, llvm::function_ref< Value(OpBuilder &, Location, ValueRange)> cond)

static void	finalizeScaffoldLoop (OpBuilder &b, Location loc, const ScaffoldLoop &loop, ValueRange yields)

static Value	emitCmpLT (OpBuilder &b, Location loc, Value l, Value r)

static Value	emitCmpEQ (OpBuilder &b, Location loc, Value l, Value r)

static Value	emitSelect (OpBuilder &b, Location loc, Value c, Value t, Value f)

static SmallVector< Value >	getPrimalResults (OpBuilder &b, Location loc, OpName forOp, ValueRange state, MGradientUtilsReverse *gutils)

static Value	cloneSlot (OpBuilder &builder, Location loc, Value base, Value buf, Value slot, int64_t budget)

static Value	allocCloneSlots (OpBuilder &b, Location loc, int64_t budget, Value proto, ClonableTypeInterface iface)

static void	freeCloneSlots (OpBuilder &builder, Value ref, Value buf, int64_t budget)

static SmallVector< Value >	cacheBinomial (OpName forOp, int64_t budget, MGradientUtilsReverse *gutils)

static void	hoistPlaceholdersBefore (Operation root, Operation before)

static LogicalResult	reverseBinomial (OpName forOp, int64_t budget, OpBuilder &builder, MGradientUtilsReverse *gutils, SmallVector< Value > caches, ArrayRef< bool > operandsActive, ArrayRef< Value > incomingGradients)

static PeriodicSchedule	getStaticPeriodicSchedule (OpName forOp)

static CacheCreator	makeSegmentCacheCreator (MGradientUtilsReverse gutils, Operation anchor)

static void	primeSegmentGradients (OpBuilder &builder, MGradientUtilsReverse gutils, Block origBody, Block *revLoopBody, ArrayRef< bool > operandsActive)

static void	primeStepGradients (OpBuilder &builder, MGradientUtilsReverse gutils, Block origBody, Block *revOuterBody, ArrayRef< bool > operandsActive)

static void	materializeSegmentValues (OpBuilder &builder, Location loc, PeriodicSchedule &sched)

static Value	segmentLength (OpBuilder &builder, Location loc, const PeriodicSchedule &sched, Value base)

static Value	segmentBase (OpBuilder &builder, Location loc, const PeriodicSchedule &sched, Value index)

static void	materializeDynamicSchedule (OpBuilder &builder, Location loc, OpName forOp, MGradientUtilsReverse *gutils, PeriodicSchedule &sched)

static SmallVector< Value >	cachePeriodic (OpName forOp, Operation op, MGradientUtilsReverse gutils)

static LogicalResult	reversePeriodic (OpName forOp, Operation op, OpBuilder &builder, MGradientUtilsReverse gutils, SmallVector< Value > caches, ArrayRef< bool > operandsActive, ArrayRef< Value > incomingGradients)

static std::optional< SmallVector< Value > >	tryCacheValues (OpName forOp, Operation op, MGradientUtilsReverse gutils)

static std::optional< LogicalResult >	tryCreateReverseModeAdjoint (OpName forOp, Operation op, OpBuilder &builder, MGradientUtilsReverse gutils, SmallVector< Value > caches, ArrayRef< bool > operandsActive, ArrayRef< Value > incomingGradients)

Detailed Description

template<typename FinalClass, typename OpName>
struct mlir::enzyme::LoopCheckpointing< FinalClass, OpName >

Definition at line 37 of file LoopCheckpointing.h.

Member Typedef Documentation

◆ CacheCreator

template<typename FinalClass , typename OpName >

using mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::CacheCreator = std::function<std::pair<Value, Value>(Type)>

Definition at line 1374 of file LoopCheckpointing.h.

Member Function Documentation

◆ allocCloneSlots()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::allocCloneSlots	(	OpBuilder &	b,
		Location	loc,
		int64_t	budget,
		Value	proto,
		ClonableTypeInterface	iface )

inlinestatic

Definition at line 592 of file LoopCheckpointing.h.

References clone().

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cacheBinomial().

◆ binomialCheckpointingAttrName()

template<typename FinalClass , typename OpName >

static StringRef mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::binomialCheckpointingAttrName ( )

inlinestatic

Definition at line 85 of file LoopCheckpointing.h.

◆ cacheBinomial()

template<typename FinalClass , typename OpName >

static SmallVector< Value > mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cacheBinomial	(	OpName	forOp,
		int64_t	budget,
		MGradientUtilsReverse *	gutils )

inlinestatic

Definition at line 656 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::allocCloneSlots(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cloneSlot(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::copyBlockWithoutTerminator(), mlir::enzyme::MGradientUtils::erase(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::BinomialCacheLayout::get(), mlir::enzyme::MGradientUtils::getNewFromOriginal(), mlir::enzyme::MGradientUtilsReverse::initAndPushCache(), mlir::enzyme::MGradientUtils::invertPointerM(), mlir::enzyme::MGradientUtils::originalToNewFnOps, and mlir::enzyme::MGradientUtils::replaceOrigOpWith().

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::tryCacheValues().

◆ cachePeriodic()

template<typename FinalClass , typename OpName >

static SmallVector< Value > mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cachePeriodic	(	OpName	forOp,
		Operation *	op,
		MGradientUtilsReverse *	gutils )

inlinestatic

Definition at line 1525 of file LoopCheckpointing.h.

References clone(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::copyBlockWithoutTerminator(), mlir::enzyme::MGradientUtils::erase(), mlir::enzyme::MGradientUtils::getNewFromOriginal(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::hoistPlaceholdersBefore(), mlir::enzyme::MGradientUtilsReverse::initAndPushCache(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::isDynamic(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::materializeDynamicSchedule(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::numItersV, mlir::enzyme::MGradientUtils::originalToNewFnOps, mlir::enzyme::MGradientUtils::replaceOrigOpWith(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::startV, and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::stepV.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::tryCacheValues().

◆ castToType()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::castToType	(	OpBuilder &	builder,
		Location	loc,
		Value	v,
		Type	targetType )

inlinestatic

Definition at line 243 of file LoopCheckpointing.h.

◆ checkpointBufferType()

template<typename FinalClass , typename OpName >

static MemRefType mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointBufferType	(	int64_t	budget,
		Type	t )

inlinestatic

Definition at line 251 of file LoopCheckpointing.h.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::createStore().

◆ checkpointingEnabled()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointingEnabled ( OpName forOp )

inlinestatic

Definition at line 114 of file LoopCheckpointing.h.

◆ checkpointPeriodAttrName()

template<typename FinalClass , typename OpName >

static StringRef mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointPeriodAttrName ( )

inlinestatic

Definition at line 92 of file LoopCheckpointing.h.

◆ checkpointRow()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointRow	(	OpBuilder &	b,
		Location	loc,
		Value	buf,
		Value	slot,
		MemRefType	rowTy )

inlinestatic

Definition at line 262 of file LoopCheckpointing.h.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::loadCheckpoint(), and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::storeCheckpoint().

◆ cloneSlot()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cloneSlot	(	OpBuilder &	builder,
		Location	loc,
		Value	base,
		Value	buf,
		Value	slot,
		int64_t	budget )

inlinestatic

Definition at line 564 of file LoopCheckpointing.h.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cacheBinomial(), and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reverseBinomial().

◆ copyBlockWithoutTerminator()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::copyBlockWithoutTerminator	(	OpBuilder &	builder,
		Block *	b,
		MGradientUtilsReverse *	gutils,
		IRMapping &	mapping )

inlinestatic

Definition at line 337 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::isDefinedInBlock(), mlir::enzyme::MGradientUtils::originalToNewFn, and mlir::enzyme::MGradientUtils::originalToNewFnOps.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cacheBinomial(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cachePeriodic(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reverseBinomial(), and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reversePeriodic().

◆ createScaffoldForLoop()

template<typename FinalClass , typename OpName >

static ScaffoldLoop mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::createScaffoldForLoop	(	OpBuilder &	b,
		Location	loc,
		Value	lb,
		Value	ub,
		Value	step,
		ValueRange	inits )

inlinestatic

Definition at line 504 of file LoopCheckpointing.h.

◆ createScaffoldWhileLoop()

template<typename FinalClass , typename OpName >

static ScaffoldLoop mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::createScaffoldWhileLoop	(	OpBuilder &	b,
		Location	loc,
		ValueRange	inits,
		llvm::function_ref< Value(OpBuilder &, Location, ValueRange)>	cond )

inlinestatic

Definition at line 518 of file LoopCheckpointing.h.

◆ createStore()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::createStore	(	OpBuilder &	b,
		Location	loc,
		int64_t	budget,
		Type	valTy )

inlinestatic

Definition at line 457 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointBufferType().

◆ destroyStore()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::destroyStore	(	OpBuilder &	b,
		Location	loc,
		Value	store )

inlinestatic

Definition at line 475 of file LoopCheckpointing.h.

◆ emitAdd()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitAdd	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 213 of file LoopCheckpointing.h.

◆ emitCmpEQ()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitCmpEQ	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 545 of file LoopCheckpointing.h.

◆ emitCmpLT()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitCmpLT	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 541 of file LoopCheckpointing.h.

◆ emitConst()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitConst	(	OpBuilder &	b,
		Location	loc,
		int64_t	v )

inlinestatic

Definition at line 209 of file LoopCheckpointing.h.

◆ emitDivU()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitDivU	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 225 of file LoopCheckpointing.h.

◆ emitMin()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitMin	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 229 of file LoopCheckpointing.h.

◆ emitMul()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitMul	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 221 of file LoopCheckpointing.h.

◆ emitSelect()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitSelect	(	OpBuilder &	b,
		Location	loc,
		Value	c,
		Value	t,
		Value	f )

inlinestatic

Definition at line 549 of file LoopCheckpointing.h.

◆ emitSub()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::emitSub	(	OpBuilder &	b,
		Location	loc,
		Value	l,
		Value	r )

inlinestatic

Definition at line 217 of file LoopCheckpointing.h.

◆ enableCheckpointingAttrName()

template<typename FinalClass , typename OpName >

static StringRef mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::enableCheckpointingAttrName ( )

inlinestatic

Definition at line 81 of file LoopCheckpointing.h.

◆ finalizeScaffoldLoop()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::finalizeScaffoldLoop	(	OpBuilder &	b,
		Location	loc,
		const ScaffoldLoop &	loop,
		ValueRange	yields )

inlinestatic

Definition at line 535 of file LoopCheckpointing.h.

◆ freeCloneSlots()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::freeCloneSlots	(	OpBuilder &	builder,
		Value	ref,
		Value	buf,
		int64_t	budget )

inlinestatic

Definition at line 623 of file LoopCheckpointing.h.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reverseBinomial().

◆ getBodyBlock()

template<typename FinalClass , typename OpName >

static Block * mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getBodyBlock ( OpName forOp )

inlinestatic

Definition at line 169 of file LoopCheckpointing.h.

◆ getCarriedResults()

template<typename FinalClass , typename OpName >

static ResultRange mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getCarriedResults ( Operation * loop )

inlinestatic

Definition at line 194 of file LoopCheckpointing.h.

◆ getCarriedTerminatorOperands()

template<typename FinalClass , typename OpName >

static OperandRange mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getCarriedTerminatorOperands ( Operation * term )

inlinestatic

Definition at line 187 of file LoopCheckpointing.h.

◆ getCheckpointBudget()

template<typename FinalClass , typename OpName >

static std::optional< int64_t > mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getCheckpointBudget ( OpName forOp )

inlinestatic

Definition at line 146 of file LoopCheckpointing.h.

◆ getIndexLikeType()

template<typename FinalClass , typename OpName >

static Type mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getIndexLikeType ( OpBuilder & builder )

inlinestatic

Definition at line 205 of file LoopCheckpointing.h.

◆ getInductionVar()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getInductionVar ( OpName forOp )

inlinestatic

Definition at line 163 of file LoopCheckpointing.h.

◆ getInits()

template<typename FinalClass , typename OpName >

static ValueRange mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getInits ( OpName forOp )

inlinestatic

Definition at line 161 of file LoopCheckpointing.h.

◆ getNumIterationsValue()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getNumIterationsValue	(	OpBuilder &	builder,
		Location	loc,
		OpName	forOp,
		MGradientUtilsReverse *	gutils )

inlinestatic

Definition at line 233 of file LoopCheckpointing.h.

◆ getNumRegionIterArgs()

template<typename FinalClass , typename OpName >

static size_t mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getNumRegionIterArgs ( OpName forOp )

inlinestatic

Definition at line 179 of file LoopCheckpointing.h.

◆ getPrimalResults()

template<typename FinalClass , typename OpName >

static SmallVector< Value > mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getPrimalResults	(	OpBuilder &	b,
		Location	loc,
		OpName	forOp,
		ValueRange	state,
		MGradientUtilsReverse *	gutils )

inlinestatic

Definition at line 558 of file LoopCheckpointing.h.

◆ getStaticPeriodicSchedule()

template<typename FinalClass , typename OpName >

static PeriodicSchedule mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::getStaticPeriodicSchedule ( OpName forOp )

inlinestatic

Definition at line 1329 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nInner, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nOuter, and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::trailingIters.

◆ hasBinomialAttr()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::hasBinomialAttr ( OpName forOp )

inlinestatic

Definition at line 110 of file LoopCheckpointing.h.

◆ hoistPlaceholdersBefore()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::hoistPlaceholdersBefore	(	Operation *	root,
		Operation *	before )

inlinestatic

Definition at line 884 of file LoopCheckpointing.h.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cachePeriodic().

◆ isDefinedInBlock()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::isDefinedInBlock	(	Block *	b,
		Value	v )

inlinestatic

Definition at line 326 of file LoopCheckpointing.h.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::copyBlockWithoutTerminator().

◆ loadCheckpoint()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::loadCheckpoint	(	OpBuilder &	b,
		Location	loc,
		Value	buf,
		Value	slot,
		Type	valTy )

inlinestatic

Definition at line 296 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointRow().

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::loadSlot().

◆ loadSlot()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::loadSlot	(	OpBuilder &	b,
		Location	loc,
		Value	store,
		Value	slot,
		Type	valTy )

inlinestatic

Definition at line 470 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::loadCheckpoint().

◆ makeSegmentCacheCreator()

template<typename FinalClass , typename OpName >

static CacheCreator mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::makeSegmentCacheCreator	(	MGradientUtilsReverse *	gutils,
		Operation *	anchor )

inlinestatic

Definition at line 1376 of file LoopCheckpointing.h.

◆ materializeDynamicSchedule()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::materializeDynamicSchedule	(	OpBuilder &	builder,
		Location	loc,
		OpName	forOp,
		MGradientUtilsReverse *	gutils,
		PeriodicSchedule &	sched )

inlinestatic

Definition at line 1511 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::numItersV, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::startV, and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::stepV.

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cachePeriodic().

◆ materializeSegmentValues()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::materializeSegmentValues	(	OpBuilder &	builder,
		Location	loc,
		PeriodicSchedule &	sched )

inlinestatic

Definition at line 1453 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::isDynamic(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nInner, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nInnerV, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nOuter, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nOuterV, and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::numItersV.

◆ needsBinomialCheckpointing()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::needsBinomialCheckpointing ( OpName forOp )

inlinestatic

Definition at line 141 of file LoopCheckpointing.h.

◆ needsCheckpointing()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::needsCheckpointing ( OpName forOp )

inlinestatic

Definition at line 126 of file LoopCheckpointing.h.

◆ preserveAttributesButCheckpointing()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::preserveAttributesButCheckpointing	(	Operation *	newOp,
		Operation *	oldOp )

inlinestatic

Definition at line 99 of file LoopCheckpointing.h.

◆ primeSegmentGradients()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::primeSegmentGradients	(	OpBuilder &	builder,
		MGradientUtilsReverse *	gutils,
		Block *	origBody,
		Block *	revLoopBody,
		ArrayRef< bool >	operandsActive )

inlinestatic

Definition at line 1412 of file LoopCheckpointing.h.

References mlir::enzyme::localizeGradients(), and mlir::enzyme::MDiffeGradientUtils::zeroDiffe().

◆ primeStepGradients()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::primeStepGradients	(	OpBuilder &	builder,
		MGradientUtilsReverse *	gutils,
		Block *	origBody,
		Block *	revOuterBody,
		ArrayRef< bool >	operandsActive )

inlinestatic

Definition at line 1431 of file LoopCheckpointing.h.

References mlir::enzyme::localizeGradients(), and mlir::enzyme::MDiffeGradientUtils::zeroDiffe().

◆ publishesCopiedSeeds()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::publishesCopiedSeeds ( )

inlinestatic

Definition at line 321 of file LoopCheckpointing.h.

◆ reverseBinomial()

template<typename FinalClass , typename OpName >

static LogicalResult mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reverseBinomial	(	OpName	forOp,
		int64_t	budget,
		OpBuilder &	builder,
		MGradientUtilsReverse *	gutils,
		SmallVector< Value >	caches,
		ArrayRef< bool >	operandsActive,
		ArrayRef< Value >	incomingGradients )

inlinestatic

Definition at line 895 of file LoopCheckpointing.h.

References mlir::enzyme::MGradientUtilsReverse::addToDiffe(), cached, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cloneSlot(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::copyBlockWithoutTerminator(), mlir::enzyme::MDiffeGradientUtils::diffe(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::freeCloneSlots(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::BinomialCacheLayout::get(), mlir::enzyme::MGradientUtils::getNewFromOriginal(), mlir::enzyme::MGradientUtils::getShadowType(), mlir::enzyme::MGradientUtils::invertedPointers, mlir::enzyme::MGradientUtils::isConstantValue(), mlir::enzyme::MGradientUtils::Logic, mlir::enzyme::MGradientUtilsReverse::popCache(), mlir::enzyme::MEnzymeLogic::visitChild(), and mlir::enzyme::MDiffeGradientUtils::zeroDiffe().

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::tryCreateReverseModeAdjoint().

◆ reversePeriodic()

template<typename FinalClass , typename OpName >

static LogicalResult mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reversePeriodic	(	OpName	forOp,
		Operation *	op,
		OpBuilder &	builder,
		MGradientUtilsReverse *	gutils,
		SmallVector< Value >	caches,
		ArrayRef< bool >	operandsActive,
		ArrayRef< Value >	incomingGradients )

inlinestatic

Definition at line 1646 of file LoopCheckpointing.h.

References mlir::enzyme::MGradientUtilsReverse::addToDiffe(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::copyBlockWithoutTerminator(), mlir::enzyme::MDiffeGradientUtils::diffe(), mlir::enzyme::MGradientUtils::isConstantValue(), mlir::enzyme::MGradientUtils::Logic, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::numItersV, mlir::enzyme::MGradientUtilsReverse::popCache(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::startV, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::stepV, mlir::enzyme::MEnzymeLogic::visitChild(), and mlir::enzyme::MDiffeGradientUtils::zeroDiffe().

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::tryCreateReverseModeAdjoint().

◆ segmentBase()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::segmentBase	(	OpBuilder &	builder,
		Location	loc,
		const PeriodicSchedule &	sched,
		Value	index )

inlinestatic

Definition at line 1503 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nInnerV.

◆ segmentLength()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::segmentLength	(	OpBuilder &	builder,
		Location	loc,
		const PeriodicSchedule &	sched,
		Value	base )

inlinestatic

Definition at line 1475 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::hasTrailing(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::isDynamic(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nInner, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nInnerV, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::nOuter, mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::numItersV, and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::PeriodicSchedule::trailingIters.

◆ setInsertionPointToBodyEnd()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::setInsertionPointToBodyEnd	(	OpBuilder &	builder,
		Block *	block )

inlinestatic

Definition at line 175 of file LoopCheckpointing.h.

◆ splitOutsideRefs()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::splitOutsideRefs	(	OpName	forOp,
		SmallVectorImpl< Value > &	mutableRefs,
		SmallVectorImpl< Value > &	immutableRefs )

inlinestatic

Definition at line 365 of file LoopCheckpointing.h.

◆ storeCheckpoint()

template<typename FinalClass , typename OpName >

static void mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::storeCheckpoint	(	OpBuilder &	b,
		Location	loc,
		Value	buf,
		Value	slot,
		Value	val )

inlinestatic

Definition at line 283 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::checkpointRow().

Referenced by mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::storeSlot().

◆ storesAreLoopCarried()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::storesAreLoopCarried ( )

inlinestatic

Definition at line 455 of file LoopCheckpointing.h.

◆ storeSlot()

template<typename FinalClass , typename OpName >

static Value mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::storeSlot	(	OpBuilder &	b,
		Location	loc,
		Value	store,
		Value	slot,
		Value	val )

inlinestatic

Definition at line 464 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::storeCheckpoint().

◆ supportsDynamicPeriodic()

template<typename FinalClass , typename OpName >

static bool mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::supportsDynamicPeriodic ( )

inlinestatic

Definition at line 124 of file LoopCheckpointing.h.

◆ tryCacheValues()

template<typename FinalClass , typename OpName >

static std::optional< SmallVector< Value > > mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::tryCacheValues	(	OpName	forOp,
		Operation *	op,
		MGradientUtilsReverse *	gutils )

inlinestatic

Definition at line 1830 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cacheBinomial(), mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::cachePeriodic(), mlir::enzyme::MGradientUtils::getNewFromOriginal(), and mlir::enzyme::MGradientUtilsReverse::initAndPushCache().

◆ tryCreateReverseModeAdjoint()

template<typename FinalClass , typename OpName >

static std::optional< LogicalResult > mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::tryCreateReverseModeAdjoint	(	OpName	forOp,
		Operation *	op,
		OpBuilder &	builder,
		MGradientUtilsReverse *	gutils,
		SmallVector< Value >	caches,
		ArrayRef< bool >	operandsActive,
		ArrayRef< Value >	incomingGradients )

inlinestatic

Definition at line 1882 of file LoopCheckpointing.h.

References mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reverseBinomial(), and mlir::enzyme::LoopCheckpointing< FinalClass, OpName >::reversePeriodic().

The documentation for this struct was generated from the following file:

MLIR/Implementations/LoopCheckpointing.h

Classes

Public Types

Static Public Member Functions

Detailed Description

Member Typedef Documentation

◆ CacheCreator

Member Function Documentation

◆ allocCloneSlots()

◆ binomialCheckpointingAttrName()

◆ cacheBinomial()

◆ cachePeriodic()

◆ castToType()

◆ checkpointBufferType()

◆ checkpointingEnabled()

◆ checkpointPeriodAttrName()

◆ checkpointRow()

◆ cloneSlot()

◆ copyBlockWithoutTerminator()

◆ createScaffoldForLoop()

◆ createScaffoldWhileLoop()

◆ createStore()

◆ destroyStore()

◆ emitAdd()

◆ emitCmpEQ()

◆ emitCmpLT()

◆ emitConst()

◆ emitDivU()

◆ emitMin()

◆ emitMul()

◆ emitSelect()

◆ emitSub()

◆ enableCheckpointingAttrName()

◆ finalizeScaffoldLoop()

◆ freeCloneSlots()

◆ getBodyBlock()

◆ getCarriedResults()

◆ getCarriedTerminatorOperands()

◆ getCheckpointBudget()

◆ getIndexLikeType()

◆ getInductionVar()

◆ getInits()

◆ getNumIterationsValue()

◆ getNumRegionIterArgs()

◆ getPrimalResults()

◆ getStaticPeriodicSchedule()

◆ hasBinomialAttr()

◆ hoistPlaceholdersBefore()

◆ isDefinedInBlock()

◆ loadCheckpoint()

◆ loadSlot()

◆ makeSegmentCacheCreator()

◆ materializeDynamicSchedule()

◆ materializeSegmentValues()

◆ needsBinomialCheckpointing()

◆ needsCheckpointing()

◆ preserveAttributesButCheckpointing()

◆ primeSegmentGradients()

◆ primeStepGradients()

◆ publishesCopiedSeeds()

◆ reverseBinomial()

◆ reversePeriodic()

◆ segmentBase()

◆ segmentLength()

◆ setInsertionPointToBodyEnd()

◆ splitOutsideRefs()

◆ storeCheckpoint()

◆ storesAreLoopCarried()

◆ storeSlot()

◆ supportsDynamicPeriodic()

◆ tryCacheValues()

◆ tryCreateReverseModeAdjoint()