Backend-defined execution hierarchy nodes shared by layouts and tile calls. More...

Namespaces
namespace	attr
	Function specific attribute names.

namespace	op_attr

namespace	tile
	Dispatch interface for lowering tile instruction Calls.

namespace	transform

Classes
class	Axis

class	AxisNode

struct	CallFFIKernelAttr
	Launch metadata for call_ffi_kernel. More...

class	ComposeLayout

class	ComposeLayoutNode

class	Dependency
	Managed reference to DependencyObj. More...

class	DependencyObj
	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	ExecScope

class	ExecScopeNode

class	Function
	Managed reference to FunctionNode. More...

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

class	IndexMap

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

class	Iter

class	IterNode

class	Layout

class	LayoutNode

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	SBlockDependenceInfo
	Managed reference to SBlockDependenceInfoObj. More...

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

class	SBlockScope
	Managed reference to SBlockScopeObj. More...

class	SBlockScopeObj
	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	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	StmtExprMutator
	Mutator that recursively mutates stmts and exprs on them. More...

class	StmtExprMutatorWithAnalyzer
	IRMutator with an analyzer context. More...

class	StmtExprVisitor
	Native visitor for TIRx statements and their expression operands. More...

class	StmtExprVisitorWithAnalyzer

class	StmtFunctor
	Type-dispatched statement functor with a caller-selected signature. More...

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

class	StmtSRef
	Managed reference to StmtSRefObj. More...

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

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

class	TensorAxisHash

class	TensorMapType
	Managed reference to TensorMapTypeNode. More...

class	TensorMapTypeNode
	The type of tensor map. More...

class	TensorType
	Managed reference to TensorTypeNode. More...

class	TensorTypeNode
	Structural type of a TIRx tensor variable. More...

class	TensorVar
	Checked zero-state view over an ordinary VarNode with TensorType. More...

class	TileLayout

class	TileLayoutNode

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

using	TensorAxis = std::pair< TensorVar, 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	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.

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	TVMStructFieldKind : int { kDLTensorAddr , kDLTensorData , kDLTensorShape , kDLTensorStrides , kDLTensorNDim , kDLTensorTypeCode , kDLTensorTypeBits , kDLTensorTypeLanes , kDLTensorByteOffset , kDLTensorDeviceId , kDLTensorDeviceType , kDLTensorKindBound_ , kTVMValueContent , kTVMFFIAnyTypeIndex , kTVMFFIAnyZeroPadding , kTVMFFIAnyUnionValue , kTVMValueKindBound_ , kInt64ArrayElem }

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

ffi::Array< ffi::Array< TensorRegion > >	GetSBlockAccessRegion (const s_tir::SBlock &block, const ffi::Map< Var, TensorVar > &tensor_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< TensorRegion > >	GetSBlockReadWriteRegion (const s_tir::SBlock &block, const ffi::Map< Var, TensorVar > &tensor_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< TensorVar, ffi::Optional< Stmt > >	DetectTensorAccessLCA (const Function &func)
	Detect the lowest common ancestor(LCA) of tensor access, including both high-level access(TensorLoad, TensorStore) and low-level access(Load, Store and opaque access). The LCA may be a For loop or a Block.

const s_tir::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 `StmtSRefObj::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.

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

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

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

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

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

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

const FunctionNode *	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 Function.

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

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

TensorVar	GetTensorVar (const VarNode *var)
	Recover a checked tensor view from an ordinary VarNode pointer.

ffi::ObjectPtr< TensorTypeNode >	CopyTensorType (const TensorVar &var)

TensorVar	RebuildTensorVar (const TensorVar &var, ffi::ObjectPtr< TensorTypeNode > type, ffi::Optional< ffi::String > name=std::nullopt)

TVM_TIRX_DLL TensorVar	decl_tensor (ffi::Array< PrimExpr > shape, PrimType dtype=PrimType::Float(32), ffi::String name="tensor", ffi::String storage_scope="", Location loc=UnknownLoc())
	Construct a new tensor given shape, and dtype.

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

TVM_TIRX_DLL TensorLoad	MakeTensorLoad (TensorVar tensor, ffi::Array< PrimExpr > indices, Location loc=UnknownLoc())
	Construct a TensorLoad from a TensorVar.

TVM_TIRX_DLL TensorRegion	FullTensorRegion (TensorVar tensor)
	Select the entire tensor.

TVM_TIRX_DLL TensorRegion	TensorRegionFromPoint (TensorVar tensor, ffi::Array< PrimExpr > indices)
	Construct unit or vector-lane ranges from point indices.

TVM_TIRX_DLL Function	RenewDef (Function func)
	Renew variable definitions and remap their uses in a function.

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

TVM_TIRX_DLL const Op &	call_extern_op ()
	Call an external C function with argument and result types supplied by the caller.

TVM_TIRX_DLL const Op &	call_pure_extern_op ()
	Call a pure external C function with types supplied by the caller.

TVM_TIRX_DLL const Op &	call_llvm_intrin_op ()
	Invoke an LLVM intrinsic.

TVM_TIRX_DLL const Op &	call_llvm_pure_intrin_op ()
	Invoke a pure LLVM intrinsic.

TVM_TIRX_DLL const Op &	call_spirv_pure_glsl450_op ()
	Invoke a pure GLSL450 SPIR-V intrinsic.

TVM_TIRX_DLL const Op &	abi_field_get_op ()
	Read a runtime structure field.

TVM_TIRX_DLL const Op &	abi_field_set_op ()
	Write a runtime structure field.

TVM_TIRX_DLL const Op &	throw_last_error_op ()
	Raise the last runtime error.

TVM_TIRX_DLL const Op &	stack_alloca_op ()
	Allocate stack storage for runtime values.

TVM_TIRX_DLL const Op &	stack_make_shape_op ()
	Construct a stack-allocated shape tuple.

TVM_TIRX_DLL const Op &	stack_make_dltensor_op ()
	Construct a stack-allocated DLTensor.

TVM_TIRX_DLL const Op &	call_packed_op ()
	Invoke a runtime packed function.

TVM_TIRX_DLL const Op &	call_ffi_kernel_op ()
	Launch a kernel using the packed-function argument convention.

TVM_TIRX_DLL const Op &	call_cpacked_op ()
	Invoke a C packed function.

TVM_TIRX_DLL const Op &	call_packed_lowered_op ()
	Invoke a packed function using an explicit argument stack.

TVM_TIRX_DLL const Op &	call_cpacked_lowered_op ()
	Invoke a C packed function using an explicit argument stack.

TVM_TIRX_DLL const Op &	kernel_replace_point_op ()
	Marker where a transform should replace generated kernel initialization.

TVM_TIRX_DLL const Op &	ignore_loop_partition_op ()
	Annotate a predicate not be considered as target condition of loop partition.

TVM_TIRX_DLL const Op &	cpu_parallel_barrier_op ()
	Synchronize all workers in the current CPU parallel launch.

TVM_TIRX_DLL const Op &	gpu_thread_return_op ()
	Return from a GPU thread without returning a function value.

TVM_TIRX_DLL PrimExpr	gpu_thread_return (Location loc=UnknownLoc())
	Return from a GPU thread without returning a function value.

TVM_TIRX_DLL const Op &	gpu_thread_filter_op ()
	Mark an opaque predicate that selects one thread on an active-set axis.

TVM_TIRX_DLL const Op &	gpu_active_thread_selector_op ()
	Analysis-only active-thread selector.

TVM_TIRX_DLL const Op &	gpu_thread_invariant_op ()
	Mark a condition to be thread invariant.

TVM_TIRX_DLL const Op &	gpu_storage_sync_op ()
	Synchronize accesses in a storage scope.

TVM_TIRX_DLL const Op &	gpu_warp_shuffle_op ()
	Read the value from a selected lane in the warp.

TVM_TIRX_DLL const Op &	gpu_warp_shuffle_up_op ()
	Read the value from a lower lane in the warp.

TVM_TIRX_DLL const Op &	gpu_warp_shuffle_down_op ()
	Read the value from a higher lane in the warp.

TVM_TIRX_DLL const Op &	gpu_warp_shuffle_xor_op ()
	Read the value from the lane selected by an XOR mask.

TVM_TIRX_DLL const Op &	gpu_warp_activemask_op ()
	Return the bit mask of active lanes in the warp.

TVM_TIRX_DLL const Op &	gpu_thread_allreduce_op ()
	Cross-thread reduction with an explicit typed combiner and identities.

TVM_TIRX_DLL const Op &	gpu_dp4a_op ()
	Dot product of two int8x4 vectors and add an optional accumulator.

TVM_TIRX_DLL PrimExpr	logaddexp (PrimExpr a, PrimExpr b, Location loc=UnknownLoc())
	Compute log(exp(a) + exp(b)).

	TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED (logaddexp)

TVM_TIRX_DLL PrimExpr	fast_erf_float_expr (PrimExpr arg, int bits)
	Fast_erf_float expression from Eigen.

TVM_TIRX_DLL const Op &	alloc_tensor_op ()
	Allocate a tensor with exactly one TensorType in ty_args.

TVM_TIRX_DLL const Op &	decl_tensor_op ()
	Declare a tensor view with exactly one TensorType in ty_args.

TVM_TIRX_DLL const Op &	reinterpret_op ()
	Reinterpret the bits of a value as an explicit primitive or pointer type.

TVM_TIRX_DLL Expr	reinterpret (Type target_ty, Expr value, Location loc=UnknownLoc())
	perform reinterpret cast value to type.

PrimExpr	reinterpret (PrimType t, PrimExpr value, Location loc=UnknownLoc())

Expr	ConstHandle (int64_t value, Location loc=UnknownLoc())
	Construct an opaque pointer from its integer payload.

TVM_TIRX_DLL const Op &	address_of_op ()
	Return the address of a tensor element or addressable variable.

TVM_TIRX_DLL const Op &	isnullptr_op ()
	Test whether a pointer is null.

TVM_TIRX_DLL const Op &	prefetch_op ()
	Prefetch a memory address.

TVM_TIRX_DLL const Op &	ptr_byte_offset_op ()
	Cast a handle to a typed pointer after adding a byte offset.

TVM_TIRX_DLL const Op &	static_handle_op ()
	Create a function local static handle that iniitalizes to nullptr.

TVM_TIRX_DLL const Op &	handle_add_byte_offset_op ()
	Add a byte offset to an opaque pointer.

TVM_TIRX_DLL const Op &	atomic_add_op ()
	atomic add instruction, corresponding e.g. to atomicAdd in CUDA.

TVM_TIRX_DLL const Op &	assume_aligned_op ()
	Assume a tensor's base address has the given constant byte alignment.

TVM_TIRX_DLL const Op &	undef_op ()
	Returns an initialized but arbitrary value.

TVM_TIRX_DLL const Op &	masked_load_op ()
	Masked tensor load.

TVM_TIRX_DLL const Op &	masked_store_op ()
	Masked tensor store.

TVM_TIRX_DLL const Op &	tensor_data_ptr_op ()
	Project the physical pointer associated with a TensorVar definition.

TVM_TIRX_DLL const Op &	launch_thread_op ()
	Thread launch region: operands are a nonempty StringImm tag and a scalar signed or unsigned integer extent wider than one bit. The sole body parameter is a fresh thread-index PrimVar matching the extent type. Tags starting with vthread denote virtual threads. There are no attrs or results.

TVM_TIRX_DLL const Op &	device_entry_op ()
	Mark a user-facing device entry containing device scope definitions. Takes no operands, body parameters, attributes or results.

TVM_TIRX_DLL const Op &	device_context_op ()
	Supply lexical device context for allocation and packed-call lowering. Operands are integer device type and device ID; there are no body parameters, attributes or results. The region does not change the active runtime device.

TVM_TIRX_DLL const Op &	compute_scope_op ()
	Outline the body as a CPU compute helper named by a StringImm operand. There are no body parameters, attributes or results.

TVM_TIRX_DLL const Op &	parallel_launch_op ()
	Launch a CPU worker team around parallel loops and team barriers. Takes no operands, body parameters, attributes or results.

TVM_TIRX_DLL const Op &	vector_high_op ()
	Get the high level half of the vector.

TVM_TIRX_DLL const Op &	vector_low_op ()
	Get the low-level half of the vector.

TVM_TIRX_DLL const Op &	vector_combine_op ()
	Concat two vectors.

TVM_TIRX_DLL const Op &	get_active_lane_mask_op ()
	Calculate a predicate mask given an upper bound (limit) and a current value (base).

bool	IsNoOp (const Stmt &stmt)
	Whether stmt is undefined, an integer Evaluate, or an empty sequence.

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

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

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.

TVM_TIRX_DLL void	RegisterExecScope (const ffi::String &backend, const ffi::String &name, int rank)
	Register a backend's execution hierarchy. Smaller rank means wider scope.

TVM_TIRX_DLL bool	ScopeHigher (const ExecScope &a, const ExecScope &b)
	Compare scopes within the same backend's hierarchy.

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

DLDataType	DefaultIndexType ()

Variables
constexpr int	kPSUMMaxElemPerBank = 512

constexpr int	kPSUMBankNum = 8

Detailed Description

Backend-defined execution hierarchy nodes shared by layouts and tile calls.

Typedef Documentation

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

◆ TensorAxis

using tvm::tirx::TensorAxis = typedef std::pair<TensorVar, int>

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

◆ TVectorizable

using tvm::tirx::TVectorizable = typedef bool

Whether the op is overloaded for vector form.

Enumeration Type Documentation

◆ 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

◆ TVMStructFieldKind

enum tvm::tirx::TVMStructFieldKind : int

Enumerator
kDLTensorAddr
kDLTensorData
kDLTensorShape
kDLTensorStrides
kDLTensorNDim
kDLTensorTypeCode
kDLTensorTypeBits
kDLTensorTypeLanes
kDLTensorByteOffset
kDLTensorDeviceId
kDLTensorDeviceType
kDLTensorKindBound_
kTVMValueContent
kTVMFFIAnyTypeIndex
kTVMFFIAnyZeroPadding
kTVMFFIAnyUnionValue
kTVMValueKindBound_
kInt64ArrayElem

Function Documentation

◆ abi_field_get_op()

TVM_TIRX_DLL const Op & tvm::tirx::abi_field_get_op ( )

Read a runtime structure field.

Arguments, in order:

args[0]: arr, The array.
args[1]: index, The index.
args[2]: field, The field index.

◆ abi_field_set_op()

TVM_TIRX_DLL const Op & tvm::tirx::abi_field_set_op ( )

Write a runtime structure field.

Arguments, in order:

args[0]: arr, The array.
args[1]: index, The index.
args[2]: field, The field index.
args[3]: value, The value to use.

◆ address_of_op()

TVM_TIRX_DLL const Op & tvm::tirx::address_of_op ( )

Return the address of a tensor element or addressable variable.

Arguments, in order:

args[0]: obj, The referenced object.

◆ alloc_tensor_op()

TVM_TIRX_DLL const Op & tvm::tirx::alloc_tensor_op ( )

Allocate a tensor with exactly one TensorType in ty_args.

The optional args[0] is a Tuple of producer-owned placement coordinates. DictAttrs directly holds allocation annotations. The TensorType owns all shape, dtype, storage scope and access/layout metadata, and is the result type.

◆ assume_aligned_op()

TVM_TIRX_DLL const Op & tvm::tirx::assume_aligned_op ( )

Assume a tensor's base address has the given constant byte alignment.

This leaf operation carries a compiler fact, without checking or changing the address. Alignment must be a power of two between 1 and 2^27 bytes.

◆ atomic_add_op()

TVM_TIRX_DLL const Op & tvm::tirx::atomic_add_op ( )

atomic add instruction, corresponding e.g. to atomicAdd in CUDA.

Arguments, in order:

args[0]: ptr, The pointer.
args[1]: value, The value to use.

◆ CalculateConstantBytes()

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

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

Parameters

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

◆ CalculateExprComplexity()

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

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

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

Parameters

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

◆ call_cpacked_lowered_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_cpacked_lowered_op ( )

Invoke a C packed function using an explicit argument stack.

Arguments, in order:

args[0]: func_name, The function name.
args[1]: args_stack, The argument stack.
args[2]: begin, The start index.
args[3]: end, The end index.

◆ call_cpacked_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_cpacked_op ( )

Invoke a C packed function.

Arguments, in order:

args[0]: func_name, The function name.
args[1...]: args, trailing Expr operands.

◆ call_extern_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_extern_op ( )

Call an external C function with argument and result types supplied by the caller.

Arguments, in order:

args[0]: func_name, The function name.
args[1...]: args, trailing Expr operands.

◆ call_ffi_kernel_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_ffi_kernel_op ( )

Launch a kernel using the packed-function argument convention.

Arguments are the kernel symbol, kernel operands, then launch values. CallFFIKernelAttr::launch_params describes the launch-value suffix. Host backends may consume this call directly or lower it to call_packed.

◆ call_llvm_intrin_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_llvm_intrin_op ( )

Invoke an LLVM intrinsic.

Arguments, in order:

args[0]: intrin_id, The intrinsic identifier.
args[1...]: args, trailing Expr operands.

◆ call_llvm_pure_intrin_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_llvm_pure_intrin_op ( )

Invoke a pure LLVM intrinsic.

Arguments, in order:

args[0]: intrin_id, The intrinsic identifier.
args[1...]: args, trailing Expr operands.

◆ call_packed_lowered_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_packed_lowered_op ( )

Invoke a packed function using an explicit argument stack.

Arguments, in order:

args[0]: func_name, The function name.
args[1]: args_stack, The argument stack.
args[2]: begin, The start index.
args[3]: end, The end index.

◆ call_packed_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_packed_op ( )

Invoke a runtime packed function.

Arguments, in order:

args[0]: func_name, The function name.
args[1...]: args, trailing Expr operands.

◆ call_pure_extern_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_pure_extern_op ( )

Call a pure external C function with types supplied by the caller.

Arguments, in order:

args[0]: func_name, The function name.
args[1...]: args, trailing Expr operands.

◆ call_spirv_pure_glsl450_op()

TVM_TIRX_DLL const Op & tvm::tirx::call_spirv_pure_glsl450_op ( )

Invoke a pure GLSL450 SPIR-V intrinsic.

Arguments, in order:

args[0]: intrin_id, The intrinsic identifier.
args[1...]: args, trailing PrimExpr operands.

◆ compute_scope_op()

TVM_TIRX_DLL const Op & tvm::tirx::compute_scope_op ( )

Outline the body as a CPU compute helper named by a StringImm operand. There are no body parameters, attributes or results.

◆ ConstHandle()

Expr tvm::tirx::ConstHandle	(	int64_t	value,
		Location	loc = `UnknownLoc()`
	)

inline

Construct an opaque pointer from its integer payload.

◆ 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

◆ CopyTensorType()

ffi::ObjectPtr< TensorTypeNode > tvm::tirx::CopyTensorType ( const TensorVar & var )

inline

◆ cpu_parallel_barrier_op()

TVM_TIRX_DLL const Op & tvm::tirx::cpu_parallel_barrier_op ( )

Synchronize all workers in the current CPU parallel launch.

Every worker must reach this operation. It must be outside partitioned parallel loops, whose iteration counts can differ between workers.

◆ decl_tensor()

TVM_TIRX_DLL TensorVar tvm::tirx::decl_tensor	(	ffi::Array< PrimExpr >	shape,
		PrimType	dtype = `PrimType::Float(32)`,
		ffi::String	name = `"tensor"`,
		ffi::String	storage_scope = `""`,
		Location	loc = `UnknownLoc()`
	)

Construct a new tensor given shape, and dtype.

Parameters

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

Returns: The created tensor.

See also: TensorVar for complete constructor.

◆ decl_tensor_op()

TVM_TIRX_DLL const Op & tvm::tirx::decl_tensor_op ( )

Declare a tensor view with exactly one TensorType in ty_args.

args[0] is the existing physical pointer backing the view. There are no attributes. The TensorType is the result type; no memory is allocated.

◆ DefaultIndexPrimType()

PrimType tvm::tirx::DefaultIndexPrimType ( )

inline

if TVM_INDEX_DEFAULT_I64 is set, return int64, otherwise return int32

◆ DefaultIndexType()

DLDataType tvm::tirx::DefaultIndexType ( )

inline

◆ DetectTensorAccessLCA()

ffi::Map< TensorVar, ffi::Optional< Stmt > > tvm::tirx::DetectTensorAccessLCA ( const Function & func )

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

Parameters

func	The Function to be detected.

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

◆ device_context_op()

TVM_TIRX_DLL const Op & tvm::tirx::device_context_op ( )

Supply lexical device context for allocation and packed-call lowering. Operands are integer device type and device ID; there are no body parameters, attributes or results. The region does not change the active runtime device.

◆ device_entry_op()

TVM_TIRX_DLL const Op & tvm::tirx::device_entry_op ( )

Mark a user-facing device entry containing device scope definitions. Takes no operands, body parameters, attributes or results.

◆ fast_erf_float_expr()

TVM_TIRX_DLL PrimExpr tvm::tirx::fast_erf_float_expr	(	PrimExpr	arg,
		int	bits
	)

Fast_erf_float expression from Eigen.

Parameters

arg	The input expression.
bits	The number of bits in the type.

Returns: The constructed expression.

◆ FindAnchorBlock()

const s_tir::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 FunctionNode * 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 Function.

Parameters

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

Returns: The entry function.

◆ FullTensorRegion()

TVM_TIRX_DLL TensorRegion tvm::tirx::FullTensorRegion ( TensorVar tensor )

Select the entire tensor.

◆ get_active_lane_mask_op()

TVM_TIRX_DLL const Op & tvm::tirx::get_active_lane_mask_op ( )

Calculate a predicate mask given an upper bound (limit) and a current value (base).

Arguments, in order:

args[0]: base, The base value.
args[1]: limit, The limit value.

◆ GetSBlockAccessRegion()

ffi::Array< ffi::Array< TensorRegion > > tvm::tirx::GetSBlockAccessRegion	(	const s_tir::SBlock &	block,
		const ffi::Map< Var, TensorVar > &	tensor_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
tensor_var_map	The outside tensors which may be accessed the block. It is a map from tensor var to the tensor.

Returns

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

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

◆ GetSBlockReadWriteRegion()

ffi::Array< ffi::Array< TensorRegion > > tvm::tirx::GetSBlockReadWriteRegion	(	const s_tir::SBlock &	block,
		const ffi::Map< Var, TensorVar > &	tensor_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
tensor_var_map	The outside tensors which may be accessed the block. It is a map from tensor var to the tensor

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

◆ GetShardingVarFromIndex()

ffi::Optional< Var > tvm::tirx::GetShardingVarFromIndex	(	PrimExpr	index,
		ffi::Map< Var, Range >	var_range,
		const sym::Analyzer &	analyzer
	)

Suppose we want to shard a tensor along a specific dimension, we need to know how to rewrite the access index of the tensor. To make it simple, we only support the case that the access can be rewritten by changing the extent of an iter var.

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

◆ GetTensorVar()

TensorVar tvm::tirx::GetTensorVar ( const VarNode * var )

inline

Recover a checked tensor view from an ordinary VarNode pointer.

◆ gpu_active_thread_selector_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_active_thread_selector_op ( )

Analysis-only active-thread selector.

gpu_active_thread_selector(var, pred) denotes the unique value of var in the current active domain for which pred is true. It is used only inside ExecContext/DispatchContext metadata, for predicates such as ptx.elect_sync() whose selected lane cannot be inferred structurally.

◆ gpu_dp4a_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_dp4a_op ( )

Dot product of two int8x4 vectors and add an optional accumulator.

Arguments, in order:

args[0]: vec1, The first input vector.
args[1]: vec2, The second input vector.
args[2]: acc, The accumulator.

◆ gpu_storage_sync_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_storage_sync_op ( )

Synchronize accesses in a storage scope.

Arguments, in order:

args[0]: storage_scope, The storage scope.

◆ gpu_thread_allreduce_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_thread_allreduce_op ( )

Cross-thread reduction with an explicit typed combiner and identities.

void gpu_thread_allreduce(LambdaExpr combine, Expr identity, Expr values, PrimExpr predicate, Expr destinations, Expr thread_axes);

For N values, combine binds lhs[0:N] followed by rhs[0:N] and returns an N-element Tuple, or a scalar when N is one. Identity, values, destinations and thread_axes may each be a scalar or an explicit Tuple of fields. Each value, identity, pair of parameters and result have the same primitive type. Inactive inputs are replaced by their identities. Destinations are N tensor loads at index zero of one-element result temporaries (optionally cast for boolean storage). Each result temporary must be accessed only at index zero. Thread axes are reduction thread variables or zero for simplified unit axes. Other thread indices remain fixed. The operation writes the reduced values to the destination tensors and returns void.

◆ gpu_thread_filter_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_thread_filter_op ( )

Mark an opaque predicate that selects one thread on an active-set axis.

Arguments, in order:

args[0]: var, The CUDA-index-bound thread-axis variable.
args[1]: pred, The runtime predicate used as an If condition. Canonical thread predicates can be used directly without this wrapper.

◆ gpu_thread_invariant_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_thread_invariant_op ( )

Mark a condition to be thread invariant.

Arguments, in order:

args[0]: cond, The condition.

◆ gpu_thread_return()

TVM_TIRX_DLL PrimExpr tvm::tirx::gpu_thread_return ( Location loc = UnknownLoc() )

Return from a GPU thread without returning a function value.

Parameters

loc	The location of this operation in the source.

Returns: The thread return expression.

◆ gpu_thread_return_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_thread_return_op ( )

Return from a GPU thread without returning a function value.

◆ gpu_warp_activemask_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_warp_activemask_op ( )

Return the bit mask of active lanes in the warp.

◆ gpu_warp_shuffle_down_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_warp_shuffle_down_op ( )

Read the value from a higher lane in the warp.

Arguments, in order:

args[0]: mask, The mask.
args[1]: value, The value to use.
args[2]: offset, The offset.
args[3]: width, The width.
args[4]: warp_size, The number of threads per warp.

◆ gpu_warp_shuffle_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_warp_shuffle_op ( )

Read the value from a selected lane in the warp.

Arguments, in order:

args[0]: mask, The mask.
args[1]: value, The value to use.
args[2]: warp_id, The warp identifier.
args[3]: width, The width.
args[4]: warp_size, The number of threads per warp.

◆ gpu_warp_shuffle_up_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_warp_shuffle_up_op ( )

Read the value from a lower lane in the warp.

Arguments, in order:

args[0]: mask, The mask.
args[1]: value, The value to use.
args[2]: offset, The offset.
args[3]: width, The width.
args[4]: warp_size, The number of threads per warp.

◆ gpu_warp_shuffle_xor_op()

TVM_TIRX_DLL const Op & tvm::tirx::gpu_warp_shuffle_xor_op ( )

Read the value from the lane selected by an XOR mask.

Arguments, in order:

args[0]: mask, The mask.
args[1]: value, The value to use.
args[2]: lane_mask, The lane mask.
args[3]: width, The width.
args[4]: warp_size, The number of threads per warp.

◆ handle_add_byte_offset_op()

TVM_TIRX_DLL const Op & tvm::tirx::handle_add_byte_offset_op ( )

Add a byte offset to an opaque pointer.

Arguments, in order:

args[0]: handle, The handle.
args[1]: offset, The offset.

◆ ignore_loop_partition_op()

TVM_TIRX_DLL const Op & tvm::tirx::ignore_loop_partition_op ( )

Annotate a predicate not be considered as target condition of loop partition.

Arguments, in order:

args[0]: predicate, The predicate.

◆ IsNoOp()

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

inline

Whether stmt is undefined, an integer Evaluate, or an empty sequence.

◆ isnullptr_op()

TVM_TIRX_DLL const Op & tvm::tirx::isnullptr_op ( )

Test whether a pointer is null.

Arguments, in order:

args[0]: x, The input value.

◆ kernel_replace_point_op()

TVM_TIRX_DLL const Op & tvm::tirx::kernel_replace_point_op ( )

Marker where a transform should replace generated kernel initialization.

◆ launch_thread_op()

TVM_TIRX_DLL const Op & tvm::tirx::launch_thread_op ( )

Thread launch region: operands are a nonempty StringImm tag and a scalar signed or unsigned integer extent wider than one bit. The sole body parameter is a fresh thread-index PrimVar matching the extent type. Tags starting with vthread denote virtual threads. There are no attrs or results.

◆ logaddexp()

TVM_TIRX_DLL PrimExpr tvm::tirx::logaddexp	(	PrimExpr	a,
		PrimExpr	b,
		Location	loc = `UnknownLoc()`
	)

Compute log(exp(a) + exp(b)).

Parameters

a	Left operand.
b	Right operand.
loc	The location of this operation in the source.

Returns: The result expression.

◆ MakeTensorLoad()

TVM_TIRX_DLL TensorLoad tvm::tirx::MakeTensorLoad	(	TensorVar	tensor,
		ffi::Array< PrimExpr >	indices,
		Location	loc = `UnknownLoc()`
	)

Construct a TensorLoad from a TensorVar.

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

◆ masked_load_op()

TVM_TIRX_DLL const Op & tvm::tirx::masked_load_op ( )

Masked tensor load.

Arguments, in order:

args[0]: tensor, The tensor.
args[1]: index, The index.
args[2...]: args, trailing PrimExpr operands.

◆ masked_store_op()

TVM_TIRX_DLL const Op & tvm::tirx::masked_store_op ( )

Masked tensor store.

Arguments, in order:

args[0]: tensor, The tensor.
args[1]: value, The value to use.
args[2]: index, The index.
args[3...]: args, trailing PrimExpr operands.

◆ operator!=()

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

inline

◆ operator==()

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

inline

◆ parallel_launch_op()

TVM_TIRX_DLL const Op & tvm::tirx::parallel_launch_op ( )

Launch a CPU worker team around parallel loops and team barriers. Takes no operands, body parameters, attributes or results.

◆ prefetch_op()

TVM_TIRX_DLL const Op & tvm::tirx::prefetch_op ( )

Prefetch a memory address.

Arguments, in order:

args[0]: ptr, The pointer.
args[1]: rw, The read/write mode.
args[2]: locality, The locality hint.
args[3]: cache_type, The cache policy.

◆ ptr_byte_offset_op()

TVM_TIRX_DLL const Op & tvm::tirx::ptr_byte_offset_op ( )

Cast a handle to a typed pointer after adding a byte offset.

Arguments, in order:

args[0]: data, Base pointer.
args[1]: byte_offset, Offset in bytes. The result type specifies the pointed-to element type.

◆ RebuildTensorVar()

TensorVar tvm::tirx::RebuildTensorVar	(	const TensorVar &	var,
		ffi::ObjectPtr< TensorTypeNode >	type,
		ffi::Optional< ffi::String >	name = `std::nullopt`
	)

inline

◆ RegisterExecScope()

TVM_TIRX_DLL void tvm::tirx::RegisterExecScope	(	const ffi::String &	backend,
		const ffi::String &	name,
		int	rank
	)

Register a backend's execution hierarchy. Smaller rank means wider scope.

◆ reinterpret() [1/2]

PrimExpr tvm::tirx::reinterpret	(	PrimType	t,
		PrimExpr	value,
		Location	loc = `UnknownLoc()`
	)

inline

◆ reinterpret() [2/2]

TVM_TIRX_DLL Expr tvm::tirx::reinterpret	(	Type	target_ty,
		Expr	value,
		Location	loc = `UnknownLoc()`
	)

perform reinterpret cast value to type.

Parameters

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

Returns: The result expression.

Note: This function may return value if the type is the same.

◆ reinterpret_op()

TVM_TIRX_DLL const Op & tvm::tirx::reinterpret_op ( )

Reinterpret the bits of a value as an explicit primitive or pointer type.

Arguments, in order:

args[0]: x, The input value.

◆ RenewDef()

TVM_TIRX_DLL Function tvm::tirx::RenewDef ( Function func )

Renew variable definitions and remap their uses in a function.

Definition regions follow the registered structural hooks. Free variables and function attributes are preserved, and the function signature is rebuilt.

Parameters

func	The function to renew.

Returns: The function with fresh variable definitions.

◆ ScopeHigher()

TVM_TIRX_DLL bool tvm::tirx::ScopeHigher	(	const ExecScope &	a,
		const ExecScope &	b
	)

Compare scopes within the same backend's hierarchy.

◆ 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 StmtSRefObj::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 s_tir::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

◆ Specialize()

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

Specialize parameters of Function.

Parameters

func	The Function 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:

from __future__ import annotations
 
@Ts.function
def mem_copy(A: T.Tensor((m, n), "float32"), B: T.Tensor((m, n), "float32"),
             m: T.int32, n: T.int32) -> None:
    for i, j in T.grid(m, n):
        with Ts.sblock():
            vi, vj = Ts.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

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

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

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

◆ stack_alloca_op()

TVM_TIRX_DLL const Op & tvm::tirx::stack_alloca_op ( )

Allocate stack storage for runtime values.

Arguments, in order:

args[0]: dtype_str, The data type name.
args[1]: num, The number of entries.

◆ stack_make_dltensor_op()

TVM_TIRX_DLL const Op & tvm::tirx::stack_make_dltensor_op ( )

Construct a stack-allocated DLTensor.

Arguments, in order:

args[0]: data, The input data.
args[1]: shape, The shape.
args[2]: strides, The strides.
args[3]: ndim, The number of dimensions.
args[4]: arr_dtype, The array data type.
args[5]: elem_offset, The element offset.

◆ stack_make_shape_op()

TVM_TIRX_DLL const Op & tvm::tirx::stack_make_shape_op ( )

Construct a stack-allocated shape tuple.

Arguments, in order:

args[0...]: args, trailing IntExpr operands.

◆ static_handle_op()

TVM_TIRX_DLL const Op & tvm::tirx::static_handle_op ( )

Create a function local static handle that iniitalizes to nullptr.

◆ SubstituteWithDataTypeLegalization() [1/2]

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

Substitution may change the data type of the expression.

Returns: The result.

◆ SubstituteWithDataTypeLegalization() [2/2]

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

This statement overload legalizes only core TIRX nodes. Dtype-changing substitutions must be applied after lowering dialect blocks: structural traversal of extension statements does not legalize their iterator domains or bindings. Type-preserving mappings continue to traverse schedulable blocks structurally before lowering.

Returns: The result.

◆ tensor_data_ptr_op()

TVM_TIRX_DLL const Op & tvm::tirx::tensor_data_ptr_op ( )

Project the physical pointer associated with a TensorVar definition.

Argument: tensor, the TensorVar whose physical pointer is projected. The result pointer type is derived from the TensorType dtype and storage scope of the sole TensorVar argument. This operation is consumed by TIRx lowering and code generation.

◆ TensorRegionFromPoint()

TVM_TIRX_DLL TensorRegion tvm::tirx::TensorRegionFromPoint	(	TensorVar	tensor,
		ffi::Array< PrimExpr >	indices
	)

Construct unit or vector-lane ranges from point indices.

◆ TensorWithOffsetAlignment()

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

Creates a TIR tensor for the provided parameters.

Parameters

shape	shape of the tensor
dtype	data type
name	tensor 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 tensor

◆ throw_last_error_op()

TVM_TIRX_DLL const Op & tvm::tirx::throw_last_error_op ( )

Raise the last runtime error.

◆ TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED()

tvm::tirx::TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED ( logaddexp )

◆ undef_op()

TVM_TIRX_DLL const Op & tvm::tirx::undef_op ( )

Returns an initialized but arbitrary value.

◆ vector_combine_op()

TVM_TIRX_DLL const Op & tvm::tirx::vector_combine_op ( )

Concat two vectors.

Arguments, in order:

args[0]: vec1, The first input vector.
args[1]: vec2, The second input vector.

◆ vector_high_op()

TVM_TIRX_DLL const Op & tvm::tirx::vector_high_op ( )

Get the high level half of the vector.

Arguments, in order:

args[0]: vec, The input vector.

◆ vector_low_op()

TVM_TIRX_DLL const Op & tvm::tirx::vector_low_op ( )

Get the low-level half of the vector.

Arguments, in order:

args[0]: vec, The input vector.

◆ VerifyMemory()

TVM_TIRX_DLL bool tvm::tirx::VerifyMemory ( const Function & 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.

◆ VerifyWellFormed() [1/2]

TVM_TIRX_DLL bool tvm::tirx::VerifyWellFormed	(	const Function &	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.
Tensor loads have consistent source, indices, and result types.
Loop control, tile-operation categories, and scope IDs are valid.

Dialect statements require their dialect-specific verifier.

Parameters

func	The Function 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.

◆ VerifyWellFormed() [2/2]

TVM_TIRX_DLL 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 Function (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.

Variable Documentation

◆ kPSUMBankNum

constexpr int tvm::tirx::kPSUMBankNum = 8

constexpr

◆ kPSUMMaxElemPerBank

constexpr int tvm::tirx::kPSUMMaxElemPerBank = 512

constexpr

Namespaces

Classes

Typedefs

Enumerations

Functions

Variables

Detailed Description

Typedef Documentation

◆ FAxisFuser

◆ FAxisSplitter

◆ FLegalize

◆ FLowerIntrinsic

◆ TensorAxis

◆ TGlobalSymbol

◆ TIRVarAxis

◆ TVectorizable

Enumeration Type Documentation

◆ DepKind

◆ TVMStructFieldKind

Function Documentation

◆ abi_field_get_op()

◆ abi_field_set_op()

◆ address_of_op()

◆ alloc_tensor_op()

◆ assume_aligned_op()

◆ atomic_add_op()

◆ CalculateConstantBytes()

◆ CalculateExprComplexity()

◆ CalculateWorkspaceBytes()

◆ call_cpacked_lowered_op()

◆ call_cpacked_op()

◆ call_extern_op()

◆ call_ffi_kernel_op()

◆ call_llvm_intrin_op()

◆ call_llvm_pure_intrin_op()

◆ call_packed_lowered_op()

◆ call_packed_op()

◆ call_pure_extern_op()

◆ call_spirv_pure_glsl450_op()

◆ compute_scope_op()

◆ ConstHandle()

◆ ContainsNode()

◆ CopyTensorType()

◆ cpu_parallel_barrier_op()

◆ decl_tensor()

◆ decl_tensor_op()

◆ DefaultIndexPrimType()

◆ DefaultIndexType()

◆ DetectTensorAccessLCA()

◆ device_context_op()

◆ device_entry_op()

◆ fast_erf_float_expr()

◆ FindAnchorBlock()

◆ FindEntryFunc()

◆ FullTensorRegion()

◆ get_active_lane_mask_op()

◆ GetSBlockAccessRegion()

◆ GetSBlockReadWriteRegion()

◆ GetShardingVarFromIndex()

◆ GetTensorVar()

◆ gpu_active_thread_selector_op()

◆ gpu_dp4a_op()

◆ gpu_storage_sync_op()

◆ gpu_thread_allreduce_op()

◆ gpu_thread_filter_op()

◆ gpu_thread_invariant_op()

◆ gpu_thread_return()

◆ gpu_thread_return_op()

◆ gpu_warp_activemask_op()

◆ gpu_warp_shuffle_down_op()

◆ gpu_warp_shuffle_op()

◆ gpu_warp_shuffle_up_op()

◆ gpu_warp_shuffle_xor_op()

◆ handle_add_byte_offset_op()

◆ ignore_loop_partition_op()

◆ IsNoOp()

◆ isnullptr_op()

◆ kernel_replace_point_op()

◆ launch_thread_op()

◆ logaddexp()