[Bug][TVMScript] Printed TVMScript Cannot Be Parsed #14328
Description
In the following script where we try to print out the IRModule, the printed out script will have T_equal[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] == True which cannot be parsed back to IRModule.
CC @cyx-6
import tvm
from tvm.script import tir as T
@tvm.script.ir_module
class Expected:
@T.prim_func
def equal(
rxplaceholder: T.Buffer((T.int64(3), T.int64(3)), "bool"),
T_equal: T.Buffer((T.int64(3), T.int64(3)), "bool"),
):
T.func_attr({"tir.noalias": True})
# with T.block("root"):
for ax0, ax1 in T.grid(T.int64(3), T.int64(3)):
with T.block("T_equal"):
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
T.reads(rxplaceholder[v_ax0, v_ax1])
T.writes(T_equal[v_ax0, v_ax1])
T_equal[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] == T.bool(True)
if __name__ == "__main__":
mod = Expected
mod.show()