[Relax] Fix end2end tuning for relax graph (#261)

* fix for relax * lint fix
microsoft · Dec 8, 2024 · 803e172 · 803e172
1 parent 5e33ce1
commit 803e172
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/bitblas/relax/transform/apply_fast_tuning.py b/bitblas/relax/transform/apply_fast_tuning.py
@@ -137,10 +137,12 @@ def transform_module(  # pylint: disable=missing-function-docstring
                         updated_functions[g_var] = sch.mod["main"].with_attr("tir.is_scheduled", 1)
                         continue
 
-                if check_func_with_dynamic(func):
+                specalized_function = func.with_attr("global_symbol", g_var.name_hint)
+
+                if check_func_with_dynamic(specalized_function):
 
                     dispatch_mod = fast_tune_with_dynamic_range(
-                        func,
+                        specalized_function,
                         target=target,
                         topk=self.topk,
                         parallel_build=self.parallel_build,
@@ -161,7 +163,7 @@ def transform_module(  # pylint: disable=missing-function-docstring
                 else:
                     # otherwise is static shape analysis
                     _, best = fast_tune(
-                        func,
+                        specalized_function,
                         target=target,
                         topk=self.topk,
                         parallel_build=self.parallel_build,