apache / apache/tvm

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

Open
#17,561 1 comment 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.