axis_group_graph.h File Reference

#include <tvm/arith/iter_affine_map.h>
#include <tvm/relax/distributed/type.h>
#include <tvm/relax/expr.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt_functor.h>
#include <algorithm>
#include <limits>
#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::BufferAxisHash
class	tvm::tirx::BufferAxisGraphExtractor
	Construct an axis group graph from a PrimFunc. Two buffer 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
namespace	tvm::relax
namespace	tvm::relax::distributed

Typedefs
using	tvm::tirx::TIRVarAxis = std::pair< Var, int >
using	tvm::tirx::BufferAxis = std::pair< BufferVar, int >

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