check stage pipeline

tlc-pack · Jul 13, 2021 · e7b42c0 · e7b42c0
1 parent 94f8252
commit e7b42c0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/tir/schedule/primitive/fuse_split.cc b/src/tir/schedule/primitive/fuse_split.cc
@@ -314,6 +314,7 @@ Array<StmtSRef> Split(ScheduleState self, const StmtSRef& loop_sref,
   // - The execution order has not changed. (The block executes with the same args and the same
   // order with before.
   // Step 1. Check correctness
+  GetScopeRootAndCheckStagePipeline(self, loop_sref);
   const auto* loop = loop_sref->StmtAs<ForNode>();
   if (loop == nullptr) {
     throw NotLoopError(self->mod, loop_sref->stmt->GetTypeKey());
@@ -410,6 +411,7 @@ StmtSRef Fuse(ScheduleState self, const Array<StmtSRef>& loop_srefs) {
   const ForNode* outer_loop = nullptr;
   arith::Analyzer analyzer;
   // Step 1. check correctness
+  GetScopeRootAndCheckStagePipeline(self, loop_srefs[0]);
   for (const StmtSRef& sref : loop_srefs) {
     const auto* loop = sref->StmtAs<ForNode>();
     if (loop == nullptr) {