[Unity] MetaScheduleApplyDatabase using workload from records

src/relax/transform/meta_schedule.cc

@@ -126,15 +126,19 @@ Pass MetaScheduleApplyDatabase(Optional<String> work_dir, bool enable_warning =
         if (Optional<meta_schedule::TuningRecord> opt_record =
                 database->QueryTuningRecord(tir_mod, target, gv->name_hint)) {
           meta_schedule::TuningRecord record = opt_record.value();
-          tir::Schedule sch =
-              tir::Schedule::Traced(tir_mod, /*seed=*/-1, /*debug_mask=*/0,
-                                    /*error_render_level=*/tir::ScheduleErrorRenderLevel::kDetail);
+          tir::Schedule sch{nullptr};
           if (!mod_eq_structural->Equal(tir_mod, record->workload->mod)) {
             // When the database lookup succeeds while structural equality check fails,
             // it implies that the anchor block based equality has been used during tuning.
             // The trace in the record cannot directly be applied to this query module.
+            sch = tir::Schedule::Traced(
+                tir_mod, /*seed=*/-1, /*debug_mask=*/0,
+                /*error_render_level=*/tir::ScheduleErrorRenderLevel::kDetail);
             meta_schedule::ScheduleUsingAnchorTrace(sch, record->trace, target);
           } else {
+            sch = tir::Schedule::Traced(
+                record->workload->mod, /*seed=*/-1, /*debug_mask=*/0,
+                /*error_render_level=*/tir::ScheduleErrorRenderLevel::kDetail);
             record->trace->ApplyToSchedule(sch, /*remove_postproc=*/false);
           }
           IRModule new_mod = sch->mod();

tests/python/relax/test_transform_meta_schedule_apply_database.py

@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import tvm
+import tvm.testing
+from tvm import tir
+from tvm import meta_schedule as ms
+from tvm import relax
+from tvm.script import ir as I, tir as T
+
+target = tvm.target.Target("llvm --num-cores=16")
+
+
+def test_apply_to_func_with_different_block_name():
+    @I.ir_module
+    class RecordModule:
+        @T.prim_func
+        def main(A: T.Buffer((2,), "float32"), B: T.Buffer((2,), "float32")):
+            T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
+            for i in T.serial(2):
+                with T.block("block"):
+                    vi = T.axis.spatial(2, i)
+                    B[vi] = A[vi]
+
+    @I.ir_module
+    class BlockRenamedModule:
+        @T.prim_func
+        def main(A: T.Buffer((2,), "float32"), B: T.Buffer((2,), "float32")):
+            T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
+            for i in T.serial(2):
+                with T.block("renamed_block"):
+                    vi = T.axis.spatial(2, i)
+                    B[vi] = A[vi]
+
+    @I.ir_module
+    class Expected:
+        @T.prim_func
+        def main(A: T.Buffer((2,), "float32"), B: T.Buffer((2,), "float32")):
+            T.func_attr(
+                {
+                    "tir.is_scheduled": T.bool(True),
+                    "global_symbol": "main",
+                    "tir.noalias": T.bool(True),
+                }
+            )
+            for i in T.serial(2):
+                with T.block("renamed_block"):
+                    vi = T.axis.spatial(2, i)
+                    B[vi] = A[vi]
+
+    def create_trace(mod: tvm.IRModule):
+        sch = tir.Schedule(mod)
+        _ = sch.get_block("block")
+        return sch.trace
+
+    db = ms.database.create(kind="memory")
+    db.commit_workload(RecordModule)
+    db.commit_tuning_record(
+        ms.database.TuningRecord(
+            create_trace(RecordModule), ms.database.Workload(RecordModule), [0.0], target
+        )
+    )
+
+    with db, target:
+        mod = relax.transform.MetaScheduleApplyDatabase()(BlockRenamedModule)
+    tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()