[REFACTOR][IR] Move declared call results into TIRX builders

python/tvm/ir/expr.py

@@ -451,9 +451,8 @@ def __init__(self, *args, **kwargs):
 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 @@ def __init__(
         # 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):

python/tvm/tirx/script/builder/ir.py

@@ -43,6 +43,7 @@
 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 @@
 # 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]

python/tvm/tirx/script/parser/parser.py

@@ -31,12 +31,24 @@
 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
 
 
+@dispatch.register(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
+
+
+@dispatch.register(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.