tvm.tir

tvm.tir#

Namespace for Tensor-level IR

class tvm.tir.Add(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Add node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Allocate(buffer_var: Var, dtype: str, extents: List[PrimExpr], condition: PrimExpr, body: Stmt, annotations: Mapping[str, Object] | None = None, span: Span | None = None)[源代码]#

Allocate node.

参数:

buffer_var (Var) -- The buffer variable.
dtype (str) -- The data type of the buffer.
extents (list of Expr) -- The extents of the allocate
condition (PrimExpr) -- The condition.
body (Stmt) -- The body statement.
annotations (Optional[Mapping[str, Object]]) -- Additional annotation hints
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.AllocateConst(buffer_var: Var, dtype: str, extents: List[PrimExpr], data_or_idx: NDArray | int, body: Stmt, annotations: Mapping[str, Object] | None = None, span: Span | None = None)[源代码]#

Allocate constant node.

参数:

buffer_var (Var) -- The buffer variable.
dtype (str) -- The data type of the buffer.
extents (list of Expr) -- The extents of the allocate
data_or_idx (Union[NDArray, int]) -- If an NDArray, this is the const data associated with the constant. If an integer, this is the index into the "constants" attribute of the IRModule that contains the AllocateConst.
body (Stmt) -- The body statement.
annotations (Optional[Mapping[str, Object]]) -- Additional annotations about the allocation.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.And(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

And node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Any(span: Span | None = None)[源代码]#

Any node.

spanOptional[Span]: The location of this expression in the source code.

class tvm.tir.AssertStmt(condition: PrimExpr, message: PrimExpr, body: Stmt, span: Span | None = None)[源代码]#

AssertStmt node.

参数:

condition (PrimExpr) -- The assert condition.
message (PrimExpr) -- The error message.
body (tvm.tir.Stmt) -- The body statement.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.AttrStmt(node: Object, attr_key: str, value: PrimExpr, body: Stmt, span: Span | None = None)[源代码]#

AttrStmt node.

参数:

node (Object) -- The node to annotate the attribute
attr_key (str) -- Attribute type key.
value (PrimExpr) -- The value of the attribute
body (Stmt) -- The body statement.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.BijectiveLayout[源代码]#

Bijective mapping for two layouts (src-layout and dst-layout). It provides shape and index conversion between each other.

Do not construct directly, use bijective_layout instead. See the documentation of bijective_layout for more details.

参数:

src_layout (str or Layout) -- source layout.
dst_layout (str or Layout) -- destination layout.

参见

bijective_layout: Declare a layout

backward_index(index)[源代码]#

Given the indices of the dst-layout, infer the src index.

参数:: index (Array of Expr) -- The indices in dst-layout.
返回:: src_index -- The inferred indices in src-layout.
返回类型:: Array of Expr

backward_shape(shape)[源代码]#

Given the shape of the dst-layout, infer the src shape.

参数:: shape (Array of Expr) -- The shape in dst-layout.
返回:: src_shape -- The inferred shape in src-layout.
返回类型:: Array of Expr

forward_index(index)[源代码]#

Given the indices of the src-layout, infer the dst index.

参数:: index (Array of Expr) -- The indices in src-layout.
返回:: dst_index -- The inferred indices in dst-layout.
返回类型:: Array of Expr

forward_shape(shape)[源代码]#

Given the shape of the src-layout, infer the dst shape.

参数:: shape (Array of Expr) -- The shape in src-layout.
返回:: dst_shape -- The inferred shape in dst-layout.
返回类型:: Array of Expr

class tvm.tir.Block(iter_vars: List[IterVar], reads: List[BufferRegion], writes: List[BufferRegion], name_hint: str, body: Stmt, init: Stmt | None = None, alloc_buffers: List[Buffer] | None = None, match_buffers: List[MatchBufferRegion] | None = None, annotations: Mapping[str, Object] | None = None, span: Span | None = None)[源代码]#

Block node.

参数:

iter_vars (List[IterVar]) -- The block Variable.
reads (List[BufferRegion]) -- The read buffer regions of the block.
writes (List[BufferRegion]) -- The write buffer regions of the block.
name_hint (str) -- the name_hint of the block.
body (Stmt) -- The body of the block.
init (Optional[Stmt]) -- The init block of the reduction block
alloc_buffers (Optional[list[Buffer]]) -- The buffer allocations
match_buffers (Optional[List[MatchBufferRegion]]) -- The subregion buffer match
annotations (Optional[Mapping[str, Object]]) -- Additional annotation hints.
span (Optional[Span]) -- The location of this block in the source code.

class tvm.tir.BlockDependenceInfo(mod: IRModule | PrimFunc)[源代码]#

An object that helps build and query block level dependences using the 2 core objects BlockScope and StmtSRef

The data structures exposed are: 1) sref2scope: Mapping from the srefs to its corresponding BlockScope 2) stmt2ref: Mapping from blocks to corresponding StmtSRefs

Note that this object does not store SRefs to loops as the purpose is only to expose block level dependences. This provides the advantage that the scope block (parent block) for a given block sref can be directly accessed as sref->parent

get_block_scope(block_sref: StmtSRef) → BlockScope[源代码]#

Get the BlockScope correpsonding to the block sref

参数:: block_sref (StmtSRef) -- The block sref to be retrieved
返回:: scope -- The corresponding BlockScope
返回类型:: StmtSRef

get_sref(block: Block) → StmtSRef | None[源代码]#

Return the corresponding sref that points to the block

参数:: stmt (Block) -- The block for which the sref is to be retrived
返回:: sref -- The corresponding sref
返回类型:: StmtSRef

class tvm.tir.BlockRealize(iter_values: List[PrimExpr], predicate: PrimExpr | bool, block: Block, span: Span | None = None)[源代码]#

BlockRealize node.

参数:

iter_values (List[PrimExpr]) -- The binding values of the block var.
predicate (Union[PrimExpr, bool]) -- The predicate of the block.
block (Block) -- The block to realize
span (Optional[Span]) -- The location of this block_realize in the source code.

class tvm.tir.Broadcast(value: PrimExpr, lanes: PrimExpr, span: Span | None = None)[源代码]#

Broadcast node.

参数:

value (PrimExpr) -- The value of the expression.
lanes (PrimExpr) -- The lanes of the expression.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Buffer[源代码]#

Symbolic data buffer in TVM.

Buffer provide a way to represent data layout specialization of data structure in TVM.

Do not construct directly, use decl_buffer() instead. See the documentation of decl_buffer() for more details.

参见

decl_buffer: Declare a buffer

access_ptr(access_mask, ptr_type='handle', content_lanes=1, offset=0, extent=None)[源代码]#

Get an access pointer to the head of buffer.

This is the recommended method to get buffer data ptress when interacting with external functions.

参数:

access_mask (int) -- The access pattern MASK. Indicate whether the access will read or write to the data content.
ptr_type (str, optional) -- The data type of the result pointer. Do not specify unless we want to cast pointer to specific type.
content_lanes (int, optional) -- The number of lanes for the data type. This value is greater than one for vector types.
offset (Expr, optional) -- The offset of pointer. We can use it to offset by the number of elements from the address of ptr.
extent (Expr, optional) -- The extent of pointer.

示例

# Get access ptr for read
buffer.access_ptr("r")
# Get access ptr for read/write with bitmask
buffer.access_ptr(Buffer.READ | Buffer.WRITE)
# Get access ptr for read/write with str flag
buffer.access_ptr("rw")
# Get access ptr for read with offset
buffer.access_ptr("r", offset = 100)
# Get access ptr for read with extent
buffer.access_ptr("r", extent = 100)

get_flattened_buffer()[源代码]#

Generate a Buffer that is a flattened version of this buffer.

返回:: flattened -- The corresponding flat buffer.
返回类型:: Buffer

offset_of(indices)[源代码]#

Determine the offset of the provided indices in the flattened buffer.

参数:: indices (Union[PrimExpr, List[PrimExpr]]) -- The indices of the element in the original buffer.
返回:: flattened_indices -- The offset indices of the element in the flattened buffer.
返回类型:: List[PrimExpr]

scope()[源代码]#: Return the storage scope associated with this buffer. :returns: scope -- The storage scope associated with this buffer. :rtype: str

vload(begin, dtype=None, predicate=None)[源代码]#

Generate an Expr that loads dtype from begin index.

参数:

begin (Array of Expr) -- The beginning index in unit of Buffer.dtype
dtype (str) -- The data type to be loaded, can be vector type which have lanes that is multiple of Buffer.dtype
predicate (Optional[PrimExpr]) -- A vector mask of boolean values indicating which lanes of a vector are to be loaded. The number lanes of the mask must be equal to the number of lanes being loaded.

返回:

load -- The corresponding load expression.

返回类型:

Expr

vstore(begin, value, predicate=None)[源代码]#

Generate a Stmt that store value into begin index.

参数:

begin (Array of Expr) -- The beginning index in unit of Buffer.dtype
value (Expr) -- The value to be stored.
predicate (Optional[PrimExpr]) -- A vector mask of boolean values indicating which lanes of a vector are to be stored. The number lanes of the mask must be equal to the number of lanes in value.

返回:

store -- The corresponding store stmt.

返回类型:

Stmt

class tvm.tir.BufferLoad(buffer: Buffer, indices: List[PrimExpr], predicate: PrimExpr | None = None, span: Span | None = None)[源代码]#

Buffer load node.

参数:

buffer (Buffer) -- The buffer to be loaded.
indices (List[PrimExpr]) -- The buffer indices to load values from.
span (Optional[Span]) -- The location of this expression in the source code.
predicate (Optional[PrimExpr]) -- A vector mask of boolean values indicating which lanes of a vector are to be loaded. The number lanes of the mask must be equal to the number of lanes being loaded.

class tvm.tir.BufferRealize(buffer: Buffer, bounds: List[Range], condition: PrimExpr, body: Stmt, span: Span | None = None)[源代码]#

Buffer realize node.

参数:

buffer (Buffer) -- The buffer.
bounds (List[Range]) -- The value we to be stored.
condition (PrimExpr) -- The realize condition.
body (Stmt) -- The body of the statement.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.BufferRegion(buffer: Buffer, region: List[Range])[源代码]#

BufferRegion node.

参数:

buffer (Buffer) -- The buffer of the buffer region
region (List[Range]) -- The region array of the buffer region

class tvm.tir.BufferStore(buffer: Buffer, value: PrimExpr, indices: List[PrimExpr], predicate: PrimExpr | None = None, span: Span | None = None)[源代码]#

Buffer store node.

参数:

buffer (Buffer) -- The buffer.
value (PrimExpr) -- The value we to be stored.
indices (List[PrimExpr]) -- The indices location to be stored.
predicate (Optional[PrimExpr]) -- A vector mask of boolean values indicating which lanes of a vector are to be stored. The number lanes of the mask must be equal to the number of lanes in value.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.Call(dtype: str, op: Op | str, args: List[PrimExpr], span: Span | None = None)[源代码]#

Call node.

参数:

dtype (str) -- The return data type
op (Union[Op, str]) -- The function to be called, or the name to the global tvm.Op
args (list of Expr) -- The input arguments to the call
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.CallEffectKind[源代码]#: Possible kinds of Call effects.

class tvm.tir.Cast(dtype, value, span: Span | None = None)[源代码]#

Cast expression.

参数:

dtype (str) -- The data type
value (PrimExpr) -- The value of the function.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.CommReducer(lhs: List[Var], rhs: List[Var], result: List[PrimExpr], identity_element: List[PrimExpr], span: Span | None = None)[源代码]#

Commutative reduce operator

参数:

lhs (List[Var]) -- The left arguments of the reducer.
rhs (List[Var]) -- The right arguments of the reducer.
result (List[PrimExpr]) -- The reduction results.
identity_element (List[PrimExpr]) -- The identity elements.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.DataProducer[源代码]#

class tvm.tir.DeclBuffer(buffer: Buffer, body: Stmt, span: Span | None = None)[源代码]#

DeclBuffer node.

参数:

buffer (Buffer) -- The buffer being declared.
body (Stmt) -- The body statement to be executed.
span (Optional[Span]) -- The location of this DeclBuffer in the source code.

class tvm.tir.Div(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Div node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.EQ(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

EQ node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Evaluate(value: PrimExpr, span: Span | None = None)[源代码]#

Evaluate node.

参数:

value (PrimExpr) -- The expression to be evaluated.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.FloatImm(dtype: str, value: float, span: Span | None = None)[源代码]#

Float constant.

参数:

dtype (str) -- The data type
value (float) -- The constant value.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.FloorDiv(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

FloorDiv node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.FloorMod(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

FloorMod node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.For(loop_var: Var, min: PrimExpr, extent: PrimExpr, kind: ForKind, body: Stmt, thread_binding: IterVar | None = None, annotations: Mapping[str, Object] | None = None, span: Span | None = None)[源代码]#

For node.

参数:

loop_var (Var) -- The loop variable.
min (PrimExpr) -- The beginning value.
extent (PrimExpr) -- The length of the loop.
kind (ForKind) -- The type of the for.
body (Stmt) -- The body statement.
thread_binding (Optional[tir.IterVar]) -- The thread this loop binds to. Only valid if kind is ThreadBinding
annotations (Optional[Mapping[str, Object]]) -- Additional annotation hints.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.ForKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[源代码]#: The kind of the for loop.

备注

ForKind can change the control flow semantics of the loop and need to be considered in all TIR passes.

class tvm.tir.GE(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

GE node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.GT(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

GT node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.IfThenElse(condition: PrimExpr, then_case: Stmt, else_case: Stmt | None, span: Span | None = None)[源代码]#

IfThenElse node.

参数:

condition (PrimExpr) -- The expression
then_case (Stmt) -- The statement to execute if condition is true.
else_case (Optional[Stmt]) -- The statement to execute if condition is false.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.IndexMap(initial_indices, final_indices, inverse_index_map)[源代码]#

A mapping from multi-dimensional indices to another set of multi-dimensional indices

参数:

initial_indices (List[Var]) -- Variables representing the indices prior to remapping.
final_indices (List[PrimExpr]) -- Expressions defining the indices after remapping.
inverse_index_map (Union[Callable, Optional[IndexMap]]) -- The optional pre-defined inverse index map. When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. Otherwise, the inverse index map will be computed on the fly. It is the user's responsibility to ensure the correctness of the pre-defined inverse index map.

static from_func(mapping_function: Callable, ndim: int | None = None, inverse_index_map: Callable | IndexMap | None = None, *, index_dtype: str = 'int64')[源代码]#

Create an index map from a function

参数:

mapping_function (Callable) -- The function to map from source indices to target indices. The function should accept tir.Var parameters and return a either a tir.PrimExpr, or a list of tir.PrimExpr. Returning a tir.PrimExpr is equivalent to returning a list of length 1 containing that tir.PrimExpr.
ndim (Optional[int]) -- The dimensionality of the buffer to which this transformation should be applied. If mapping_function uses variadic argument *args, ndim must be specified. If mapping_function does not use variadic arguments, ndim is optional.
inverse_index_map (Union[Callable, Optional[IndexMap]]) -- The optional pre-defined inverse index map. When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. Otherwise, the inverse index map will be computed on the fly. It is the user's responsibility to ensure the correctness of the pre-defined inverse index map.

返回:

index_map -- Returns an IndexMap representing the mapping_function.

返回类型:

IndexMap

static from_func_with_separators(mapping_function: Callable, ndim: int | None = None, inverse_index_map: Callable | IndexMap | None = None, *, index_dtype: str = 'int64')[源代码]#

Create an index map from a function

参数:

mapping_function (Callable) -- The function to map from source indices to target indices. The function should accept tir.Var parameters and return either a tir.PrimExpr or a list. Each element of the returned list should be either a tir.PrimExpr or the object IndexMap.AXIS_SEPARATOR. Returning a tir.PrimExpr is equivalent to returning a list of length 1 containing that tir.PrimExpr.
ndim (Optional[int]) -- The dimensionality of the buffer to which this transformation should be applied. If mapping_function uses variadic argument *args, ndim must be specified. If mapping_function does not use variadic arguments, ndim is optional.
inverse_index_map (Union[Callable, Optional[IndexMap]]) -- The optional pre-defined inverse index map. When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. Otherwise, the inverse index map will be computed on the fly. It is the user's responsibility to ensure the correctness of the pre-defined inverse index map.
index_dtype (str) -- The default index dtype to use for input iters in the mapping function.

返回:

ret -- Returns a tuple whose first element is an IndexMap representing the mapping_function, and whose second index is a list of indices at which IndexMap.AXIS_SEPARATOR occurred.

返回类型:

Tuple[IndexMap, List[int]]

inverse(shape: List[Range | PrimExpr]) → IndexMap[源代码]#

Return the inverse of the map

Throws an error if the function is not bijective.

参数:: shape (List[Union[Range,PrimExpr]]) -- The region over which the inverse should be determined. Used for validating that the mapping is bijective over this range.
返回:: inverse -- The inverse
返回类型:: IndexMap

is_equivalent_to(other_map: IndexMap) → bool[源代码]#

Return if the index maps are equivalent.

参数:: other_map (IndexMap) -- The IndexMap to which the comparison should be made.
返回:: is_equivalent -- True if the two mappings represent the same transformation, otherwise False
返回类型:: bool

map_indices(indices: List[PrimExpr]) → List[PrimExpr][源代码]#

Apply the index map to a set of indices

参数:: indices (List[PrimExpr]) -- The indices to be mapped
返回:: result -- The mapped indices
返回类型:: List[PrimExpr]

map_ndarray(arr_src: NDArray) → NDArray[源代码]#

Apply thie index map to transform the layout of the input NDArray

参数:: arr_src (runtime.NDArray) -- The NDArray to be transformed
返回:: arr_dst -- The transformed NDArray
返回类型:: runtime.NDArray

map_shape(shape: List[PrimExpr]) → List[PrimExpr][源代码]#

Apply the index map to a buffer shape

参数:: shape (List[PrimExpr]) -- The buffer shape to be mapped
返回:: result -- The mapped shape
返回类型:: List[PrimExpr]

non_surjective_inverse(shape: List[Range | PrimExpr]) → Tuple[IndexMap, PrimExpr][源代码]#

Return the inverse of the map

Can be applied to transformations that introduce padding.

参数:: shape (List[Union[Range,PrimExpr]]) -- The region over which the inverse should be determined. Used for determining the predicate.
返回:: result -- The inverse, and a predicate for which the inverse maps to a valid index in the input range.
返回类型:: Tuple[IndexMap, PrimExpr]

示例

index_map = IndexMap.from_func(lambda i: [i//4, i%4])
inverse_map, predicate = index_map.non_surjective_inverse([14])
assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"

class tvm.tir.IntImm(dtype: str, value: int, span: Span | None = None)[源代码]#

Int constant.

参数:

dtype (str) -- The data type
value (int) -- The constant value.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.IterVar(dom: Range, var: Var | str, iter_type: int, thread_tag: str = '', span: Span | None = None)[源代码]#

Represent iteration variable.

IterVar represents axis iterations in the computation.

参数:

dom (Range) -- The domain of the iteration.
var (Union[Var, str]) -- The internal variable that is used for iteration.
iter_type (int) -- The iteration type.
thread_tag (str) -- The thread type tag.
span (Optional[Span]) -- The location of this expression in the source code.

参见

te.thread_axis: Create thread axis IterVar.
te.reduce_axis: Create reduce axis IterVar.

class tvm.tir.LE(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

LE node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.LT(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

LT node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Layout[源代码]#

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

参见

layout: Declare a layout

factor_of(axis)[源代码]#

Get the factor size of the subordinate axis.

参数:: axis (str) -- The axis name, need to be [a-z,A-Z]
返回:: factor -- the size of the subordinate-axis of axis (if axis is a primal-axis), or the size of axis itself (if axis is a subordinate-axis). Return -1 if axis is not in the layout.
返回类型:: int

index_of(axis)[源代码]#

Get the index of an axis

参数:: axis (str) -- The axis name, need to be [a-z,A-Z]
返回:: index -- The index of the axis, -1 if not found.
返回类型:: int

class tvm.tir.Let(var: Var, value: PrimExpr, body: PrimExpr, span: Span | None = None)[源代码]#

Let node.

参数:

var (Var) -- The variable in the binding.
value (PrimExpr) -- The value in to be bound.
body (PrimExpr) -- The body expression.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.LetStmt(var: Var, value: PrimExpr, body: Stmt, span: Span | None = None)[源代码]#

LetStmt node.

参数:

var (Var) -- The variable in the binding.
value (PrimExpr) -- The value in to be bound.
body (Stmt) -- The body statement.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.MatchBufferRegion(buffer: Buffer, source: BufferRegion)[源代码]#

MatchBufferRegion node.

参数:

buffer (Buffer) -- The target buffer
source (BufferRegion) -- The region of source buffer

class tvm.tir.Max(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Max node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Min(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Min node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Mod(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Mod node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Mul(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Mul node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.NE(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

NE node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Not(a: PrimExpr, span: Span | None = None)[源代码]#

Not node.

参数:

a (PrimExpr) -- The input value
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Or(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Or node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Prefetch(buffer: Buffer, bounds: List[Range], span: Span | None = None)[源代码]#

Prefetch node.

参数:

buffer (Buffer) -- The buffer to be prefetched.
bounds (List[Range]) -- The bounds to be prefetched.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.PrimFunc(params, body, ret_type=None, buffer_map=None, attrs=None, span=None)[源代码]#

A function declaration expression.

参数:

params (List[Union[tvm.tir.Var, tvm.tir.Buffer]]) -- List of input parameters to the function.
body (tvm.tir.Stmt) -- The body of the function.
ret_type (tvm.ir.Type) -- The return type annotation of the function.
buffer_map (Map[tvm.tir.Var, tvm.tir.Buffer]) -- The buffer binding map.
attrs (Optional[tvm.Attrs]) -- Attributes of the function, can be None
span (Optional[Span]) -- The location of this itervar in the source code.

specialize(param_map: Mapping[Var, PrimExpr | Buffer])[源代码]#

Specialize parameters of PrimFunc

参数:: param_map (Mapping[Var, Union[PrimExpr, Buffer]]) -- The mapping from function params to the instance

示例

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.block():
            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: tir.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})

The specialized function:

@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.block():
            vi, vj = T.axis.remap("SS", [i, j])
            B[vi, vj] = A[vi, vj]

返回:: func -- The new function with parameter specialized
返回类型:: PrimFunc

with_body(new_body, span=None)[源代码]#

Create a new PrimFunc with the same set signatures but a new body.

参数:

new_body (Stmt) -- The new body.
span (Optional[Span]) -- The location of this itervar in the source code.

返回:

new_func -- The created new function.

返回类型:

PrimFunc

class tvm.tir.ProducerLoad(producer: DataProducer, indices: List[PrimExpr], span: Span | None = None)[源代码]#

Producer load node.

参数:

producer (DataProducer) -- The buffer to be loaded.
indices (List[PrimExpr]) -- The buffer indices.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.ProducerRealize(producer: DataProducer, bounds: List[Range], condition: PrimExpr, body: Stmt, storage_scope: str = '', span: Span | None = None)[源代码]#

ProducerRealize node.

参数:

producer (DataProducer) -- The data producer.
bounds (List[Range]) -- The bound of realize
condition (PrimExpr) -- The realize condition.
body (Stmt) -- The realize body
storage_scope (str) -- The storage scope associated with this realization
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.ProducerStore(producer: DataProducer, value: PrimExpr, indices: List[PrimExpr], span: Span | None = None)[源代码]#

ProducerStore node.

参数:

producer (DataProducer) -- The data producer.
value (PrimExpr) -- The value to be stored.
indices (list of Expr) -- The index arguments of the store.
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.Ramp(base: PrimExpr, stride: PrimExpr, lanes: PrimExpr, span: Span | None = None)[源代码]#

Ramp node.

参数:

base (PrimExpr) -- The base expression.
stride (PrimExpr) -- The stride of the ramp.
lanes (PrimExpr) -- The lanes of the expression.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Reduce(combiner: CommReducer, src: List[PrimExpr], rdom: List[IterVar], condition: PrimExpr, value_index: int, init: List[PrimExpr] | None = None, span: Span | None = None)[源代码]#

Reduce node.

参数:

combiner (CommReducer) -- The combiner.
src (list of Expr) -- The source expression.
rdom (list of IterVar) -- The iteration domain
condition (PrimExpr) -- The reduce condition.
value_index (int) -- The value index.
init (list of Expr) -- The initial value for output. This can be an int, float or ProducerLoad
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Select(condition: PrimExpr, true_value: PrimExpr, false_value: PrimExpr, span: Span | None = None)[源代码]#

Select node.

备注

Select may compute both true_value and false_value. Use tvm.tir.if_then_else instead if you want to get a conditional expression that only evaluates the correct branch.

参数:

condition (PrimExpr) -- The condition expression.
true_value (PrimExpr) -- The value to take when condition is true.
false_value (PrimExpr) -- The value to take when condition is false.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.SeqStmt(seq: List[Stmt], span: Span | None = None)[源代码]#

Sequence of statements.

参数:

seq (List[Stmt]) -- The statements
span (Optional[Span]) -- The location of the stmt in the source code.

class tvm.tir.Shuffle(vectors: List[PrimExpr], indices: List[PrimExpr], span: Span | None = None)[源代码]#

Shuffle node.

参数:

vectors (List[PrimExpr]) -- The vectors
indices (List[PrimExpr]) -- The indices
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.SizeVar(name: str, dtype: str | Type, span: Span | None = None)[源代码]#

Symbolic variable to represent a tensor index size: which is greater or equal to zero.

参数:

name (str) -- The name
dtype (Union[str, ir.Type]) -- The data type
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Stmt[源代码]#: Base class of all the statements.

class tvm.tir.StringImm(value: str, span: Span | None = None)[源代码]#

String constant.

参数:

value (str) -- The value of the function.
span (Optional[Span]) -- The location of this expression in the source code.

class tvm.tir.Sub(a: PrimExpr, b: PrimExpr, span: Span | None = None)[源代码]#

Sub node.

参数:

a (PrimExpr) -- The left hand operand.
b (PrimExpr) -- The right hand operand.
span (Optional[Span]) -- The location of this expression in the source code.

tvm.tir.TVMBackendAllocWorkspace(device_type, device_id, nbytes, dtype_code_hint, dtype_bits_hint)[源代码]#

Backend function to allocate temporal workspace

参数:

device_type (int) -- The device type which the space will be allocated.
device_id (int) -- The device id which the space will be allocated.
nbytes (int) -- The size of the space requested.
dtype_code_hint (int) -- The type code of the array elements. Only used in certain backends such as OpenGL.
dtype_bits_hint (int) -- The type bits of the array elements. Only used in certain backends such as OpenGL.

返回:

call -- The call expression.

返回类型:

tvm.tir

目录

tvm.tir#