From 2f2d5d439ae255c88ee4bb162490dd24eb13c6c7 Mon Sep 17 00:00:00 2001 From: Hongyi Jin <3231950289@qq.com> Date: Mon, 6 Mar 2023 21:43:49 -0500 Subject: [PATCH] [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. --- python/tvm/meta_schedule/tune_context.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index 38a46ebe757e..f15976b1cc47 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -44,7 +44,8 @@ def _normalize_mod(mod: Union[PrimFunc, IRModule]) -> IRModule: """Normalize the input to an IRModule""" if isinstance(mod, PrimFunc): - mod = mod.with_attr("global_symbol", "main") + if not (mod.attrs and "global_symbol" in mod.attrs): + mod = mod.with_attr("global_symbol", "main") mod = mod.with_attr("tir.noalias", True) mod = IRModule({"main": mod}) if not isinstance(mod, IRModule):