Revert "Skip applying sch_rule when both ann and sch_rule are defined"

This reverts commit 4915c6a.
Lucien0 · Apr 13, 2022 · 8a48fb5 · 8a48fb5
1 parent daea033
commit 8a48fb5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc
@@ -145,7 +145,7 @@ class PostOrderApplyNode : public SpaceGeneratorNode {
                        << ann.value();
         }
 
-        if ((ann.defined() && sch_rule.defined()) || (!has_schedule_rule && !sch_rule.defined()) ||
+        if ((!sch_rule.defined() && !has_schedule_rule) ||
             (ann.defined() && ann.value() == "None")) {
           stack.emplace_back(sch, blocks);
           continue;

diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc
@@ -2085,7 +2085,10 @@ Optional<TensorizeInfo> GetTensorizeLoopMapping(const tir::ScheduleState& self,
 
       if (i_desc != n_desc_vars - 1 && iter_types[i] == IterVarType::kCommReduce) continue;
 
-      if (int_block_extent->value == int_desc_extent->value) {
+      // if (int_block_extent->value == int_desc_extent->value) {
+      if (!tir::UsesVar(r, [&block_loop_vars](const tir::VarNode* var) {
+            return block_loop_vars.count(var);
+          })) {
         block_loop = block_loops[i];
         LOG(INFO) << "Selected " << i << " th block loop " << block_loops[i]->loop_var << ", "
                   << block_loop->extent;