[Bug] v0.18.0 tir.UnrollLoop uses unbounded memory during compilation
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 19h
- Merged PRs (30d)
- 111
Description
## Steps to Reproduce:
```python
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
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 const&, tvm::Target)>::AssignTypedLambda(tvm::$_5, std::__cxx11::basic_string, std::allocator >)::{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 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::AssignTypedLambda(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
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the reproducer and the tir.UnrollLoop pass, then inspect tvm/src/tir/transforms/unsupported_dtype_legalize.cc around line 483 and the compilation path reached by tvm.build. Compare compilation with and without tir.UnrollLoop, and verify that the reproducer no longer consumes unbounded memory or crashes while preserving the expected validation behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100