apache / apache/tvm

[Bug] [Relax] Build fails when applying `dlight.gpu.GeneralReduction` to `R.nn.group_norm` with dynamic shapes and `R.reshape`

Open
#17,531 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 19h
Merged PRs (30d)
111

Description

### Actual behavior
When building the TVMScript below using `dlight.gpu.GeneralReduction()`, the build fails with the following error:
`InternalError: Check failed: (!divisor.is_const(0)) is false: Find divide by zero`

```python
from tvm.script import ir as I
from tvm.script import tir as T
from tvm.script import relax as R

@I.ir_module
class Module:
@R.function
def reshape_norm(
inp_0: R.Tensor((1, 512, "w", "h"), dtype="float16"),
inp_1: R.Tensor((512,), dtype="float16"),
inp_2: R.Tensor((512,), dtype="float16")
)-> R.Tensor((1, 512, "w * h"), dtype="float16"):
w = T.int64()
h = T.int64()
with R.dataflow():
lv = R.reshape(inp_0, R.shape([1, 512, w * h]))
lv1 = R.nn.group_norm(data = lv, gamma = inp_1, beta = inp_2, num_groups=32, channel_axis=1, axes=[2], epsilon=9.9999999999999995e-07, center=True, scale=True)
R.output(lv1)
return lv1
```

- **However, if I modify the input tensor `inp_0` and the output tensor shape to `(1, 512, "n")` and remove the `R.reshape` operation, the build completes successfully without errors.**
- It works well with other `dlight` schedules. If I remove `dl.gpu.GeneralReduction()`, the build also completes with other dlight schedules.

### Environment

- TVM Version: v0.18.0
- Commit Hash: 30b7b1c75

### Steps to reproduce

```python
import tvm
from tvm import relax
import tvm.dlight as dl

@tvm.transform.module_pass(opt_level=0)
def dynshape_build_pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
seq = tvm.transform.Sequential(
[
relax.backend.DispatchSampling(),
relax.backend.DispatchSortScan(),
relax.transform.LegalizeOps(),
dl.ApplyDefaultSchedule(
dl.gpu.Matmul(),
dl.gpu.GEMV(),
dl.gpu.Reduction(),
dl.gpu.GeneralReduction(),
dl.gpu.Fallback(),
),
relax.transform.RewriteDataflowReshape(),
relax.transform.ToNonDataflow(),
relax.transform.RemovePurityChecking(),
relax.transform.CallTIRRewrite(),
relax.transform.StaticPlanBlockMemory(),
relax.transform.RewriteCUDAGraph(),
relax.transform.LowerAllocTensor(),
relax.transform.KillAfterLastUse(),
relax.transform.LowerRuntimeBuiltin(),
relax.transform.ComputePrimValue(),
relax.transform.VMShapeLower(),
relax.transform.AttachGlobalSymbol(),
],
)
mod = seq(mod)
return mod

# `Module` as TVMScript in 'Actual behavior'
mod = Module
mod = relax.get_pipeline()(mod)
target = tvm.target.Target("cuda")
ex = relax.build(mod, target=target, pipeline=dynshape_build_pipeline)
```

cc @junrushao

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by running the supplied Python reproducer with the dynshape_build_pipeline and the cuda target, then trace the failure through dl.gpu.GeneralReduction() when R.reshape and dynamic shapes are present. Done means the original TVMScript builds successfully with GeneralReduction enabled, while the existing successful no-reshape case remains unaffected.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.