[Bug] [Relax] `adaptive_avg_pool2d` compilation crashes on CUDA when output size does not evenly divide input
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
## Reproduction
```python
import tvm
from tvm import relax
import tvm.relax.op as R
from tvm.s_tir import dlight
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorStructInfo((1, 3, 8, 8), "float32"))
with bb.function("main", [x]):
with bb.dataflow():
out = bb.emit(R.nn.adaptive_avg_pool2d(x, output_size=(3, 3))) # 8 % 3 != 0
gv = bb.emit_output(out)
bb.emit_func_output(gv)
mod = bb.get()
pipeline = tvm.ir.transform.Sequential([
relax.transform.LegalizeOps(),
dlight.ApplyDefaultSchedule(dlight.gpu.Fallback()),
])
with tvm.target.Target("cuda"):
mod_cuda = pipeline(mod)
tvm.relax.build(mod_cuda, target="cuda") # RuntimeError
```
## Error message
```
RuntimeError: Memory verification failed with the following errors:
Variable `x` is directly accessed by host memory
(it is not contained in a thread environment or in the function arguments)
```
## Expected behavior
`adaptive_avg_pool2d` should compile and run on CUDA for any valid output size, as it does on CPU (`llvm`).
## Trigger condition
The crash occurs when the output size does not evenly divide the input spatial dimensions. Even divisions compile and run correctly.
| Input spatial | Output size | H % oh == 0? | CUDA |
|:---:|:---:|:---:|:---:|
| 8×8 | 4×4 | Yes | OK |
| 8×8 | 2×2 | Yes | OK |
| 8×8 | 1×1 | Yes | OK |
| 8×8 | 3×3 | No | **Crash** |
| 8×8 | 5×5 | No | **Crash** |
| 16×16 | 3×3 | No | **Crash** |
| 7×7 | 3×3 | No | **Crash** |
| 7×7 | 1×1 | Yes | OK |
All cases above compile and run correctly on CPU (`llvm`).
## Analysis
When the pool window is non-uniform (i.e., different output positions require different window sizes), the TIR produced by `topi.nn.adaptive_pool` contains conditional branches with variable-extent reduction loops. After `dlight.gpu.Fallback()` scheduling, these loops are not bound to GPU threads, causing the memory verifier to reject the generated code.
## Environment
- TVM: main branch (commit `0b0afd8dd`)
- Target: `cuda`
- Python: 3.11
- OS: Ubuntu Linux
cc @junrushao
Contributor guide
No contributing guide indexed for this repository
Research direction
Reproduce the issue with the provided Python script, then inspect the TIR from topi.nn.adaptive_pool after Relax LegalizeOps and dlight.gpu.Fallback(). Trace the variable-extent reduction loops and their GPU thread binding, using the memory verifier failure as the guide. Done means adaptive_avg_pool2d compiles and runs on CUDA for both evenly and non-evenly dividing output sizes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100