[Bug] v0.18.0 tir.UnrollLoop uses unbounded memory during compilation

[talha-ahsan]

Steps to Reproduce:

import tvm
from tvm import tir
from tvm.tir.analysis.analysis import verify_well_formed, verify_memory

from tvm.script import tir as T

@T.prim_func(private=True)
def main(h1: T.handle, h2: T.handle, h3: T.handle, h4: T.handle):
    N = T.uint32(is_size_var=True)
    bi3 = T.match_buffer(h1, (N, N), "int32")
    bu3 = T.match_buffer(h2, (N, N), "uint32")
    bb3 = T.match_buffer(h3, (N, N), "bool")
    bf3 = T.match_buffer(h4, (N, N))
    for b45 in T.unroll(T.max(T.uint32(1618984980), T.uint32(1747359052)), T.max(T.uint32(1618984980), T.uint32(1747359052)) + T.uint32(1077590990)):
        bf3[T.float32(0.7505693304727864), b45] = T.float32(0.88788715142083252)

func = main
mod = tvm.ir.IRModule({'main': func})
if not verify_well_formed(mod) and verify_memory(func):
    print("Validation failed")
else: 
    print("Beginning Compilation")
    with tvm.transform.PassContext(opt_level=4):
        nopt_mod = tvm.build(mod)
    print("Success!")

Expected Behavior:

Successful Compilation or a reason for why the compilation target is invalid

Reality: Continuous compilation, with slowly increasing memory usage until a machine crash occurs

When tested with tvm.transform.PassContext(opt_level=4, disabled_pass=["tir.UnrollLoop"]) the compilation fails with exit code 1 and the following error message:

Beginning Compilation
Traceback (most recent call last):
  File "reprod_without_unroll.py", line 24, in <module>
    nopt_mod = tvm.build(mod)
  File "/tvm/python/tvm/driver/build_module.py", line 297, in build
    rt_mod_host = _driver_ffi.tir_to_runtime(annotated_mods, target_host)
  File "/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 245, in __call__
    raise_last_ffi_error()
  File "/tvm/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error
    raise py_err
tvm.error.InternalError: Traceback (most recent call last):
  11: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::runtime::Module (tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)>::AssignTypedLambda<tvm::$_5>(tvm::$_5, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  10: tvm::TIRToRuntime(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target const&)
  9: tvm::SplitMixedModule(tvm::IRModule, tvm::Target const&, tvm::Target const&)
  8: tvm::ApplyPasses(tvm::IRModule, tvm::transform::Sequential)
  7: tvm::transform::Pass::operator()(tvm::IRModule) const
  6: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  5: tvm::transform::SequentialNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  4: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  3: tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  2: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::tir::PrimFunc (tvm::tir::PrimFunc, tvm::IRModule, tvm::transform::PassContext)>::AssignTypedLambda<tvm::tir::transform::FP8StorageLegalize()::$_3>(tvm::tir::transform::FP8StorageLegalize()::$_3)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  1: tvm::tir::StorageLegalizer::Legalize(tvm::tir::PrimFunc)
  0: _ZN3tvm7runtime6detail
  File "/tvm/src/tir/transforms/unsupported_dtype_legalize.cc", line 483
InternalError: Check failed: func->buffer_map.size() == 0 (4 vs. 0) : This pass must be called after MakePackedAPI

[talha-ahsan]

Oh some Device/Environment Information if needed:

OS: Ubuntu 20.04 LTS Python: 3.10.4 TVM: v0.18.0 built from source with CPU only, no GPU usage in place

Please let me know if there's any additional information required, or if this is intended behavior for large loops. :)