[Bug] Crash when applying transform_layout with axis separator

[quic-sanirudh]

When we apply transform_layout on a write buffer generated by cache_write with axis separators, we see an error that the sref pointed by the block does not exist in stmt2ref. The below code reproduces the error:

Expected behavior

The code should work fine without errors and apply the layout transform with axis separators

Actual behavior

Code crashes with error Check failed: (it != self_->stmt2ref.end()) is false: InternalError: The sref points to a statement that does not exist in stmt2ref

Environment

Tested on Upstream TVM main with Ubuntu 18.04

Steps to reproduce

import tvm
import tvm.contrib.hexagon
import tvm.testing
from tvm.script import tir as T


@T.prim_func
def add(
    p0: T.Buffer((T.int64(33), T.int64(128)), "float32"),
    p1: T.Buffer((T.int64(33), T.int64(128)), "float32"),
    T_add: T.Buffer((T.int64(33), T.int64(128)), "float32"),
):
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    # with T.block("root"):
    for ax0, ax1 in T.grid(T.int64(33), T.int64(128)):
        with T.block("T_add"):
            v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
            T.reads(p0[v_ax0, v_ax1], p1[v_ax0, v_ax1])
            T.writes(T_add[v_ax0, v_ax1])
            T_add[v_ax0, v_ax1] = p0[v_ax0, v_ax1] + p1[v_ax0, v_ax1]


def transform_fn(x, y):
    return [x // 32, y, tvm.te.AXIS_SEPARATOR, x % 32]


def main():
    sch = tvm.tir.Schedule(add, debug_mask="all")
    block_rv = sch.get_block("T_add")
    sch.cache_write(block_rv, 0, "global")
    sch.transform_layout(block_rv, ("write", 0), transform_fn, pad_value=0.0)


if __name__ == "__main__":
    main()

Triage

• needs-triage

[quic-sanirudh]

cc @Lunderberg @vinx13

[quic-sanirudh]

The issue is fixed with #14023, hence closing this issue.