nautasolva opened a new issue, #16889: URL: https://github.com/apache/tvm/issues/16889
When used on a block with a init statement, blockize creates a separate init block that is not discoverable by any means. This hinders further scheduling, like tensorizing the init block. ### Expected behavior When using `blockize` on a loop that contains an init statement, the init is moved to a new block `<block>_init` that should be discoverable with `get_block` or `get_children_blocks` on the newly created outer block. ### Actual behavior Init block exists in the TIR module but does not seem to be registered by the schedule. `get_block("<block_init>)"` fails with `InternalError: Check failed: (it != self_->stmt2ref.end()) is false` <details><summary>Stacktrace</summary> <p> Traceback (most recent call last): File "/home/dev/tvm_upstream/../tvm/playground/blockize_init_bug.py", line 31, in <module> a_init = sch.get_block("A_init") File "/home/dev/tvm_upstream/python/tvm/tir/schedule/_type_checker.py", line 340, in wrap return func(*args, **kwargs) File "/home/dev/tvm_upstream/python/tvm/tir/schedule/schedule.py", line 499, in get_block return _ffi_api.ScheduleGetBlock( # type: ignore # pylint: disable=no-member File "/home/dev/tvm_upstream/python/tvm/_ffi/_ctypes/packed_func.py", line 239, in __call__ raise_last_ffi_error() File "/home/dev/tvm_upstream/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error raise py_err File "/home/dev/tvm_upstream/src/tir/schedule/traced_schedule.cc", line 128, in tvm::tir::TracedScheduleNode::GetBlock(tvm::runtime::String const&, tvm::runtime::Optional<tvm::runtime::String> const&) BlockRV result = ConcreteScheduleNode::GetBlock(name, func_name); File "/home/dev/tvm_upstream/src/tir/schedule/concrete_schedule.cc", line 321, in tvm::tir::ConcreteScheduleNode::GetBlock(tvm::runtime::String const&, tvm::runtime::Optional<tvm::runtime::String> const&) Array<StmtSRef> blocks = tir::GetBlocks(this->state_, name, gv); File "/home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc", line 46, in tvm::tir::GetBlocks(tvm::tir::ScheduleState const&, tvm::runtime::String const&, tvm::GlobalVar const&) finder(prim_func->body); File "/home/dev/tvm_upstream/src/tir/ir/stmt_functor.cc", line 142, in tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::BlockNode const*) this->VisitStmt(op->init.value()); File "/home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc", line 29, in VisitStmt_ void VisitStmt_(const BlockNode* block) override { File "/home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc", line 32, in VisitStmt_ ICHECK(it != self_->stmt2ref.end()); tvm.error.InternalError: Traceback (most recent call last): 5: tvm::tir::TracedScheduleNode::GetBlock(tvm::runtime::String const&, tvm::runtime::Optional<tvm::runtime::String> const&) at /home/dev/tvm_upstream/src/tir/schedule/traced_schedule.cc:128 4: tvm::tir::ConcreteScheduleNode::GetBlock(tvm::runtime::String const&, tvm::runtime::Optional<tvm::runtime::String> const&) at /home/dev/tvm_upstream/src/tir/schedule/concrete_schedule.cc:321 3: tvm::tir::GetBlocks(tvm::tir::ScheduleState const&, tvm::runtime::String const&, tvm::GlobalVar const&) at /home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc:46 2: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::BlockNode const*) at /home/dev/tvm_upstream/src/tir/ir/stmt_functor.cc:142 1: VisitStmt_ at /home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc:29 0: VisitStmt_ at /home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc:32 File "/home/dev/tvm_upstream/src/tir/schedule/primitive/get_block_loop.cc", line 32 </p> </details> ### Environment Reproducible on main (d4056ca79571d4265a12beeedd1b1565953df936) ### Steps to reproduce ```python import tvm from tvm.script import ir as I from tvm.script import tir as T @I.ir_module class Module: @T.prim_func def main(): # with T.block("root"): A_sum = T.alloc_buffer((1,), "float32") A = T.alloc_buffer((1, 16), "float32") for nn, ff in T.grid(1, 16): with T.block("A"): v_nn, v_ff = T.axis.remap("SR", [nn, ff]) T.reads(A[v_nn, v_ff]) T.writes(A_sum[v_nn]) with T.init(): A_sum[v_nn] = T.float32(0) A_sum[v_nn] = A_sum[v_nn] + A[v_nn, v_ff] sch = tvm.tir.Schedule(Module) a = sch.get_block("A") loop_n, loop_f = sch.get_loops(a) sch.blockize(loop_f) print(sch.mod) # <-- A_init exists a_init = sch.get_block("A_init") # <-- fails with InternalError: Check failed: (it != self_->stmt2ref.end()) is false ``` ### Triage * tir:schedule -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org