zacharywhitley opened a new pull request, #20391:
URL: https://github.com/apache/tvm/pull/20391
## Summary
TVM's Metal runtime allocates every device buffer with
`MTLResourceStorageModePrivate`. On Apple Silicon at graph scale
(~11+ chained producer-consumer dispatches with certain fusion
patterns, tracked in #20157) this exposes a driver-level correctness
failure: outputs are silently wrong, `cosine(LLVM, Metal)` drops to
~0.75-0.95 on the repro.
Switching the allocation to `MTLResourceStorageModeShared` (unified
memory on UMA — same physical bytes, stronger cache-coherency
guarantees) closes the gate to numeric equivalence with LLVM.
Adds an env-flag opt-in so callers can select the storage mode
without a rebuild. Default remains `Private` for perf; this is a
correctness escape hatch, not a default change.
## Verification (macOS 15.5 / Apple M2 Max)
| Workload | `Private` (default) |
`TVM_METAL_STORAGE_MODE=shared` |
| ------------------------------ | ------------------: |
------------------------------: |
| WaveNet-style Relax subgraph | cosine = 0.7556 |
**cosine = 1.000** |
| Synthetic N=11 (gated tanh×σ) | cosine = 0.9465 |
**cosine = 1.000** |
| Synthetic N=20 | cosine = 0.7064 |
**cosine = 1.000** |
| MobileNetV2 native Metal | argmax = 915 (✗) | **argmax = 470,
5.3 ms warm** |
All match LLVM to fp32 noise (`max_delta ≤ 2e-5`).
## What this doesn't do
Not a root-cause fix. On UMA, `Private` and `Shared` refer to the same
physical memory; the difference is driver-side cache coherency. Twelve
TVM-side hypotheses were ruled out via C++ instrumentation before
converging on the storage-mode workaround: Relax memory reuse,
`KillAfterLastUse`, TIR storage rewrite, per-encoder barriers,
per-dispatch `waitUntilCompleted`, buffer zero-init, MSL codegen (MSL
is byte-identical between working and broken sizes). The bug survives
all software-side synchronization and shows up only under `Private`
mode, pointing at Metal driver behavior rather than TVM codegen.
## Suggested follow-ups
1. Benchmark `Shared` vs `Private` on representative GPU workloads to
decide whether the macOS default should flip.
2. Consider promoting from env-flag to a target-level config option
(`Target({"kind": "metal", "storage_mode": "shared"})`) once perf
is characterized.
3. File a companion Apple Metal report with the reproducer for
driver-side investigation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
