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 298b94aff371be8eaf053def2474b6468dae1b9e Author: Tianqi Chen <[email protected]> AuthorDate: Sun Sep 20 22:40:56 2026 +0000 [TVMScript] Reserve module identities before signature construction --- include/tvm/script/ir_builder/ir/ir.h | 3 +++ python/tvm/script/ir_builder/ir/__init__.py | 1 + python/tvm/script/ir_builder/ir/ir.py | 5 +++++ src/script/ir_builder/ir/ir.cc | 15 +++++++++++++-- src/tirx/script/builder/frame.cc | 3 ++- 5 files changed, 24 insertions(+), 3 deletions(-) diff --git a/include/tvm/script/ir_builder/ir/ir.h b/include/tvm/script/ir_builder/ir/ir.h index 761e2a995f..9abbdc2b07 100644 --- a/include/tvm/script/ir_builder/ir/ir.h +++ b/include/tvm/script/ir_builder/ir/ir.h @@ -36,6 +36,9 @@ namespace ir { */ TVM_DLL IRModuleFrame IRModule(); +/*! \brief Reserve a stable function identity before evaluating its signature. */ +TVM_DLL GlobalVar ReserveFunction(const ffi::String& func_name); + /*! * \brief Declare a Function without given the specific function implementation. * \note It is usually used in cross-function call. And we can specify the function by `DefFunction` diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/ir/__init__.py index 2b3785a349..04a973c307 100644 --- a/python/tvm/script/ir_builder/ir/__init__.py +++ b/python/tvm/script/ir_builder/ir/__init__.py @@ -20,6 +20,7 @@ from ..protocol import MISSING, _at, expression_args from .frame import IRModuleFrame from .ir import ( + reserve_function, decl_function, def_function, ir_module, diff --git a/python/tvm/script/ir_builder/ir/ir.py b/python/tvm/script/ir_builder/ir/ir.py index e2810ce698..ea37fe643d 100644 --- a/python/tvm/script/ir_builder/ir/ir.py +++ b/python/tvm/script/ir_builder/ir/ir.py @@ -66,6 +66,11 @@ def ir_module() -> IRModuleFrame: return _ffi_api.IRModule() # type: ignore[attr-defined] # pylint: disable=no-member +def reserve_function(func_name: str) -> GlobalVar: + """Reserve a module identity before a declaration frame evaluates its signature.""" + return _ffi_api.ReserveFunction(func_name) + + def decl_function(func_name: str, func_signature: BaseFunc) -> GlobalVar: """Declare a Function without given the specific function implementation. diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc index fd9f3274b2..da0cc81819 100644 --- a/src/script/ir_builder/ir/ir.cc +++ b/src/script/ir_builder/ir/ir.cc @@ -70,12 +70,22 @@ inline ffi::Optional<Type> GetGlobalVarType(const BaseFunc& func) { return std::nullopt; } -GlobalVar DeclFunction(const ffi::String& func_name, const BaseFunc& func_signature) { +GlobalVar ReserveFunction(const ffi::String& func_name) { IRModuleFrame frame = FindModuleFrame(); TVM_FFI_CHECK(!frame->global_var_map.count(func_name), ValueError) << "function " << func_name << " already exists"; + GlobalVar gv(func_name); + frame->global_var_map.Set(func_name, gv); + return gv; +} - GlobalVar gv = GlobalVar(func_name); +GlobalVar DeclFunction(const ffi::String& func_name, const BaseFunc& func_signature) { + IRModuleFrame frame = FindModuleFrame(); + GlobalVar gv = frame->global_var_map.count(func_name) + ? frame->global_var_map.at(func_name) + : GlobalVar(func_name); + TVM_FFI_CHECK(!frame->functions.count(gv), ValueError) + << "function " << func_name << " already exists"; if (auto ty = GetGlobalVarType(func_signature)) { gv->ty = ty.value(); } else { @@ -192,6 +202,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() .def("script.ir_builder.ir.IRModule", IRModule) + .def("script.ir_builder.ir.ReserveFunction", ReserveFunction) .def("script.ir_builder.ir.DeclFunction", DeclFunction) .def("script.ir_builder.ir.DefFunction", DefFunction) .def("script.ir_builder.ir.ModuleAttrs", ModuleAttrs) diff --git a/src/tirx/script/builder/frame.cc b/src/tirx/script/builder/frame.cc index 5c008ec1a5..bef985e5dc 100644 --- a/src/tirx/script/builder/frame.cc +++ b/src/tirx/script/builder/frame.cc @@ -195,7 +195,8 @@ void PrimFuncFrameNode::ExitWithScope() { "function scope, if it's defined in a Module"; const ir::IRModuleFrame& frame = opt_frame.value(); const ffi::String& func_name = name.value_or(""); - if (!frame->global_var_map.count(func_name)) { + if (!frame->global_var_map.count(func_name) || + !frame->functions.count(frame->global_var_map.at(func_name))) { // Case. First time visiting the function. global_var = ir::DeclFunction(func_name, func); }
