[Bug] [Bug][Relax] StaticPlanBlockMemory can emit non-dominating storage across if branches for reshape chains
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
A well-formed Relax function with an if expression should allocate storage in a scope that dominates all uses of that storage.
For the program below, both branches perform:
reshape(x, [32]) -> reshape(flat, [4, 8])
The function is mathematically an identity function. It should return the input tensor for both cond=True and cond=False.
### Actual behavior
The official LLVM target-default Relax pipeline builds the module, but running the VM with cond=False crashes with a native segfault.
Observed locally:
official target-default: native segfault
skip StaticPlanBlockMemory: output matches, max_abs = 0.0
Pass-level evidence points to StaticPlanBlockMemory. After CallTIRRewrite, the two branches contain separate branch-local builtin.alloc_tensor calls. After StaticPlanBlockMemory, the then branch allocates
storage, but the else branch uses the same storage variable without a dominating alloc_storage:
if cond:
storage = R.memory.alloc_storage(...)
alloc = R.memory.alloc_tensor(storage, ...)
...
else:
storage: R.Object
alloc2 = R.memory.alloc_tensor(storage, ...)
...
This makes the generated memory IR invalid for the cond=False path.
### Environment
OS: Linux x86_64
Python: 3.10.12
TVM version: 0.26.dev1
TVM commit: 5a8dae4d95c55c8fec9246a607a28c3ff54ffe05
Target: llvm
Relax VM exec_mode: bytecode
### Steps to reproduce
import numpy as np
import tvm
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def main(
cond: R.Tensor((), dtype="bool"),
x: R.Tensor((4, 8), dtype="float32"),
) -> R.Tensor((4, 8), dtype="float32"):
if cond:
flat_t: R.Tensor((32,), dtype="float32") = R.reshape(x, R.shape([32]))
out_t: R.Tensor((4, 8), dtype="float32") = R.reshape(flat_t, R.shape([4, 8]))
out: R.Tensor((4, 8), dtype="float32") = out_t
else:
flat_f: R.Tensor((32,), dtype="float32") = R.reshape(x, R.shape([32]))
out_f: R.Tensor((4, 8), dtype="float32") = R.reshape(flat_f, R.shape([4, 8]))
out: R.Tensor((4, 8), dtype="float32") = out_f
return out
target = tvm.target.Target("llvm")
exe = relax.build(
Module,
target=target,
relax_pipeline=relax.get_default_pipeline(target),
exec_mode="bytecode",
)
vm = relax.VirtualMachine(exe, tvm.cpu())
x_np = np.arange(32, dtype="float32").reshape(4, 8) * np.float32(0.125)
x = tvm.runtime.tensor(x_np, tvm.cpu())
print("cond=True")
print(vm["main"](tvm.runtime.tensor(np.array(True, dtype="bool"), tvm.cpu()), x).numpy())
print("cond=False")
# This call segfaults locally.
print(vm["main"](tvm.runtime.tensor(np.array(False, dtype="bool"), tvm.cpu()), x).numpy())
Observed crash:
Segfault encountered
runtime/memory/memory_manager.cc: StorageObj::AllocTensor
runtime/vm/builtin.cc: alloc_tensor builtin
runtime/vm/vm.cc: VM RunLoop
I also ran an ablation that keeps the rest of the CPU pipeline but skips only StaticPlanBlockMemory. The same input builds and runs correctly with max_abs=0.0, which suggests that the issue is in storage
planning for branch-local reshape chains.
### Triage
- needs-triage
- type: bug
- relax
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by running the provided Python reproduction through the official Relax pipeline and compare it with the ablation that skips StaticPlanBlockMemory. Inspect StaticPlanBlockMemory output after CallTIRRewrite, focusing on alloc_storage and alloc_tensor across both if branches. Done means both cond=True and cond=False execute without a segfault and return the identity result.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100