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 ec242df27a4ee72d412dcf27c756ddf0e2891d81 Author: Tianqi Chen <[email protected]> AuthorDate: Mon Sep 21 02:46:25 2026 +0000 Resolve concrete callables through registered builder policies --- python/tvm/script/ir_builder/protocol.py | 14 ++++++++++++++ python/tvm/script/parser_v2/transform.py | 7 +++++++ 2 files changed, 21 insertions(+) diff --git a/python/tvm/script/ir_builder/protocol.py b/python/tvm/script/ir_builder/protocol.py index 98975d3bd7..35e19104ce 100644 --- a/python/tvm/script/ir_builder/protocol.py +++ b/python/tvm/script/ir_builder/protocol.py @@ -179,3 +179,17 @@ def select_lazy(operation, condition, true_value, false_value): return true_value() if condition else false_value() return operation(condition, true_value(), false_value()) + +def register_call_kind(builder, value_type, adapter): + """Register a concrete callable's construction policy on its builder namespace.""" + policies = dict(getattr(builder, "__tvm_call_kinds__", {})) + policies[value_type] = adapter + builder.__tvm_call_kinds__ = policies + + +def callee(builder, value): + """Resolve an already-evaluated callable without altering global call behavior.""" + for value_type, adapter in getattr(builder, "__tvm_call_kinds__", {}).items(): + if isinstance(value, value_type): + return adapter(value) + return value diff --git a/python/tvm/script/parser_v2/transform.py b/python/tvm/script/parser_v2/transform.py index 80bd0fb1ec..674afc823b 100644 --- a/python/tvm/script/parser_v2/transform.py +++ b/python/tvm/script/parser_v2/transform.py @@ -214,6 +214,13 @@ class Transformer(ast.NodeTransformer): child.format_spec = self._format_spec(child.format_spec) else: self._expression_children(node, declarations) + if isinstance(node, ast.Call): + node.func = self._call( + self.infrastructure_name, + "callee", + [self._name(self.builder_name, original), node.func], + original, + ) if isinstance(node, ast.Starred) or ( isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store) ):
