apache / apache/tvm

[Bug] Segfault in `tvm.compile` (Relax→TIR, CUDA target) inside `tir::transform::InjectPTXLDG32` / `PTXRewriter::VisitStmt_(BufferStore)` when compiling `torch.export` model returning `(tril, triu)` tuple

Open
#18,612 1 comment 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

## Summary

I am seeing a **hard segfault** (no Python exception) during `tvm.compile(...)` for a CUDA target. The crash consistently occurs inside the TIR pass:

* `tvm::tir::transform::InjectPTXLDG32(bool)`
* `tvm::tir::PTXRewriter::VisitStmt_(tvm::tir::BufferStoreNode const*)`
* `tvm::tir::BufferStore::BufferStore(...)`

The input IRModule is produced by converting a PyTorch `torch.export` program using `tvm.relax.frontend.torch.from_exported_program`. The PyTorch model is intentionally small (`Linear(4,4)`) and returns a **tuple of tensors**: `(torch.tril(x), torch.triu(x))`.

This looks like a bug in the `InjectPTXLDG32` rewrite logic, or an unsafe assumption in the pass leading to a null/invalid `BufferStore` construction.

---

## Environment

From the repro output:

* TVM version: `0.22.0`
* TVM commit: `9dbf3f22ff6f44962472f9af310fda368ca85ef2`
* LLVM: `17.0.6`
* Python: `3.10.16` (from stack paths)
* NumPy: `2.2.6`
* PyTorch: `2.9.0+cu128`
* CUDA GPU: `NVIDIA RTX A6000` (sm_86)

Target string used:

```
cuda -keys=cuda,gpu -arch=sm_86 -max_num_threads=1024 -thread_warp_size=32
```

---

## Minimal Repro Script

```python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import torch
import torch.nn as nn
import tvm
from tvm import tir

def print_env_info():
print("==== Environment Info ====")
print("TVM version:", getattr(tvm, "__version__", "unknown"))
try:
print("TVM git commit:", tvm.support.libinfo().get("GIT_COMMIT_HASH", "unknown"))
except Exception:
print("TVM git commit: unknown")
try:
print("TVM LLVM version:", tvm.support.libinfo().get("LLVM_VERSION", "unknown"))
except Exception:
print("TVM LLVM version: unknown")
print("Python (numpy) version:", np.__version__)
print("PyTorch version:", torch.__version__)
print("CUDA available (torch):", torch.cuda.is_available())
if torch.cuda.is_available():
try:
print("CUDA device:", torch.cuda.get_device_name(0))
except Exception:
pass
print("==========================\n")

class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(4, 4)

def forward(self, x):
x = self.linear(x)
return torch.tril(x), torch.triu(x)

def export_to_relax(mod: nn.Module, x: torch.Tensor) -> tvm.IRModule:
mod = mod.to("cpu").eval()
x = x.to("cpu")
ep = torch.export.export(mod, (x,))
from tvm.relax.frontend.torch import from_exported_program
return from_exported_program(ep)

def main():
print_env_info()

if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for this repro, but torch.cuda.is_available() is False")

target_str = "cuda -keys=cuda,gpu -arch=sm_86 -max_num_threads=1024 -thread_warp_size=32"
target = tvm.target.Target(target_str)

relax_pipeline = "default"
tir_pipeline = "default"

model = MyModel()
x = torch.zeros((1, 4), dtype=torch.float32)

print("[repro] exporting torch -> relax ...")
ir_mod = export_to_relax(model, x)

disabled_pass = [
"DeadCodeElimination",
"CanonicalizeBindings",
"Simplify",
"UnrollLoop",
"VectorizeLoop",
"StorageRewrite",
"RemoveNoOp",
"LoopPartition",
]

pass_config = {
"relax.FuseOps.max_depth": 2,
"relax.lift_transform_params.consume_params": 1,
"tir.disable_storage_rewrite": 1,
"tir.disable_vectorize": 1,
"tir.instrument_bound_checkers": 1,
"tir.merge_static_smem": 1,
"tir.noalias": 1,
"tir.ptx_ldg32": 1,
"tir.use_async_copy": 1,
}

pc_kwargs = {
"opt_level": 3,
"disabled_pass": disabled_pass,
"config": pass_config,
}

print("[repro] target:", target)
print("[repro] relax_pipeline:", relax_pipeline)
print("[repro] tir_pipeline:", tir_pipeline)
print("[repro] opt_level:", pc_kwargs["opt_level"])
print("[repro] disabled_pass:", disabled_pass)
print("[repro] PassContext.config keys:", sorted(pass_config.keys()))
print("[repro] compiling with tvm.compile ...")

with tvm.transform.PassContext(**pc_kwargs):
_ = tvm.compile(
ir_mod,
target=target,
relax_pipeline=relax_pipeline,
tir_pipeline=tir_pipeline,
)

print("[repro] compile finished (no crash).")

if __name__ == "__main__":
main()
```

---

## Actual Behavior

Segfault during compilation:

```text
!!!!!!! Segfault encountered !!!!!!!
...
tvm::tir::BufferStore::BufferStore(...)
tvm::tir::PTXRewriter::VisitStmt_(tvm::tir::BufferStoreNode const*)
...
tvm::tir::transform::InjectPTXLDG32(bool)
Segmentation fault (core dumped)
```

This is a hard crash (core dumped), not a recoverable error.

---

## Expected Behavior

`tvm.compile(...)` should either:

* successfully compile the module, or
* raise a normal Python exception / diagnostic if some pass config is invalid,

but it should **not segfault**.

### Triage

* needs-triage
* bug

Contributor guide

No contributing guide indexed for this repository

Research direction

Reproduce the crash with the supplied torch.export model, CUDA target, and PassContext configuration, then inspect InjectPTXLDG32 and PTXRewriter::VisitStmt_(BufferStore) using the reported stack trace as the entry point. Done means tvm.compile completes successfully or reports a normal diagnostic instead of segfaulting; add a regression test if the relevant test location can be identified.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.