[TIR][Schedule] fix replace (apache#323)

spectrometerHBH · Siyuan Feng · commit bcd15d620cf8 · 2021-07-04T05:20:06.000Z
diff --git a/src/tir/schedule/schedule_replace.cc b/src/tir/schedule/schedule_replace.cc
@@ -414,7 +414,7 @@ class ChildReplacer : private StmtMutator {
         // Case 2. stmt is BlockRealize, src_stmt is Block
         if (realize->block.get() == src_stmt_) {
           ObjectPtr<BlockRealizeNode> new_realize = make_object<BlockRealizeNode>(*realize);
-          new_realize->block = GetRef<Block>(static_cast<const BlockNode*>(src_stmt_));
+          new_realize->block = GetRef<Block>(static_cast<const BlockNode*>(tgt_stmt_.get()));
           new_stmt = BlockRealize(std::move(new_realize));
         }
       }

Original file line number	Diff line number	Diff line change
`@@ -414,7 +414,7 @@ class ChildReplacer : private StmtMutator {`
`414`	`414`	`// Case 2. stmt is BlockRealize, src_stmt is Block`
`415`	`415`	`if (realize->block.get() == src_stmt_) {`
`416`	`416`	`ObjectPtr<BlockRealizeNode> new_realize = make_object<BlockRealizeNode>(*realize);`
`417`		`- new_realize->block = GetRef<Block>(static_cast<const BlockNode*>(src_stmt_));`
	`417`	`+ new_realize->block = GetRef<Block>(static_cast<const BlockNode*>(tgt_stmt_.get()));`
`418`	`418`	`new_stmt = BlockRealize(std::move(new_realize));`
`419`	`419`	`}`
`420`	`420`	`}`