[Bug] s_tir.LowerInitBlock skips T.init() when symbolic reduction extent is 0, returning uninitialized memory #20974

Expected behavior

When compiling and executing a reduction whose reduction extent is 0 at runtime (either via tvm.relax.frontend.onnx.from_onnx on an ONNX model with a zero-length dimension, or via dynamic R.prod / R.matmul / R.sum in Relax), TVM should initialize the output buffer to the reduction identity element (1.0 for ReduceProd / prod, 0.0 for MatMul / ReduceSum), matching NumPy and the ONNX specification.

Actual behavior

s_tir.transform.LowerInitBlock (src/s_tir/transform/lower_init_block.cc, InitBlockLower::DoLowering, lines 64–81) lowers SBlockNode::init into a branch if (v_k == 0) { init_body } inside the reduction loop for (k = 0; k < K; ++k).

Whenever a symbolic or zero-extent reduction bound K == 0 at runtime:

1. In tvm.relax.frontend.onnx.from_onnx, zero-sized dimensions are converted to symbolic tir.Var("v"), so ReduceProd(float32[2, 0], axes=[1]) compiles to a loop for (i1 = 0; i1 < v; ++i1) with if (i1 == 0) out[i0] = 1.0f inside the loop. At runtime v == 0, the loop executes zero iterations, T.init() is completely skipped, and TVM returns [0.0, 0.0] (or uninitialized memory) instead of [1.0, 1.0].
2. In R.matmul (python/tvm/relax/transform/legalize_ops/linear_algebra.py line 59 only special-cases static isinstance(..., IntImm) zero dims), a dynamic (2, k) @ (k, 3) called with k = 0 skips T.init() and returns uninitialized heap floats (e.g., 4.36e+27) instead of 0.0.

Environment

• OS: Windows 11 (10.0.26200)
• TVM version: 0.27.0.post1 (also verified in main src/s_tir/transform/lower_init_block.cc)
• Python: 3.12.0, onnx: 1.23.2, Target: llvm

Steps to reproduce

import numpy as np
import onnx
from onnx import TensorProto, helper
import onnx.reference
import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

# 1) ONNX ReduceProd over zero-extent axis (shape [2, 0], axes=[1])
node = helper.make_node("ReduceProd", ["X", "axes"], ["Y"], keepdims=0)
graph = helper.make_graph(
    [node],
    "reduce_prod_zero_k",
    [helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 0])],
    [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2])],
    [helper.make_tensor("axes", TensorProto.INT64, [1], [1])],
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
onnx.checker.check_model(model, full_check=True)

x_np = np.empty((2, 0), dtype=np.float32)
ref_out = onnx.reference.ReferenceEvaluator(model).run(None, {"X": x_np})[0]

mod = relax.transform.DecomposeOpsForInference()(from_onnx(model))
exe = tvm.compile(mod, target="llvm")
vm = relax.VirtualMachine(exe, tvm.cpu())
tvm_out = vm["main"](tvm.runtime.tensor(x_np)).numpy()

print("ONNX ReferenceEvaluator:", ref_out)  # [1. 1.]
print("TVM Relax LLVM:         ", tvm_out)  # [0. 0.]  <-- T.init() skipped!

Triage

• needs-triage
• tir:transform
• relax