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 fd6884ded4 [REFACTOR][IR] Move declared call results into TIRX
builders (#20382)
fd6884ded4 is described below
commit fd6884ded431680cb20c0e6a3edb4ac9b1ae86f8
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 17 19:43:39 2026 -0400
[REFACTOR][IR] Move declared call results into TIRX builders (#20382)
Keep shared Call construction independent of callee signatures: omitted
return types remain missing so Relax normalization performs argument
validation and dependent result deduction.
Resolve declared TIRX results in the parser's builder path, preserving
ordinary module calls and exact pointer returns through the existing
dialect token hooks.
---
python/tvm/ir/expr.py | 20 ++++----------------
python/tvm/tirx/script/builder/ir.py | 16 ++++++++++++++++
python/tvm/tirx/script/parser/parser.py | 14 +++++++++++++-
3 files changed, 33 insertions(+), 17 deletions(-)
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 5b61fc73ef..7d27eed70c 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -451,9 +451,8 @@ class TensorLoad(_CallableExprWithOp):
class Call(_CallableExprWithOp):
"""Core function call node.
- When ``ret_ty`` is omitted, use the callee signature's declared return type
- if available, or a missing type otherwise. Argument-dependent signatures
- retain a missing type for subsequent normalization.
+ When ``ret_ty`` is omitted, use a missing type for subsequent
normalization.
+ Builders may supply a known result type explicitly.
"""
op: Expr
@@ -474,25 +473,14 @@ class Call(_CallableExprWithOp):
# pylint: disable=import-outside-toplevel
from .attrs import DictAttrs
from .op import Op
- from .type import PointerType, PrimType, TupleType, Type
+ from .type import PointerType, PrimType, Type
if isinstance(op, str):
op = Op.get(op)
if attrs is not None and isinstance(attrs, dict):
attrs = DictAttrs(attrs)
if ret_ty is None:
- # Reuse a declared signature without invoking dialect-specific
inference.
- signature = getattr(op, "ty", None)
- ret_ty = getattr(signature, "ret_type", None)
- if not isinstance(ret_ty, Type):
- ret_ty = getattr(signature, "ret", None)
- # Rich signatures may specialize their result using arguments.
- # Reuse only fixed shared scalar, pointer, or void results
here.
- is_fixed_result = isinstance(ret_ty, PrimType | PointerType)
or (
- isinstance(ret_ty, TupleType) and not ret_ty.fields
- )
- if not is_fixed_result or getattr(signature, "derive_func",
None) is not None:
- ret_ty = Type.missing()
+ ret_ty = Type.missing()
if isinstance(ret_ty, str) and ret_ty == "handle":
ret_ty = PointerType(PrimType("void"))
elif ret_ty is not None and not isinstance(ret_ty, Type):
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index f9be5ec39d..9572cac8b3 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -43,6 +43,7 @@ from tvm.ir.prim import _ffi_api as _prim_ffi_api
from tvm.runtime import convert
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.ir import meta_var
+from tvm.script.ir_builder.ir.frame import IRModuleFrame
from tvm.target import Target
# pylint: disable=unused-import
@@ -101,6 +102,21 @@ from .external_kernel import call_kernel
# pylint: enable=unused-import
+def _call_global(func: ir.GlobalVar, *args: Expr) -> Call:
+ """Build a TIRX call using the declared function's exact result type."""
+ if IRBuilder.is_in_scope():
+ for module_frame in reversed(list(IRBuilder.current().frames)):
+ if isinstance(module_frame, IRModuleFrame) and func in
module_frame.functions:
+ declaration = module_frame.functions[func]
+ if isinstance(declaration, tir.PrimFunc):
+ # The Relax-facing signature may erase pointer results to
Any.
+ return Call(func, args, ret_ty=declaration.ret_type)
+ break
+ if isinstance(func.ty, ir.FuncType):
+ return Call(func, args, ret_ty=func.ty.ret_type)
+ return Call(func, args)
+
+
def cast(value, dtype, span=None):
"""Cast an expression to the requested data type."""
return _prim_ffi_api._cast(dtype, value, span) # type:
ignore[attr-defined]
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index 753e46c413..407c30352f 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -31,12 +31,24 @@ from tvm.script.parser._core import Parser,
collect_signature_type_vars, dispatc
from tvm.script.parser.core.doc import from_doc
from tvm.tirx import Buffer, IterVar, Layout, buffer_data, is_buffer_var
from tvm.tirx.script import builder as T
-from tvm.tirx.script.builder.ir import name_meta_class_value
+from tvm.tirx.script.builder.ir import _call_global, name_meta_class_value
from .entry import _OptionalAnnotation, inline
from .entry import constexpr as _constexpr_sentinel
[email protected](token="tirx", type_name="enter_token")
+def enter_token(self: Parser) -> dict[str, Any]:
+ context = {"GlobalVar.__call__": GlobalVar.__call__}
+ GlobalVar.__call__ = _call_global
+ return context
+
+
[email protected](token="tirx", type_name="exit_token")
+def exit_token(self: Parser, context: dict[str, Any]) -> None:
+ GlobalVar.__call__ = context["GlobalVar.__call__"]
+
+
def slice_buffer_from_region(br: TensorRegion) -> Buffer:
"""Create a matched DeclBuffer from a TensorRegion.
