mlir::enzyme::batchutils Namespace Reference

Classes
struct	BatchCacheKey
struct	BatchDiffCacheKey

Functions
template<typename SourceOp >
BatchDiffCacheKey	createDiffCacheKey (SourceOp uop, FunctionOpInterface fn)
template<typename SourceOp , std::enable_if_t< llvm::is_one_of< SourceOp, ForwardDiffOp, AutoDiffOp >::value, bool > = true>
SmallVector< MemoryEffects::EffectInstance >	findCallerEffects (SourceOp callerOp, FunctionOpInterface innerFnOp, const SmallVector< MemoryEffects::EffectInstance > &innerEffects)
template<typename SourceOp , std::enable_if_t< llvm::is_one_of< SourceOp, ForwardDiffOp, AutoDiffOp >::value, bool > = true>
llvm::SmallVector< SourceOp, 2 >	pruneGradDefs (BatchDiffCacheKey &key, SmallVector< SourceOp > &allDiffs)
template<typename SourceOp , std::enable_if_t< llvm::is_one_of< SourceOp, ForwardDiffOp, AutoDiffOp >::value, bool > = true>
llvm::SmallVector< SourceOp >	pruneMemoryEffects (SymbolTableCollection &symbolTable, BatchDiffCacheKey &key, SmallVector< SourceOp > &prunedSources, DenseMap< SourceOp, SmallVector< MemoryEffects::EffectInstance > > &callerEffectMap, llvm::DenseMap< FunctionOpInterface, SmallVector< MemoryEffects::EffectInstance > > &innerEffectCache)
mlir::TensorType	applyBatchSizes (mlir::Type Ty, llvm::ArrayRef< int64_t > batchSizes)
LogicalResult	handleCallOp (func::CallOp callOp, OpBuilder &builder, IRMapping &mapper, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
void	batchCloneBlock (OpBuilder &builder, Block *blk, IRMapping &mapper, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache, bool withoutTerminator)
void	batchCloneRegion (OpBuilder &builder, Region *src, Region *dest, IRMapping &mapper, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
FunctionOpInterface	batchCloneFunction (OpBuilder &builder, FunctionOpInterface F, Twine name, llvm::ArrayRef< int64_t > batchSizes, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
FunctionOpInterface	batchOperationWithoutInsertingCallOp (OpBuilder &builder, T CI, FunctionOpInterface fn, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	batchOperation (SymbolTableCollection &symbolTable, OpBuilder &builder, T CI, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	batchOperation (SymbolTableCollection &symbolTable, PatternRewriter &rewriter, T CI, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	batchOperation (OpBuilder &builder, T CI, FunctionOpInterface fn, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
template<typename T >
LogicalResult	batchOperation (PatternRewriter &rewriter, T CI, FunctionOpInterface fn, std::map< BatchCacheKey, FunctionOpInterface > &batchedFunctionCache)
void	batchOperationInline (PatternRewriter &rewriter, enzyme::BatchOp batchOp, FunctionOpInterface func)

Function Documentation

applyBatchSizes()

mlir::TensorType mlir::enzyme::batchutils::applyBatchSizes	(	mlir::Type	Ty,
		llvm::ArrayRef< int64_t >	batchSizes )

Definition at line 41 of file EnzymeBatchPass.cpp.

Referenced by batchCloneBlock(), batchCloneFunction(), and batchCloneRegion().

batchCloneBlock()

void mlir::enzyme::batchutils::batchCloneBlock	(	OpBuilder &	builder,
		Block *	blk,
		IRMapping &	mapper,
		llvm::ArrayRef< int64_t >	batchSizes,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache,
		bool	withoutTerminator )

Definition at line 100 of file EnzymeBatchPass.cpp.

References applyBatchSizes(), batchCloneRegion(), and handleCallOp().

Referenced by batchCloneRegion(), and batchOperationInline().

batchCloneFunction()

FunctionOpInterface mlir::enzyme::batchutils::batchCloneFunction	(	OpBuilder &	builder,
		FunctionOpInterface	F,
		Twine	name,
		llvm::ArrayRef< int64_t >	batchSizes,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 189 of file EnzymeBatchPass.cpp.

References applyBatchSizes(), and batchCloneRegion().

Referenced by batchOperationWithoutInsertingCallOp(), and handleCallOp().

batchCloneRegion()

void mlir::enzyme::batchutils::batchCloneRegion	(	OpBuilder &	builder,
		Region *	src,
		Region *	dest,
		IRMapping &	mapper,
		llvm::ArrayRef< int64_t >	batchSizes,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 163 of file EnzymeBatchPass.cpp.

References applyBatchSizes(), and batchCloneBlock().

Referenced by batchCloneBlock(), and batchCloneFunction().

batchOperation() [1/4]

template<typename T >

LogicalResult mlir::enzyme::batchutils::batchOperation	(	OpBuilder &	builder,
		T	CI,
		FunctionOpInterface	fn,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 80 of file EnzymeBatchPass.h.

References batchOperationWithoutInsertingCallOp().

batchOperation() [2/4]

template<typename T >

LogicalResult mlir::enzyme::batchutils::batchOperation	(	PatternRewriter &	rewriter,
		T	CI,
		FunctionOpInterface	fn,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 101 of file EnzymeBatchPass.h.

References batchOperationWithoutInsertingCallOp().

batchOperation() [3/4]

template<typename T >

LogicalResult mlir::enzyme::batchutils::batchOperation	(	SymbolTableCollection &	symbolTable,
		OpBuilder &	builder,
		T	CI,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 61 of file EnzymeBatchPass.h.

References batchOperation().

Referenced by batchOperation(), and batchOperation().

batchOperation() [4/4]

template<typename T >

LogicalResult mlir::enzyme::batchutils::batchOperation	(	SymbolTableCollection &	symbolTable,
		PatternRewriter &	rewriter,
		T	CI,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 71 of file EnzymeBatchPass.h.

References batchOperation().

batchOperationInline()

void mlir::enzyme::batchutils::batchOperationInline ( PatternRewriter & rewriter, enzyme::BatchOp batchOp, FunctionOpInterface func )

inline

Definition at line 121 of file EnzymeBatchPass.h.

References batchCloneBlock().

batchOperationWithoutInsertingCallOp()

template<typename T >

FunctionOpInterface mlir::enzyme::batchutils::batchOperationWithoutInsertingCallOp	(	OpBuilder &	builder,
		T	CI,
		FunctionOpInterface	fn,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 239 of file EnzymeBatchPass.cpp.

References batchCloneFunction().

Referenced by batchOperation(), and batchOperation().

createDiffCacheKey()

template<typename SourceOp >

BatchDiffCacheKey mlir::enzyme::batchutils::createDiffCacheKey	(	SourceOp	uop,
		FunctionOpInterface	fn )

Definition at line 65 of file EnzymeBatchDiffPass.h.

findCallerEffects()

template<typename SourceOp , std::enable_if_t< llvm::is_one_of< SourceOp, ForwardDiffOp, AutoDiffOp >::value, bool > = true>

SmallVector< MemoryEffects::EffectInstance > mlir::enzyme::batchutils::findCallerEffects	(	SourceOp	callerOp,
		FunctionOpInterface	innerFnOp,
		const SmallVector< MemoryEffects::EffectInstance > &	innerEffects )

Definition at line 101 of file EnzymeBatchDiffPass.h.

References mlir::enzyme::oputils::getEffectOfVal().

Referenced by pruneMemoryEffects().

handleCallOp()

LogicalResult mlir::enzyme::batchutils::handleCallOp	(	func::CallOp	callOp,
		OpBuilder &	builder,
		IRMapping &	mapper,
		llvm::ArrayRef< int64_t >	batchSizes,
		std::map< BatchCacheKey, FunctionOpInterface > &	batchedFunctionCache )

Definition at line 54 of file EnzymeBatchPass.cpp.

References batchCloneFunction().

Referenced by batchCloneBlock().

pruneGradDefs()

template<typename SourceOp , std::enable_if_t< llvm::is_one_of< SourceOp, ForwardDiffOp, AutoDiffOp >::value, bool > = true>

llvm::SmallVector< SourceOp, 2 > mlir::enzyme::batchutils::pruneGradDefs	(	BatchDiffCacheKey &	key,
		SmallVector< SourceOp > &	allDiffs )

Definition at line 184 of file EnzymeBatchDiffPass.h.

pruneMemoryEffects()

template<typename SourceOp , std::enable_if_t< llvm::is_one_of< SourceOp, ForwardDiffOp, AutoDiffOp >::value, bool > = true>

llvm::SmallVector< SourceOp > mlir::enzyme::batchutils::pruneMemoryEffects	(	SymbolTableCollection &	symbolTable,
		BatchDiffCacheKey &	key,
		SmallVector< SourceOp > &	prunedSources,
		DenseMap< SourceOp, SmallVector< MemoryEffects::EffectInstance > > &	callerEffectMap,
		llvm::DenseMap< FunctionOpInterface, SmallVector< MemoryEffects::EffectInstance > > &	innerEffectCache )

Definition at line 225 of file EnzymeBatchDiffPass.h.

References mlir::enzyme::oputils::collectFnEffects(), mlir::enzyme::oputils::collectOpEffects(), findCallerEffects(), and mlir::enzyme::oputils::mayAlias().