apache / apache/tvm

[Bug][Metal] Metal codegen produces silently-wrong output on a HiFi-GAN-style Relax graph (cosine 0.756 vs LLVM 1.0000)

Open
#20,157 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

## Environment

- TVM: `apache-tvm==0.25.0.post1` (pip wheel)
- OS: macOS 15.5, arm64 (Apple Silicon, 12-core Apple Silicon Mac)
- Python: 3.13
- Backends compiled into the wheel: LLVM (host), Metal. Vulkan and
CUDA codegen absent (verified — `tvm.compile(mod, Target("vulkan"))`
raises `Cannot find global function target.build.vulkan`).

## Expected behaviour

Compiling the same Relax IR to `Target("llvm", ...)` and
`Target("metal", host="llvm")` and executing under the Relax VM with
identical inputs should produce numerically equivalent outputs within
`float32` rounding tolerance.

## Actual behaviour

LLVM produces the expected output. Metal produces output that is
uncorrelated with LLVM:

- `cosine(LLVM, Metal) = 0.755631`
- `max_abs_delta = 4.0493e+00`
- `RMS(Metal - LLVM) = 76% of RMS(LLVM)`

The bug is **deterministic** across fresh Python subprocesses — a
single `python3 repro.py` invocation reproduces byte-exact-identical
wrong Metal output every time (verified across 8 consecutive fresh
subprocesses, all producing identical SHA-256 of the output tensor).

## Reproducer

Repo: **https://github.com/zacharywhitley/tvm-metal-codegen-bug-repro**

Contents (all with published SHA-256 in the repo README):

- `minimum_repro_synthetic.onnx` (~3.4 MB, 174 nodes, opset 17) —
a 174-node HiFi-GAN-style Relax subgraph. All float initializers
are i.i.d. `N(0, 0.1)` samples with a fixed seed. Graph topology,
node names, initializer names/dtypes/shapes, and value_info are
exactly those of a HiFi-GAN-style ONNX decoder subgraph extracted
with `onnx.utils.extract_model`; only float weight values are
synthetic.
- `expected_output_synthetic.f32` (589 824 bytes) — LLVM oracle output
for the deterministic input, raw fp32 little-endian, shape
`[1, 144, 1024]`.
- `expected_output_synthetic.npy` — same, NumPy format.
- `repro.py` — self-contained driver.

To reproduce:

```
git clone https://github.com/zacharywhitley/tvm-metal-codegen-bug-repro
cd tvm-metal-codegen-bug-repro
pip install "apache-tvm==0.25.0.post1" "onnx>=1.16"
python3 repro.py
```

Op inventory (post-simplification, from `onnx.load`):

- 49× Slice
- 36× Mul
- 34× Conv (1-D)
- 19× Add
- 13× Tanh
- 13× Sigmoid
- 4× Split
- 3× Sub
- 3× Concat
- (174 total)

No `ConvTranspose` in this minimum subgraph — an important isolation
point since we previously suspected the Metal ConvTranspose lowering
(see "Isolation performed" below).

## Reproduction

```python
import numpy as np, onnx, tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

model = onnx.load("minimum_repro_synthetic.onnx")
mod = from_onnx(model, keep_params_in_input=False)

# Deterministic inputs
ch = np.arange(128, dtype=np.float32).reshape(-1, 1)
f = np.arange(1024, dtype=np.float32).reshape(1, -1)
z_p = (np.sin(2 * np.pi * (ch * 3 + f) / 128) * 0.5).reshape(1, 128, 1024).astype("float32")
y_mask = np.ones((1, 1, 1024), dtype=np.float32)

def run(target, dev):
ex = tvm.compile(mod, target)
vm = relax.VirtualMachine(ex, dev)
out = vm["main"](tvm.runtime.tensor(z_p, dev),
tvm.runtime.tensor(y_mask, dev))
dev.sync()
return (out.numpy() if hasattr(out, "numpy") else out[0].numpy()).ravel()

wf_llvm = run(tvm.target.Target({"kind":"llvm","mtriple":"arm64-apple-darwin"}),
tvm.runtime.cpu())
wf_metal = run(tvm.target.Target({"kind":"metal"},
host={"kind":"llvm","mtriple":"arm64-apple-darwin"}),
tvm.metal(0))

def cos(a, b):
return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

print("LLVM vs METAL cosine:", cos(wf_llvm, wf_metal))
# Expected: ~1.0. Observed: 0.756.
```

Recommend running LLVM and Metal in separate Python subprocesses —
see "Additional observations" for why. The attached `repro.py`
handles this automatically (default `--fresh-subprocess` mode) and
prints a correctness table.

## Isolation performed

- **Not a single-op ConvTranspose1D bug.** A hand-written Relax
module with a single `nn.conv1d_transpose` at the same shape the
parent decoder uses (`[1,128,1024] × [128,64,16] stride=8 pad=8`)
produces byte-exact LLVM ↔ Metal agreement (`cosine=1.0000000`,
`max_abs_delta=0.0e+00`).
- **Not the first three flow blocks.** The exact same op pattern
(WaveNet-like `Conv → gated (tanh × sigmoid) → Conv`) is stacked
three times in the prefix of this graph without triggering the
divergence. Only when the fourth flow block's first `in_layers.1`
Conv is added does Metal diverge from LLVM. Bisection table:

| Anchor node index | Nodes | Fresh-proc cosine | In-process cosine |
|---|---:|---:|---:|
| flow.6/post/Conv (1 flow block) | 49 | 1.000000 | 1.000000 |
| flow.4/post/Conv (2 blocks) | 102 | – | 1.000000 |
| flow.2/post/Conv (3 blocks) | 155 | – | 1.000000 |
| flow.0/pre/Conv (4 blocks, pre only) | 162 | – | 1.000000 |
| flow.0/enc/in_layers.0/Conv | 164 | – | 1.000000 |
| flow.0/enc/res_skip_layers.0/Conv | 170 | **1.000000** | **0.905** *(in-proc, run ≥ 2)* |
| flow.0/enc/Slice_1_output_0 | 171 | 1.000000 | 0.914 |
| flow.0/enc/Add_output_0 | 172 | 1.000000 | 1.000000 |
| flow.0/enc/Mul_output_0 (mask × add) | 173 | 1.000000 | 0.832 |
| **flow.0/enc/in_layers.1/Conv** *(this repro)* | **174** | **~0.80** | **~0.80** |
| flow.0/enc/res_skip_layers.1/Conv | 180 | 0.794 | 0.706 |
| flow.0/post/Conv (4 blocks) | 208 | 0.645 | 0.645 |
| /dec/conv_pre/Conv | 215 | 0.648 | 0.648 |
| /dec/ups.0/ConvTranspose | 217 | 0.046 | 0.046 |
| waveform (full 410-node decoder) | 410 | 0.05 – 0.11 | 0.01 – 0.11 |

Cosine numbers in the bisection table are from the original trained
weights. With synthetic `N(0, 0.1)` weights the divergence pattern
is qualitatively the same (correct through node 170; wrong at node
174) with slightly different magnitudes.

## Additional observations (in-process state carryover)

A related pattern surfaced during bisection that may narrow the
search space:

- On a strictly smaller subgraph (170 nodes; extracted at
`/flow/flows.0/enc/res_skip_layers.0/Conv_output_0`), a fresh
Python subprocess produces `cosine=1.0000000` on the first Metal
compile — correct.
- Running Metal a **second time** on the same graph within the same
process produces `cosine=0.905` — deterministically wrong, and
byte-exact-identical across further in-process runs.
- The 174-node graph in this reproducer produces `cosine=0.756`
(synthetic weights) or `cosine=0.802` (trained weights) on the
**first** Metal compile in a fresh process.

Reading: at least one component of the failure appears to be state
that persists across `tvm.compile(mod, Target("metal"))` invocations
within a single process (compilation cache, device kernel cache, or a
static-scoped IR mutation). A large-enough graph triggers the bug
cold; a smaller graph only triggers it once that state has been
perturbed by a prior Metal compile.

The attached `repro.py --in-process` toggles this mode for
comparison.

## Ask

- Is this a known-fixed issue in `main` / 0.26?
- Which Relax/TIR pass is the most likely culprit given the minimum
reproducer's op profile (Conv + Slice + gated activation chain)
and the in-process state-carryover observation? A pointer to
which pass to bisect against would let me narrow further.
- Is the compilation-cache / device-kernel-cache lifecycle expected
to be process-scoped and safely mutable across `tvm.compile(...)`
calls?

Happy to reduce further if pointed at a specific pass or op class,
share intermediate outputs at any bisection anchor, or run against a
`tvm` build with candidate patches.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with repro.py and the 174-node synthetic ONNX graph, running LLVM and Metal in separate processes as described. Bisect the Relax/TIR compilation path around flow.0/enc/in_layers.1/Conv at node 174, comparing the relevant passes and compilation-cache behavior against the 170-node case. Done means Metal output matches LLVM within float32 rounding tolerance across fresh and repeated compilations.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.