Construct an axis group graph from a tirx::Function. Two tensor axis are connected if they are accessed by the same index. More...

#include <axis_group_graph.h>

Inheritance diagram for tvm::tirx::TensorAxisGraphExtractor:

Collaboration diagram for tvm::tirx::TensorAxisGraphExtractor:

Public Member Functions
void	DFSGraph (TensorAxis cur, std::unordered_set< TensorAxis, TensorAxisHash > visited, std::vector< TensorAxis > tensor_axis_group)

Public Member Functions inherited from tvm::s_tir::StmtExprVisitor
virtual ffi::Optional< VisitInterrupt >	Visit_ (const SBlockRealizeNode *op)

Public Member Functions inherited from tvm::tirx::StmtExprVisitor
virtual ffi::Optional< VisitInterrupt >	Visit_ (const BindNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const RegionStmtNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const IfNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const ForNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const WhileNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const ReturnNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const BreakNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const ContinueNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const AssertStmtNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const SeqStmtNode *op)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const EvaluateNode *op)

ffi::Optional< VisitInterrupt >	Visit_ (const VarNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const OpaqueExprNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const TupleNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const TupleGetItemNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const CallNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const prim::LetNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const prim::RampNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const prim::BroadcastNode *op) override

ffi::Optional< VisitInterrupt >	Visit_ (const prim::ShuffleNode *op) override

Public Member Functions inherited from tvm::ExprVisitor
virtual ffi::Optional< VisitInterrupt >	Visit_ (const TensorRegionNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const GlobalVarNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const GenericConstNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const DataTypeImmNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const IntImmNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const FloatImmNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const OpNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const StringImmNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::CastNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::AddNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::LShiftNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::RShiftNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::BitwiseAndNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::BitwiseOrNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::BitwiseXorNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::BitwiseNotNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::SubNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::MulNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::DivNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::ModNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::FloorDivNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::FloorModNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::MinNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::MaxNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::EQNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::NENode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::LTNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::LENode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::GTNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::GENode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::AndNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::OrNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::NotNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit_ (const prim::SelectNode *node)

virtual ffi::Optional< VisitInterrupt >	Visit (ffi::AnyView value)
	Construct a visitor with the core expression hooks.

Public Member Functions inherited from tvm::ObjectVisitor
	ObjectVisitor ()
	Construct a visitor using structural fallback for every value.

	ObjectVisitor (const ObjectVisitor &other)=delete

ObjectVisitor &	operator= (const ObjectVisitor &other)=delete

Expected< ffi::Optional< VisitInterrupt > >	VisitExpected (ffi::AnyView value) noexcept
	Adapt the virtual throwing entrypoint to the structural ABI.

Static Public Member Functions
static std::vector< std::vector< TIRVarAxis > >	GetTIRVarAxisGraph (const tirx::Function &function)

Static Public Member Functions inherited from tvm::s_tir::StmtExprVisitor
static ffi::Optional< VisitInterrupt >	VisitBlock (tirx::StmtExprVisitor visitor, const SBlockNode op)

static ffi::Optional< VisitInterrupt >	VisitBlockRealize (tirx::StmtExprVisitor visitor, const SBlockRealizeNode op)

Additional Inherited Members
Protected Types inherited from tvm::ObjectVisitor
using	VTable = ObjectFunctor< ffi::Optional< VisitInterrupt >(const ffi::Object , ObjectVisitor )>
	Exact native dispatch table with owning interrupt results.

Protected Member Functions inherited from tvm::s_tir::StmtExprVisitor
	StmtExprVisitor (const VTable *vtable)

Protected Member Functions inherited from tvm::tirx::StmtExprVisitor
	StmtExprVisitor (const VTable *vtable)

Protected Member Functions inherited from tvm::ExprVisitor
	ExprVisitor (const VTable *vtable)
	Construct a visitor with an extended native dispatch table.

Protected Member Functions inherited from tvm::ObjectVisitor
	ObjectVisitor (const VTable *vtable)
	Construct a visitor with a finalized native dispatch table.

Static Protected Member Functions inherited from tvm::s_tir::StmtExprVisitor
static void	InitVTable (VTable *vtable)

Static Protected Member Functions inherited from tvm::tirx::StmtExprVisitor
static void	InitVTable (VTable *vtable)

Static Protected Member Functions inherited from tvm::ExprVisitor
static void	InitVTable (VTable *vtable)
	Register core expression hooks in a fresh mutable table.

Static Protected Member Functions inherited from tvm::ObjectVisitor
static void	InitVTable (VTable *vtable)
	Initialize the base registrations in a fresh mutable table.

template<typename Self , typename Node >
static void	SetDispatch (VTable *vtable)
	Register an exact node type in a fresh table initialized by its parent. Derived node types require separate registrations.

Detailed Description

Construct an axis group graph from a tirx::Function. Two tensor axis are connected if they are accessed by the same index.

Member Function Documentation

◆ DFSGraph()

void tvm::tirx::TensorAxisGraphExtractor::DFSGraph	(	TensorAxis	cur,
		std::unordered_set< TensorAxis, TensorAxisHash > *	visited,
		std::vector< TensorAxis > *	tensor_axis_group
	)

inline

◆ GetTIRVarAxisGraph()

static std::vector< std::vector< TIRVarAxis > > tvm::tirx::TensorAxisGraphExtractor::GetTIRVarAxisGraph ( const tirx::Function & function )

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()