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

[XYZboom]

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

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