apache / apache/tvm

[Bug] relax.build with batch_norm on CUDA fails with 'buffer v_red occurred before its declaration' for specific input shape

Open
#20,036 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.