[Bug] relax.build with batch_norm on CUDA fails with 'buffer v_red occurred before its declaration' for specific input shape
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
`relax.build(mod, target="cuda")` should compile successfully for any valid Relax IR module, including `relax.op.nn.batch_norm` with any valid input shape. The compiled module should execute and produce correct batch normalization results.
### Actual behavior
`relax.build` crashes with an `InternalError` during CUDA codegen:
```
tvm.error.InternalError: Check failed: (it != info_map_.end()) is false: Load/Store of buffer v_red (0x9c5db70) occurred before its declaration.
```
### Environment
- **OS**: Linux (x86_64)
- **GPU**: NVIDIA GeForce RTX 3080 Ti (12GB VRAM)
- **CUDA driver**: 580.76.05
- **TVM version**: 0.25.0.post1 (installed via pip)
- **Target**: `cuda` (GPU compilation)
- **Python**: 3.12
### Steps to reproduce
```python
from tvm import relax
bb = relax.BlockBuilder()
v = relax.Var("v", relax.TensorStructInfo(
shape=relax.ShapeExpr([1, 2, 1]), dtype="float32"))
with bb.function("f", [v]):
bn = bb.emit(relax.op.nn.batch_norm(v,
gamma=relax.op.ones(relax.ShapeExpr([2]), dtype="float32"),
beta=relax.op.zeros(relax.ShapeExpr([2]), dtype="float32"),
moving_mean=relax.op.zeros(relax.ShapeExpr([2]), dtype="float32"),
moving_var=relax.op.ones(relax.ShapeExpr([2]), dtype="float32"),
axis=1))
out = bb.emit(relax.TupleGetItem(bn, 0))
bb.emit_func_output(out)
mod = bb.get()
ex = relax.build(mod, target="cuda")
```
### Triage
* needs-triage
* bug
* CUDA
* codegen
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the provided Relax IR reproducer and trace relax.build(mod, target="cuda") through CUDA codegen, focusing on relax.op.nn.batch_norm and the reported v_red buffer declaration error. Verify the fix by compiling the reproducer successfully and executing the compiled module to confirm correct batch-normalization results.
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
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100