axis_group_graph.h File Reference

#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/stmt_functor.h>
#include <tvm/sym/iter_affine_map.h>
#include <tvm/tirx/function.h>
#include <algorithm>
#include <limits>
#include <optional>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

Include dependency graph for axis_group_graph.h:

Go to the source code of this file.

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

Namespaces
namespace	tvm
	An object that builds and maintains block scope and StmtSref mapping for Dependence analysis.
namespace	tvm::tirx
	Backend-defined execution hierarchy nodes shared by layouts and tile calls.
namespace	tvm::relax
namespace	tvm::relax::distributed

Typedefs
using	tvm::tirx::TIRVarAxis = std::pair< Var, int >
using	tvm::tirx::TensorAxis = std::pair< TensorVar, int >

Functions
ffi::Optional< Var >	tvm::tirx::GetShardingVarFromIndex (PrimExpr index, ffi::Map< Var, Range > var_range, const sym::Analyzer &analyzer)
	Suppose we want to shard a tensor along a specific dimension, we need to know how to rewrite the access index of the tensor. To make it simple, we only support the case that the access can be rewritten by changing the extent of an iter var.