[Bug][MetaSchedule][CUDA] tune_relax for minimal conv2d aborts during candidate generation on Windows

[Nanmur]

[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:

[task_scheduler.cc:193] TaskScheduler picks Task #0: "conv2d"

The process then exits with:

LASTEXITCODE=-1073740940

On Windows this corresponds to 0xC0000374, commonly reported as heap corruption.

Environment

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.

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:

[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:

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).