tvm_tirx.tirx.script.alloc_tensor

Contents

tvm_tirx.tirx.script.alloc_tensor#

tvm_tirx.tirx.script.alloc_tensor(placement=None, *, ty_args, attrs=None, annotations=None, loc=UNKNOWN_LOC, ty=None)#

Allocate the tensor described by one explicit Tensor type argument.

Optional placement is an ordinary tuple operand. Allocation annotations belong to attrs; annotations accepts a Python dictionary whose scalar values are normalized to typed IR constants. Under an active builder, emit one binding and return its tensor variable in an AlreadyEmitted receipt. Parsed assignments consume this receipt automatically; programmatic code uses T.bind_(receipt) or receipt.value to obtain the variable without emitting another binding. Outside a builder return the allocation Call. Construct ir.Call("tirx.alloc_tensor", ...) directly when a pure expression is needed inside a builder.

Tensor allocation conveniences such as T.alloc_local and T.alloc_shared use the same receipt convention. Scalar allocation and declaration helpers instead return AlreadyEmitted[TensorLoad] in an active builder, consumed through the same binding protocol.

Examples

Parsed TVMScript uses an ordinary assignment. The allocation creates one binding, and subsequent tensor accesses use its variable.

from tvm_tirx import ir
from tvm_tirx.script import tirx as T
from tvm_tirx.script.ir_builder import IRBuilder

@T.function(private=True)
def example():
    tensor = T.alloc_tensor(ty_args=[T.Tensor((4,), "int32", scope="local")])
    tensor[0] = 1

tensor_type = T.Tensor((4,), "int32", scope="local")
with IRBuilder() as builder:
    raw = ir.Call("tirx.alloc_tensor", [], ty_args=[tensor_type])
    receipt = T.alloc_tensor(ty_args=[tensor_type])
    tensor = T.bind_(receipt, name="tensor")
    assert tensor.same_as(receipt.value)
binding = builder.get()
assert isinstance(binding, ir.Bind)
assert binding.var.same_as(tensor)
assert not binding.value.same_as(raw)
Parameters:
Return type:

Call | AlreadyEmitted[Var]

Contents