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 Expr &e) override
	Visitor to Exprs, can be overriden to do recursive changes to Exprs.

void	VisitExpr_ (const TensorLoadNode *op) override

void	VisitExpr_ (const BufferRegionNode *op) override

void	VisitExpr_ (const VarNode *op) override

void	VisitExpr_ (const TensorLoadNode *op) override

void	VisitExpr_ (const OpaqueExprNode *op) override

void	VisitExpr_ (const BufferRegionNode *op) override

void	VisitExpr_ (const TupleNode *op) override

void	VisitExpr_ (const TupleGetItemNode *op) override

void	VisitExpr_ (const prim::LetNode *op) override

void	VisitExpr_ (const CallNode *op) override

void	VisitExpr_ (const prim::AddNode *op) override

void	VisitExpr_ (const prim::SubNode *op) override

void	VisitExpr_ (const prim::MulNode *op) override

void	VisitExpr_ (const prim::DivNode *op) override

void	VisitExpr_ (const prim::ModNode *op) override

void	VisitExpr_ (const prim::FloorDivNode *op) override

void	VisitExpr_ (const prim::FloorModNode *op) override

void	VisitExpr_ (const prim::MinNode *op) override

void	VisitExpr_ (const prim::MaxNode *op) override

void	VisitExpr_ (const prim::EQNode *op) override

void	VisitExpr_ (const prim::NENode *op) override

void	VisitExpr_ (const prim::LTNode *op) override

void	VisitExpr_ (const prim::LENode *op) override

void	VisitExpr_ (const prim::GTNode *op) override

void	VisitExpr_ (const prim::GENode *op) override

void	VisitExpr_ (const prim::AndNode *op) override

void	VisitExpr_ (const prim::OrNode *op) override

void	VisitExpr_ (const prim::CastNode *op) override

void	VisitExpr_ (const prim::NotNode *op) override

void	VisitExpr_ (const prim::SelectNode *op) override

void	VisitExpr_ (const prim::RampNode *op) override

void	VisitExpr_ (const prim::BroadcastNode *op) override

void	VisitExpr_ (const prim::ShuffleNode *op) override

void	VisitExpr_ (const IntImmNode *op) override

void	VisitExpr_ (const FloatImmNode *op) override

void	VisitExpr_ (const prim::StringImmNode *op) override

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

void	VisitExpr_ (const TensorLoadNode *op) override

void	VisitExpr_ (const OpaqueExprNode *op) override

void	VisitExpr_ (const BufferRegionNode *op) override

void	VisitExpr_ (const TupleNode *op) override

void	VisitExpr_ (const TupleGetItemNode *op) override

void	VisitExpr_ (const prim::LetNode *op) override

void	VisitExpr_ (const CallNode *op) override

void	VisitExpr_ (const prim::AddNode *op) override

void	VisitExpr_ (const prim::SubNode *op) override

void	VisitExpr_ (const prim::MulNode *op) override

void	VisitExpr_ (const prim::DivNode *op) override

void	VisitExpr_ (const prim::ModNode *op) override

void	VisitExpr_ (const prim::FloorDivNode *op) override

void	VisitExpr_ (const prim::FloorModNode *op) override

void	VisitExpr_ (const prim::MinNode *op) override

void	VisitExpr_ (const prim::MaxNode *op) override

void	VisitExpr_ (const prim::EQNode *op) override

void	VisitExpr_ (const prim::NENode *op) override

void	VisitExpr_ (const prim::LTNode *op) override

void	VisitExpr_ (const prim::LENode *op) override

void	VisitExpr_ (const prim::GTNode *op) override

void	VisitExpr_ (const prim::GENode *op) override

void	VisitExpr_ (const prim::AndNode *op) override

void	VisitExpr_ (const prim::OrNode *op) override

void	VisitExpr_ (const prim::CastNode *op) override

void	VisitExpr_ (const prim::NotNode *op) override

void	VisitExpr_ (const prim::SelectNode *op) override

void	VisitExpr_ (const prim::RampNode *op) override

void	VisitExpr_ (const prim::BroadcastNode *op) override

void	VisitExpr_ (const prim::ShuffleNode *op) override

void	VisitExpr_ (const IntImmNode *op) override

void	VisitExpr_ (const FloatImmNode *op) override

void	VisitExpr_ (const prim::StringImmNode *op) override

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

virtual void	VisitBufferUse (const BufferVar &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.

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 ReturnNode *op) override

void	VisitStmt_ (const BreakNode *op) override

void	VisitStmt_ (const ContinueNode *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

void	VisitStmt_ (const ScopeIdDefStmtNode *op) override

void	VisitStmt_ (const tirx::TilePrimitiveCallNode *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()