This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvmscript-generic-parser-builder in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 04273c52a6e867a1f8fc13468924e47c1cfbc514 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Sep 21 02:46:05 2026 +0000 Register dialect construction policies for global function calls --- python/tvm/relax/script/builder_v2/__init__.py | 13 +++++++++++++ python/tvm/tirx/script/builder_v2/__init__.py | 9 +++++++++ 2 files changed, 22 insertions(+) diff --git a/python/tvm/relax/script/builder_v2/__init__.py b/python/tvm/relax/script/builder_v2/__init__.py index 34cd2124fc..9725b8acb5 100644 --- a/python/tvm/relax/script/builder_v2/__init__.py +++ b/python/tvm/relax/script/builder_v2/__init__.py @@ -19,6 +19,8 @@ # pylint: disable=wildcard-import,redefined-builtin,invalid-name import builtins as _python import numbers as _numbers +import sys as _sys +from functools import partial as _partial import tvm_ffi as _ffi @@ -445,3 +447,14 @@ def select(condition, true_value, false_value): __all__ += ["logical_and", "logical_not", "logical_or", "select"] + + +def _call_global(function, *args): + return _relax.Call(function, [_relax.utils.convert_to_expr(value) for value in args]) + + +def _global_callee(function): + return _partial(_call_global, function) + + +_protocol.register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee) diff --git a/python/tvm/tirx/script/builder_v2/__init__.py b/python/tvm/tirx/script/builder_v2/__init__.py index 9e28d89d70..f12acf0198 100644 --- a/python/tvm/tirx/script/builder_v2/__init__.py +++ b/python/tvm/tirx/script/builder_v2/__init__.py @@ -17,6 +17,7 @@ """Concrete TIRx construction operations over the shared native IRBuilder stack.""" import builtins as _python +import sys as _sys from functools import partial as _partial from functools import wraps as _wraps @@ -30,6 +31,7 @@ from tvm.script.ir_builder.base import IRBuilderFrame as _NativeFrame from tvm.script.ir_builder.protocol import MISSING as _MISSING from tvm.script.ir_builder.protocol import at as _at from tvm.script.ir_builder.protocol import expression_args as _expression_args +from tvm.script.ir_builder.protocol import register_call_kind as _register_call_kind from tvm.script.ir_builder.protocol import register_declaration as _register_declaration from tvm.script.ir_builder.protocol import span_context as _span_context from tvm.tirx.script import builder as _T @@ -549,3 +551,10 @@ def select(condition, true_value, false_value): if not isinstance(condition, _ir.Expr): return true_value if condition else false_value return _tir.if_then_else(condition, true_value, false_value) + + +def _global_callee(function): + return _partial(_T.ir._call_global, function) + + +_register_call_kind(_sys.modules[__name__], _ir.GlobalVar, _global_callee)
