From cf231fa3854cb417ddfea41982661acf88a101ea Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 15:58:53 +0000 Subject: [PATCH 1/2] [REFACTOR][IR] Move declared call results into TIRX builders Leave shared Call results missing unless supplied explicitly so dialect normalization owns argument validation and dependent result deduction. Resolve exact TIRX declarations through the existing parser token hooks. --- python/tvm/ir/expr.py | 20 +-- python/tvm/tirx/script/builder/ir.py | 16 +++ python/tvm/tirx/script/parser/parser.py | 14 +- tests/python/relax/test_blockbuilder_core.py | 57 +++++++++ tests/python/relax/test_expr.py | 18 +++ .../tvmscript/test_tvmscript_parser_tir.py | 121 ++++++++++++++++++ 6 files changed, 229 insertions(+), 17 deletions(-) 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. diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py index 212b9d438877..a10b9196c5dc 100644 --- a/tests/python/relax/test_blockbuilder_core.py +++ b/tests/python/relax/test_blockbuilder_core.py @@ -40,6 +40,63 @@ def nop(): pass +@pytest.mark.parametrize("global_callee", [False, True]) +@pytest.mark.parametrize("ret_ty", [tvm.ir.PrimType("int32"), tvm.ir.TupleType([])]) +def test_normalize_call_checks_fixed_result_arguments(global_callee, ret_ty): + signature = rx.FuncType([tvm.ir.PrimType("int32")], ret_ty) + if global_callee: + callee = tvm.ir.GlobalVar("callee") + rx.expr._update_type(callee, signature) + else: + callee = rx.Var("callee", signature) + arg = rx.Var("arg", tvm.ir.PrimType("int32")) + wrong_arg = rx.Var("arg", tvm.ir.PrimType("float32")) + bb = rx.BlockBuilder() + call = callee(arg) + assert call.ty.is_missing() + tvm.ir.assert_structural_equal(bb.normalize(call).ty, ret_ty) + with pytest.raises(ValueError, match="Number of arguments and parameters mismatch"): + bb.normalize(callee()) + with pytest.raises(ValueError, match="type mismatch"): + bb.normalize(callee(wrong_arg)) + # An explicit result remains an override of normalization's type deduction. + override = tvm.ir.Call(callee, [], ret_ty=ret_ty) + tvm.ir.assert_structural_equal(bb.normalize(override).ty, ret_ty) + + +def test_normalize_call_opaque_and_dependent_results(): + scalar = tvm.ir.PrimType("int32") + opaque = rx.Var("opaque", rx.FuncType.opaque_func(ret=scalar)) + bb = rx.BlockBuilder() + tvm.ir.assert_structural_equal(bb.normalize(opaque()).ty, scalar) + + seen = [] + + @tvm.register_global_func("test.call_builder.derive", override=True) + def derive(call, _ctx): + seen.append(call) + return call.args[0].ty + + custom = rx.Var( + "custom", + rx.FuncType.opaque_func(derive_func=tvm.ir.EnvFunc.get("test.call_builder.derive")), + ) + arg = rx.Var("arg", tvm.ir.PrimType("float32")) + call = custom(arg) + assert call.ty.is_missing() + tvm.ir.assert_structural_equal(bb.normalize(call).ty, arg.ty) + assert len(seen) == 1 + + n = tirx.Var("n", "int64") + dependent = rx.Var( + "dependent", rx.FuncType([rx.TensorType([n], "float32")], rx.TensorType([n + 1], "float32")) + ) + tensor = rx.Var("tensor", rx.TensorType([5], "float32")) + call = dependent(tensor) + assert call.ty.is_missing() + tvm.ir.assert_structural_equal(bb.normalize(call).ty, rx.TensorType([6], "float32")) + + def test_block_builder(): m = tirx.Var("m", "int64") n = tirx.Var("n", "int64") diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py index fc53813afed2..29aeeacd1ebf 100644 --- a/tests/python/relax/test_expr.py +++ b/tests/python/relax/test_expr.py @@ -360,6 +360,24 @@ def test_call(): assert call.args[0].same_as(arg) +@pytest.mark.parametrize("signature_type", [tvm.ir.FuncType, rx.FuncType]) +@pytest.mark.parametrize( + "ret_ty", + [ + tvm.ir.PrimType("int32"), + tvm.ir.PointerType(tvm.ir.PrimType("float32")), + tvm.ir.TupleType([]), + ], +) +def test_call_result_is_explicit(signature_type, ret_ty): + gv = tvm.ir.GlobalVar("callee") + rx.expr._update_type(gv, signature_type([], ret_ty)) + assert tvm.ir.Call(gv, []).ty.is_missing() + assert gv().ty.is_missing() + tvm.ir.assert_structural_equal(tvm.ir.Call(gv, [], ret_ty=ret_ty).ty, ret_ty) + assert tvm.ir.Call(gv, [], ret_ty=tvm.ir.Type.missing()).ty.is_missing() + + def test_call_accepts_core_expr_operator(): """relax.Call aliases the core ir.Call constructor.""" dtype = tvm.ir.PrimType("int32") diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py index a4baa9b4568a..b847b4110640 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_tir.py +++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py @@ -24,6 +24,127 @@ from tvm.script.parser import tirx as T +def test_declared_global_call_results(): + original_call = ir.GlobalVar.__call__ + mod = tvm.script.from_source( + """ +@I.ir_module +class Module: + @T.prim_func + def main(p: T.handle("float32")): + Module.noop() + value = Module.scalar() + T.int32(1) + data = Module.pointer(p) + buffer = T.decl_buffer((1,), "float32", data=data) + buffer[0] = T.Cast("float32", value) + + @T.prim_func + def noop(): + T.evaluate(0) + + @T.prim_func + def scalar() -> T.int32: + return 1 + + @T.prim_func + def pointer(p: T.handle("float32")) -> T.handle("float32"): + return p +""" + ) + assert ir.GlobalVar.__call__ is original_call + calls = {} + + def collect(node, visitor): + if isinstance(node.op, ir.GlobalVar): + calls[node.op.name_hint] = node + visitor.default_visit(node) + + tvm_ffi.structural_visit(mod["main"].body, [(ir.Call, collect)]) + for name in ["noop", "scalar", "pointer"]: + tvm.ir.assert_structural_equal(calls[name].ty, mod[name].ret_type) + tvm.ir.assert_structural_equal(mod, tvm.script.from_source(mod.script())) + assert mod.get_global_var("scalar")().ty.is_missing() + + +def test_global_call_static_signature(): + callee = ir.GlobalVar("callee") + tvm.relax.expr._update_type(callee, ir.FuncType([], ir.PrimType("int32"))) + func = tvm.script.from_source( + """ +@T.prim_func +def main() -> T.int32: + return callee() +""", + extra_vars={"callee": callee, "T": T}, + ) + tvm.ir.assert_structural_equal(func.body.value.ty, ir.PrimType("int32")) + assert callee().ty.is_missing() + + +def test_mixed_module_global_call_adapters(): + original_call = ir.GlobalVar.__call__ + mod = tvm.script.from_source( + """ +@I.ir_module +class Module: + @T.prim_func + def scalar() -> T.int32: + return 1 + + @R.function + def pair(x: R.Tuple(R.Tensor((1,), "float32"), R.Tensor((1,), "float32"))): + return x + + @R.function + def main(x: R.Tensor((1,), "float32")): + result = Module.pair((x, x)) + return result +""" + ) + assert ir.GlobalVar.__call__ is original_call + call = mod["main"].body.blocks[0].bindings[-1].value + assert isinstance(call, ir.Call) + assert isinstance(call.args[0], ir.Tuple) + tvm.ir.assert_structural_equal(mod, tvm.script.from_source(mod.script())) + + +@pytest.mark.parametrize("fail_inner", [False, True]) +def test_nested_global_call_adapters(fail_inner): + from tvm.script.parser.core.diagnostics import Source + from tvm.script.parser.core.parser import Parser + + parser = Parser(Source(""), {}) + original_call = ir.GlobalVar.__call__ + callee = ir.GlobalVar("callee") + arg = tvm.relax.Var("arg", tvm.relax.TensorType([1], "float32")) + with parser.with_dispatch_token("relax"): + relax_call = ir.GlobalVar.__call__ + try: + with parser.with_dispatch_token("tirx"): + if fail_inner: + raise ValueError("unwind inner token") + except ValueError: + assert fail_inner + assert ir.GlobalVar.__call__ is relax_call + call = callee((arg, arg)) + assert isinstance(call.args[0], ir.Tuple) + assert call.ty.is_missing() + assert ir.GlobalVar.__call__ is original_call + + +def test_global_call_adapter_restored_after_parser_error(): + original_call = ir.GlobalVar.__call__ + with pytest.raises(tvm.error.DiagnosticError): + tvm.script.from_source( + """ +@T.prim_func +def main(): + undefined_function() +""" + ) + assert ir.GlobalVar.__call__ is original_call + + def test_tir_buffer_proxy(): buffer_0 = T.Buffer((128, 128), "float32") assert ( From c4a66f73a394caa27cf1ced246630cd6d252e757 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Thu, 17 Sep 2026 22:19:32 +0000 Subject: [PATCH 2/2] [REFACTOR][IR] Keep call builder patch focused on implementation --- tests/python/relax/test_blockbuilder_core.py | 57 --------- tests/python/relax/test_expr.py | 18 --- .../tvmscript/test_tvmscript_parser_tir.py | 121 ------------------ 3 files changed, 196 deletions(-) diff --git a/tests/python/relax/test_blockbuilder_core.py b/tests/python/relax/test_blockbuilder_core.py index a10b9196c5dc..212b9d438877 100644 --- a/tests/python/relax/test_blockbuilder_core.py +++ b/tests/python/relax/test_blockbuilder_core.py @@ -40,63 +40,6 @@ def nop(): pass -@pytest.mark.parametrize("global_callee", [False, True]) -@pytest.mark.parametrize("ret_ty", [tvm.ir.PrimType("int32"), tvm.ir.TupleType([])]) -def test_normalize_call_checks_fixed_result_arguments(global_callee, ret_ty): - signature = rx.FuncType([tvm.ir.PrimType("int32")], ret_ty) - if global_callee: - callee = tvm.ir.GlobalVar("callee") - rx.expr._update_type(callee, signature) - else: - callee = rx.Var("callee", signature) - arg = rx.Var("arg", tvm.ir.PrimType("int32")) - wrong_arg = rx.Var("arg", tvm.ir.PrimType("float32")) - bb = rx.BlockBuilder() - call = callee(arg) - assert call.ty.is_missing() - tvm.ir.assert_structural_equal(bb.normalize(call).ty, ret_ty) - with pytest.raises(ValueError, match="Number of arguments and parameters mismatch"): - bb.normalize(callee()) - with pytest.raises(ValueError, match="type mismatch"): - bb.normalize(callee(wrong_arg)) - # An explicit result remains an override of normalization's type deduction. - override = tvm.ir.Call(callee, [], ret_ty=ret_ty) - tvm.ir.assert_structural_equal(bb.normalize(override).ty, ret_ty) - - -def test_normalize_call_opaque_and_dependent_results(): - scalar = tvm.ir.PrimType("int32") - opaque = rx.Var("opaque", rx.FuncType.opaque_func(ret=scalar)) - bb = rx.BlockBuilder() - tvm.ir.assert_structural_equal(bb.normalize(opaque()).ty, scalar) - - seen = [] - - @tvm.register_global_func("test.call_builder.derive", override=True) - def derive(call, _ctx): - seen.append(call) - return call.args[0].ty - - custom = rx.Var( - "custom", - rx.FuncType.opaque_func(derive_func=tvm.ir.EnvFunc.get("test.call_builder.derive")), - ) - arg = rx.Var("arg", tvm.ir.PrimType("float32")) - call = custom(arg) - assert call.ty.is_missing() - tvm.ir.assert_structural_equal(bb.normalize(call).ty, arg.ty) - assert len(seen) == 1 - - n = tirx.Var("n", "int64") - dependent = rx.Var( - "dependent", rx.FuncType([rx.TensorType([n], "float32")], rx.TensorType([n + 1], "float32")) - ) - tensor = rx.Var("tensor", rx.TensorType([5], "float32")) - call = dependent(tensor) - assert call.ty.is_missing() - tvm.ir.assert_structural_equal(bb.normalize(call).ty, rx.TensorType([6], "float32")) - - def test_block_builder(): m = tirx.Var("m", "int64") n = tirx.Var("n", "int64") diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py index 29aeeacd1ebf..fc53813afed2 100644 --- a/tests/python/relax/test_expr.py +++ b/tests/python/relax/test_expr.py @@ -360,24 +360,6 @@ def test_call(): assert call.args[0].same_as(arg) -@pytest.mark.parametrize("signature_type", [tvm.ir.FuncType, rx.FuncType]) -@pytest.mark.parametrize( - "ret_ty", - [ - tvm.ir.PrimType("int32"), - tvm.ir.PointerType(tvm.ir.PrimType("float32")), - tvm.ir.TupleType([]), - ], -) -def test_call_result_is_explicit(signature_type, ret_ty): - gv = tvm.ir.GlobalVar("callee") - rx.expr._update_type(gv, signature_type([], ret_ty)) - assert tvm.ir.Call(gv, []).ty.is_missing() - assert gv().ty.is_missing() - tvm.ir.assert_structural_equal(tvm.ir.Call(gv, [], ret_ty=ret_ty).ty, ret_ty) - assert tvm.ir.Call(gv, [], ret_ty=tvm.ir.Type.missing()).ty.is_missing() - - def test_call_accepts_core_expr_operator(): """relax.Call aliases the core ir.Call constructor.""" dtype = tvm.ir.PrimType("int32") diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py index b847b4110640..a4baa9b4568a 100644 --- a/tests/python/tvmscript/test_tvmscript_parser_tir.py +++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py @@ -24,127 +24,6 @@ from tvm.script.parser import tirx as T -def test_declared_global_call_results(): - original_call = ir.GlobalVar.__call__ - mod = tvm.script.from_source( - """ -@I.ir_module -class Module: - @T.prim_func - def main(p: T.handle("float32")): - Module.noop() - value = Module.scalar() + T.int32(1) - data = Module.pointer(p) - buffer = T.decl_buffer((1,), "float32", data=data) - buffer[0] = T.Cast("float32", value) - - @T.prim_func - def noop(): - T.evaluate(0) - - @T.prim_func - def scalar() -> T.int32: - return 1 - - @T.prim_func - def pointer(p: T.handle("float32")) -> T.handle("float32"): - return p -""" - ) - assert ir.GlobalVar.__call__ is original_call - calls = {} - - def collect(node, visitor): - if isinstance(node.op, ir.GlobalVar): - calls[node.op.name_hint] = node - visitor.default_visit(node) - - tvm_ffi.structural_visit(mod["main"].body, [(ir.Call, collect)]) - for name in ["noop", "scalar", "pointer"]: - tvm.ir.assert_structural_equal(calls[name].ty, mod[name].ret_type) - tvm.ir.assert_structural_equal(mod, tvm.script.from_source(mod.script())) - assert mod.get_global_var("scalar")().ty.is_missing() - - -def test_global_call_static_signature(): - callee = ir.GlobalVar("callee") - tvm.relax.expr._update_type(callee, ir.FuncType([], ir.PrimType("int32"))) - func = tvm.script.from_source( - """ -@T.prim_func -def main() -> T.int32: - return callee() -""", - extra_vars={"callee": callee, "T": T}, - ) - tvm.ir.assert_structural_equal(func.body.value.ty, ir.PrimType("int32")) - assert callee().ty.is_missing() - - -def test_mixed_module_global_call_adapters(): - original_call = ir.GlobalVar.__call__ - mod = tvm.script.from_source( - """ -@I.ir_module -class Module: - @T.prim_func - def scalar() -> T.int32: - return 1 - - @R.function - def pair(x: R.Tuple(R.Tensor((1,), "float32"), R.Tensor((1,), "float32"))): - return x - - @R.function - def main(x: R.Tensor((1,), "float32")): - result = Module.pair((x, x)) - return result -""" - ) - assert ir.GlobalVar.__call__ is original_call - call = mod["main"].body.blocks[0].bindings[-1].value - assert isinstance(call, ir.Call) - assert isinstance(call.args[0], ir.Tuple) - tvm.ir.assert_structural_equal(mod, tvm.script.from_source(mod.script())) - - -@pytest.mark.parametrize("fail_inner", [False, True]) -def test_nested_global_call_adapters(fail_inner): - from tvm.script.parser.core.diagnostics import Source - from tvm.script.parser.core.parser import Parser - - parser = Parser(Source(""), {}) - original_call = ir.GlobalVar.__call__ - callee = ir.GlobalVar("callee") - arg = tvm.relax.Var("arg", tvm.relax.TensorType([1], "float32")) - with parser.with_dispatch_token("relax"): - relax_call = ir.GlobalVar.__call__ - try: - with parser.with_dispatch_token("tirx"): - if fail_inner: - raise ValueError("unwind inner token") - except ValueError: - assert fail_inner - assert ir.GlobalVar.__call__ is relax_call - call = callee((arg, arg)) - assert isinstance(call.args[0], ir.Tuple) - assert call.ty.is_missing() - assert ir.GlobalVar.__call__ is original_call - - -def test_global_call_adapter_restored_after_parser_error(): - original_call = ir.GlobalVar.__call__ - with pytest.raises(tvm.error.DiagnosticError): - tvm.script.from_source( - """ -@T.prim_func -def main(): - undefined_function() -""" - ) - assert ir.GlobalVar.__call__ is original_call - - def test_tir_buffer_proxy(): buffer_0 = T.Buffer((128, 128), "float32") assert (