[Bug] [Relax][CUDA] Default Relax CUDA build emits unscheduled TIR for R.power and fails memory verification
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Description
I found that a minimal Relax program containing only `R.power(x, 2.0)` fails to build for CUDA with the default Relax build pipeline.
The failure happens during TIR build / memory verification. The generated TIR directly accesses buffers from host code and has no GPU thread binding:
```text
RuntimeError: Memory verification failed with the following errors:
Variable `T_power` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.
Variable `x` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.
Did you forget to bind?
```
However, if I explicitly apply `relax.transform.LegalizeOps()` followed by `tir.transform.DefaultGPUSchedule()` under the CUDA target context, the same module builds successfully. This suggests that the default CUDA Relax build pipeline legalizes `R.power` to TIR but does not apply the required GPU scheduling / thread binding before memory verification.
### Actual behavior
The default CUDA Relax build fails:
```
default_cuda_build: FAILED
```
The failure is:
```
RuntimeError: Memory verification failed with the following errors:
Variable `T_power` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.
Variable `x` is directly accessed by host memory (it is not contained in a thread environment or in the function arguments.
Did you forget to bind?
```
The generated TIR has no thread binding:
```
@T.prim_func
def power(
x: T.Buffer((T.int64(1), T.int64(2), T.int64(1), T.int64(1)), "float32"),
T_power: T.Buffer((T.int64(1), T.int64(2), T.int64(1), T.int64(1)), "float32"),
):
T.func_attr({
"target": T.target({
"arch": "sm_86",
"keys": ["cuda", "gpu"],
"kind": "cuda",
"max_num_threads": 1024,
"thread_warp_size": 32
}),
"tir.noalias": True
})
for ax1 in range(2):
T_power_1 = T.Buffer((T.int64(2),), data=T_power.data)
x_1 = T.Buffer((T.int64(2),), data=x.data)
T_power_1[ax1] = T.pow(x_1[ax1], T.float32(2.0))
```
Explicit LegalizeOps() alone still fails:
```
legalize_ops_then_cuda_build: FAILED
```
But applying DefaultGPUSchedule() after LegalizeOps() under the CUDA target context succeeds:
```
legalize_ops_default_gpu_schedule_then_cuda_build: OK
```
The scheduled TIR contains GPU thread binding:
```
for ax0_ax1_ax2_ax3_fused_0 in T.thread_binding(T.int64(1), thread="blockIdx.x"):
for ax0_ax1_ax2_ax3_fused_1 in T.thread_binding(T.int64(2), thread="threadIdx.x"):
...
```
### Environment
TVM: 0.23.0
LLVM: 17.0.6
Python: 3.10.16 (from stack paths)
NumPy: 2.2.6
### Steps to reproduce
```python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import platform
import traceback
import tvm
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R
@I.ir_module
class PowerModule:
@R.function
def main(
x: R.Tensor((1, 2, 1, 1), dtype="float32")
) -> R.Tensor((1, 2, 1, 1), dtype="float32"):
with R.dataflow():
y: R.Tensor((1, 2, 1, 1), dtype="float32") = R.power(
x, R.const(2.0, "float32")
)
R.output(y)
return y
def build_relax_module(mod, target):
if hasattr(tvm, "compile"):
return tvm.compile(mod, target=target)
return relax.build(mod, target=target)
def print_header(title):
print("\n" + "=" * 100)
print(title)
print("=" * 100)
def print_env(target):
print_header("Environment")
print("python:", sys.version.replace("\n", " "))
print("platform:", platform.platform())
print("tvm version:", getattr(tvm, "__version__", ""))
print("tvm path:", getattr(tvm, "__file__", ""))
print("TVM_CUDA_TARGET env:", os.environ.get("TVM_CUDA_TARGET", ""))
try:
dev = tvm.cuda(0)
print("tvm.cuda(0).exist:", dev.exist)
print("tvm.cuda(0):", dev)
except Exception as e:
print("tvm.cuda(0): failed:", repr(e))
print("target:", target)
print("has tvm.compile:", hasattr(tvm, "compile"))
print("has relax.transform.LegalizeOps:", hasattr(relax.transform, "LegalizeOps"))
print("has tir.transform.DefaultGPUSchedule:", hasattr(tvm.tir.transform, "DefaultGPUSchedule"))
def print_exception(e):
print(type(e).__name__, repr(e))
print("\nTraceback:")
traceback.print_exc()
def try_build_case(name, mod, target):
print_header(name)
print("IRModule:")
print(mod)
print("\nTarget:")
print(target)
try:
build_relax_module(mod, target)
print("\n[BUILD] OK")
return True
except Exception as e:
print("\n[BUILD] FAILED")
print_exception(e)
return False
def try_transform_case(name, transform_fn, mod):
print_header(name)
try:
out = transform_fn(mod)
print("[TRANSFORM] OK")
print("\nTransformed IRModule:")
print(out)
return out
except Exception as e:
print("[TRANSFORM] FAILED")
print_exception(e)
return None
def _schedule_with_target(mod, target):
with target:
return tvm.tir.transform.DefaultGPUSchedule()(mod)
def main():
target = tvm.target.Target(
"cuda -keys=cuda,gpu -arch=sm_86 -max_num_threads=1024 -thread_warp_size=32"
)
print_env(target)
print_header("Original Relax IR")
print(PowerModule)
results = {}
results["default_cuda_build"] = try_build_case(
"CASE 1: default CUDA Relax build",
PowerModule,
target,
)
legalized_mod = try_transform_case(
"CASE 2A: relax.transform.LegalizeOps()",
lambda mod: relax.transform.LegalizeOps()(mod),
PowerModule,
)
if legalized_mod is not None:
results["legalize_ops_then_cuda_build"] = try_build_case(
"CASE 2B: build after LegalizeOps",
legalized_mod,
target,
)
else:
results["legalize_ops_then_cuda_build"] = False
scheduled_mod = None
if legalized_mod is not None:
scheduled_mod = try_transform_case(
"CASE 3A: DefaultGPUSchedule after LegalizeOps under CUDA target context",
lambda mod: _schedule_with_target(mod, target),
legalized_mod,
)
if scheduled_mod is not None:
results["legalize_ops_default_gpu_schedule_then_cuda_build"] = try_build_case(
"CASE 3B: build after LegalizeOps + DefaultGPUSchedule",
scheduled_mod,
target,
)
else:
results["legalize_ops_default_gpu_schedule_then_cuda_build"] = False
print_header("Summary")
for k, v in results.items():
print(f"{k}: {'OK' if v else 'FAILED'}")
if __name__ == "__main__":
main()
```
### Triage
* needs-triage
* bug
cc @junrushao
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the provided PowerModule reproduction and run build_relax_module for the default CUDA Relax build, then compare the LegalizeOps and DefaultGPUSchedule cases under the CUDA target context. Trace the Relax build pipeline around these entry points to determine why scheduling is omitted; done means the default build succeeds without memory-verification errors and produces GPU thread binding.
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
- 48/100