Construct an axis group graph from a PrimFunc. Two buffer axis are connected if they are accessed by the same index. More...

#include <axis_group_graph.h>

Inheritance diagram for tvm::tirx::BufferAxisGraphExtractor:

Collaboration diagram for tvm::tirx::BufferAxisGraphExtractor:

Public Member Functions
void	DFSGraph (BufferAxis cur, std::unordered_set< BufferAxis, BufferAxisHash > visited, std::vector< BufferAxis > buffer_axis_group)

Static Public Member Functions
static std::vector< std::vector< TIRVarAxis > >	GetTIRVarAxisGraph (const PrimFunc &prim_func)

Additional Inherited Members
Protected Member Functions inherited from tvm::tirx::StmtExprVisitor
void	VisitExpr (const PrimExpr &e) override
	Visitor to Exprs, can be overriden to do recursive changes to Exprs. More...

void	VisitExpr_ (const BufferLoadNode *op) override

void	VisitExpr_ (const VarNode *op) override

void	VisitExpr_ (const SizeVarNode *op) override

void	VisitExpr_ (const BufferLoadNode *op) override

void	VisitExpr_ (const ProducerLoadNode *op) override

void	VisitExpr_ (const LetNode *op) override

void	VisitExpr_ (const CallNode *op) override

void	VisitExpr_ (const AddNode *op) override

void	VisitExpr_ (const SubNode *op) override

void	VisitExpr_ (const MulNode *op) override

void	VisitExpr_ (const DivNode *op) override

void	VisitExpr_ (const ModNode *op) override

void	VisitExpr_ (const FloorDivNode *op) override

void	VisitExpr_ (const FloorModNode *op) override

void	VisitExpr_ (const MinNode *op) override

void	VisitExpr_ (const MaxNode *op) override

void	VisitExpr_ (const EQNode *op) override

void	VisitExpr_ (const NENode *op) override

void	VisitExpr_ (const LTNode *op) override

void	VisitExpr_ (const LENode *op) override

void	VisitExpr_ (const GTNode *op) override

void	VisitExpr_ (const GENode *op) override

void	VisitExpr_ (const AndNode *op) override

void	VisitExpr_ (const OrNode *op) override

void	VisitExpr_ (const ReduceNode *op) override

void	VisitExpr_ (const CastNode *op) override

void	VisitExpr_ (const NotNode *op) override

void	VisitExpr_ (const SelectNode *op) override

void	VisitExpr_ (const RampNode *op) override

void	VisitExpr_ (const BroadcastNode *op) override

void	VisitExpr_ (const ShuffleNode *op) override

void	VisitExpr_ (const IntImmNode *op) override

void	VisitExpr_ (const FloatImmNode *op) override

void	VisitExpr_ (const StringImmNode *op) override

Protected Member Functions inherited from tvm::tirx::ExprVisitor
void	VisitExpr_ (const VarNode *op) override

void	VisitExpr_ (const SizeVarNode *op) override

void	VisitExpr_ (const BufferLoadNode *op) override

void	VisitExpr_ (const ProducerLoadNode *op) override

void	VisitExpr_ (const LetNode *op) override

void	VisitExpr_ (const CallNode *op) override

void	VisitExpr_ (const AddNode *op) override

void	VisitExpr_ (const SubNode *op) override

void	VisitExpr_ (const MulNode *op) override

void	VisitExpr_ (const DivNode *op) override

void	VisitExpr_ (const ModNode *op) override

void	VisitExpr_ (const FloorDivNode *op) override

void	VisitExpr_ (const FloorModNode *op) override

void	VisitExpr_ (const MinNode *op) override

void	VisitExpr_ (const MaxNode *op) override

void	VisitExpr_ (const EQNode *op) override

void	VisitExpr_ (const NENode *op) override

void	VisitExpr_ (const LTNode *op) override

void	VisitExpr_ (const LENode *op) override

void	VisitExpr_ (const GTNode *op) override

void	VisitExpr_ (const GENode *op) override

void	VisitExpr_ (const AndNode *op) override

void	VisitExpr_ (const OrNode *op) override

void	VisitExpr_ (const ReduceNode *op) override

void	VisitExpr_ (const CastNode *op) override

void	VisitExpr_ (const NotNode *op) override

void	VisitExpr_ (const SelectNode *op) override

void	VisitExpr_ (const RampNode *op) override

void	VisitExpr_ (const BroadcastNode *op) override

void	VisitExpr_ (const ShuffleNode *op) override

void	VisitExpr_ (const IntImmNode *op) override

void	VisitExpr_ (const FloatImmNode *op) override

void	VisitExpr_ (const StringImmNode *op) override

Protected Member Functions inherited from tvm::tirx::StmtVisitor
virtual void	VisitBufferDef (const Buffer &buffer, bool alloc_data)
	Visit buffer at definition site (AllocBuffer, DeclBuffer, SBlock alloc_buffers). Visits buffer shape, strides, elem_offset via VisitExpr. More...

virtual void	VisitBufferUse (const Buffer &buffer)
	Visit buffer at use site (BufferStore, BufferLoad, SBlock reads/writes). By default, this is a no-op, as buffer fields (shape, strides, elem_offset) are visited at their definition site. More...

void	VisitStmt_ (const BindNode *op) override

void	VisitStmt_ (const AttrStmtNode *op) override

void	VisitStmt_ (const IfThenElseNode *op) override

void	VisitStmt_ (const ForNode *op) override

void	VisitStmt_ (const WhileNode *op) override

void	VisitStmt_ (const AllocBufferNode *op) override

void	VisitStmt_ (const DeclBufferNode *op) override

void	VisitStmt_ (const BufferStoreNode *op) override

void	VisitStmt_ (const AssertStmtNode *op) override

void	VisitStmt_ (const SeqStmtNode *op) override

void	VisitStmt_ (const EvaluateNode *op) override

void	VisitStmt_ (const SBlockNode *op) override

void	VisitStmt_ (const SBlockRealizeNode *op) override

Detailed Description

Construct an axis group graph from a PrimFunc. Two buffer axis are connected if they are accessed by the same index.

Member Function Documentation

◆ DFSGraph()

void tvm::tirx::BufferAxisGraphExtractor::DFSGraph	(	BufferAxis	cur,
		std::unordered_set< BufferAxis, BufferAxisHash > *	visited,
		std::vector< BufferAxis > *	buffer_axis_group
	)

inline

◆ GetTIRVarAxisGraph()

static std::vector<std::vector<TIRVarAxis> > tvm::tirx::BufferAxisGraphExtractor::GetTIRVarAxisGraph ( const PrimFunc & prim_func )

inlinestatic

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

include/tvm/relax/distributed/axis_group_graph.h

Public Member Functions

Static Public Member Functions

Additional Inherited Members

Detailed Description

Member Function Documentation

◆ DFSGraph()

◆ GetTIRVarAxisGraph()