Skip to content

Commit

Permalink
[MetaSchedule] preseve global_symbol attached to function after apply…
Browse files Browse the repository at this point in the history
…ing 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.
  • Loading branch information
jinhongyii authored Mar 7, 2023
1 parent 012d6a7 commit 2f2d5d4
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion python/tvm/meta_schedule/tune_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 2f2d5d4

Please sign in to comment.