apache / apache/tvm

[Bug] Segfault in `tvm.compile` on **LLVM (CPU) target** when `tir.ptx_ldg32=1`: unexpectedly runs `tir::transform::InjectPTXLDG32` / `PTXRewriter` and crashes in `BufferStore`

Open
#18,617 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

## Summary

`tvm.compile` segfaults when compiling a Relax module (imported from `torch.export`) **even with a pure CPU target** (`llvm -keys=cpu`). The crash occurs inside the PTX-specific pass:

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

This is unexpected because the target is **LLVM/CPU**, yet the compilation pipeline still enters a PTX rewriting pass. Removing `tir.ptx_ldg32` from PassContext avoids the crash.

This suggests either:

1. `tir.ptx_ldg32` enables a PTX-only pass without checking whether the target is CUDA/PTX-capable, or
2. `InjectPTXLDG32` lacks a defensive early-exit / target predicate and can crash on non-PTX code paths.

---

## Environment

From the repro output:

* TVM: `0.22.0`
* Commit: `9dbf3f22ff6f44962472f9af310fda368ca85ef2`
* LLVM: `17.0.6`
* Python: `3.10.16` (from stack paths)
* NumPy: `2.2.6`
* PyTorch: `2.9.0+cu128`

Target used in repro:

```
llvm -keys=cpu -mtriple=x86_64-unknown-linux-gnu
```

---

## Reproduction Steps

1. Convert a small PyTorch module to Relax via `torch.export.export` + `tvm.relax.frontend.torch.from_exported_program`.
2. Call `tvm.compile` under a PassContext with `config={"tir.ptx_ldg32": 1}`.
3. Observe segfault during compilation.

---

## Minimal Repro Script

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

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

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

class BranchNet(nn.Module):
def __init__(self, k: int):
super().__init__()
self.conv1 = nn.Conv2d(1, 16, k, 1)
self.conv2 = nn.Conv2d(16, 32, 3, 1)
self.pool = nn.MaxPool2d(2)

s1 = 28 - k + 1
s2 = s1 - 2
sp = s2 // 2
self.fc = nn.Linear(32 * sp * sp, 10)

def forward(self, x):
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = self.pool(x)
x = x.reshape(x.shape[0], -1)
return self.fc(x)

class M(nn.Module):
def __init__(self):
super().__init__()
self.b1 = BranchNet(3)
self.b2 = BranchNet(5)
self.b3 = BranchNet(7)
self.out = nn.Linear(30, 10)

def forward(self, x):
a = self.b1(x)
b = self.b2(x)
c = self.b3(x)
y = self.out(torch.cat([a, b, c], dim=1))
return F.log_softmax(y, dim=1)

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

target = tvm.target.Target("llvm -keys=cpu -mtriple=x86_64-unknown-linux-gnu")
tir_pipeline = tir.get_default_tir_pipeline(target) # explicit_default
relax_pipeline = "default"

x = torch.rand(1, 1, 28, 28, dtype=torch.float32)
ir_mod = export_to_relax(M(), x)

pc = {
"opt_level": 0,
"disabled_pass": ["LoopPartition"],
"config": {
"tir.ptx_ldg32": 1,
},
}

print("[repro] target:", target)
print("[repro] tir_pipeline: explicit_default")
print("[repro] compiling with tvm.compile ...")
with tvm.transform.PassContext(**pc):
tvm.compile(ir_mod, target=target, relax_pipeline=relax_pipeline, tir_pipeline=tir_pipeline)

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

---

## Actual Behavior

Segfault during compilation. Stack trace shows PTX rewrite pass even though target is LLVM:

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

---

## Expected Behavior

On a CPU/LLVM target:

1. Setting `tir.ptx_ldg32=1` should either be ignored (no-op) or rejected with a clear error message, and
2. PTX-specific passes such as `InjectPTXLDG32` should **not run** for non-CUDA targets, and
3. TVM should never segfault; failures should be surfaced as Python exceptions with diagnostics.

### Triage

Please refer to the list of label tags [here](https://github.com/apache/tvm/wiki/Issue-Triage-Labels) to find the relevant tags and add them below in a bullet format (example below).

* needs-triage
* bug

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by running the minimal reproduction with the LLVM target and explicit tir.get_default_tir_pipeline(target). Trace the named InjectPTXLDG32 and PTXRewriter entry points to determine why they run with tir.ptx_ldg32 on a CPU target. Done means the PTX pass is skipped or rejected clearly for non-PTX targets and compilation no longer segfaults; the issue names no source file or test.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers
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.