From abfeb923d78c13bfdd9815a95919e0ce0aac7c6a Mon Sep 17 00:00:00 2001 From: Zachary Whitley Date: Sat, 19 Sep 2026 17:12:34 -0400 Subject: [PATCH] [Metal] Add TVM_METAL_STORAGE_MODE opt-in for correctness at scale (#20157) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 around 0.75-0.95 on the repro. Switching 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. Verified fixes on macOS 15.5 / Apple M2 Max: - WaveNet-style Relax subgraph: cosine 0.7556 -> 1.000 - Synthetic N=11 gated activations: cosine 0.9465 -> 1.000 - Synthetic N=20 gated activations: cosine 0.7064 -> 1.000 - MobileNetV2 Metal: argmax 915 (buggy) -> 470 (gold), 5.3 ms warm inference Twelve TVM-side hypotheses ruled out via prior instrumentation (Relax memory reuse, KillAfterLastUse, TIR storage rewrite, per-encoder barriers, per-dispatch waitUntilCompleted, buffer zero-init, Metal MSL codegen — MSL is byte-identical between working and broken sizes). The bug survives all software-side synchronization and shows up only with Private mode, pointing at Metal driver behavior rather than TVM codegen. Follow-ups: - Benchmark Shared vs Private on typical GPU workloads to decide whether the macOS default should flip. - Consider promoting from env-flag to a target-level config option ({"kind": "metal", "storage_mode": "shared"}) once perf is characterized. --- src/backend/metal/runtime/metal_device_api.mm | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/backend/metal/runtime/metal_device_api.mm b/src/backend/metal/runtime/metal_device_api.mm index 2a44d98f109f..b8026fb84caa 100644 --- a/src/backend/metal/runtime/metal_device_api.mm +++ b/src/backend/metal/runtime/metal_device_api.mm @@ -20,6 +20,9 @@ /*! * \file metal_device_api.mm */ +#include +#include + #include #include #include @@ -184,15 +187,14 @@ int GetWarpSize(id dev) { id buf; AUTORELEASEPOOL { id dev = GetDevice(device); - // GPU memory only + // GPU memory only. Callers can override via TVM_METAL_STORAGE_MODE + // ("shared" or "managed") when Private-mode exposes a driver-level + // correctness issue on their workload — see apache/tvm#20157. MTLResourceOptions storage_mode = MTLResourceStorageModePrivate; - /* - #if TARGET_OS_IPHONE - storage_mode = MTLResourceStorageModeShared; - #else - storage_mode = MTLResourceStorageModeManaged; - #endif - */ + if (const char* env = std::getenv("TVM_METAL_STORAGE_MODE")) { + if (std::strcmp(env, "shared") == 0) storage_mode = MTLResourceStorageModeShared; + if (std::strcmp(env, "managed") == 0) storage_mode = MTLResourceStorageModeManaged; + } buf = [dev newBufferWithLength:nbytes options:storage_mode]; TVM_FFI_ICHECK(buf != nil); };