Tile Primitive Authoring API#
Kernel code calls tile primitives through Tx.tile:
from tvm.script import tirx as Tx
Tx.tile.copy(dst, src)
Tx.tile.wg.gemm_async(accum, lhs, rhs)
The scope namespaces bind an operation to an execution scope.
Tx.tile.copy(...) uses thread scope, while Tx.tile.warp.copy(...),
Tx.tile.wg.copy(...), and Tx.tile.cta.copy(...) request wider
cooperation. The dialect also exposes those scope objects at its root: for
example, Tx.cta.copy(...)
and Tx.tile.cta.copy(...) construct the same call. This documentation uses
the explicit Tx.tile form consistently. See the
programming guide for the model, primitive catalog,
and dispatch configuration.
Tile primitive shorthand namespace for TIRx TVMScript.
- tvm.tirx.script.tile.compose_op(workspace: dict[str, Var] | None = None, dispatch: str | None = None, **kwargs) → ComposeOpFrame
Compose a TIRx op.
- Parameters:
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator
- Returns:
res – The result ComposeOpFrame.
- Return type:
frame.ComposeOpFrame
Scope namespaces#
- tvm.tirx.script.tile.thread = <Tx.thread: thread-scope tile primitives>
Bind a cooperation scope to every tile primitive reached through it.
Tx.cluster/Tx.cta/Tx.wg(warpgroup) /Tx.warpare the instances exposed on theTxsurface. Attribute access resolves a tile-primitive op name against the publicTxsurface (registered and dynamic ops alike) and binds this namespace’s scope, soTx.warp.copy(dst, src)emits a copy at warp scope andTx.cta.sum(out, x)reduces at CTA scope. A bareTx.copy(...)(no namespace prefix) stays at the defaultthreadscope.
- tvm.tirx.script.tile.warp = <Tx.warp: warp-scope tile primitives>
Bind a cooperation scope to every tile primitive reached through it.
Tx.cluster/Tx.cta/Tx.wg(warpgroup) /Tx.warpare the instances exposed on theTxsurface. Attribute access resolves a tile-primitive op name against the publicTxsurface (registered and dynamic ops alike) and binds this namespace’s scope, soTx.warp.copy(dst, src)emits a copy at warp scope andTx.cta.sum(out, x)reduces at CTA scope. A bareTx.copy(...)(no namespace prefix) stays at the defaultthreadscope.
- tvm.tirx.script.tile.wg = <Tx.wg: warpgroup-scope tile primitives>
Bind a cooperation scope to every tile primitive reached through it.
Tx.cluster/Tx.cta/Tx.wg(warpgroup) /Tx.warpare the instances exposed on theTxsurface. Attribute access resolves a tile-primitive op name against the publicTxsurface (registered and dynamic ops alike) and binds this namespace’s scope, soTx.warp.copy(dst, src)emits a copy at warp scope andTx.cta.sum(out, x)reduces at CTA scope. A bareTx.copy(...)(no namespace prefix) stays at the defaultthreadscope.
- tvm.tirx.script.tile.warpgroup = <Tx.warpgroup: warpgroup-scope tile primitives>
Bind a cooperation scope to every tile primitive reached through it.
Tx.cluster/Tx.cta/Tx.wg(warpgroup) /Tx.warpare the instances exposed on theTxsurface. Attribute access resolves a tile-primitive op name against the publicTxsurface (registered and dynamic ops alike) and binds this namespace’s scope, soTx.warp.copy(dst, src)emits a copy at warp scope andTx.cta.sum(out, x)reduces at CTA scope. A bareTx.copy(...)(no namespace prefix) stays at the defaultthreadscope.
- tvm.tirx.script.tile.cta = <Tx.cta: cta-scope tile primitives>
Bind a cooperation scope to every tile primitive reached through it.
Tx.cluster/Tx.cta/Tx.wg(warpgroup) /Tx.warpare the instances exposed on theTxsurface. Attribute access resolves a tile-primitive op name against the publicTxsurface (registered and dynamic ops alike) and binds this namespace’s scope, soTx.warp.copy(dst, src)emits a copy at warp scope andTx.cta.sum(out, x)reduces at CTA scope. A bareTx.copy(...)(no namespace prefix) stays at the defaultthreadscope.
- tvm.tirx.script.tile.cluster = <Tx.cluster: cluster-scope tile primitives>
Bind a cooperation scope to every tile primitive reached through it.
Tx.cluster/Tx.cta/Tx.wg(warpgroup) /Tx.warpare the instances exposed on theTxsurface. Attribute access resolves a tile-primitive op name against the publicTxsurface (registered and dynamic ops alike) and binds this namespace’s scope, soTx.warp.copy(dst, src)emits a copy at warp scope andTx.cta.sum(out, x)reduces at CTA scope. A bareTx.copy(...)(no namespace prefix) stays at the defaultthreadscope.
Operations#
ScopedOp objects are callable but are not ordinary Python functions, so
they are listed explicitly. Autodoc follows each wrapper to its underlying
builder function, whose signature includes an internal scope parameter.
Kernel authors do not pass that parameter directly: choose a scope through
Tx.tile.thread / warp / wg / cta / cluster. Bound scope
objects expose a generic (*args, **kwargs) Python signature and forward to
the same underlying operation.
- tvm.tirx.script.tile.add(dst: BufferRegion | Var, src1: BufferRegion | Var | FloatImm, src2: BufferRegion | Var | FloatImm, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Add data from src1 and src2, store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for add result.
src1 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 1, or float.
src2 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 2, or float.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.binary_chain(output: BufferRegion | Var, data: BufferRegion | Var, operand0: BufferRegion | Var | FloatImm, operand1: BufferRegion | Var | FloatImm, op0: str | Op, op1: str | Op, reverse1: bool = False, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Chain multiple binary operations together.
- if not reverse1:
output = (operand0 op0 data) op1 operand1
- else:
output = operand1 op1 (operand0 op0 data)
- Parameters:
output (Union[BufferRegion, Buffer]) – The destination buffer region for the result.
data (Union[BufferRegion, Buffer]) – The input data to operate on.
operand0 (Union[BufferRegion, Buffer, FloatImm]) – The first operand to combine with data.
operand1 (Union[BufferRegion, Buffer, FloatImm]) – The second operand to use in chained operation.
op0 (Union[str, Op]) – The first binary operation to perform.
op1 (Union[str, Op]) – The second binary operation to perform.
reverse1 (bool) – Whether to reverse the order of the second binary operation.
workspace (Dict[str, Buffer]) – The workspace of the operator.
config (Dict[str, Any]) – The scheduler configuration.
- tvm.tirx.script.tile.binary_reduce(binary_output: BufferRegion | Var, reduce_output: BufferRegion | Var, binary_input1: BufferRegion | Var | FloatImm, binary_input2: BufferRegion | Var | FloatImm, binary_op: str | Op, reduce_op: str | Op, reduce_axes: int | tuple[int] = -1, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Combine a binary operation with a reduction operation.
- Parameters:
binary_output (Union[BufferRegion, Buffer]) – The destination buffer region for binary operation result.
reduce_output (Union[BufferRegion, Buffer]) – The destination buffer region for reduction result.
binary_input1 (Union[BufferRegion, Buffer, FloatImm]) – The first source input for binary operation.
binary_input2 (Union[BufferRegion, Buffer, FloatImm]) – The second source input for binary operation.
binary_op (Union[str, Op]) – The binary operation to perform.
reduce_op (Union[str, Op]) – The reduction operation to perform.
reduce_axes (Union[int, Tuple[int]]) – The axes to reduce over.
workspace (Dict[str, Buffer]) – The workspace of the operator.
config (Dict[str, Any]) – The scheduler configuration.
- tvm.tirx.script.tile.cast(dst, src=None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Cast — overloaded.
cast(value, dtype)— expression-level cast: returnsT.cast(value, dtype). Also acceptscast(value, dtype=...)as a kwarg form.cast(dst, src, workspace=..., dispatch=...)— buffer-level Cast operator.
- tvm.tirx.script.tile.copy(dst: BufferRegion | Var, src: BufferRegion | Var, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Copy data from src to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region.
src (Union[BufferRegion, Buffer]) – The source buffer region.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.copy_async(dst: BufferRegion | Var, src: BufferRegion | Var, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
- tvm.tirx.script.tile.exp(dst: BufferRegion | Var, src: BufferRegion | Var | None = None, bias: BufferRegion | Var | FloatImm | None = None, scale: FloatImm | None = None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Exponentiate all elements in src and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for exp result. When src is omitted, also used as the source (in-place).
src (Union[BufferRegion, Buffer], optional) – The source buffer region. If omitted, dst is used (in-place).
bias (Optional[Union[BufferRegion, Buffer, FloatImm]]) – The bias of the exp src. Only supported on Trn.
scale (Optional[FloatImm]) – The scale of the exp src. Only supported on Trn.
workspace (Dict[str, Buffer]) – The workspace of the operator.
- tvm.tirx.script.tile.exp2(dst: BufferRegion | Var, src: BufferRegion | Var | None = None, bias: BufferRegion | Var | FloatImm | None = None, scale: FloatImm | None = None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Compute base-2 exponential (2^x) of all elements in src and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for exp2 result. When src is omitted, also used as the source (in-place).
src (Union[BufferRegion, Buffer], optional) – The source buffer region. If omitted, dst is used (in-place).
bias (Optional[Union[BufferRegion, Buffer, FloatImm]]) – The bias of the exp2 src.
scale (Optional[FloatImm]) – The scale of the exp2 src.
workspace (Dict[str, Buffer]) – The workspace of the operator.
- tvm.tirx.script.tile.log2(dst: BufferRegion | Var, src: BufferRegion | Var | None = None, bias: BufferRegion | Var | FloatImm | None = None, scale: FloatImm | None = None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Compute base-2 logarithm of all elements in src and store to dst.
- tvm.tirx.script.tile.fdiv(dst: BufferRegion | Var, src1: BufferRegion | Var, src2: BufferRegion | Var | FloatImm, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
(Float) Div data from src2 to src1, store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for div result.
src1 (Union[BufferRegion, Buffer]) – The source buffer region 1.
src2 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 2, or float.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.fill(dst: BufferRegion | Var, value: Expr, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Fill the buffer region with the value.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region.
value (Expr) – The value to be filled.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.fma(dst: BufferRegion | Var, src: BufferRegion | Var, scale: BufferRegion | Var | Expr, bias: BufferRegion | Var | Expr, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Fused multiply-add: dst = src * scale + bias.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region.
src (Union[BufferRegion, Buffer]) – The input buffer region.
scale (Union[BufferRegion, Buffer, Expr]) – The scale factor (buffer region or scalar).
bias (Union[BufferRegion, Buffer, Expr]) – The bias term (buffer region or scalar).
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.gemm(D: BufferRegion | Var, A: BufferRegion | Var, B: BufferRegion | Var, C: BufferRegion | Var, transpose_A: bool = False, transpose_B: bool = False, alpha: Expr = 1.0, beta: Expr = 0.0, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
General matrix multiplication.
D = A * B * alpha + C * beta
- Parameters:
D (Union[BufferRegion, Buffer]) – The buffer of matrix D.
A (Union[BufferRegion, Buffer]) – The buffer of matrix A.
B (Union[BufferRegion, Buffer]) – The buffer of matrix B.
C (Union[BufferRegion, Buffer]) – The buffer of matrix C.
transpose_A (bool) – Whether to transpose A.
transpose_B (bool) – Whether to transpose B.
alpha (Expr) – The scalar alpha.
beta (Expr) – The scalar beta.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.gemm_async(C: BufferRegion | Var, A: BufferRegion | Var, B: BufferRegion | Var, SFA: BufferRegion | Var | None = None, SFB: BufferRegion | Var | None = None, transA: bool = False, transB: bool = False, accum: bool = False, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
General matrix multiplication asynchronously.
- Parameters:
C (Union[BufferRegion, Buffer]) – The buffer of matrix C.
A (Union[BufferRegion, Buffer]) – The buffer of matrix A.
B (Union[BufferRegion, Buffer]) – The buffer of matrix B.
SFA (Optional[Union[BufferRegion, Buffer]]) – The scale factor buffer for matrix A (block-scaled MMA only).
SFB (Optional[Union[BufferRegion, Buffer]]) – The scale factor buffer for matrix B (block-scaled MMA only).
transA (bool) – False if A is K-major (MxK), True if A is MN-major (KxM).
transB (bool) – False if B is K-major (NxK), True if B is MN-major (KxN).
accum (bool) – Whether C is accumulated. C = A * B if accum is False, otherwise C += A * B.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.max(dst, src=None, axes: int | tuple[int] = -1, accum: bool = False, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Max — overloaded.
max(a, b)— expression: returnstirx.max(a, b).max(dst, src, axes=, accum=)— reduction operator over buffers.
- tvm.tirx.script.tile.maximum(dst: BufferRegion | Var, src1: BufferRegion | Var | FloatImm, src2: BufferRegion | Var | FloatImm, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Maximum all elements in src1 and src2 and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for maximum result.
src1 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 1, or float.
src2 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 2, or float.
workspace (Dict[str, Buffer]) – The workspace of the operator.
- tvm.tirx.script.tile.memset(dst: BufferRegion | Var, value: Expr, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Set all elements in dst to value.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for memset.
value (Expr) – The value to be set.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.min(dst, src=None, axes: int | tuple[int] = -1, accum: bool = False, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Min — overloaded.
min(a, b)— expression: returnstirx.min(a, b).min(dst, src, axes=, accum=)— reduction operator over buffers.
- tvm.tirx.script.tile.minimum(dst: BufferRegion | Var, src1: BufferRegion | Var | FloatImm, src2: BufferRegion | Var | FloatImm, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Minimum all elements in src1 and src2 and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for minimum result.
src1 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 1, or float.
src2 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 2, or float.
workspace (Dict[str, Buffer]) – The workspace of the operator.
- tvm.tirx.script.tile.mul(dst: BufferRegion | Var, src1: BufferRegion | Var | FloatImm, src2: BufferRegion | Var | FloatImm, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Multiply data from src1 and src2, store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for mul result.
src1 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 1, or float.
src2 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 2, or float.
workspace (Dict[str, Buffer]) – The workspace of the operator.
- tvm.tirx.script.tile.permute_layout(dst: BufferRegion | Var, src: BufferRegion | Var, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Move data so the buffer’s bytes are arranged under a different layout.
Logical shape is preserved (
dst.shape == src.shape); only the byte placement changes (dst.layout != src.layout).dstandsrcmay alias the same SMEM (in-place) or be two distinct buffers.- Parameters:
dst (Union[BufferRegion, Buffer]) – Destination view (carries the target layout).
src (Union[BufferRegion, Buffer]) – Source view (carries the current layout).
workspace (Dict[str, Buffer]) – Optional workspace for the operator.
dispatch (Optional[str]) – Force a specific dispatch variant by name.
- tvm.tirx.script.tile.reciprocal(dst: BufferRegion | Var, src: BufferRegion | Var | None = None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Reciprocal all elements in src and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for reciprocal result. When src is omitted, also used as the source (in-place).
src (Union[BufferRegion, Buffer], optional) – The source buffer region. If omitted, dst is used (in-place).
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.reduce_negate(output: BufferRegion | Var, input: BufferRegion | Var, reduce_op: str | Op, reduce_axes: int | tuple[int] = -1, accum: bool = False, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Negate the result of a reduction operation.
- Parameters:
output (Union[BufferRegion, Buffer]) – The destination buffer region for the negated reduction result.
input (Union[BufferRegion, Buffer]) – The input buffer region to reduce.
reduce_axes (Union[int, Tuple[int]]) – The axes to reduce over.
accum (bool) – Whether to accumulate the result into the output.
reduce_op (Union[str, Op]) – The reduction operation to perform before negation.
workspace (Dict[str, Buffer]) – The workspace of the operator.
config (Dict[str, Any]) – The scheduler configuration.
- tvm.tirx.script.tile.select(dst: BufferRegion | Var, true_value: BufferRegion | Var | FloatImm, false_value: BufferRegion | Var | FloatImm, pred: LambdaExpr | Callable[[...], Expr], scope: ExecScope | None = None)
Select between two values based on a predicate.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for the result.
true_value (Union[BufferRegion, Buffer, FloatImm]) – The value to select if the predicate is true.
false_value (Union[BufferRegion, Buffer, FloatImm]) – The value to select if the predicate is false.
pred (Union[LambdaExpr, Callable[..., Expr]]) – The predicate to evaluate. The callable should take the same number of arguments as the dimensions of the destination buffer.
- tvm.tirx.script.tile.silu(dst: BufferRegion | Var, src: BufferRegion | Var, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Compute SiLU (x * sigmoid(x)) for all elements in src and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for SiLU result.
src (Union[BufferRegion, Buffer]) – The source buffer region.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.sqrt(dst: BufferRegion | Var, src: BufferRegion | Var | None = None, bias: BufferRegion | Var | FloatImm | None = None, scale: FloatImm | None = None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Sqrt all elements in src and store to dst.
dst = sqrt(src * scale + bias) (if scale or bias are provided)
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for sqrt result. When src is omitted, also used as the source (in-place).
src (Union[BufferRegion, Buffer], optional) – The source buffer region. If omitted, dst is used (in-place).
bias (Optional[Union[BufferRegion, Buffer, FloatImm]]) – The bias of the sqrt src. Only supported on Trn.
scale (Optional[FloatImm]) – The scale of the sqrt src. Only supported on Trn.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.sub(dst: BufferRegion | Var, src1: BufferRegion | Var, src2: BufferRegion | Var | FloatImm, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Sub data from src2 to src1, store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for sub result.
src1 (Union[BufferRegion, Buffer]) – The source buffer region 1.
src2 (Union[BufferRegion, Buffer, FloatImm]) – The source buffer region 2, or float.
workspace (Dict[str, Buffer]) – The workspace of the operator.
- tvm.tirx.script.tile.sum(dst: BufferRegion | Var, src: BufferRegion | Var, axes: int | tuple[int] = -1, accum: bool = False, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Sum all elements in src and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for sum result.
src (Union[BufferRegion, Buffer]) – The source buffer region.
accum (bool) – Whether dst is accumulated.
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.
- tvm.tirx.script.tile.unary_reduce(unary_output: BufferRegion | Var, reduce_output: BufferRegion | Var, unary_input: BufferRegion | Var, unary_op: str | Op, reduce_op: str | Op, bias: BufferRegion | Var | FloatImm | None = None, scale: FloatImm | None = None, reduce_axes: int | tuple[int] = -1, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Combine a unary operation with a reduction operation.
- Parameters:
unary_output (Union[BufferRegion, Buffer]) – The destination buffer region for unary operation result.
reduce_output (Union[BufferRegion, Buffer]) – The destination buffer region for reduction result.
unary_input (Union[BufferRegion, Buffer]) – The source input for unary operation.
reduce_op (Union[str, Op]) – The reduction operation to perform.
bias (Optional[Union[BufferRegion, Buffer, FloatImm]]) – The bias to apply before unary operation.
scale (Optional[FloatImm]) – The scale to apply before unary operation.
reduce_axes (Union[int, Tuple[int]]) – The axes to reduce over.
workspace (Dict[str, Buffer]) – The workspace of the operator.
config (Dict[str, Any]) – The scheduler configuration.
- tvm.tirx.script.tile.zero(dst: BufferRegion | Var, src: BufferRegion | Var | None = None, workspace: dict[str, Var] | None = None, dispatch: str | None = None, scope: ExecScope | None = None, **kwargs)
Zero out all elements in src and store to dst.
- Parameters:
dst (Union[BufferRegion, Buffer]) – The destination buffer region for zero result. When src is omitted, also used as the source (in-place).
src (Union[BufferRegion, Buffer], optional) – The source buffer region. If omitted, dst is used (in-place).
workspace (Optional[Dict[str, Buffer]]) – The workspace of the operator.