This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new b693ebeb57 [TIRx] Add scalar builder initialization and assignment
(#20724)
b693ebeb57 is described below
commit b693ebeb57ab9144d244a1318c20146f90799955
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 12:45:04 2026 -0400
[TIRx] Add scalar builder initialization and assignment (#20724)
Extend native local_scalar with optional initialization and in-place
naming while preserving its existing storage representation and
allocation annotations.
Expose assign for indexed tensor loads and one-element tensor storage by
reusing parser scalar-store behavior, keeping destination checks and
statement emission consistent.
---
python/tvm/tirx/script/ir_builder/ir.py | 59 +++++++++++++++++++--
.../tirx/script/test_tirx_script_ir_builder.py | 61 ++++++++++++++++++++++
.../python/tirx/script/test_tirx_script_tensors.py | 3 --
3 files changed, 116 insertions(+), 7 deletions(-)
diff --git a/python/tvm/tirx/script/ir_builder/ir.py
b/python/tvm/tirx/script/ir_builder/ir.py
index e1ea6455dd..d4ffb16a95 100644
--- a/python/tvm/tirx/script/ir_builder/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -30,7 +30,7 @@ from tvm.ir import Range, TensorRegion, Type
from tvm.ir.location import UNKNOWN_LOC, Location
from tvm.ir.type import _init_type_api, _normalize_type
from tvm.runtime import MemoryScope
-from tvm.script.ir_builder.base import MISSING
+from tvm.script.ir_builder.base import MISSING, AlreadyEmitted, IRBuilder
from tvm.script.ir_builder.ir import Lambda, meta_var
from tvm.script.parser.protocol_registry import (
register_mutable_decl as _register_mutable_decl,
@@ -411,10 +411,60 @@ def shared_scalar(
@_register_mutable_decl("tirx.local_scalar")
def local_scalar(
- dtype: str = "float32", *, annotations: dict[str, Any] | None = None
+ dtype: str = "float32",
+ init: Any = None,
+ *,
+ name: str | None = None,
+ annotations: dict[str, Any] | None = None,
) -> _ir.TensorLoad:
- """Allocate a zero-dimensional tensor in local memory."""
- return alloc_scalar(dtype=dtype, scope="local", annotations=annotations)
+ """Allocate scalar local storage, optionally naming and initializing it.
+
+ Parameters
+ ----------
+ dtype : str
+ The scalar data type.
+ init : Any
+ Initial value. None leaves the storage uninitialized; otherwise one
+ tensor store is emitted immediately after the allocation.
+ name : str | None
+ Optional name for the allocated tensor variable, set in place.
+ annotations : dict[str, Any] | None
+ Annotations passed through to the allocation.
+
+ Returns
+ -------
+ scalar : TensorLoad
+ A load from index zero of the one-element local tensor, with no layout.
+ """
+ scalar = alloc_scalar(dtype=dtype, scope="local", annotations=annotations)
+ if name is not None:
+ IRBuilder.name(name, scalar.source)
+ if init is not None:
+ assign(scalar, init)
+ return scalar
+
+
+def assign(dst: _ir.TensorLoad | _ir.Var, value: Any) ->
AlreadyEmitted[_ir.Stmt]:
+ """Store a value into scalar storage and return its emission receipt.
+
+ Parameters
+ ----------
+ dst : TensorLoad | Var
+ A tensor load identifying the destination tensor and indices, or a
+ tensor variable with constant shape (1,). Immutable scalar variables
+ and whole multi-element tensors are not assignment destinations.
+ value : Any
+ The scalar value to store, converted by the ordinary tensor-store API.
+
+ Returns
+ -------
+ result : AlreadyEmitted[Stmt]
+ Receipt for the exact emitted store. Parser consumption does not emit
+ it again. Active builder source locations are preserved.
+ """
+ from .parser_protocol import set_mutable_cell_ # pylint:
disable=import-outside-toplevel
+
+ return set_mutable_cell_(dst, value)
def _is_meta_class_instance(value: Any) -> bool:
@@ -1037,6 +1087,7 @@ __all__ = [
"alloc_scalar",
"alloc_shared",
"alloc_tensor",
+ "assign",
"bf16",
"bfloat16",
"boolean",
diff --git a/tests/python/tirx/script/test_tirx_script_ir_builder.py
b/tests/python/tirx/script/test_tirx_script_ir_builder.py
index 75afd04f8e..005c3f4fc6 100644
--- a/tests/python/tirx/script/test_tirx_script_ir_builder.py
+++ b/tests/python/tirx/script/test_tirx_script_ir_builder.py
@@ -476,6 +476,67 @@ def test_concrete_mutable_scalar(declare):
assert scalar.source.elem_offset == 2
+def test_local_scalar_initialization_and_assign():
+ from tvm.script import tirx as TS
+ from tvm.script.ir_builder.base import AlreadyEmitted
+
+ assert TS.assign is T.assign
+ loc = tvm.ir.SourceLoc(tvm.ir.SourceName("scalar.py"), 3, 1, 3, 20)
+ annotations = {"alignment": 16}
+ with IRBuilder() as ib:
+ with T.function():
+ with ib.with_loc(loc):
+ scalar = T.local_scalar("int32", 2, name="counter",
annotations=annotations)
+ uninitialized = T.local_scalar()
+ receipt = T.assign(scalar, 3)
+ assert isinstance(receipt, AlreadyEmitted)
+ T.emit_(receipt)
+ source_receipt = T.assign(scalar.source, 4)
+ global_tensor = T.bind(T.alloc_tensor(ty_args=[T.Tensor((2,),
"int32")]))
+ indexed_receipt = T.assign(global_tensor[1], 5)
+ for invalid in (tirx.Var("immutable", "int32"), scalar + 1,
global_tensor):
+ with pytest.raises(TypeError, match="scalar storage"):
+ T.assign(invalid, 6)
+ statements = list(ib.get().body.seq)
+ stores = [stmt for stmt in statements if isinstance(stmt,
tvm.ir.TensorStore)]
+ allocations = [stmt for stmt in statements if _is_tensor_binding(stmt,
"tirx.alloc_tensor")]
+ assert len(allocations) == 3
+ assert len(stores) == 4
+ assert statements[0].same_as(allocations[0])
+ assert statements[1].same_as(stores[0])
+ assert allocations[0].var.same_as(scalar.source)
+ assert allocations[0].value.attrs["alignment"] == annotations["alignment"]
+ assert scalar.source.name == "counter"
+ assert scalar.source.ty.dtype == "int32"
+ assert scalar.source.ty.scope == tvm.runtime.MemoryScope("local")
+ assert scalar.source.ty.layout is None
+ assert list(scalar.source.ty.shape) == [1]
+ assert uninitialized.source.ty.dtype == "float32"
+ assert all(store.dest.same_as(scalar.source) for store in stores[:3])
+ assert stores[1].same_as(receipt.value)
+ assert stores[2].same_as(source_receipt.value)
+ assert stores[3].same_as(indexed_receipt.value)
+ assert stores[3].dest.same_as(global_tensor)
+ assert list(stores[3].indices) == [1]
+ assert [store.value.value for store in stores] == [2, 3, 4, 5]
+ assert all(store.loc.same_as(loc) for store in stores)
+
+ @TS.function
+ def parsed():
+ scalar = TS.local_scalar("int32", 1)
+ TS.assign(scalar, 2)
+ scalar = 3
+
+ parsed_stores = [stmt for stmt in parsed.body.seq if isinstance(stmt,
tvm.ir.TensorStore)]
+ assert len(parsed.body.seq) == 4
+ assert [store.value.value for store in parsed_stores] == [1, 2, 3]
+ assert all(store.dest.same_as(parsed.body.seq[0].var) for store in
parsed_stores)
+ assert all(isinstance(store.loc, tvm.ir.SourceLoc) for store in
parsed_stores)
+ assert [store.loc.start_line for store in parsed_stores] == list(
+ range(parsed_stores[0].loc.start_line, parsed_stores[0].loc.start_line
+ 3)
+ )
+
+
@pytest.mark.skipif(not tvm.runtime.enabled("llvm"), reason="LLVM is not
enabled")
@pytest.mark.parametrize(
"kind",
diff --git a/tests/python/tirx/script/test_tirx_script_tensors.py
b/tests/python/tirx/script/test_tirx_script_tensors.py
index b36f91299a..8562bcd910 100644
--- a/tests/python/tirx/script/test_tirx_script_tensors.py
+++ b/tests/python/tirx/script/test_tirx_script_tensors.py
@@ -306,9 +306,6 @@ def test_alloc_apis_reject_name_argument():
with pytest.raises(TypeError):
T.alloc_tensor(ty_args=[T.Tensor((1,), "int32")], name="buf")
- with pytest.raises(TypeError):
- T.local_scalar("int32", name="idx")
-
def test_tensor():
# fmt: off
