[Bug] `relax.build` crashes with a `ScheduleError` during the dlight GPU scheduling pass

[XYZboom]

Expected behavior

relax.build(mod, target="cuda") should compile successfully for any valid Relax IR module containing relax.op.nn.conv2d with valid input shapes. The compiled module should execute and produce correct convolution results.

Actual behavior

relax.build crashes with a ScheduleError during the dlight GPU scheduling pass:

tvm.s_tir.schedule.schedule.ScheduleError: ScheduleError: An error occurred in the schedule primitive 'bind'.

Environment

• OS: Linux (x86_64, conda environment)
• 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([2, 1, 4, 20]), dtype="float32"))
w = relax.Var("w", relax.TensorStructInfo(
    shape=relax.ShapeExpr([1, 1, 3, 1]), dtype="float32"))
with bb.function("f", [v, w]):
    c = bb.emit(relax.op.nn.conv2d(v, w,
        strides=[1, 1], padding=[0, 0], dilation=[1, 1], groups=1))
    bb.emit_func_output(c)
mod = bb.get()
relax.build(mod, target="cuda")

error log:

Traceback (most recent call last):
  File "/root/autodl-tmp/data/maybeBug/tvm_dlight_reduction_conv2d/main.py", line 31, in <module>
    relax.build(mod, target="cuda")
  File "/root/miniconda3/lib/python3.12/site-packages/tvm/relax/vm_build.py", line 270, in build
    mod = relax_pipeline(mod)
          ^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/lib/python3.12/site-packages/tvm/ir/transform.py", line 171, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/tvm_ffi/cython/function.pxi", line 968, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in TVMFFIPyCallback(void*, TVMFFIAny const*, int, TVMFFIAny*)
  File "<unknown>", line 0, in TVMFFIPyCallManager::ForwardPyErrorToFFI()
  File "<unknown>", line 0, in TVMFFICyErrorSetRaisedFromPyError(_object*)
  File "/root/miniconda3/lib/python3.12/site-packages/tvm/relax/backend/cuda/pipeline.py", line 87, in _pipeline
    mod = seq(mod)

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/ir/transform.py", line 171, in __call__
    return _ffi_transform_api.RunPass(self, mod)

  File "python/tvm_ffi/cython/function.pxi", line 968, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in TVMFFIPyCallback(void*, TVMFFIAny const*, int, TVMFFIAny*)
  File "<unknown>", line 0, in TVMFFIPyCallManager::ForwardPyErrorToFFI()
  File "<unknown>", line 0, in TVMFFICyErrorSetRaisedFromPyError(_object*)
  File "/root/miniconda3/lib/python3.12/site-packages/tvm/ir/transform.py", line 238, in _pass_func
    return inst.transform_module(mod, ctx)

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/s_tir/dlight/base/transform.py", line 70, in transform_module
    sch = _apply_rules(func, target, self.rules, tunable=False)

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/s_tir/dlight/base/transform.py", line 88, in _apply_rules
    space = rule.apply(func, target, tunable)

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/s_tir/dlight/gpu/reduction.py", line 111, in apply
    self._sch_inner_spatial(

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/s_tir/dlight/gpu/reduction.py", line 285, in _sch_inner_spatial
    sch.bind(s, "threadIdx.x")

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/s_tir/schedule/_type_checker.py", line 368, in wrap
    return func(*args, **kwargs)

  File "/root/miniconda3/lib/python3.12/site-packages/tvm/s_tir/schedule/schedule.py", line 1302, in bind
    _ffi_api.ScheduleBind(self, loop, thread_axis)  # type: ignore # pylint: disable=no-member

  File "python/tvm_ffi/cython/function.pxi", line 968, in tvm_ffi.core.Function.__call__
  File "/project/src/s_tir/schedule/concrete_schedule.cc", line 652, in virtual void tvm::s_tir::ConcreteScheduleNode::Bind(const tvm::s_tir::LoopRV&, const tvm::ffi::String&)
tvm.s_tir.schedule.schedule.ScheduleError: ScheduleError: An error occurred in the schedule primitive 'bind'.

The IR with diagnostic is:
# from tvm.script import ir as I
# from tvm.script import tirx as T
# from tvm.tirx.layout import Axis

@I.ir_module
class Module:
    @T.prim_func(private=True, s_tir=True)
    def main(var_v: T.handle, var_w: T.handle, var_conv2d_nchw: T.handle):
        T.func_attr({"op_pattern": 4, "tirx.noalias": True})
        v = T.match_buffer(var_v, (T.int64(2), T.int64(1), T.int64(4), T.int64(20)))
        w = T.match_buffer(var_w, (T.int64(1), T.int64(1), T.int64(3), T.int64(1)))
        conv2d_nchw = T.match_buffer(var_conv2d_nchw, (T.int64(2), T.int64(1), T.int64(2), T.int64(20)))
        with T.sblock("root"):
            T.reads()
            T.writes()
            conv2d_nchw_rf_local = T.sblock_alloc_buffer((T.int64(16), T.int64(2), T.int64(1), T.int64(2), T.int64(20)), scope="local")
            for ax0_ax1_ax2_fused_0 in T.thread_binding(T.int64(8), thread="blockIdx.x"):
                for ax0_ax1_ax2_fused_1 in T.thread_binding(T.int64(10), thread="threadIdx.x"):
                    for ax3_fused_1 in T.thread_binding(T.int64(16), thread="threadIdx.y"):
                        with T.sblock("conv2d_nchw_rf_init"):
                            vax3_fused_1 = T.axis.spatial(T.int64(16), ax3_fused_1)
                            v0 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_fused_0 * T.int64(10) + ax0_ax1_ax2_fused_1) // T.int64(40))
                            v1 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_fused_0 * T.int64(10) + ax0_ax1_ax2_fused_1) % T.int64(40) // T.int64(20))
                            v2 = T.axis.spatial(T.int64(20), (ax0_ax1_ax2_fused_0 * T.int64(10) + ax0_ax1_ax2_fused_1) % T.int64(20))
                            T.reads()
                            T.writes(conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2])
                            conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2] = T.float32(0.0)
                        for ax3_fused_0 in range(T.int64(1)):
                            for u in range(1):
                                with T.sblock("conv2d_nchw_rf_update"):
                                    vax3_fused_1 = T.axis.spatial(T.int64(16), ax3_fused_1)
                                    v0 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_fused_0 * T.int64(10) + ax0_ax1_ax2_fused_1) // T.int64(40))
                                    v1 = T.axis.spatial(T.int64(2), (ax0_ax1_ax2_fused_0 * T.int64(10) + ax0_ax1_ax2_fused_1) % T.int64(40) // T.int64(20))
                                    v2 = T.axis.spatial(T.int64(20), (ax0_ax1_ax2_fused_0 * T.int64(10) + ax0_ax1_ax2_fused_1) % T.int64(20))
                                    vax3_fused_0 = T.axis.reduce(T.int64(1), ax3_fused_0)
                                    T.where(ax3_fused_0 * T.int64(16) + ax3_fused_1 < T.int64(3))
                                    T.reads(conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2], v[v0, T.int64(0), v1 + (vax3_fused_0 * T.int64(16) + vax3_fused_1), v2], w[T.int64(0), T.int64(0), vax3_fused_0 * T.int64(16) + vax3_fused_1, T.int64(0)])
                                    T.writes(conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2])
                                    conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2] = conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2] + v[v0, T.int64(0), v1 + (vax3_fused_0 * T.int64(16) + vax3_fused_1), v2] * w[T.int64(0), T.int64(0), vax3_fused_0 * T.int64(16) + vax3_fused_1, T.int64(0)]
                # tirx.For#0
                for ax1_ax2_ax3_fused in range(T.int64(10)):
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                    for ax0 in range(T.int64(16)):
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                        # tirx.SBlock#1
                        with T.sblock("conv2d_nchw"):
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            vax3_fused_1 = T.axis.reduce(T.int64(16), ax0)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            v0 = T.axis.spatial(T.int64(2), T.Add(ax0_ax1_ax2_fused_0 // T.int64(4), T.int64(0)))
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            v1 = T.axis.spatial(T.int64(2), T.Add(ax0_ax1_ax2_fused_0 // T.int64(2) - ax0_ax1_ax2_fused_0 // T.int64(4) * T.int64(2), T.int64(0)))
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            v2 = T.axis.spatial(T.int64(20), ax0_ax1_ax2_fused_0 % T.int64(2) * T.int64(10) + ax1_ax2_ax3_fused % T.int64(10))
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            T.reads(conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2])
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            T.writes(conv2d_nchw[v0, T.int64(0), v1, v2])
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            with T.init():
                            ^^^^^^^^^^^^^^
                                conv2d_nchw[v0, T.int64(0), v1, v2] = T.float32(0.0)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                            conv2d_nchw[v0, T.int64(0), v1, v2] = conv2d_nchw[v0, T.int64(0), v1, v2] + conv2d_nchw_rf_local[vax3_fused_1, v0, T.int64(0), v1, v2]
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Error message: The queried subtree root tirx.For#0 in SRef tree does not have compact dataflow, because its child block tirx.SBlock#1 on SRef tree is neither a local complete block nor a local reduction block.
It violates condition #1 as a local complete block.
Definition of a local complete block:
1) All block vars are data parallel
2) Local Dominant: the block is the only writer of its output, dominating the reader of its output buffers under a given subtree
3) No overlap between the buffers the block reads and writes
It violates condition #2 as a local reduction block.
Definition of a reduction block:
1) The block has the `init` statement
2) All the block bindings are quasi-affine expressions
3) All block vars are either data parallel block vars or reduction block vars
4) Local Dominant: the block is the only writer of its output, dominating the reader of its output buffers under a given subtree
5) The reduction block vars are not used to index the output buffers

Triage

• bug
• CUDA
• dlight
• reduction
• conv2d
• codegen
• needs-triage