mlir::enzyme Namespace Reference

Namespaces
namespace	batchutils
namespace	detail
namespace	oputils
namespace	transforms

Classes
class	ActivityAnalyzer
	Helper class to analyze the differential activity. More...
class	ActivityPrinterConfig
class	AliasAnalysis
	This analysis implements interprocedural alias analysis. More...
class	AliasClassLattice
class	BackwardActivityAnnotationAnalysis
class	BackwardOriginsLattice
class	BackwardOriginsMap
struct	CacheInfo
	Information about a cache, each cache init should have one corresponding push and pop. More...
class	DenseActivityAnnotationAnalysis
class	DenseBackwardActivityAnnotationAnalysis
struct	ForLikeEnzymeOpsRemover
class	ForwardActivityAnnotationAnalysis
class	ForwardOriginsLattice
class	ForwardOriginsMap
struct	IfLikeEnzymeOpsRemover
class	MapOfSetsLattice
class	MDiffeGradientUtils
class	MEnzymeLogic
class	MFnTypeInfo
class	MGradientUtils
class	MGradientUtilsReverse
class	MTypeAnalysis
class	MTypeResults
class	OriginalClasses
	Alias classes for freshly created, e.g., allocated values. More...
class	PointsToPointerAnalysis
class	PointsToSets
class	SetLattice
class	SparseSetLattice
class	Utils

Typedefs
using	ValueOriginSet = SetLattice<OriginAttr>
using	AliasClassSet = SetLattice<DistinctAttr>
	A set of alias class identifiers to be treated as a single union.
typedef void	buildReturnFunction(OpBuilder &, mlir::Block *)

Enumerations
enum class	LoopCacheType { TENSOR , MEMREF }

Functions
void	runActivityAnnotations (FunctionOpInterface callee, ArrayRef< enzyme::Activity > argActivities={}, const ActivityPrinterConfig &config=ActivityPrinterConfig())
void	runDataFlowActivityAnalysis (FunctionOpInterface callee, ArrayRef< enzyme::Activity > argumentActivity, bool print=false, bool verbose=false, bool annotate=false)
bool	sortAttributes (Attribute a, Attribute b)
	Used when serializing to ensure a consistent order.
bool	sortArraysLexicographic (ArrayAttr a, ArrayAttr b)
template<typename OpTy >
void	registerAutoDiffUsingControlFlowInterface (MLIRContext &context)
template<typename OpTy >
void	registerAutoDiffUsingBranchInterface (MLIRContext &context)
template<typename OpTy >
void	registerAutoDiffUsingRegionTerminatorInterface (MLIRContext &context)
template<typename OpTy >
void	registerAutoDiffUsingReturnInterface (MLIRContext &context)
template<typename OpTy , int... storedvals>
void	registerAutoDiffUsingMemoryIdentityInterface (MLIRContext &context)
template<typename OpTy >
void	registerAutoDiffUsingAllocationInterface (MLIRContext &context)
void	registerAffineDialectAutoDiffInterface (DialectRegistry &registry)
void	registerArithDialectAutoDiffInterface (DialectRegistry &registry)
void	registerBuiltinDialectAutoDiffInterface (DialectRegistry &registry)
void	registerLLVMDialectAutoDiffInterface (DialectRegistry &registry)
void	registerLLVMExtDialectAutoDiffInterface (DialectRegistry &registry)
void	registerNVVMDialectAutoDiffInterface (DialectRegistry &registry)
void	registerMemRefDialectAutoDiffInterface (DialectRegistry &registry)
void	registerComplexDialectAutoDiffInterface (DialectRegistry &registry)
void	registerSCFDialectAutoDiffInterface (DialectRegistry &registry)
void	registerCFDialectAutoDiffInterface (DialectRegistry &registry)
void	registerLinalgDialectAutoDiffInterface (DialectRegistry &registry)
void	registerMathDialectAutoDiffInterface (DialectRegistry &registry)
void	registerFuncDialectAutoDiffInterface (DialectRegistry &registry)
void	registerTensorDialectAutoDiffInterface (DialectRegistry &registry)
void	registerEnzymeDialectAutoDiffInterface (DialectRegistry &registry)
void	registerCoreDialectAutodiffInterfaces (DialectRegistry &registry)
mlir::TypedAttr	getConstantAttr (mlir::Type type, llvm::StringRef value)
bool	inlineAutodiffOp (enzyme::AutoDiffOp &op, RewriterBase &rewriter, SymbolTableCollection &symbolTable)
void	minCutCache (Block *forward, Block *reverse, SmallVector< CacheInfo > &caches, PatternRewriter &rewriter, const IRMapping &fwdrevmap, Operation *lastFwd=nullptr)
static LoopCacheType	getCacheType (Operation *op)
static bool	hasMinCut (Operation *op)
static bool	hasLICM (Operation *op)
void	localizeGradients (OpBuilder &builder, MGradientUtilsReverse *gutils, Block *fwd)
void	removalBlockExplore (Block *block, IRMapping &mapping, PatternRewriter &rewriter, llvm::SetVector< Value > &gradients, llvm::MapVector< Value, CacheInfo > &caches)
bool	blockCmp (mlir::Block *a, mlir::Block *b)
bool	opCmp (mlir::Operation *a, mlir::Operation *b)
bool	regionCmp (mlir::Region *a, mlir::Region *b)
bool	valueCmp (mlir::Value a, mlir::Value b)
Type	getConcatType (Value val, int64_t width)
Value	getConcatValue (OpBuilder &builder, Location loc, ArrayRef< Value > argList)
Value	getExtractValue (OpBuilder &builder, Location loc, Type argTy, Value val, int64_t index)
void	computeAffineIndices (OpBuilder &builder, Location loc, AffineMap map, ValueRange operands, SmallVectorImpl< Value > &indices)

Variables
constexpr llvm::StringLiteral	undefinedSetString = "<undefined>"
constexpr llvm::StringLiteral	unknownSetString = "<unknown>"
static constexpr llvm::StringLiteral	kPreserveCacheAttrName = "preserve_cache"

Typedef Documentation

AliasClassSet

using mlir::enzyme::AliasClassSet = SetLattice<DistinctAttr>

A set of alias class identifiers to be treated as a single union.

May be marked as "unknown", which is a conservative pessimistic state, or as "undefined", which is a "not-yet-analyzed" initial state. Undefined state is different from an empty alias set.

Definition at line 47 of file DataFlowAliasAnalysis.h.

buildReturnFunction

typedef void mlir::enzyme::buildReturnFunction(OpBuilder &, mlir::Block *)

Definition at line 13 of file EnzymeLogic.h.

ValueOriginSet

using mlir::enzyme::ValueOriginSet = SetLattice<OriginAttr>

Definition at line 17 of file ActivityAnnotations.h.

Enumeration Type Documentation

LoopCacheType

enum class mlir::enzyme::LoopCacheType

strong

Enumerator
TENSOR
MEMREF

Definition at line 72 of file RemovalUtils.h.

Function Documentation

blockCmp()

bool mlir::enzyme::blockCmp	(	mlir::Block *	a,
		mlir::Block *	b )

Referenced by valueCmp().

computeAffineIndices()

void mlir::enzyme::computeAffineIndices	(	OpBuilder &	builder,
		Location	loc,
		AffineMap	map,
		ValueRange	operands,
		SmallVectorImpl< Value > &	indices )

Definition at line 171 of file Utils.cpp.

getCacheType()

static LoopCacheType mlir::enzyme::getCacheType ( Operation * op )

static

Definition at line 74 of file RemovalUtils.h.

References MEMREF, and TENSOR.

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

getConcatType()

Type mlir::enzyme::getConcatType	(	Value	val,
		int64_t	width )

Definition at line 133 of file Utils.cpp.

Referenced by getConcatValue().

getConcatValue()

Value mlir::enzyme::getConcatValue	(	OpBuilder &	builder,
		Location	loc,
		ArrayRef< Value >	argList )

Definition at line 155 of file Utils.cpp.

References getConcatType().

Referenced by mlir::enzyme::detail::memoryIdentityForwardHandler().

getConstantAttr()

mlir::TypedAttr mlir::enzyme::getConstantAttr	(	mlir::Type	type,
		llvm::StringRef	value )

Definition at line 26 of file CoreDialectsAutoDiffImplementations.cpp.

getExtractValue()

Value mlir::enzyme::getExtractValue	(	OpBuilder &	builder,
		Location	loc,
		Type	argTy,
		Value	val,
		int64_t	index )

Definition at line 163 of file Utils.cpp.

Referenced by mlir::enzyme::detail::memoryIdentityForwardHandler().

hasLICM()

static bool mlir::enzyme::hasLICM ( Operation * op )

static

Definition at line 86 of file RemovalUtils.h.

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

hasMinCut()

static bool mlir::enzyme::hasMinCut ( Operation * op )

static

Definition at line 82 of file RemovalUtils.h.

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

inlineAutodiffOp()

bool mlir::enzyme::inlineAutodiffOp	(	enzyme::AutoDiffOp &	op,
		RewriterBase &	rewriter,
		SymbolTableCollection &	symbolTable )

Definition at line 380 of file InlineEnzymeRegions.cpp.

localizeGradients()

void mlir::enzyme::localizeGradients	(	OpBuilder &	builder,
		MGradientUtilsReverse *	gutils,
		Block *	fwd )

Definition at line 32 of file RemovalUtils.cpp.

References mlir::enzyme::MDiffeGradientUtils::getDifferential(), and mlir::enzyme::MGradientUtils::isConstantValue().

minCutCache()

void mlir::enzyme::minCutCache	(	Block *	forward,
		Block *	reverse,
		SmallVector< CacheInfo > &	caches,
		PatternRewriter &	rewriter,
		const IRMapping &	fwdrevmap,
		Operation *	lastFwd = nullptr )

Definition at line 427 of file RemovalUtils.cpp.

References annotate_ops(), bfs(), computeRankOfType(), computeSizeOfType(), DebugGraphviz(), dump(), filterGraph(), findCommonAncestor(), mlir::enzyme::CacheInfo::initOp, isMovable(), mlir::enzyme::CacheInfo::popOp, and mlir::enzyme::CacheInfo::pushOp.

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

opCmp()

bool mlir::enzyme::opCmp	(	mlir::Operation *	a,
		mlir::Operation *	b )

Referenced by valueCmp().

regionCmp()

bool mlir::enzyme::regionCmp	(	mlir::Region *	a,
		mlir::Region *	b )

registerAffineDialectAutoDiffInterface()

void mlir::enzyme::registerAffineDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 777 of file AffineAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerArithDialectAutoDiffInterface()

void mlir::enzyme::registerArithDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 108 of file ArithAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerAutoDiffUsingAllocationInterface()

template<typename OpTy >

void mlir::enzyme::registerAutoDiffUsingAllocationInterface

(

MLIRContext &

context

)

Definition at line 243 of file CoreDialectsAutoDiffImplementations.h.

registerAutoDiffUsingBranchInterface()

template<typename OpTy >

void mlir::enzyme::registerAutoDiffUsingBranchInterface

(

MLIRContext &

context

)

Definition at line 214 of file CoreDialectsAutoDiffImplementations.h.

registerAutoDiffUsingControlFlowInterface()

template<typename OpTy >

void mlir::enzyme::registerAutoDiffUsingControlFlowInterface

(

MLIRContext &

context

)

Definition at line 208 of file CoreDialectsAutoDiffImplementations.h.

registerAutoDiffUsingMemoryIdentityInterface()

template<typename OpTy , int... storedvals>

void mlir::enzyme::registerAutoDiffUsingMemoryIdentityInterface

(

MLIRContext &

context

)

Definition at line 237 of file CoreDialectsAutoDiffImplementations.h.

registerAutoDiffUsingRegionTerminatorInterface()

template<typename OpTy >

void mlir::enzyme::registerAutoDiffUsingRegionTerminatorInterface

(

MLIRContext &

context

)

Definition at line 221 of file CoreDialectsAutoDiffImplementations.h.

registerAutoDiffUsingReturnInterface()

template<typename OpTy >

void mlir::enzyme::registerAutoDiffUsingReturnInterface

(

MLIRContext &

context

)

Definition at line 229 of file CoreDialectsAutoDiffImplementations.h.

registerBuiltinDialectAutoDiffInterface()

void mlir::enzyme::registerBuiltinDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 336 of file BuiltinAutoDiffTypeInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerCFDialectAutoDiffInterface()

void mlir::enzyme::registerCFDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 36 of file CFAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerComplexDialectAutoDiffInterface()

void mlir::enzyme::registerComplexDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 78 of file ComplexAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerCoreDialectAutodiffInterfaces()

void mlir::enzyme::registerCoreDialectAutodiffInterfaces

(

DialectRegistry &

registry

)

Definition at line 466 of file CoreDialectsAutoDiffImplementations.cpp.

References registerAffineDialectAutoDiffInterface(), registerArithDialectAutoDiffInterface(), registerBuiltinDialectAutoDiffInterface(), registerCFDialectAutoDiffInterface(), registerComplexDialectAutoDiffInterface(), registerEnzymeDialectAutoDiffInterface(), registerFuncDialectAutoDiffInterface(), registerLinalgDialectAutoDiffInterface(), registerLLVMDialectAutoDiffInterface(), registerLLVMExtDialectAutoDiffInterface(), registerMathDialectAutoDiffInterface(), registerMemRefDialectAutoDiffInterface(), registerNVVMDialectAutoDiffInterface(), registerSCFDialectAutoDiffInterface(), and registerTensorDialectAutoDiffInterface().

Referenced by main().

registerEnzymeDialectAutoDiffInterface()

void mlir::enzyme::registerEnzymeDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 32 of file EnzymeAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerFuncDialectAutoDiffInterface()

void mlir::enzyme::registerFuncDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 257 of file FuncAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerLinalgDialectAutoDiffInterface()

void mlir::enzyme::registerLinalgDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 412 of file LinalgAutoDiffOpInterfaceImpl.cpp.

References attachAllInterfaces().

Referenced by registerCoreDialectAutodiffInterfaces().

registerLLVMDialectAutoDiffInterface()

void mlir::enzyme::registerLLVMDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 365 of file LLVMAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerLLVMExtDialectAutoDiffInterface()

void mlir::enzyme::registerLLVMExtDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 31 of file LLVMExtAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerMathDialectAutoDiffInterface()

void mlir::enzyme::registerMathDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 34 of file MathAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerMemRefDialectAutoDiffInterface()

void mlir::enzyme::registerMemRefDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 292 of file MemRefAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerNVVMDialectAutoDiffInterface()

void mlir::enzyme::registerNVVMDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 29 of file NVVMAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerSCFDialectAutoDiffInterface()

void mlir::enzyme::registerSCFDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 1264 of file SCFAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

registerTensorDialectAutoDiffInterface()

void mlir::enzyme::registerTensorDialectAutoDiffInterface

(

DialectRegistry &

registry

)

Definition at line 118 of file ArithAutoDiffOpInterfaceImpl.cpp.

Referenced by registerCoreDialectAutodiffInterfaces().

removalBlockExplore()

void mlir::enzyme::removalBlockExplore	(	Block *	block,
		IRMapping &	mapping,
		PatternRewriter &	rewriter,
		llvm::SetVector< Value > &	gradients,
		llvm::MapVector< Value, CacheInfo > &	caches )

Definition at line 81 of file RemovalUtils.cpp.

References getOp, mlir::enzyme::CacheInfo::merge(), mlir::enzyme::CacheInfo::popOp, and mlir::enzyme::CacheInfo::pushedValue().

Referenced by mlir::enzyme::IfLikeEnzymeOpsRemover< FinalClass, OpName >::removeEnzymeOps().

runActivityAnnotations()

void mlir::enzyme::runActivityAnnotations	(	FunctionOpInterface	callee,
		ArrayRef< enzyme::Activity >	argActivities = {},
		const ActivityPrinterConfig &	config = ActivityPrinterConfig() )

Definition at line 991 of file ActivityAnnotations.cpp.

References mlir::enzyme::ActivityPrinterConfig::annotate, mlir::enzyme::ForwardOriginsMap::getOrigins(), mlir::enzyme::MapOfSetsLattice< KeyT, ElementT >::join(), mlir::enzyme::BackwardOriginsLattice::meet(), mlir::enzyme::MapOfSetsLattice< KeyT, ElementT >::serialize(), traversePointsToSets(), and mlir::enzyme::ActivityPrinterConfig::verbose.

runDataFlowActivityAnalysis()

void mlir::enzyme::runDataFlowActivityAnalysis	(	FunctionOpInterface	callee,
		ArrayRef< enzyme::Activity >	argumentActivity,
		bool	print = false,
		bool	verbose = false,
		bool	annotate = false )

Definition at line 1071 of file DataFlowActivityAnalysis.cpp.

References ValueActivity::getActiveVal(), and printActivityAnalysisResults().

sortArraysLexicographic()

bool mlir::enzyme::sortArraysLexicographic	(	ArrayAttr	a,
		ArrayAttr	b )

Definition at line 43 of file DataFlowLattice.cpp.

References sortAttributes().

sortAttributes()

bool mlir::enzyme::sortAttributes	(	Attribute	a,
		Attribute	b )

Used when serializing to ensure a consistent order.

Definition at line 35 of file DataFlowLattice.cpp.

Referenced by sortArraysLexicographic().

valueCmp()

bool mlir::enzyme::valueCmp	(	mlir::Value	a,
		mlir::Value	b )

Definition at line 104 of file Utils.cpp.

References blockCmp(), and opCmp().

Variable Documentation

kPreserveCacheAttrName

llvm::StringLiteral mlir::enzyme::kPreserveCacheAttrName = "preserve_cache"

staticconstexpr

Definition at line 28 of file RemovalUtils.h.

Referenced by mlir::enzyme::IfLikeEnzymeOpsRemover< FinalClass, OpName >::removeEnzymeOps().

undefinedSetString

llvm::StringLiteral mlir::enzyme::undefinedSetString = "<undefined>"

constexpr

Definition at line 38 of file DataFlowLattice.h.

Referenced by deserializeAliasSummary(), deserializePointsTo(), mlir::enzyme::PointsToSets::print(), and mlir::enzyme::SetLattice< ValueT >::print().

unknownSetString

llvm::StringLiteral mlir::enzyme::unknownSetString = "<unknown>"

constexpr

Definition at line 39 of file DataFlowLattice.h.

Referenced by deserializeAliasSummary(), deserializePointsTo(), mlir::enzyme::PointsToSets::print(), and mlir::enzyme::SetLattice< ValueT >::print().