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

[zacharywhitley]

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

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.