diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py index 5b61fc73ef24..7d27eed70c14 100644 --- a/python/tvm/ir/expr.py +++ b/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): diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index f9be5ec39d75..9572cac8b386 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/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] diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 753e46c41319..407c30352fc3 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.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.