Namespaces
	attr
	PrimFunc specific attribute names.

	builtin
	Collection of builtin intrinsics as ops.

	callback

	transform

Classes
class	BufferAxisHash

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	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	SLayout
	Managed reference to SLayoutNode. More...

class	SBijectiveLayoutNode

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	SBlockDependenceInfoNode
	An object that helps build and query block level dependences using the 2 core objects SBlockScope and StmtSRef. More...

class	SBlockDependenceInfo
	Managed reference to SBlockDependenceInfoNode. More...

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

class	StmtSRef
	Managed reference to StmtSRefNode. More...

class	SRefTreeCreator

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	Dependency
	Managed reference to DependencyNode. 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	SBlockScope
	Managed reference to SBlockScopeNode. More...

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

class	PipelineNode

class	Pipeline

class	CopyPipelineNode

class	CopyPipeline

class	BufferNode
	Node to represent a buffer. More...

class	Buffer
	Buffer is a symbolic n-darray structure. It is a composition of primitive symbolic types, used to specify the memory layout of the Tensor used in program input. More...

class	DataProducerNode
	Base node for data producers. More...

class	DataProducer
	Managed reference to DataProducerNode. More...

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

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

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	ExecContext
	Per-program-point ExecContext: active set + scope kind + split. More...

class	ScopeIdDefNode

class	ScopeIdDef

class	ScopeIdDefVerifier

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

class	ExecScopeNode

class	ExecScope

class	StringImmNode
	ffi::String constants, only used in asserts. More...

class	StringImm
	Managed reference to StringImmNode. More...

class	CastNode
	Cast value from one data type to another. More...

class	Cast
	Managed reference to CastNode. More...

class	BinaryOpNode
	Base template to implement binary ops. More...

class	AddNode
	a + b More...

class	Add
	Managed reference to AddNode. More...

class	SubNode
	a - b More...

class	Sub
	Managed reference to SubNode. More...

class	MulNode
	a * b More...

class	Mul
	Managed reference to MulNode. More...

class	DivNode
	a / b in the C semnatics. More...

class	Div
	Managed reference to DivNode. More...

class	ModNode
	a % b in the C semnatics. More...

class	Mod
	Managed reference to ModNode. More...

class	FloorDivNode
	Floor division, floor(a/b) More...

class	FloorDiv
	Managed reference to FloorDivNode. More...

class	FloorModNode
	The remainder of the floordiv. More...

class	FloorMod
	Managed reference to FloorModNode. More...

class	MinNode
	min(a, b) More...

class	Min
	Managed reference to MinNode. More...

class	MaxNode
	max(a, b) More...

class	Max
	Managed reference to MaxNode. More...

class	CmpOpNode
	Base template to implement comparison ops. More...

class	EQNode
	a == b More...

class	EQ
	Managed reference to EQNode. More...

class	NENode
	a != b More...

class	NE
	Managed reference to NENode. More...

class	LTNode
	a < b More...

class	LT
	Managed reference to LTNode. More...

struct	LENode
	a <= b More...

class	LE
	Managed reference to LENode. More...

class	GTNode
	a > b More...

class	GT
	Managed reference to GTNode. More...

class	GENode
	a >= b More...

class	GE
	Managed reference to GENode. More...

class	AndNode
	a && b More...

class	And
	Managed reference to AndNode. More...

class	OrNode
	a \|\| b More...

class	Or
	Managed reference to OrNode. More...

class	NotNode
	!a More...

class	Not
	Managed reference to NotNode. More...

class	SelectNode
	return true_value if condition is true, otherwise return false_value. More...

class	Select
	Managed reference to SelectNode. More...

class	BufferLoadNode
	Load value from the high dimension buffer. More...

class	BufferLoad
	Managed reference to BufferLoadNode. More...

class	ProducerLoadNode
	Load value from the result produced by the producer. More...

class	ProducerLoad
	Managed reference to ProducerLoadNode. More...

class	RampNode
	Construct a vector with lanes elements where its i-th element equals base + i * stride. This is useful to construct a index for a continuous vector load. More...

class	Ramp
	Managed reference to RampNode. More...

class	BroadcastNode
	Create a vector where all the elements are value. More...

class	Broadcast
	Managed reference to BroadcastNode. More...

class	LetNode
	Let binding. Bind var to value then evaluate body. More...

class	Let
	Managed reference to LetNode. More...

class	CallNode
	Call node. More...

class	Call
	Managed reference to CallNode. More...

class	ShuffleNode
	Shuffle instruction. vec = concat(vectors) result = (vec[indices[0]], vec[indices[1]] ...) More...

class	Shuffle
	Managed reference to ShuffleNode. More...

class	CommReducerNode
	A commutative reducer node to represent a commutative binary operator with identity element. More...

class	CommReducer
	Managed reference to CommReducerNode. More...

class	ReduceNode
	Reduction operator. More...

class	Reduce
	Managed reference to ReduceNode. 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 PrimExpr &n, Args...)>

class	ExprVisitor
	ExprVisitor. More...

class	ExprMutator
	ExprMutator that mutates expressions. More...

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

class	PrimFunc
	Managed reference to PrimFuncNode. More...

class	TensorIntrinNode
	Tensor intrinsics for tensorization. More...

class	TensorIntrin
	Managed reference to TensorIntrinNode. More...

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

class	IndexMap

class	AxisAttrMap

class	LayoutNode

class	Layout

class	AxisNode

class	Axis

class	AxisRegEntry

class	IterNode

class	Iter

class	TileLayoutNode

class	TileLayout

class	SwizzleLayoutNode

class	SwizzleLayout

class	ComposeLayoutNode

class	ComposeLayout

class	PredicateNode

class	Predicate

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

class	Stmt
	Container of all statements. More...

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

class	Bind
	Managed reference to BindNode. 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	AttrStmt
	Managed reference to AttrStmtNode. More...

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

class	AssertStmt
	Managed reference to AssertStmtNode. More...

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

class	BufferStore
	Managed reference to BufferStoreNode. More...

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

class	DeclBuffer
	Managed reference to DeclBufferNode. More...

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

class	AllocBuffer
	Managed reference to AllocBufferNode. More...

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

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

class	Evaluate
	Managed reference to EvaluateNode. More...

class	SeqStmt
	Sequence statement. More...

class	IfThenElseNode
	IfThenElse statement. More...

class	IfThenElse
	Managed reference to IfThenElseNode. More...

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

class	For
	Managed reference to ForNode. More...

class	WhileNode
	A While loop. More...

class	While
	Managed reference to WhileNode. More...

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

class	Break
	Managed reference to BreakNode. More...

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

class	Continue
	Managed reference to ContinueNode. More...

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

class	BufferRegion
	Managed reference to BufferRegionNode. 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	MatchBufferRegion
	Managed reference to MatchBufferRegionNode. More...

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

class	SBlock
	Managed reference to SBlockNode. More...

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

class	SBlockRealize
	Managed reference to BlockRealizeNode. More...

class	ExecScopeStmtNode
	A statement that annotates the execution scope for its body. More...

class	ExecScopeStmt
	Managed reference to ExecScopeStmtNode. More...

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

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

class	StmtVisitor
	StmtVisitor. More...

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

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

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

class	ScheduleContextNode
	The context information of the kernel required by op schedule. More...

class	ScheduleContext
	Managed reference to ScheduleContextNode. More...

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

class	DispatchContext
	Managed reference to DispatchContextNode. More...

class	TilePrimitiveCallNode
	TIRX TilePrimitiveCall stmt. More...

class	TilePrimitiveCall
	Managed reference to TilePrimitiveCallNode. More...

class	VarNode
	A variable node in the IR. More...

class	Var
	a named variable in TIR More...

class	SizeVarNode
	A variable node represent a tensor index size, whose value must be non-negative. More...

class	SizeVar
	a named variable represents a tensor index size More...

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

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

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

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

using	IntImmNode = tvm::IntImmNode

using	FloatImmNode = tvm::FloatImmNode

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. More...

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

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

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

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

using	TScriptDtypePrintLocation = Integer

using	TCallEffectKind = Integer
	Use integer to record the kind. More...

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

using	FOpScheduler = ffi::TypedFunction< Stmt(tvm::Op, ffi::Array< ffi::ObjectRef >, ScheduleContext)>
	The type of the function that schedules a TIRX operator. More...

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	BufferType : int { kDefault = 1 , kAutoBroadcast = 2 }
	buffer type More...

enum class	ScopeKind : int { kWorld = 0 , kKernel = 1 , kCluster = 2 , kCta = 3 , kWarpgroup = 4 , kWarp = 5 , kThread = 6 }
	The target execution scope kind of an ExecScopeStmt. 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 , kOpaque = 3 , kOpaque = kUpdateState , kUnrolled = 5 , kUnrolled = 3 , kVectorized = 6 , kVectorized = 2 , 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, 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. More...

ffi::Array< ffi::Array< BufferRegion > >	GetSBlockAccessRegion (const SBlock &block, const ffi::Map< Var, Buffer > &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. More...

ffi::Array< ffi::Array< BufferRegion > >	GetSBlockReadWriteRegion (const SBlock &block, const ffi::Map< Var, Buffer > &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. More...

ffi::Map< Buffer, 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. More...

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. More...

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. More...

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. More...

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

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

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

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

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

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

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

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) More...

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

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

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

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

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

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

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. More...

DataType	DefaultIndexType ()
	if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32 More...

Buffer	decl_buffer (ffi::Array< PrimExpr > shape, DataType dtype=DataType::Float(32), ffi::String name="buffer", ffi::String storage_scope="", ffi::Optional< ffi::Array< IntImm >> axis_separators=std::nullopt, Span span=Span())
	Construct a new buffer given shape, and dtype. More...

tirx::Buffer	BufferWithOffsetAlignment (ffi::Array< PrimExpr > shape, DataType dtype, std::string name, int data_alignment, int offset_factor, bool compact, std::string memory_scope="")
	Creates TIR Buffer for provided parameters. More...

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

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`. More...

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

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. More...

std::string	ScopeKindToString (ScopeKind kind)
	Convert a ScopeKind to its string name (e.g. kKernel -> "kernel"). More...

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

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

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

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 ==). More...

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

template<typename K , typename V >
std::unordered_map< K, V >	as_unordered_map (const ffi::Map< K, V > &dmap)

PrimFunc	Specialize (PrimFunc func, const ffi::Map< Var, ffi::Variant< Buffer, PrimExpr >> &param_map)
	Specialize parameters of PrimFunc. More...

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

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

template<typename ValueType , typename = typename std::enable_if<std::is_pod<ValueType>::value>::type>
PrimExpr	make_const (DataType t, ValueType value, Span span=Span())
	Make a const value with certain data type. More...

PrimExpr	make_zero (DataType t, Span span=Span())
	Make a const zero expr. More...

PrimExpr	const_true (int lanes=1, Span span=Span())
	Make a constant true expression. More...

PrimExpr	const_false (int lanes=1, Span span=Span())
	Make a constant false expression. More...

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

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

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

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

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

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

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

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

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. More...

bool	is_positive_const (const PrimExpr &a)

bool	is_negative_const (const PrimExpr &a)

template<typename ValueType >
PrimExpr	MakeConstScalar (DataType t, ValueType value, Span span=Span())

template<>
PrimExpr	MakeConstScalar (DataType t, bool value, Span span)

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

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

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 More...

void	PostOrderVisit (const ffi::ObjectRef &node, std::function< void(const ffi::ObjectRef &)> fvisit)
	Recursively visit the ir in post DFS order node, apply fvisit Each node is guaranteed to be visited only once. More...

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

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

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

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

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

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

template<typename Obj , typename Expr , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>
auto	Substitute (Obj &&obj, const std::unordered_map< const VarNode *, Expr > &vmap)
	Substitute the vars specified by vmap. More...

template<typename Obj , typename Expr , typename Hasher , typename EqualityChecker , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>
auto	Substitute (Obj &&obj, const std::unordered_map< Var, Expr, Hasher, EqualityChecker > &vmap)
	Substitute the vars specified by vmap. More...

template<typename Obj , typename Expr , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>
auto	Substitute (Obj &&obj, const std::unordered_map< IterVar, Expr > &iter_vmap)
	Substitute the vars specified by vmap. More...

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

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

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

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. More...

const Op &	cast ()
	See pesudo code below: More...

const Op &	permute_dims ()
	See pesudo code below: More...

const Op &	copy ()
	See pesudo code below: More...

const Op &	copy_async ()
	See pesudo code below: More...

const Op &	fill ()
	See pesudo code below: More...

const Op &	gemm ()
	See pesudo code below: More...

const Op &	gemm_async ()
	See pesudo code below: More...

const Op &	zero ()

const Op &	sqrt ()

const Op &	exp ()

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 &	tvm_kernel_replace_point ()
	See pesudo code below: More...

const char *	IterVarType2String (IterVarType t)

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

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<Buffer, 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.

◆ FloatImmNode

using tvm::tirx::FloatImmNode = typedef tvm::FloatImmNode

◆ FLowerIntrinsic

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

The intrinsic lowering function for given op.

◆ FOpScheduler

using tvm::tirx::FOpScheduler = typedef ffi::TypedFunction<Stmt(tvm::Op, ffi::Array<ffi::ObjectRef>, ScheduleContext)>

The type of the function that schedules a TIRX operator.

Parameters

op	The operator.
args	The arguments.
context	The schedule context.

◆ IntImmNode

using tvm::tirx::IntImmNode = typedef tvm::IntImmNode

◆ Region

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

◆ TCallEffectKind

using tvm::tirx::TCallEffectKind = typedef Integer

Use integer to record the kind.

◆ 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>

◆ TScriptDtypePrintLocation

using tvm::tirx::TScriptDtypePrintLocation = typedef Integer

◆ 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.

Enumeration Type Documentation

◆ BufferType

enum tvm::tirx::BufferType : int

buffer type

Enumerator
kDefault
kAutoBroadcast

◆ CallEffectKind

enum 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 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 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
kOpaque
kOpaque	Opaque function, cannot make any assumption.
kUnrolled	The execution is unrolled.
kUnrolled	The loop body must be unrolled.
kVectorized	The loop is vectorized.
kVectorized	Vector SIMD loop. The loop body will be vectorized.
kParallelized	The loop is parallelized.
kTensorized	Marks boundary of tensorization intrinsic.

◆ ScopeBinding

enum 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; T.filter(var, lo, hi) 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 tvm::tirx::ScopeKind : int

strong

The target execution scope kind of an ExecScopeStmt.

Replaces the string-keyed name of ExecScope. One value per user-facing with T.<kind>(): construct, plus kWorld for the cross-kernel root scope used by axe-layout's pid axis. Ordered from coarsest to finest; smaller integer = wider scope, so ScopeKindHigher is a plain <.

Enumerator
kWorld
kKernel
kCluster
kCta
kWarpgroup
kWarp
kThread

◆ ScriptDtypePrintLocation

enum 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.

◆ as_unordered_map()

template<typename K , typename V >

std::unordered_map<K, V> tvm::tirx::as_unordered_map ( const ffi::Map< K, V > & dmap )

inline

◆ binary_chain()

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

◆ binary_reduce()

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

◆ BufferWithOffsetAlignment()

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

Creates TIR Buffer for 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.
compact	If the statement has already bound to a compact buffer.
memory_scope	memory scope of the buffer

◆ CalculateConstantBytes()

size_t tvm::tirx::CalculateConstantBytes	(	const PrimFunc &	func,
		const Integer &	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,
		const Integer &	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)

◆ const_false()

PrimExpr tvm::tirx::const_false	(	int	lanes = `1`,
		Span	span = `Span()`
	)

inline

Make a constant false expression.

Parameters

lanes	The number of lanes in the bool
span	The location of this operation in the source.

Returns: The result expression.

◆ const_true()

PrimExpr tvm::tirx::const_true	(	int	lanes = `1`,
		Span	span = `Span()`
	)

inline

Make a constant true expression.

Parameters

lanes	The number of lanes in the bool
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)

◆ decl_buffer()

Buffer tvm::tirx::decl_buffer	(	ffi::Array< PrimExpr >	shape,
		DataType	dtype = `DataType::Float(32)`,
		ffi::String	name = `"buffer"`,
		ffi::String	storage_scope = `""`,
		ffi::Optional< ffi::Array< IntImm >>	axis_separators = `std::nullopt`,
		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
axis_separators	Divisions defining the groups of axes that will be flattened together.
span	The location of this object in the source code.

Returns: The created buffer.

See also: Buffer for complete constructor.

◆ DefaultIndexType()

DataType tvm::tirx::DefaultIndexType ( )

inline

if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32

◆ DetectBufferAccessLCA()

ffi::Map<Buffer, 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 ( )

◆ 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.

◆ 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(Buffer A, Buffer B, Buffer C, Buffer 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)

◆ GetSBlockAccessRegion()

ffi::Array<ffi::Array<BufferRegion> > tvm::tirx::GetSBlockAccessRegion	(	const SBlock &	block,
		const ffi::Map< Var, Buffer > &	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, Buffer > &	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,
		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 T.kernel() 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,
		const DataType &	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

◆ make_const()

template<typename ValueType , typename = typename std::enable_if<std::is_pod<ValueType>::value>::type>

PrimExpr tvm::tirx::make_const	(	DataType	t,
		ValueType	value,
		Span	span = `Span()`
	)

inline

Make a const value with certain data type.

Parameters

t	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.

◆ make_zero()

PrimExpr tvm::tirx::make_zero	(	DataType	t,
		Span	span = `Span()`
	)

inline

Make a const zero expr.

Parameters

t	The target type.
span	The location of this operation in the source.

Returns: the result expression.

◆ MakeConstScalar() [1/2]

template<>

PrimExpr tvm::tirx::MakeConstScalar	(	DataType	t,
		bool	value,
		Span	span
	)

inline

◆ MakeConstScalar() [2/2]

template<typename ValueType >

PrimExpr tvm::tirx::MakeConstScalar	(	DataType	t,
		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<<() [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
	)

◆ permute_dims()

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

See pesudo code below:

Tx.permute_dims(BufferRegion buffer, List order)

◆ PostOrderVisit()

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

Recursively visit the ir in post DFS order node, apply fvisit Each node is guaranteed to be visited only once.

Parameters

node	The ir 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 the IR in pre DFS order node, apply fvisit. If fvisit returns false, it won't visit the children of the node.

Parameters

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

◆ 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. kKernel -> "kernel").

◆ 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

◆ Specialize()

PrimFunc tvm::tirx::Specialize	(	PrimFunc	func,
		const ffi::Map< Var, ffi::Variant< Buffer, PrimExpr >> &	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/10]

template<typename T >

ffi::Array<T> tvm::tirx::Substitute	(	const ffi::Array< T > &	arr,
		std::function< ffi::Optional< PrimExpr >(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/10]

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/10]

Range tvm::tirx::Substitute	(	const Range &	range,
		std::function< ffi::Optional< PrimExpr >(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/10]

template<typename Obj , typename Expr , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>

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.

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() [5/10]

template<typename Obj >

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const ffi::Map< Var, PrimExpr > &	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/10]

template<typename Obj , typename Expr , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const std::unordered_map< const VarNode *, Expr > &	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/10]

template<typename Obj , typename Expr , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const std::unordered_map< IterVar, Expr > &	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() [8/10]

template<typename Obj , typename Expr , typename Hasher , typename EqualityChecker , typename = std::enable_if_t<std::is_base_of_v<PrimExpr, Expr>>>

auto tvm::tirx::Substitute	(	Obj &&	obj,
		const std::unordered_map< Var, Expr, 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() [9/10]

PrimExpr tvm::tirx::Substitute	(	PrimExpr	expr,
		std::function< ffi::Optional< PrimExpr >(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() [10/10]

Stmt tvm::tirx::Substitute	(	Stmt	stmt,
		std::function< ffi::Optional< PrimExpr >(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 ( )

◆ tvm_kernel_replace_point()

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

See pesudo code below:

tvm_kernel_replace_point()

◆ TypeAnnotation()

PrimExpr tvm::tirx::TypeAnnotation	(	DataType	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

◆ FloatImmNode

◆ FLowerIntrinsic

◆ FOpScheduler

◆ IntImmNode

◆ Region

◆ TCallEffectKind

◆ TGlobalSymbol

◆ TIRVarAxis

◆ TScriptDtypePrintLocation

◆ TScriptPrinterName

◆ TVectorizable

Enumeration Type Documentation

◆ BufferType

◆ CallEffectKind

◆ DepKind

◆ ForKind

◆ IterVarType

◆ ScopeBinding

◆ ScopeKind

◆ ScriptDtypePrintLocation

Function Documentation

◆ add()

◆ as_const_int()

◆ as_unordered_map()

◆ binary_chain()

◆ binary_reduce()

◆ BufferWithOffsetAlignment()

◆ CalculateConstantBytes()

◆ CalculateExprComplexity()

◆ CalculateWorkspaceBytes()

◆ cast()

◆ const_false()

◆ const_true()

◆ ContainsNode()

◆ copy()

◆ copy_async()

◆ decl_buffer()

◆ DefaultIndexType()

◆ DetectBufferAccessLCA()

◆ EncodeSplitSide()

◆ exp()

◆ fdiv()

◆ fill()

◆ FilterNarrow()

◆ FindAnchorBlock()

◆ FindEntryFunc()

◆ foldl()

◆ ForKind2String()

◆ gemm()

◆ gemm_async()

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

◆ is_one()

◆ is_positive_const()

◆ is_zero()

◆ IsPointerType()

◆ IterVarType2String()