Skip to content

Commit 2f2d5d4

Browse files
authored
[MetaSchedule] preseve global_symbol attached to function after applying MS (#14219)
Right now seems a global symbol main will be attached to every fused prim_func after Meta Schedule tuning. Although the problem get hidden because we will re-attach global symbols, this is a behavior that need to be fixed. Now the transformed prim func preserves the original global symbol. If it doesn't exist, attach global symbol main.
1 parent 012d6a7 commit 2f2d5d4

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

python/tvm/meta_schedule/tune_context.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@
4444
def _normalize_mod(mod: Union[PrimFunc, IRModule]) -> IRModule:
4545
"""Normalize the input to an IRModule"""
4646
if isinstance(mod, PrimFunc):
47-
mod = mod.with_attr("global_symbol", "main")
47+
if not (mod.attrs and "global_symbol" in mod.attrs):
48+
mod = mod.with_attr("global_symbol", "main")
4849
mod = mod.with_attr("tir.noalias", True)
4950
mod = IRModule({"main": mod})
5051
if not isinstance(mod, IRModule):

0 commit comments

Comments
 (0)