Repository navigation
[Bug] DLight Matmul crashes on fused conv2d-bias-ReLU-pad consumer chain #20615
Description
Expected behavior
Applying the default CUDA DLight rules to a valid fused S-TIR function should not abort. If the Matmul rule cannot schedule the consumer chain after its cache-write stage, it should decline the function so another rule can be considered, or otherwise produce a valid schedule.
Actual behavior
For a fused conv2d -> bias add -> ReLU -> output pad function, dl.gpu.Matmul() reaches auto_inline_consumer_chain and raises:
AssertionError: There are some consumers of the cache-write stage that are not properly inlined.
The traceback ends at:
tvm/s_tir/dlight/gpu/matmul.py, in auto_inline_consumer_chain
assert len(sch.get_consumers(block)) == 0, msg
Environment
- OS: Windows 11 x86_64
- Python: 3.11.16
- TVM source commit:
1ea930e6db63c9747b9d81a12438864db3e7126e - LLVM: 22.1.8
- Reproduced with a source build configured with
USE_LLVM=ON,USE_RTTI=ON, andUSE_CUDA=OFF. CUDA runtime support is not required to reproduce because the failure occurs while applying the Python DLight schedule rule, before CUDA code generation.
Steps to reproduce
import tvm
from tvm.script import ir as I
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.s_tir import dlight as dl
@I.ir_module
class Before:
@Ts.function
def main(
data: T.Tensor((1, 8, 9, 17), "float32"),
weight: T.Tensor((4, 8, 2, 2), "float32"),
bias: T.Tensor((1, 4, 1, 1), "float32"),
output: T.Tensor((1, 4, 9, 17), "float32"),
):
Ts.func_attr({"tirx.noalias": True})
padded_input = Ts.sblock_alloc_buffer((1, 8, 9, 17))
conv = Ts.sblock_alloc_buffer((1, 4, 8, 16))
add = Ts.sblock_alloc_buffer((1, 4, 8, 16))
relu = Ts.sblock_alloc_buffer((1, 4, 8, 16))
for n, c, h, w in T.grid(1, 8, 9, 17):
with Ts.sblock("copy_input"):
vn, vc, vh, vw = Ts.axis.remap("SSSS", [n, c, h, w])
padded_input[vn, vc, vh, vw] = data[vn, vc, vh, vw]
for n, f, y, x, rc, ry, rx in T.grid(1, 4, 8, 16, 8, 2, 2):
with Ts.sblock("conv2d_nchw"):
vn, vf, vy, vx, vrc, vry, vrx = Ts.axis.remap(
"SSSSRRR", [n, f, y, x, rc, ry, rx]
)
with Ts.init():
conv[vn, vf, vy, vx] = T.float32(0)
conv[vn, vf, vy, vx] += (
padded_input[vn, vrc, vy + vry, vx + vrx]
* weight[vf, vrc, vry, vrx]
)
for n, f, y, x in T.grid(1, 4, 8, 16):
with Ts.sblock("bias_add"):
vn, vf, vy, vx = Ts.axis.remap("SSSS", [n, f, y, x])
add[vn, vf, vy, vx] = conv[vn, vf, vy, vx] + bias[vn, vf, 0, 0]
for n, f, y, x in T.grid(1, 4, 8, 16):
with Ts.sblock("relu"):
vn, vf, vy, vx = Ts.axis.remap("SSSS", [n, f, y, x])
relu[vn, vf, vy, vx] = T.max(add[vn, vf, vy, vx], T.float32(0))
for n, f, y, x in T.grid(1, 4, 9, 17):
with Ts.sblock("pad_output"):
vn, vf, vy, vx = Ts.axis.remap("SSSS", [n, f, y, x])
output[vn, vf, vy, vx] = T.if_then_else(
vy < 8 and vx < 16,
relu[vn, vf, vy, vx],
T.float32(0),
)
target = tvm.target.Target("cuda")
with target:
dl.ApplyDefaultSchedule(
dl.gpu.Matmul(),
dl.gpu.GEMV(),
dl.gpu.Reduction(),
dl.gpu.GeneralReduction(),
dl.gpu.Fallback(),
)(Before)This reproducer uses only synthetic shapes and buffers.
Triage
- needs-triage
- type: bug
Activity