[Bug][MetaSchedule][CUDA] tune_relax for minimal conv2d aborts during candidate generation on Windows
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
# [Bug][MetaSchedule][CUDA] tune_relax for minimal conv2d aborts during candidate generation on Windows
## Problem
Tuning a minimal Relax `conv2d` module for CUDA with MetaSchedule can abort the Python process on Windows before any builder result is produced.
The last TVM log line before process termination is:
```text
[task_scheduler.cc:193] TaskScheduler picks Task #0: "conv2d"
```
The process then exits with:
```text
LASTEXITCODE=-1073740940
```
On Windows this corresponds to `0xC0000374`, commonly reported as heap corruption.
## Environment
```text
OS: Windows-10-10.0.26200-SP0
Python: 3.11.14
TVM version: 0.26.dev0
TVM commit: 2fb591c5ba4d64f145ca90e946ea374a78fbba8c
Target: {"kind":"cuda","keys":["cuda","gpu"],"max_threads_per_block":1024,"arch":"sm_120","max_shared_memory_per_block":49152,"max_num_threads":1024,"thread_warp_size":32}
CUDA device available: True
```
## Minimal Reproduction
Run this script from a real `.py` file on Windows. `LocalBuilder` uses multiprocessing, so running from stdin can hide or change process behavior.
```python
from pathlib import Path
import tvm
from tvm import relax
from tvm.s_tir.meta_schedule.relax_integration import tune_relax
def make_target():
dev = tvm.cuda(0)
if not dev.exist:
raise RuntimeError("CUDA device 0 is not available")
return tvm.target.Target.from_device(dev)
def make_relax_module():
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType((1, 3, 8, 8), "float32"))
weight = relax.Var("weight", relax.TensorType((4, 3, 3, 3), "float32"))
with bb.function("main", [x, weight]):
with bb.dataflow():
conv = bb.emit(
relax.op.nn.conv2d(
x,
weight,
strides=(1, 1),
padding=(1, 1),
)
)
output = bb.emit_output(conv)
bb.emit_func_output(output)
return bb.get()
def prepare_tuning_module(mod):
for pass_func in [
relax.transform.LegalizeOps(),
relax.transform.AnnotateTIROpPattern(),
relax.transform.FuseOps(),
relax.transform.FoldConstant(),
relax.transform.FuseTIR(),
]:
mod = pass_func(mod)
return mod
if __name__ == "__main__":
target = make_target()
mod = prepare_tuning_module(make_relax_module())
tune_relax(
mod=mod,
target=target,
params=None,
work_dir=Path("ms_cuda_conv2d_repro").resolve(),
max_trials_global=1,
num_trials_per_iter=1,
max_trials_per_task=1,
builder="local",
runner="local",
strategy="evolutionary",
module_equality="ignore-tensor",
seed=0,
)
```
## Observed Logs
The lowered Relax module successfully produces a MetaSchedule task named `conv2d`. The generated design spaces contain CUDA bindings such as `blockIdx.x` and `threadIdx.x`.
The tuning log reaches candidate generation:
```text
[task_scheduler.cc:172] Initializing Task #0: "conv2d"
[task_scheduler.cc:193] TaskScheduler picks Task #0: "conv2d"
[evolutionary_search.cc:738] Generating candidates......
[evolutionary_search.cc:505] Pick-Best-From-Database summary:
Trace replay failures: 0 failure(s)
Postproc #0 [s_tir.meta_schedule.DisallowDynamicLoop]: 0 failure(s)
Postproc #1 [s_tir.meta_schedule.RewriteCooperativeFetch]: 0 failure(s)
Postproc #2 [s_tir.meta_schedule.RewriteUnboundBlock]: 0 failure(s)
Postproc #3 [s_tir.meta_schedule.RewriteParallelVectorizeUnroll]: 0 failure(s)
Postproc #4 [s_tir.meta_schedule.RewriteReductionBlock]: 0 failure(s)
Postproc #5 [s_tir.meta_schedule.VerifyGPUCode]: 0 failure(s)
Postproc #6 [s_tir.meta_schedule.RewriteTensorize]: 0 failure(s)
[evolutionary_search.cc:740] Picked top 0 candidate(s) from database
[evolutionary_search.cc:551] Sample-Init-Population summary:
Trace replay failures: 0 failure(s)
Postproc #5 [s_tir.meta_schedule.VerifyGPUCode]: 505 failure(s)
```
Shortly after this point, the Python process terminates with:
```text
LASTEXITCODE=-1073740940
```
## Expected Behavior
`tune_relax(..., max_trials_global=1)` should either produce a valid measurement candidate or return a Python/TVM diagnostic error. It should not abort the process.
## Actual Behavior
The process aborts during or immediately after MetaSchedule candidate generation for the `conv2d` task. In this run, no useful builder or runner result is produced before process termination.
## Notes
I also tried wrapping the builder to record each `BuilderInput.mod` before delegating to `LocalBuilder`. In the failing run, no builder batch was recorded, which suggests the abort happens before `TaskScheduler::SendToBuilder` receives a measurable candidate.
The issue is reproducible with the small shape above and also with a larger shape such as input `(1, 3, 48, 320)` and weight `(16, 3, 3, 3)`.
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the provided reproduction from a real .py file on Windows and inspect the candidate-generation path around tune_relax, evolutionary_search, and TaskScheduler::SendToBuilder. Use the logged VerifyGPUCode failures and the absence of builder batches to locate where the process aborts. Done means tuning returns a diagnostic error or produces a valid measurement candidate instead of terminating the Python process.
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
- 45/100