(tvm) branch main updated: Fix crash when multiple PrimFunc objects are present in IRModule (#18384)

[cbalint13]

This is an automated email from the ASF dual-hosted git repository.

cbalint13 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 31a24a4a79 Fix crash when multiple PrimFunc objects are present in 
IRModule (#18384)
31a24a4a79 is described below

commit 31a24a4a7903819468f3749dbb1ac7f48a673a8e
Author: Qingchao Shen <qingchaos...@outlook.com>
AuthorDate: Tue Oct 21 23:45:41 2025 +0800

    Fix crash when multiple PrimFunc objects are present in IRModule (#18384)
---
 python/tvm/meta_schedule/tune_context.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/python/tvm/meta_schedule/tune_context.py 
b/python/tvm/meta_schedule/tune_context.py
index 34527f409e..c3f496265a 100644
--- a/python/tvm/meta_schedule/tune_context.py
+++ b/python/tvm/meta_schedule/tune_context.py
@@ -53,9 +53,8 @@ def _normalize_mod(mod: Union[PrimFunc, IRModule]) -> 
IRModule:
     if not isinstance(mod, IRModule):
         raise TypeError(f"Expected `mod` to be PrimFunc or IRModule, but gets: 
{mod}")
     func_names = mod.get_global_vars()
-    (func_name,) = func_names
-    if len(func_names) == 1 and func_name.name_hint != "main":
-        mod = IRModule({"main": mod[func_name]})
+    if len(func_names) == 1 and func_names[0].name_hint != "main":
+        mod = IRModule({"main": mod[func_names[0]]})
     return mod