Namespaces
namespace	attr
	PrimFunc specific attribute names.

namespace	builtin
	Collection of builtin intrinsics as ops.

namespace	callback

namespace	transform

Classes
struct	ActiveSet
	Active thread set A. The source of truth is `layout`: shard = active axes with extents offset = per-axis lower bound, possibly a selector PrimExpr. More...

class	AllocBuffer
	Managed reference to AllocBufferNode. More...

class	AllocBufferNode
	Allocate a buffer and declare it in scope. More...

class	AssertStmt
	Managed reference to AssertStmtNode. More...

class	AssertStmtNode
	Assert condition, if an error occurs, return the error message. More...

class	AttrStmt
	Managed reference to AttrStmtNode. More...

class	AttrStmtNode
	Define certain auxiliary attribute for the body to be a symbolic value. This provide auxiliary information for IR passes that transforms body. More...

class	Axis

class	AxisAttrMap

class	AxisNode

struct	AxisRange
	Active slice offset + stride * [0, extent) encoded on one TileLayout axis. More...

class	AxisRegEntry

class	Bind
	Managed reference to BindNode. More...

class	BindNode
	Bind a variable to a value in the enclosing scope. More...

class	Break
	Managed reference to BreakNode. More...

class	BreakNode
	A Break in control flow. More...

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

class	BufferAxisHash

class	BufferRegion
	Managed reference to BufferRegionNode. More...

class	BufferRegionNode
	Representing a region of multi-dimensional buffer access. More...

class	BufferRegionType
	Managed reference to BufferRegionTypeNode. More...

class	BufferRegionTypeNode
	The type of a multi-dimensional buffer region expression. More...

class	BufferStore
	Managed reference to BufferStoreNode. More...

class	BufferStoreNode
	Store value to the high dimension buffer. More...

class	BufferType
	Managed reference to BufferTypeNode. More...

class	BufferTypeNode
	Structural type of a TIRx buffer variable. More...

class	BufferVar
	Checked zero-state view over an ordinary VarNode with BufferType. More...

class	ComposeLayout

class	ComposeLayoutNode

class	Continue
	Managed reference to ContinueNode. More...

class	ContinueNode
	A Continue in control flow. More...

class	DeclBuffer
	Managed reference to DeclBufferNode. More...

class	DeclBufferNode
	Declare a buffer that can be used in the body. More...

class	Dependency
	Managed reference to DependencyNode. More...

class	DependencyNode
	A tuple (src, dst, kind) representing certain types of dependency. For example, (A, B, kRAW) means block B depends on block A, and the dependency kind is read-after-write, which means block B reads the result written by block A. More...

class	DispatchContext
	Managed reference to DispatchContextNode. More...

class	DispatchContextNode
	The context information of the kernel required by op dispatch. More...

class	Evaluate
	Managed reference to EvaluateNode. More...

class	EvaluateNode
	Evaluates an expression. This is mostly used for putting a Call node into Stmt. More...

struct	ExecContext
	Per-program-point ExecContext: active set + scope kind + split. More...

class	ExecScope

class	ExecScopeNode

struct	ExecSplit
	One scope_switch split. Fields are sparse dicts keyed by active-set axis name, e.g. laneid/warpid/cta_id/wid_in_wg/wgid or factorized CTA axes such as cbx/cby/cbz. An empty map denotes the empty layout (e.g. intra under scope_kind=thread). More...

struct	ExprDeepEqual
	Compare two expressions recursively and check if they are equal to each other without var remapping. More...

class	ExprFunctor
	A dynamical functor that dispatches on in the first Expr argument. You can use this as a more powerful Visitor, since it allows you to define function signatures of Visit Function. More...

class	ExprFunctor< R(const Expr &n, Args...)>

class	ExprMutator
	ExprMutator that mutates expressions. More...

class	ExprVisitor
	ExprVisitor. More...

class	For
	Managed reference to ForNode. More...

class	ForNode
	A for loop, with possible type annotations. More...

class	IfThenElse
	Managed reference to IfThenElseNode. More...

class	IfThenElseNode
	IfThenElse statement. More...

class	IndexMap

class	IndexMapNode
	Defines a mapping between two representations of indices into a buffer. More...

class	Iter

class	IterNode

class	IterVar
	Iteration Variable, represents an iteration over an integer interval. More...

class	IterVarNode
	An iteration variable representing an iteration over a one dimensional interval. More...

class	LambdaExpr
	Managed reference to LambdaExprNode. More...

class	LambdaExprNode
	A reified Python lambda: a list of bound variables and a body over them. More...

class	Layout

class	LayoutNode

class	MatchBufferRegion
	Managed reference to MatchBufferRegionNode. More...

class	MatchBufferRegionNode
	Match introduces a constraint that the source buffer region can be remapped to the data layout specified by the buffer field. The constraint can be checked in later part of lowering (or optionally during runtime). More...

class	PrimFunc
	Managed reference to PrimFuncNode. More...

class	PrimFuncNode
	Primitive functions that contains TIR statements. More...

class	PrimVar
	Checked scalar view over a VarNode. More...

class	Return
	Managed reference to ReturnNode. More...

class	ReturnNode
	A return from the current function. More...

class	SBijectiveLayout
	Bijective function mapping for data layout transformation. Given two SLayout, SBijectiveLayout build and store the mapping rules, provides API to transform N-dimention tensor from the source indices (i0, i1, .., im) to the destination indices (j0, j1, .., jm). More...

class	SBijectiveLayoutNode

class	SBlock
	Managed reference to SBlockNode. More...

class	SBlockDependenceInfo
	Managed reference to SBlockDependenceInfoNode. More...

class	SBlockDependenceInfoNode
	An object that helps build and query block level dependences using the 2 core objects SBlockScope and StmtSRef. More...

class	SBlockNode
	A block is a basic schedule unit in TIR. More...

class	SBlockRealize
	Managed reference to BlockRealizeNode. More...

class	SBlockRealizeNode
	A block realization node represents execution of the block at the binding values. More...

class	SBlockScope
	Managed reference to SBlockScopeNode. More...

class	SBlockScopeNode
	An object with 1-to-1 correspondence with each block reference in the sref tree. This data structure is used to track the producer-consumer dependencies between blocks. For example even leaf nodes have a scope node, even though they have no dependencies. More...

class	ScopeIdDef

class	ScopeIdDefNode

class	ScopeIdDefStmt
	Managed reference to ScopeIdDefStmtNode. More...

class	ScopeIdDefStmtNode
	Standalone statement that declares a scope-id binding (e.g. cta_id, warp_id, lane_id). Carries a `ScopeIdDef` value. More...

class	ScopeIdDefVerifier

class	ScopeIdResolve
	Static resolver for ScopeIdDef values. Replaces the former ScopeIdResolveTable runtime registry with a closed-enum switch. More...

class	SeqStmt
	Sequence statement. More...

class	SeqStmtNode
	The container of seq statement. Represent a sequence of statements. More...

class	SLayout
	Managed reference to SLayoutNode. More...

class	SLayoutAxis

class	SLayoutNode
	SLayout is to describe how data is organized within an N-dimention tensor. It is composed of upper cases, lower cases and numbers, where upper case indicates a primal axis and the corresponding lower case with factor size indicates the subordinate axis. For example, NCHW16c can describe a 5-D tensor of [batch_size, channel, height, width, channel_block]. Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel). SLayout for scalar is defined, while both its name and axes have size 0. More...

class	SRefTreeCreator

class	Stmt
	Container of all statements. More...

class	StmtExprMutator
	Mutator that recursively mutates stmts and exprs on them. More...

class	StmtExprVisitor
	Visitor that recursively visit stmts and exprs on them. More...

class	StmtFunctor
	Same as ExprFunctor except it is applied on statements. More...

class	StmtFunctor< R(const Stmt &n, Args... args)>

class	StmtMutator
	StmtMutator that mutates the statements. More...

class	StmtNode
	Base node of all statements. More...

class	StmtSRef
	Managed reference to StmtSRefNode. More...

class	StmtSRefNode
	An object that refers to schedulable elements (block/for-loop) in TensorIR, aka "sref". More...

class	StmtVisitor
	StmtVisitor. More...

class	TensorIntrin
	Managed reference to TensorIntrinNode. More...

class	TensorIntrinNode
	Tensor intrinsics for tensorization. More...

class	TileLayout

class	TileLayoutNode

class	TilePrimitiveCall
	Managed reference to TilePrimitiveCallNode. More...

class	TilePrimitiveCallNode
	TIRX TilePrimitiveCall stmt. More...

class	While
	Managed reference to WhileNode. More...

class	WhileNode
	A While loop. More...

Typedefs
using	TIRVarAxis = std::pair< Var, int >

using	BufferAxis = std::pair< BufferVar, int >

using	FAxisFuser = ffi::TypedFunction< ffi::Optional< Iter >(Target, ffi::String, ffi::String, Iter)>

using	FAxisSplitter = ffi::TypedFunction< ffi::Array< Iter, void >(Target, ffi::String, Iter)>

using	AxisRegistry = AttrRegistry< AxisRegEntry, Axis >

using	TGlobalSymbol = ffi::String
	Global symbol of the op after lowering.

using	TVectorizable = bool
	Whether the op is overloaded for vector form.

using	FLowerIntrinsic = ffi::TypedFunction< PrimExpr(PrimExpr)>
	The intrinsic lowering function for given op.

using	FLegalize = ffi::TypedFunction< PrimExpr(PrimExpr)>
	The legalization function for given tirx op.

using	TScriptPrinterName = ffi::String
	The operator's name in TVMScript printer.

using	TScriptDtypePrintLocation = int64_t

using	TIRxOpCategory = ffi::String
	Broad TIRx op category.

using	TDeviceIntrinsicNamespace = ffi::String
	Device intrinsic namespace.

using	TCallEffectKind = int64_t
	Use integer to record the kind.

using	FArgSanitizer = ffi::TypedFunction< void(tvm::Op, ffi::Array< ffi::ObjectRef >)>
	The type of the function that sanitizes the arguments of a TIRX operator.

using	VarNode = tvm::VarNode

using	Var = tvm::Var

using	Region = ffi::Array< Range >

Enumerations
enum class	DepKind : int32_t { kRAW = 0 , kWAW = 1 , kWAR = 2 , kOpaque = 3 }
	Type of dependency. Right now we have 4 types of dependencies 1) Read-after-write (kRAW) 2) Write-after-write (kWAW) 3) Write-after-read (kWAR) 4) Opaque dependency (kOpaque) More...

enum class	ScopeKind : int { kCluster = 2 , kCta = 3 , kWarpgroup = 4 , kWarp = 5 , kThread = 6 }
	The target execution scope kind of a tile primitive call. More...

enum class	ScopeBinding : int { kKernelCluster = 0 , kKernelCta = 1 , kClusterCta = 2 , kCtaWarpgroup = 3 , kCtaWarp = 4 , kWarpgroupWarp = 5 , kWarpThread = 6 , kCtaThread = 7 , kWarpgroupThread = 8 , kClusterCtaPair = 9 }
	The binding between a parent scope and a child scope as used by a `ScopeIdDef`. The closed enum of valid (parent -> cur) pairs. More...

enum class	ScriptDtypePrintLocation : int { kNone = 0 , kFirst = 1 , kLast = 2 }
	Specifies that TVMScript printer prints the dtype as the first/last argument. If not specified, dtype will not be printed. More...

enum class	CallEffectKind : int { kExprAnnotation = 0 , kPure = 1 , kReadState = 2 , kUpdateState = 3 , kOpaque = kUpdateState , kSpecialCallArg = 4 , kEmbedInfo = 5 , kControlJump = 6 }
	The effect type of the call. More...

enum class	ForKind : int { kSerial = 0 , kParallel = 1 , kVectorized = 2 , kUnrolled = 3 , kThreadBinding = 4 }
	The kind of the loop. More...

enum	IterVarType : int { kDataPar = 0 , kThreadIndex = 1 , kCommReduce = 2 , kOrdered = 3 , kOpaque = 4 , kUnrolled = 5 , kVectorized = 6 , kParallelized = 7 , kTensorized = 8 }
	Type of iteration variable. Each IterVar have a specific type. More...

Functions
Var	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.

ffi::Array< ffi::Array< BufferRegion > >	GetSBlockAccessRegion (const SBlock &block, const ffi::Map< Var, BufferVar > &buffer_var_map)
	Auto detect the block access region according to its body stmt It will detect the access region as an array in order of appearance in AST.

ffi::Array< ffi::Array< BufferRegion > >	GetSBlockReadWriteRegion (const SBlock &block, const ffi::Map< Var, BufferVar > &buffer_var_map)
	Auto detect the block read/write region according to its body stmt. An opaque access will be counted as both a read and a write access.

ffi::Map< BufferVar, ffi::Optional< Stmt > >	DetectBufferAccessLCA (const PrimFunc &func)
	Detect the lowest common ancestor(LCA) of buffer access, including both high-level access(BufferLoad, BufferStore) and low-level access(Load, Store and opaque access). The LCA may be a For loop or a Block.

const tirx::SBlockNode *	FindAnchorBlock (const IRModule &mod)
	Find the "anchor block" of the given module. We define the anchor block to be the block with (1) an init statement and (2) having the biggest flops count. The latter condition is only used when there are multiple blocks with an init statement. For example, if the input module is conv2d + fused spatial blocks, conv2d is the anchor block. The input module may not contain more than one such block. For example, a module having two conv2d is not allowed as an input. However, a module created from winograd convolution has multiple blocks with an init statement (input transform, batched GEMM, and output transform). We use the second condition, the flops count, to determine that the batched GEMM block is the anchor block.

void	SetSeqIndex (std::unordered_map< const StmtNode *, StmtSRef > &stmt2ref, const Stmt &stmt, int seq_index, bool include_loops=true)
	Set the `StmtSRefNode::seq_index` field for stmt.

void	SetSeqIndexInChildren (std::unordered_map< const StmtNode , StmtSRef > &stmt2ref, const SeqStmtNode seq_stmt, bool include_loops=true)
	Update seq_index of the children of a SeqStmt.

template<class FLambda >
void	VisitPrimFuncs (const IRModule &mod, FLambda fvisit)
	Visit the PrimFuncs in the IRModule.

ffi::Array< Var >	UndefinedVars (const Stmt &stmt, const ffi::Array< Var > &defs)
	Find undefined vars in the statement.

ffi::Array< Var >	UndefinedVars (const PrimExpr &expr)
	Find undefined vars in the expression.

ffi::Array< Var >	UndefinedVars (const PrimExpr &expr, const ffi::Array< Var > &defs)
	Find undefined vars in the expression.

CallEffectKind	SideEffect (const PrimExpr &expr)
	Analyze the side effect of an expression.

bool	UsesVar (const Stmt &stmt, std::function< bool(const VarNode *)> vset_contains)
	Whether the given Stmt uses any var in the given variable set.

bool	UsesVar (const PrimExpr &expr, std::function< bool(const VarNode *)> vset_contains)
	Whether the given PrimExpr uses any var in the given variable set.

bool	VerifySSA (const PrimFunc &func)
	Verifies whether the IR stmt or Expr is in SSA form. That is: each Var is defined and assigned once(in Let/For)

bool	VerifyMemory (const PrimFunc &func)
	Verify if memory accesses are legal for a specific target device type.

size_t	CalculateExprComplexity (const PrimExpr &expr)
	Calculate the expression complexity based on number of symbols it contains.

size_t	CalculateConstantBytes (const PrimFunc &func, int64_t constant_byte_alignment)
	Calculate the constants size in bytes needed by the TIR allocates inside the TIR PrimFunc.

size_t	CalculateWorkspaceBytes (const PrimFunc &func, int64_t workspace_byte_alignment)
	Calculate the workspace size in bytes needed by the TIR allocates inside the TIR PrimFunc.

bool	VerifyWellFormed (const PrimFunc &func, bool assert_mode=true)
	Verify if the given TIR is well-formed. The verification includes:

bool	VerifyWellFormed (const IRModule &mod, bool assert_mode=true)
	Verify if the TIR in the given IRMOdule is well-formed.

const PrimFuncNode *	FindEntryFunc (const IRModule &mod, GlobalVar *result_g_var)
	Find the entry function of the given IRModule, i.e, functions marked by `tirx::attr::kIsEntryFunc`, whose name is `main` or being the only PrimeFunc.

PrimType	DefaultIndexPrimType ()
	if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32

DLDataType	DefaultIndexType ()

bool	operator== (const BufferVar &lhs, const BufferVar &rhs)

bool	operator!= (const BufferVar &lhs, const BufferVar &rhs)

BufferVar	GetBufferVar (const VarNode *var)
	Recover a checked buffer view from an ordinary VarNode pointer.

ffi::ObjectPtr< BufferTypeNode >	CopyBufferType (const BufferVar &var)

BufferVar	RebuildBufferVar (const BufferVar &var, ffi::ObjectPtr< BufferTypeNode > type, ffi::Optional< ffi::String > name=std::nullopt)

BufferVar	decl_buffer (ffi::Array< PrimExpr > shape, PrimType dtype=PrimType::Float(32), ffi::String name="buffer", ffi::String storage_scope="", Span span=Span())
	Construct a new buffer given shape, and dtype.

tirx::BufferVar	BufferWithOffsetAlignment (ffi::Array< PrimExpr > shape, PrimType dtype, std::string name, int data_alignment, int offset_factor, std::string memory_scope="")
	Creates a TIR buffer for the provided parameters.

TensorLoad	BufferLoad (BufferVar buffer, ffi::Array< PrimExpr > indices, Span span=Span())
	Construct a TensorLoad from a BufferVar.

ActiveSet	InitialActiveSet (int64_t lane_ext, int64_t warp_ext, int64_t cta_ext)
	Initial A at PrimFunc device entry: all threads active, offsets zero.

ActiveSet	InitialActiveSet (int64_t lane_ext, int64_t warp_ext, int64_t cta_ext, const std::vector< std::pair< std::string, int64_t > > &cta_axes)

bool	FilterNarrow (const ActiveSet &A, ScopeBinding binding, int64_t lo, int64_t hi, ActiveSet out, std::string err)
	Narrow A on the lane bound to `binding`.

bool	ScopeSwitch (const ActiveSet &A, ScopeKind scope_kind, ExecSplit out, std::string err)
	Factor A into (inter, intra) for target scope_kind.

ffi::Map< ffi::String, ffi::Array< PrimExpr > >	EncodeSplitSide (const std::unordered_map< std::string, AxisRange > &side)
	Encode one side of an ExecSplit (inter or intra) as the FFI map used by `DispatchContextNode::{inter, intra}`: axis name -> [extent, offset] for unit-stride axes, or [extent, offset, stride] for strided axes.

std::string	ScopeKindToString (ScopeKind kind)
	Convert a ScopeKind to its string name (e.g. kThread -> "thread").

ScopeKind	StringToScopeKind (const ffi::String &name)
	Parse a string name to a ScopeKind. FATAL if unknown.

std::pair< ffi::String, ffi::String >	ScopeBindingToStringPair (ScopeBinding binding)
	Convert a ScopeBinding to its (parent, cur) string pair.

ScopeBinding	StringPairToScopeBinding (const ffi::String &parent, const ffi::String &cur)
	Parse a (parent, cur) string pair to a ScopeBinding. FATAL if unknown.

bool	ScopeKindHigher (ScopeKind a, ScopeKind b)
	Strict-weak "a is wider than b" on scope kinds: `world > kernel > cluster > cta > warpgroup > warp > thread`. Only used by axe-layout scope-chain validity (the rest of the codebase compares scope identities with ==).

bool	ScopeNameHigher (const ffi::String &a, const ffi::String &b)
	String-keyed convenience over ScopeKindHigher. FATALs on bad name.

PrimFunc	Specialize (PrimFunc func, const ffi::Map< Var, ffi::Variant< BufferVar, Expr > > &param_map)
	Specialize parameters of PrimFunc.

IndexMap	Substitute (const IndexMap &index_map, std::function< ffi::Optional< PrimExpr >(const Var &var)> f_subst)
	Substitute variables in an index map.

bool	IsPointerType (const Type &type, DLDataType element_type)
	Check if type is a pointer to a runtime element type.

template<typename ValueType , typename = typename std::enable_if<std::is_standard_layout<ValueType>::value && std::is_trivial<ValueType>::value>::type>
PrimExpr	MakeConst (PrimType dtype, ValueType value, Span span=Span())
	Make a const value with certain data type.

Expr	ConstHandle (int64_t value, Span span=Span())
	Make a constant opaque-pointer value.

const int64_t *	as_const_int (const PrimExpr &x)
	Get x as constant int expression.

bool	is_const_int (const PrimExpr &x, int64_t value)
	Check whether x is a constant integer expression.

bool	is_no_op (const tirx::Stmt &stmt)
	Check whether stmt is nop.

bool	is_one (const PrimExpr &x)
	Check whether x is a constant integer 1.

bool	is_zero (const PrimExpr &x)
	Check whether x is a constant integer 0.

bool	is_const_int (const PrimExpr &x)
	Check whether x is an integer constant.

bool	is_const_number (const PrimExpr &x)
	Check whether x is an integer/float constant.

template<typename FReduce >
PrimExpr	foldl (FReduce freduce, PrimExpr init_value, const ffi::Array< PrimExpr > &values, Span span=Span())
	Left fold.

bool	is_const_power_of_two_integer (const PrimExpr &x, int *shift)
	Check whether x is a constant power of two If x is power of two, write the power to the shift.

bool	is_positive_const (const PrimExpr &a)

bool	is_negative_const (const PrimExpr &a)

template<typename ValueType >
PrimExpr	MakeConstScalar (PrimType dtype, ValueType value, Span span=Span())

template<>
PrimExpr	MakeConstScalar (PrimType dtype, bool value, Span span)

std::ostream &	operator<< (std::ostream &os, CallEffectKind side_effect)

PrimExpr	TypeAnnotation (PrimType dtype, Span span=Span())
	Create a type annotation expression.

std::ostream &	operator<< (std::ostream &os, ForKind kind)

const char *	ForKind2String (ForKind t)

Stmt	IRTransform (Stmt stmt, const ffi::Function &preorder, const ffi::Function &postorder, ffi::Optional< ffi::Array< ffi::String > > only_enable=std::nullopt)
	recursively visit the ir nodes in post DFS order, and transform it

void	PostOrderVisit (const ffi::ObjectRef &node, std::function< void(const ffi::ObjectRef &)> fvisit)
	Recursively visit a statement or expression in post DFS order, applying fvisit. Each node is guaranteed to be visited only once.

Stmt	Substitute (Stmt stmt, std::function< ffi::Optional< Expr >(const Var &var)> vmap)
	Substitute the var specified by vmap.

Expr	Substitute (Expr expr, std::function< ffi::Optional< Expr >(const Var &var)> vmap)
	Substitute the var specified by vmap.

PrimExpr	Substitute (PrimExpr expr, std::function< ffi::Optional< Expr >(const Var &var)> vmap)

Range	Substitute (const Range &range, std::function< ffi::Optional< Expr >(const Var &var)> vmap)
	Substitute the vars specified by vmap.

template<typename T >
ffi::Array< T >	Substitute (const ffi::Array< T > &arr, std::function< ffi::Optional< Expr >(const Var &var)> vmap)
	Substitute the var specified by vmap.

template<typename Obj >
auto	Substitute (Obj &&obj, const ffi::Map< Var, Expr > &vmap)
	Substitute the vars specified by vmap.

template<typename Obj , typename Replacement >
auto	Substitute (Obj &&obj, const ffi::Map< Var, Replacement > &vmap)
	Substitute the vars specified by vmap.

template<typename Obj , typename Replacement >
auto	Substitute (Obj &&obj, const std::unordered_map< const VarNode *, Replacement > &vmap)
	Substitute the vars specified by vmap.

template<typename Obj , typename Replacement , typename Hasher , typename EqualityChecker >
auto	Substitute (Obj &&obj, const std::unordered_map< Var, Replacement, Hasher, EqualityChecker > &vmap)
	Substitute the vars specified by vmap.

template<typename Obj , typename Replacement >
auto	Substitute (Obj &&obj, const std::unordered_map< IterVar, Replacement > &iter_vmap)
	Substitute the vars specified by vmap.

Stmt	SubstituteWithDataTypeLegalization (Stmt stmt, std::function< ffi::Optional< PrimExpr >(const Var &)> vmap)
	Substitute the var specified by vmap and legalize data types after substitution.

PrimExpr	SubstituteWithDataTypeLegalization (PrimExpr expr, std::function< ffi::Optional< PrimExpr >(const Var &)> vmap)
	Substitute the var specified by vmap and legalize data types after substitution.

void	PreOrderVisit (const ffi::ObjectRef &stmt_or_expr, const std::function< bool(const ffi::ObjectRef &)> &fvisit)
	Recursively visit a statement or expression in pre DFS order, applying fvisit. If fvisit returns false, it won't visit the children of the node.

template<typename Node , typename = std::enable_if_t<std::is_base_of_v<StmtNode, Node>>>
bool	ContainsNode (const Stmt &stmt)
	Check if the statement contains the specified node type.

const Op &	cast ()
	See pesudo code below:

const Op &	copy ()
	See pesudo code below:

const Op &	copy_async ()
	See pesudo code below:

const Op &	fill ()
	See pesudo code below:

const Op &	gemm ()
	See pesudo code below:

const Op &	gemm_async ()
	See pesudo code below:

const Op &	zero ()

const Op &	sqrt ()

const Op &	exp ()

const Op &	exp2 ()

const Op &	add ()

const Op &	sub ()

const Op &	mul ()

const Op &	fdiv ()

const Op &	minimum ()

const Op &	maximum ()

const Op &	reciprocal ()

const Op &	sum ()

const Op &	max ()

const Op &	min ()

const Op &	memset ()

const Op &	reduce_negate ()

const Op &	binary_reduce ()

const Op &	unary_reduce ()

const Op &	binary_chain ()

const Op &	select ()

const Op &	fma ()

const Op &	silu ()

const Op &	compose_op ()

const Op &	permute_layout ()

const char *	IterVarType2String (IterVarType t)

Variables
constexpr int	kWgSize = 4
	Warpgroup size in warps (hardware-fixed).

constexpr int	kPSUMMaxElemPerBank = 512

constexpr int	kPSUMBankNum = 8

Typedef Documentation

◆ AxisRegistry

using tvm::tirx::AxisRegistry = typedef AttrRegistry<AxisRegEntry, Axis>

◆ BufferAxis

using tvm::tirx::BufferAxis = typedef std::pair<BufferVar, int>

◆ FArgSanitizer

using tvm::tirx::FArgSanitizer = typedef ffi::TypedFunction<void(tvm::Op, ffi::Array<ffi::ObjectRef>)>

The type of the function that sanitizes the arguments of a TIRX operator.

Parameters

op	The operator.
args	The arguments.

◆ FAxisFuser

using tvm::tirx::FAxisFuser = typedef ffi::TypedFunction<ffi::Optional<Iter>(Target, ffi::String, ffi::String, Iter)>

◆ FAxisSplitter

using tvm::tirx::FAxisSplitter = typedef ffi::TypedFunction<ffi::Array<Iter, void>(Target, ffi::String, Iter)>

◆ FLegalize

using tvm::tirx::FLegalize = typedef ffi::TypedFunction<PrimExpr(PrimExpr)>

The legalization function for given tirx op.

◆ FLowerIntrinsic

using tvm::tirx::FLowerIntrinsic = typedef ffi::TypedFunction<PrimExpr(PrimExpr)>

The intrinsic lowering function for given op.

◆ Region

using tvm::tirx::Region = typedef ffi::Array<Range>

◆ TCallEffectKind

using tvm::tirx::TCallEffectKind = typedef int64_t

Use integer to record the kind.

◆ TDeviceIntrinsicNamespace

using tvm::tirx::TDeviceIntrinsicNamespace = typedef ffi::String

Device intrinsic namespace.

Expected values include "cuda", "ptx", "nvshmem", "nki", and "metal".

◆ TGlobalSymbol

using tvm::tirx::TGlobalSymbol = typedef ffi::String

Global symbol of the op after lowering.

◆ TIRVarAxis

using tvm::tirx::TIRVarAxis = typedef std::pair<Var, int>

◆ TIRxOpCategory

using tvm::tirx::TIRxOpCategory = typedef ffi::String

Broad TIRx op category.

Expected values:

"builtin"
"tile_primitive"
"device_intrin"

◆ TScriptDtypePrintLocation

using tvm::tirx::TScriptDtypePrintLocation = typedef int64_t

◆ TScriptPrinterName

using tvm::tirx::TScriptPrinterName = typedef ffi::String

The operator's name in TVMScript printer.

◆ TVectorizable

using tvm::tirx::TVectorizable = typedef bool

Whether the op is overloaded for vector form.

◆ Var

using tvm::tirx::Var = typedef tvm::Var

◆ VarNode

using tvm::tirx::VarNode = typedef tvm::VarNode

Enumeration Type Documentation

◆ CallEffectKind

enum class tvm::tirx::CallEffectKind : int

strong

The effect type of the call.

Enumerator
kExprAnnotation	Function corresponds to an annotation(e.g. likely) and can translate to identity.
kPure	Pure function that do not interacts with any external state.
kReadState	Function's that may read from states(e.g. RAM)
kUpdateState	Function that may read/write from states(e.g. RAM).
kOpaque	Opaque function, cannot make any assumption.
kSpecialCallArg	Special intrinsic to annotate call arguments info only valid as a direct argument to a call.
kEmbedInfo	Embed opaque information in the Expr, cannot be codegen.
kControlJump	Function that changes control flow.

◆ DepKind

enum class tvm::tirx::DepKind : int32_t

strong

Type of dependency. Right now we have 4 types of dependencies 1) Read-after-write (kRAW) 2) Write-after-write (kWAW) 3) Write-after-read (kWAR) 4) Opaque dependency (kOpaque)

Enumerator
kRAW
kWAW
kWAR
kOpaque

◆ ForKind

enum class tvm::tirx::ForKind : int

strong

The kind of the loop.

ForKind can change the control flow semantics of the loop. So the kind field needs to be considered in all TIR passes.

Enumerator
kSerial	default semantics – serial execution.
kParallel	Parallel execution on CPU.
kVectorized	Vector SIMD loop. The loop body will be vectorized.
kUnrolled	The loop body must be unrolled.
kThreadBinding	The loop variable is bound to a thread in an environment. In the final stage of lowering, the loop is simply removed and the loop variable is mapped to the corresponding context thread.

◆ IterVarType

enum tvm::tirx::IterVarType : int

Type of iteration variable. Each IterVar have a specific type.

The type of iter var can be overriden via stage.iter_var_attrs given they are compatible.

Enumerator
kDataPar	Data parallel iteration. This normally corresponds to axis of Tensor. Allow all IterVar manipulations. Note This does not mean the loop have to be executed in parallel fashion.
kThreadIndex	The IterVar itself is a thread-index of a fixed thread launching group. Note that this is already assumed to be parallelized. Disallow: split/fuse/vectorize/parallel
kCommReduce	Communicative reduction. Cannot be directly parallelized. Disallow: parallel/vectorize
kOrdered	Serial loops with loop carry dependency, the iteration must execute in order. Cannot be re-ordered. Disallow: reorder/parallel/vectorize
kOpaque	IterVar is opaque,. May not corresponds to any generated loop Disallow all IterVar manipulations and compute_at Note This is usually used to implement composite op or external op, where the
kUnrolled	The execution is unrolled.
kVectorized	The loop is vectorized.
kParallelized	The loop is parallelized.
kTensorized	Marks boundary of tensorization intrinsic.

◆ ScopeBinding

enum class tvm::tirx::ScopeBinding : int

strong

The binding between a parent scope and a child scope as used by a ScopeIdDef. The closed enum of valid (parent -> cur) pairs.

Single-axis bindings (target one ActiveSet box axis – laneid / warpid / cta_id, possibly via a warpid factor lane): kKernelCta, kClusterCta -> cta_id (flat) kCtaWarp -> warpid (flat) kCtaWarpgroup -> warpid (outer factor; warpgroup index) kWarpgroupWarp -> warpid (inner factor; warp-within-wg index) kWarpThread -> laneid (flat) kKernelCluster -> not a filter target (cluster_id by design) kClusterCtaPair -> hardware CTA pair id (cluster CTA rank % 2)

Multi-axis (flat-thread) bindings – linearize across two ActiveSet axes; a flat lo <= var and var < hi predicate cannot narrow them as a contiguous box range, so they fall back to plain predicate semantics: kCtaThread -> threadIdx.x within a CTA (laneid * warpid) kWarpgroupThread -> threadIdx.x within a warpgroup (laneid * wid_in_wg)

Enumerator
kKernelCluster
kKernelCta
kClusterCta
kCtaWarpgroup
kCtaWarp
kWarpgroupWarp
kWarpThread
kCtaThread
kWarpgroupThread
kClusterCtaPair

◆ ScopeKind

enum class tvm::tirx::ScopeKind : int

strong

The target execution scope kind of a tile primitive call.

Identifies the granularity at which an op executes (the per-call scope on a TilePrimitiveCall, e.g. Tx.warp.copy(...)). Ordered from coarsest to finest; smaller integer = wider scope, so ScopeKindHigher is a plain <.

Enumerator
kCluster
kCta
kWarpgroup
kWarp
kThread

◆ ScriptDtypePrintLocation

enum class tvm::tirx::ScriptDtypePrintLocation : int

strong

Specifies that TVMScript printer prints the dtype as the first/last argument. If not specified, dtype will not be printed.

Enumerator
kNone	Do not print dtype as an argument.
kFirst	Print dtype as the first argument.
kLast	FPrint dtype as the last argument.

Function Documentation

◆ add()

const Op & tvm::tirx::add ( )

◆ as_const_int()

const int64_t * tvm::tirx::as_const_int ( const PrimExpr & x )

inline

Get x as constant int expression.

Parameters

x	The expression

Returns: the address to the int expression, return nullptr, if x is not IntImm.

◆ binary_chain()

const Op & tvm::tirx::binary_chain ( )

◆ binary_reduce()

const Op & tvm::tirx::binary_reduce ( )

◆ BufferLoad()

TensorLoad tvm::tirx::BufferLoad	(	BufferVar	buffer,
		ffi::Array< PrimExpr >	indices,
		Span	span = `Span()`
	)

Construct a TensorLoad from a BufferVar.

This is the sole typed construction path for tirx loads. The result type is derived from the buffer element type and index lanes, and every tirx TensorLoad is required to have a BufferVar source.

◆ BufferWithOffsetAlignment()

tirx::BufferVar tvm::tirx::BufferWithOffsetAlignment	(	ffi::Array< PrimExpr >	shape,
		PrimType	dtype,
		std::string	name,
		int	data_alignment,
		int	offset_factor,
		std::string	memory_scope = `""`
	)

Creates a TIR buffer for the provided parameters.

Parameters

shape	shape of the buffer
dtype	data type
name	buffer name
data_alignment	alignment requirement of data pointer in bytes
offset_factor	Factor of elem_offset field, elem_offset is guaranteed to be multiple of offset_factor User can specify data_alignment and offset_factor to be 0 A default value will be picked.
memory_scope	memory scope of the buffer

◆ CalculateConstantBytes()

size_t tvm::tirx::CalculateConstantBytes	(	const PrimFunc &	func,
		int64_t	constant_byte_alignment
	)

Calculate the constants size in bytes needed by the TIR allocates inside the TIR PrimFunc.

Parameters

func	The TIR PrimFunc for which the constants size to be calculated
constant_byte_alignment	The byte alignment required for each constant allocated

◆ CalculateExprComplexity()

size_t tvm::tirx::CalculateExprComplexity ( const PrimExpr & expr )

Calculate the expression complexity based on number of symbols it contains.

Parameters

expr	The expr to be calculated.

◆ CalculateWorkspaceBytes()

size_t tvm::tirx::CalculateWorkspaceBytes	(	const PrimFunc &	func,
		int64_t	workspace_byte_alignment
	)

Calculate the workspace size in bytes needed by the TIR allocates inside the TIR PrimFunc.

Parameters

func	The TIR PrimFunc for which the workspace size to be calculated
workspace_byte_alignment	The byte alignment required for each tensor allocated in this workspace

◆ cast()

const Op & tvm::tirx::cast ( )

See pesudo code below:

Tx.cast(BufferRegion dst, BufferRegion src)

◆ compose_op()

const Op & tvm::tirx::compose_op ( )

◆ ConstHandle()

Expr tvm::tirx::ConstHandle	(	int64_t	value,
		Span	span = `Span()`
	)

inline

Make a constant opaque-pointer value.

Parameters

value	The integer payload to reinterpret as a handle.
span	The location of this operation in the source.

Returns: The result expression.

◆ ContainsNode()

template<typename Node , typename = std::enable_if_t<std::is_base_of_v<StmtNode, Node>>>

bool tvm::tirx::ContainsNode ( const Stmt & stmt )

Check if the statement contains the specified node type.

This utility potentially walks the entire statement, and should therefore not be used if it could otherwise be merged with another pass.

Parameters

stmt	The statement to be searched

Returns: Whether stmt contains Node

◆ copy()

const Op & tvm::tirx::copy ( )

See pesudo code below:

Tx.copy(BufferRegion dst, BufferRegion src)

◆ copy_async()

const Op & tvm::tirx::copy_async ( )

See pesudo code below:

Tx.Async.copy(BufferRegion dst, BufferRegion src)

◆ CopyBufferType()

ffi::ObjectPtr< BufferTypeNode > tvm::tirx::CopyBufferType ( const BufferVar & var )

inline

◆ decl_buffer()

BufferVar tvm::tirx::decl_buffer	(	ffi::Array< PrimExpr >	shape,
		PrimType	dtype = `PrimType::Float(32)`,
		ffi::String	name = `"buffer"`,
		ffi::String	storage_scope = `""`,
		Span	span = `Span()`
	)

Construct a new buffer given shape, and dtype.

Parameters

shape	The shape of the buffer,
dtype	The content data type.
name	The name of the buffer
storage_scope	The storage scope associated with this buffer
span	The location of this object in the source code.

Returns: The created buffer.

See also: BufferVar for complete constructor.

◆ DefaultIndexPrimType()

PrimType tvm::tirx::DefaultIndexPrimType ( )

inline

if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32

◆ DefaultIndexType()

DLDataType tvm::tirx::DefaultIndexType ( )

inline

◆ DetectBufferAccessLCA()

ffi::Map< BufferVar, ffi::Optional< Stmt > > tvm::tirx::DetectBufferAccessLCA ( const PrimFunc & func )

Detect the lowest common ancestor(LCA) of buffer access, including both high-level access(BufferLoad, BufferStore) and low-level access(Load, Store and opaque access). The LCA may be a For loop or a Block.

Parameters

func	The PrimFunc to be detected.

Returns: The Map from buffer to the LCA of all access to it. The lca is function root if the return stmt is std::nullopt.

◆ EncodeSplitSide()

ffi::Map< ffi::String, ffi::Array< PrimExpr > > tvm::tirx::EncodeSplitSide ( const std::unordered_map< std::string, AxisRange > & side )

Encode one side of an ExecSplit (inter or intra) as the FFI map used by DispatchContextNode::{inter, intra}: axis name -> [extent, offset] for unit-stride axes, or [extent, offset, stride] for strided axes.

◆ exp()

const Op & tvm::tirx::exp ( )

◆ exp2()

const Op & tvm::tirx::exp2 ( )

◆ fdiv()

const Op & tvm::tirx::fdiv ( )

◆ fill()

const Op & tvm::tirx::fill ( )

See pesudo code below:

Tx.fill(BufferRegion dst, PrimExpr value)

◆ FilterNarrow()

bool tvm::tirx::FilterNarrow	(	const ActiveSet &	A,
		ScopeBinding	binding,
		int64_t	lo,
		int64_t	hi,
		ActiveSet *	out,
		std::string *	err
	)

Narrow A on the lane bound to binding.

The ScopeBinding maps directly to which native axis (laneid/warpid/cta_id) to narrow, and for warpid whether to narrow the full axis (kCtaWarp), the outer factor (kCtaWarpgroup), or the inner factor (kWarpgroupWarp).

Bindings with no single-lane representation are conservative: cluster_id is not a filter target; flat thread ids are accepted only when the range can be represented as a rectangular lane/warp active set.

◆ FindAnchorBlock()

const tirx::SBlockNode * tvm::tirx::FindAnchorBlock ( const IRModule & mod )

Find the "anchor block" of the given module. We define the anchor block to be the block with (1) an init statement and (2) having the biggest flops count. The latter condition is only used when there are multiple blocks with an init statement. For example, if the input module is conv2d + fused spatial blocks, conv2d is the anchor block. The input module may not contain more than one such block. For example, a module having two conv2d is not allowed as an input. However, a module created from winograd convolution has multiple blocks with an init statement (input transform, batched GEMM, and output transform). We use the second condition, the flops count, to determine that the batched GEMM block is the anchor block.

Parameters

mod	The input TIR module.

Returns: The anchor block if found, nullptr otherwise.

◆ FindEntryFunc()

const PrimFuncNode * tvm::tirx::FindEntryFunc	(	const IRModule &	mod,
		GlobalVar *	result_g_var
	)

Find the entry function of the given IRModule, i.e, functions marked by tirx::attr::kIsEntryFunc, whose name is main or being the only PrimeFunc.

Parameters

mod	The IRModule to find the entry function.
result_g_var	The result GlobalVar of the entry function.

Returns: The entry function.

◆ fma()

const Op & tvm::tirx::fma ( )

◆ foldl()

template<typename FReduce >

PrimExpr tvm::tirx::foldl	(	FReduce	freduce,
		PrimExpr	init_value,
		const ffi::Array< PrimExpr > &	values,
		Span	span = `Span()`
	)

inline

Left fold.

Parameters

freduce	The reduction function.
init_value	The initial value.
values	The values to be folded.
span	The location of the fold in the source.

Returns: The result.

Template Parameters

FReduce The type of the reduction.

◆ ForKind2String()

const char * tvm::tirx::ForKind2String ( ForKind t )

inline

◆ gemm()

const Op & tvm::tirx::gemm ( )

See pesudo code below:

Tx.gemm(BufferVar A, BufferVar B, BufferVar C, BufferVar D, PrimExpr alpha, PrimExpr beta)

◆ gemm_async()

const Op & tvm::tirx::gemm_async ( )

See pesudo code below:

Tx.gemm_async(BufferRegion C, BufferRegion A, BufferRegion B, bool transA, bool transB, bool accum)

◆ GetBufferVar()

BufferVar tvm::tirx::GetBufferVar ( const VarNode * var )

inline

Recover a checked buffer view from an ordinary VarNode pointer.

◆ GetSBlockAccessRegion()

ffi::Array< ffi::Array< BufferRegion > > tvm::tirx::GetSBlockAccessRegion	(	const SBlock &	block,
		const ffi::Map< Var, BufferVar > &	buffer_var_map
	)

Auto detect the block access region according to its body stmt It will detect the access region as an array in order of appearance in AST.

Parameters

block	The block to be detected
buffer_var_map	The outside buffers which may be accessed the block. It is a map from buffer var to the buffer.

Returns

Array of access regions. There are three arrays of BufferRegion:

first: read regions
second: write regions
third: opaque regions

◆ GetSBlockReadWriteRegion()

ffi::Array< ffi::Array< BufferRegion > > tvm::tirx::GetSBlockReadWriteRegion	(	const SBlock &	block,
		const ffi::Map< Var, BufferVar > &	buffer_var_map
	)

Auto detect the block read/write region according to its body stmt. An opaque access will be counted as both a read and a write access.

Parameters

block	The block to be detected
buffer_var_map	The outside buffers which may be accessed the block. It is a map from buffer var to the buffer

Returns: An array only consisting of the read regions and write regions of the input block

◆ GetShardingVarFromIndex()

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.

Parameters

index	The access index
var_range	The range of each iter var
analyzer	The analyzer

Returns: The iter var whose extent to be changed

◆ InitialActiveSet() [1/2]

ActiveSet tvm::tirx::InitialActiveSet	(	int64_t	lane_ext,
		int64_t	warp_ext,
		int64_t	cta_ext
	)

Initial A at PrimFunc device entry: all threads active, offsets zero.

◆ InitialActiveSet() [2/2]

ActiveSet tvm::tirx::InitialActiveSet	(	int64_t	lane_ext,
		int64_t	warp_ext,
		int64_t	cta_ext,
		const std::vector< std::pair< std::string, int64_t > > &	cta_axes
	)

◆ IRTransform()

Stmt tvm::tirx::IRTransform	(	Stmt	stmt,
		const ffi::Function &	preorder,
		const ffi::Function &	postorder,
		ffi::Optional< ffi::Array< ffi::String > >	only_enable = `std::nullopt`
	)

recursively visit the ir nodes in post DFS order, and transform it

Parameters

stmt	The ir to be transformed.
preorder	The function called in before recursive mutation If preorder returns None, then the transform will proceed to recursive call. If preorder returns a not None Stmt/Expr, the transformer will simply return it and won't do further recursion.
postorder	The function called after recursive mutation. The recursive mutation result is passed to postorder for further mutation.
only_enable	List of String. If it is null, all IRNode will call preorder/postorder If it is not null, preorder/postorder will only be called when the IRNode's type key is in the list.

◆ is_const_int() [1/2]

bool tvm::tirx::is_const_int ( const PrimExpr & x )

inline

Check whether x is an integer constant.

Note: This only return true for integer types.

Returns: whether x is constant

◆ is_const_int() [2/2]

bool tvm::tirx::is_const_int	(	const PrimExpr &	x,
		int64_t	value
	)

inline

Check whether x is a constant integer expression.

Parameters

x	The input argument
value	the value to be compared against.

Returns: whether x is constant expression.

◆ is_const_number()

bool tvm::tirx::is_const_number ( const PrimExpr & x )

inline

Check whether x is an integer/float constant.

Note: This only return true for integer types.

Returns: whether x is constant

◆ is_const_power_of_two_integer()

bool tvm::tirx::is_const_power_of_two_integer	(	const PrimExpr &	x,
		int *	shift
	)

Check whether x is a constant power of two If x is power of two, write the power to the shift.

Parameters

x	The input expression.
shift	The output shift if x is power of two.

Returns: whether x is constant power of two

◆ is_negative_const()

bool tvm::tirx::is_negative_const ( const PrimExpr & a )

inline

◆ is_no_op()

bool tvm::tirx::is_no_op ( const tirx::Stmt & stmt )

inline

Check whether stmt is nop.

Parameters

stmt	The input statement

Returns: whether stmt is nop

◆ is_one()

bool tvm::tirx::is_one ( const PrimExpr & x )

inline

Check whether x is a constant integer 1.

Parameters

x	The input argument.

Note: This only return true for integer types.

Returns: whether x is constant 1

◆ is_positive_const()

bool tvm::tirx::is_positive_const ( const PrimExpr & a )

inline

◆ is_zero()

bool tvm::tirx::is_zero ( const PrimExpr & x )

inline

Check whether x is a constant integer 0.

Parameters

x	The input argument

Returns: whether x is constant 0

Note: This only return true for integer types.

◆ IsPointerType()

bool tvm::tirx::IsPointerType	(	const Type &	type,
		DLDataType	element_type
	)

inline

Check if type is a pointer to a runtime element type.

Parameters

type	The type to be checked.
element_type	The corresponding element type.

Returns: The check results

◆ IterVarType2String()

const char * tvm::tirx::IterVarType2String ( IterVarType t )

inline

◆ MakeConst()

template<typename ValueType , typename = typename std::enable_if<std::is_standard_layout<ValueType>::value && std::is_trivial<ValueType>::value>::type>

PrimExpr tvm::tirx::MakeConst	(	PrimType	dtype,
		ValueType	value,
		Span	span = `Span()`
	)

inline

Make a const value with certain data type.

Prefer direct IntImm or FloatImm construction when dtype is known to be scalar integer or floating point. This makes the compiled code more compact and efficient. Keep MakeConst for generic overload cases where dtype can be integer, floating point, or vector-valued and the caller needs its scalar/vector dispatch.

Parameters

dtype	The target type.
value	The input value

Returns: the result expression.

Template Parameters

ValueType The constant value type

Parameters

span	The location of this operation in the source.

◆ MakeConstScalar() [1/2]

template<>

PrimExpr tvm::tirx::MakeConstScalar	(	PrimType	dtype,
		bool	value,
		Span	span
	)

inline

◆ MakeConstScalar() [2/2]

template<typename ValueType >

PrimExpr tvm::tirx::MakeConstScalar	(	PrimType	dtype,
		ValueType	value,
		Span	span = `Span()`
	)

inline

◆ max()

const Op & tvm::tirx::max ( )

◆ maximum()

const Op & tvm::tirx::maximum ( )

◆ memset()

const Op & tvm::tirx::memset ( )

◆ min()

const Op & tvm::tirx::min ( )

◆ minimum()

const Op & tvm::tirx::minimum ( )

◆ mul()

const Op & tvm::tirx::mul ( )

◆ operator!=()

bool tvm::tirx::operator!=	(	const BufferVar &	lhs,
		const BufferVar &	rhs
	)

inline

◆ operator<<() [1/2]

std::ostream & tvm::tirx::operator<<	(	std::ostream &	os,
		CallEffectKind	side_effect
	)

inline

◆ operator<<() [2/2]

std::ostream & tvm::tirx::operator<<	(	std::ostream &	os,
		ForKind	kind
	)

◆ operator==()

bool tvm::tirx::operator==	(	const BufferVar &	lhs,
		const BufferVar &	rhs
	)

inline

◆ permute_layout()

const Op & tvm::tirx::permute_layout ( )

◆ PostOrderVisit()

void tvm::tirx::PostOrderVisit	(	const ffi::ObjectRef &	node,
		std::function< void(const ffi::ObjectRef &)>	fvisit
	)

Recursively visit a statement or expression in post DFS order, applying fvisit. Each node is guaranteed to be visited only once.

Parameters

node	The statement or expression to be visited.
fvisit	The visitor function to be applied.

◆ PreOrderVisit()

void tvm::tirx::PreOrderVisit	(	const ffi::ObjectRef &	stmt_or_expr,
		const std::function< bool(const ffi::ObjectRef &)> &	fvisit
	)

Recursively visit a statement or expression in pre DFS order, applying fvisit. If fvisit returns false, it won't visit the children of the node.

Parameters

stmt_or_expr	The statement or expression to be visited.
fvisit	The visitor function to be applied. If fvisit returns false, it won't visit the children of the node

◆ RebuildBufferVar()

BufferVar tvm::tirx::RebuildBufferVar	(	const BufferVar &	var,
		ffi::ObjectPtr< BufferTypeNode >	type,
		ffi::Optional< ffi::String >	name = `std::nullopt`
	)

inline

◆ reciprocal()

const Op & tvm::tirx::reciprocal ( )

◆ reduce_negate()

const Op & tvm::tirx::reduce_negate ( )

◆ ScopeBindingToStringPair()

std::pair< ffi::String, ffi::String > tvm::tirx::ScopeBindingToStringPair ( ScopeBinding binding )

Convert a ScopeBinding to its (parent, cur) string pair.

◆ ScopeKindHigher()

bool tvm::tirx::ScopeKindHigher	(	ScopeKind	a,
		ScopeKind	b
	)

inline

Strict-weak "a is wider than b" on scope kinds: world > kernel > cluster > cta > warpgroup > warp > thread. Only used by axe-layout scope-chain validity (the rest of the codebase compares scope identities with ==).

◆ ScopeKindToString()

std::string tvm::tirx::ScopeKindToString ( ScopeKind kind )

Convert a ScopeKind to its string name (e.g. kThread -> "thread").

◆ ScopeNameHigher()

bool tvm::tirx::ScopeNameHigher	(	const ffi::String &	a,
		const ffi::String &	b
	)

String-keyed convenience over ScopeKindHigher. FATALs on bad name.

◆ ScopeSwitch()

bool tvm::tirx::ScopeSwitch	(	const ActiveSet &	A,
		ScopeKind	scope_kind,
		ExecSplit *	out,
		std::string *	err
	)

Factor A into (inter, intra) for target scope_kind.

Returns false on factoring failure (warpgroup with warpid lane that crosses a warpgroup boundary unaligned) and writes reason to *err.

◆ select()

const Op & tvm::tirx::select ( )

◆ SetSeqIndex()

void tvm::tirx::SetSeqIndex	(	std::unordered_map< const StmtNode *, StmtSRef > &	stmt2ref,
		const Stmt &	stmt,
		int	seq_index,
		bool	include_loops = `true`
	)

inline

Set the StmtSRefNode::seq_index field for stmt.

Parameters

stmt2ref	The stmt2ref map to be updated with seq_index
stmt	The statement, or the realize node of the statement whose sref to be set
seq_index	The seq_index to be set
include_loops	Ignore ForNodes if this value is false

Note: The method is NOP for statements that are not schedulable, i.e. not For or SBlock

◆ SetSeqIndexInChildren()

void tvm::tirx::SetSeqIndexInChildren	(	std::unordered_map< const StmtNode *, StmtSRef > &	stmt2ref,
		const SeqStmtNode *	seq_stmt,
		bool	include_loops = `true`
	)

inline

Update seq_index of the children of a SeqStmt.

Parameters

stmt2ref	The stmt2ref map to be updated with indices
seq_stmt	The SeqStmt whose children need updating
include_loops	Ignore ForNodes if this value is false

◆ SideEffect()

CallEffectKind tvm::tirx::SideEffect ( const PrimExpr & expr )

Analyze the side effect of an expression.

Parameters

expr	The expression to be checked.

Returns: CallEffectKind, can be kPure, kReadState or kUpdateState

◆ silu()

const Op & tvm::tirx::silu ( )

◆ Specialize()

PrimFunc tvm::tirx::Specialize	(	PrimFunc	func,
		const ffi::Map< Var, ffi::Variant< BufferVar, Expr > > &	param_map
	)

Specialize parameters of PrimFunc.

Parameters

func	The PrimFunc to be specialized.
param_map	The mapping from function params to the instance.

Returns: The new function with parameter specialized.

Note: We can define a Meta TIR function with symbolic shape:

@T.prim_func
def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
    A = T.match_buffer(a, (m, n), "float32")
    B = T.match_buffer(b, (m, n), "float32")
    for i, j in T.grid(m, n):
        with T.sblock():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

Then we can make it specialized with given shapes or buffers.

a, _, m, n = mem_copy.params
func = mem_copy.specialize({a: tirx.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})

@T.prim_func
def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
    A = T.match_buffer(a, (16, 16), "float32")
    B = T.match_buffer(b, (16, 16), "float32")
    for i, j in T.grid(16, 16):
        with T.sblock():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

◆ sqrt()

const Op & tvm::tirx::sqrt ( )

◆ StringPairToScopeBinding()

ScopeBinding tvm::tirx::StringPairToScopeBinding	(	const ffi::String &	parent,
		const ffi::String &	cur
	)

Parse a (parent, cur) string pair to a ScopeBinding. FATAL if unknown.

◆ StringToScopeKind()

ScopeKind tvm::tirx::StringToScopeKind ( const ffi::String & name )

Parse a string name to a ScopeKind. FATAL if unknown.

◆ sub()

const Op & tvm::tirx::sub ( )

◆ Substitute() [1/11]

template<typename T >

ffi::Array< T > tvm::tirx::Substitute	(	const ffi::Array< T > &	arr,
		std::function< ffi::Optional< Expr >(const Var &var)>	vmap
	)

Substitute the var specified by vmap.

Parameters

arr	The array of Stmt/PrimExpr to be substituted
vmap	returns a new value if re-mapping is needed, otherwise returns nullptr.

Returns: The result.

◆ Substitute() [2/11]

IndexMap tvm::tirx::Substitute	(	const IndexMap &	index_map,
		std::function< ffi::Optional< PrimExpr >(const Var &var)>	f_subst
	)

Substitute variables in an index map.

Parameters

index_map	The index_map
f_subst	The substitution function

◆ Substitute() [3/11]

Range tvm::tirx::Substitute	(	const Range &	range,
		std::function< ffi::Optional< Expr >(const Var &var)>	vmap
	)

inline

Substitute the vars specified by vmap.

Parameters

range	The array of Stmt/PrimExpr to be substituted
vmap	returns a new value if re-mapping is needed, otherwise returns nullptr.

Returns: The modified Range.

◆ Substitute() [4/11]

Expr tvm::tirx::Substitute	(	Expr	expr,
		std::function< ffi::Optional< Expr >(const Var &var)>	vmap
	)

Substitute the var specified by vmap.

Parameters

expr	The source statement to be substituted
vmap	returns a new value if re-mapping is needed, otherwise returns nullptr.

Returns: The result.

◆ Substitute() [5/11]

template<typename Obj >

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const ffi::Map< Var, Expr > &	vmap
	)

Substitute the vars specified by vmap.

Delegates to the Substitute methods that use std::function. This overload allows braced-initialization of the Map, whereas the template<typename Expr> overload cannot.

Parameters

obj	The object in which TIR variables should be substituted
vmap	Map defining the TIR variables to be replaced

Returns: The modified object.

◆ Substitute() [6/11]

template<typename Obj , typename Replacement >

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const ffi::Map< Var, Replacement > &	vmap
	)

Substitute the vars specified by vmap.

Delegates to the Substitute methods that use std::function.

Parameters

obj	The object in which TIR variables should be substituted
vmap	Map defining the TIR variables to be replaced

Returns: The modified object.

◆ Substitute() [7/11]

template<typename Obj , typename Replacement >

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const std::unordered_map< const VarNode *, Replacement > &	vmap
	)

Substitute the vars specified by vmap.

Delegates to the Substitute methods that use std::function.

Parameters

obj	The object in which TIR variables should be substituted
vmap	Map defining the TIR variables to be replaced

Returns: The modified object.

◆ Substitute() [8/11]

template<typename Obj , typename Replacement >

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const std::unordered_map< IterVar, Replacement > &	iter_vmap
	)

Substitute the vars specified by vmap.

Delegates to the Substitute methods that use std::function.

Parameters

obj	The object in which TIR variables should be substituted
iter_vmap	Map defining the TIR variables to be replaced

Returns: The modified object.

◆ Substitute() [9/11]

template<typename Obj , typename Replacement , typename Hasher , typename EqualityChecker >

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const std::unordered_map< Var, Replacement, Hasher, EqualityChecker > &	vmap
	)

Substitute the vars specified by vmap.

Delegates to the Substitute methods that use std::function.

Parameters

obj	The object in which TIR variables should be substituted
vmap	Map defining the TIR variables to be replaced

Returns: The modified object.

◆ Substitute() [10/11]

PrimExpr tvm::tirx::Substitute	(	PrimExpr	expr,
		std::function< ffi::Optional< Expr >(const Var &var)>	vmap
	)

inline

◆ Substitute() [11/11]

Stmt tvm::tirx::Substitute	(	Stmt	stmt,
		std::function< ffi::Optional< Expr >(const Var &var)>	vmap
	)

Substitute the var specified by vmap.

Parameters

stmt	The source statement to be substituted
vmap	returns a new value if re-mapping is needed, otherwise returns nullptr.

Returns: The converted form.

◆ SubstituteWithDataTypeLegalization() [1/2]

PrimExpr tvm::tirx::SubstituteWithDataTypeLegalization	(	PrimExpr	expr,
		std::function< ffi::Optional< PrimExpr >(const Var &)>	vmap
	)

Substitute the var specified by vmap and legalize data types after substitution.

Parameters

expr	The source statement to be substituted
vmap	returns a new value if re-mapping is needed, otherwise returns nullptr.

Unlike Substitute, this allows the substitution to change the data type of the expression.

See also: Substitute

Returns: The result.

◆ SubstituteWithDataTypeLegalization() [2/2]

Stmt tvm::tirx::SubstituteWithDataTypeLegalization	(	Stmt	stmt,
		std::function< ffi::Optional< PrimExpr >(const Var &)>	vmap
	)

Substitute the var specified by vmap and legalize data types after substitution.

Parameters

stmt	The source statement to be substituted
vmap	returns a new value if re-mapping is needed, otherwise returns nullptr.

Unlike Substitute, this allows the substitution to change the data type of the expression.

See also: Substitute

Returns: The result.

◆ sum()

const Op & tvm::tirx::sum ( )

◆ TypeAnnotation()

PrimExpr tvm::tirx::TypeAnnotation	(	PrimType	dtype,
		Span	span = `Span()`
	)

Create a type annotation expression.

Parameters

dtype	The data type
span	The location of this object in the source code.

Returns: Expr a expression with dtype.

◆ unary_reduce()

const Op & tvm::tirx::unary_reduce ( )

◆ UndefinedVars() [1/3]

ffi::Array< Var > tvm::tirx::UndefinedVars ( const PrimExpr & expr )

Find undefined vars in the expression.

Parameters

expr	The expression to be checked.

Returns: Array of undefined vars.

◆ UndefinedVars() [2/3]

ffi::Array< Var > tvm::tirx::UndefinedVars	(	const PrimExpr &	expr,
		const ffi::Array< Var > &	defs
	)

Find undefined vars in the expression.

Parameters

expr	The expression to be checked.
defs	The vars that is defined.

Returns: Array of undefined vars.

◆ UndefinedVars() [3/3]

ffi::Array< Var > tvm::tirx::UndefinedVars	(	const Stmt &	stmt,
		const ffi::Array< Var > &	defs
	)

Find undefined vars in the statement.

Parameters

stmt	The statement to be checked.
defs	The vars that is defined.

Returns: Array of undefined vars.

◆ UsesVar() [1/2]

bool tvm::tirx::UsesVar	(	const PrimExpr &	expr,
		std::function< bool(const VarNode *)>	vset_contains
	)

Whether the given PrimExpr uses any var in the given variable set.

Parameters

expr	The PrimExpr to be checked.
vset_contains	The check function to see if var is in the variable set.

Returns: Whether expr uses any var in the given variable set.

◆ UsesVar() [2/2]

bool tvm::tirx::UsesVar	(	const Stmt &	stmt,
		std::function< bool(const VarNode *)>	vset_contains
	)

Whether the given Stmt uses any var in the given variable set.

Parameters

stmt	The Stmt to be checked.
vset_contains	The check function to see if a var is in the variable set.

Returns: Whether stmt uses any var in the given variable set.

◆ VerifyMemory()

bool tvm::tirx::VerifyMemory ( const PrimFunc & func )

Verify if memory accesses are legal for a specific target device type.

In the case that tgt is cuda, if not all workload is bound with threads, CPU code is generated that tries to access GPU memory, which is illegal. This pass performs verification for this case.

Parameters

func	The function to be verified.

Returns: Success of memory verification.

◆ VerifySSA()

bool tvm::tirx::VerifySSA ( const PrimFunc & func )

Verifies whether the IR stmt or Expr is in SSA form. That is: each Var is defined and assigned once(in Let/For)

Parameters

func	The function to be verified.

Returns: Whether IR is in SSA form.

Note: All passes in TIR consume and produce SSA form.

◆ VerifyWellFormed() [1/2]

bool tvm::tirx::VerifyWellFormed	(	const IRModule &	mod,
		bool	assert_mode = `true`
	)

Verify if the TIR in the given IRMOdule is well-formed.

In addition to the checks performed for each PrimFunc (see above), the following checks are performed:

The same TIR variable may not be defined in more than one function

Parameters

mod	The IRModule to be verified.
assert_mode	The indicator if it raises an error when the function is not well-formed.

Returns: Whether it is a well-formed TIR module.

◆ VerifyWellFormed() [2/2]

bool tvm::tirx::VerifyWellFormed	(	const PrimFunc &	func,
		bool	assert_mode = `true`
	)

Verify if the given TIR is well-formed. The verification includes:

All variables are defined prior to their point of use.
No variables are used outside of the scope of their definition.
Each variable has a single point of definition.
Expressions within a tirx::SBlock may not reference variables defined outside the block. For example, for a block with iter vars ‘vi, vj = T.axis.remap('SS’, [i,j]), the statement B[i,j] = A[i,j]would be ill-formed, because it uses the loop variablesiandjinstead of the block variablesviand vj`.

Parameters

func	The PrimFunc to be verified.
assert_mode	The indicator if it raises an error when the function is not well-formed.

Returns: Whether it is a well-formed TIR function.

◆ VisitPrimFuncs()

template<class FLambda >

void tvm::tirx::VisitPrimFuncs	(	const IRModule &	mod,
		FLambda	fvisit
	)

inline

Visit the PrimFuncs in the IRModule.

Template Parameters

FLambda The type of the PrimFunc visitor

Parameters

mod	The IRModule to be visited
fvisit	The visitor to the PrimFuncs in the IRModule

◆ zero()

const Op & tvm::tirx::zero ( )

Variable Documentation

◆ kPSUMBankNum

constexpr int tvm::tirx::kPSUMBankNum = 8

constexpr

◆ kPSUMMaxElemPerBank

constexpr int tvm::tirx::kPSUMMaxElemPerBank = 512

constexpr

◆ kWgSize

constexpr int tvm::tirx::kWgSize = 4

constexpr

Warpgroup size in warps (hardware-fixed).

Namespaces

Classes

Typedefs

Enumerations

Functions

Variables

Typedef Documentation

◆ AxisRegistry

◆ BufferAxis

◆ FArgSanitizer

◆ FAxisFuser

◆ FAxisSplitter

◆ FLegalize

◆ FLowerIntrinsic

◆ Region

◆ TCallEffectKind

◆ TDeviceIntrinsicNamespace

◆ TGlobalSymbol

◆ TIRVarAxis

◆ TIRxOpCategory

◆ TScriptDtypePrintLocation

◆ TScriptPrinterName

◆ TVectorizable

◆ Var

◆ VarNode

Enumeration Type Documentation

◆ CallEffectKind

◆ DepKind

◆ ForKind

◆ IterVarType

◆ ScopeBinding

◆ ScopeKind

◆ ScriptDtypePrintLocation

Function Documentation

◆ add()

◆ as_const_int()

◆ binary_chain()

◆ binary_reduce()

◆ BufferLoad()

◆ BufferWithOffsetAlignment()

◆ CalculateConstantBytes()

◆ CalculateExprComplexity()

◆ CalculateWorkspaceBytes()

◆ cast()

◆ compose_op()

◆ ConstHandle()

◆ ContainsNode()

◆ copy()

◆ copy_async()

◆ CopyBufferType()

◆ decl_buffer()

◆ DefaultIndexPrimType()

◆ DefaultIndexType()

◆ DetectBufferAccessLCA()

◆ EncodeSplitSide()

◆ exp()

◆ exp2()

◆ fdiv()

◆ fill()

◆ FilterNarrow()

◆ FindAnchorBlock()

◆ FindEntryFunc()

◆ fma()

◆ foldl()

◆ ForKind2String()

◆ gemm()

◆ gemm_async()

◆ GetBufferVar()

◆ GetSBlockAccessRegion()

◆ GetSBlockReadWriteRegion()

◆ GetShardingVarFromIndex()

◆ InitialActiveSet() [1/2]

◆ InitialActiveSet() [2/2]

◆ IRTransform()

◆ is_const_int() [1/2]

◆ is_const_int() [2/2]

◆ is_const_number()

◆ is_const_power_of_two_integer()

◆ is_negative_const()

◆ is_no_op()