apache / apache/tvm

[Bug] [Relax] `adaptive_avg_pool2d` compilation crashes on CUDA when output size does not evenly divide input

Open
#19,520 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.