Copilot commented on code in PR #19504:
URL: https://github.com/apache/tvm/pull/19504#discussion_r3180373616
##########
src/runtime/metal/metal_device_api.mm:
##########
@@ -24,12 +24,91 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/timer.h>
+
+#include <cctype>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+
#include "metal_common.h"
namespace tvm {
namespace runtime {
namespace metal {
+/*!
+ * \brief Resolve the storage mode used for device data buffers from the
+ * TVM_METAL_STORAGE_MODE environment variable (read once, cached).
+ *
+ * Accepted values (case-insensitive):
+ * - "private" (default; backward compatible) ->
MTLResourceStorageModePrivate
+ * - "shared" ->
MTLResourceStorageModeShared
+ * - "managed" (macOS discrete GPUs only) ->
MTLResourceStorageModeManaged
+ *
+ * Why an environment variable instead of a CMake flag or a runtime API:
+ * - CMake flag would require rebuilding TVM to flip behaviour, which makes
+ * downstream packaging (mlc-llm, tvm-ffi consumers, prebuilt wheels)
+ * awkward — they would need two builds.
+ * - A runtime SetStorageMode() API has the right shape but adds a public
+ * C++/Python surface to MetalWorkspace; it is also racy if any allocation
+ * can happen before the user sets the mode, so consumers would still need
+ * to set an env var or call the API extremely early.
+ * - An env var read-once-at-init is the smallest possible behaviour change:
+ * default (unset) keeps the historical Private mode, opt-in is a one-line
+ * export. This matches how TVM exposes other low-level toggles
+ * (e.g. TVM_LOG_DEBUG, TVM_NUM_THREADS).
+ *
+ * The motivating use case is zero-copy DLPack interop with MLX (which
+ * allocates Metal buffers as Shared). Two allocators on the same MTLDevice
+ * cannot share an MTLBuffer that has different page-mapping semantics, so
+ * DLPack capsules cross-imported between TVM-NDArray and mx.array currently
+ * fail. With TVM_METAL_STORAGE_MODE=shared, the two allocators agree on the
+ * storage class and the underlying MTLBuffer can be aliased without copy.
+ *
+ * Notes:
+ * - This intentionally does NOT change the staging buffer pool
+ * (StagingBufferPool::GetOrCreate) or the GPU->CPU temp buffer
+ * (GetTempBuffer); both must remain Shared because they are host-staging
+ * buffers whose entire purpose is host visibility. Changing them would
+ * break CPU<->GPU memcpy, which is unrelated to the device-data
+ * allocation that this flag governs.
+ * - Shared-mode device buffers are valid arguments to compute kernels
+ * (setBuffer:offset:atIndex:) on every Metal device family — MLX itself
+ * proves this in production. The kernel dispatch path in metal_module.mm
+ * does not need any change.
+ */
+inline MTLResourceOptions GetMetalStorageOptions() {
+ static const MTLResourceOptions cached = []() -> MTLResourceOptions {
+ const char* raw = std::getenv("TVM_METAL_STORAGE_MODE");
+ if (raw == nullptr || raw[0] == '\0') {
+ return MTLResourceStorageModePrivate;
+ }
+ // Lowercase a small bounded copy for case-insensitive comparison.
+ std::string v(raw);
Review Comment:
The comment says this lowercases “a small bounded copy”, but `std::string
v(raw)` copies the full env-var value (which can be quite large). Either bound
the copied length (e.g., cap to a small max) or adjust the comment to avoid
claiming bounded behavior.
##########
src/runtime/metal/metal_device_api.mm:
##########
@@ -418,10 +493,24 @@ int GetWarpSize(id<MTLDevice> dev) {
result.Set("free_syncs", static_cast<int64_t>(p.free_syncs));
return result;
})
- .def("metal.ResetProfileCounters", [](int device_id) {
- auto* ws = MetalWorkspace::Global();
- ws->CastStreamOrGetDefault(nullptr, device_id)->profile.Reset();
- });
+ .def("metal.ResetProfileCounters",
+ [](int device_id) {
+ auto* ws = MetalWorkspace::Global();
+ ws->CastStreamOrGetDefault(nullptr, device_id)->profile.Reset();
+ })
+ .def("metal.GetStorageMode",
+ []() -> ffi::String {
+ // Return a stable string for parity tests; mirrors the env-var
+ // names accepted by GetMetalStorageOptions().
+ MTLResourceOptions opts = GetMetalStorageOptions();
+ // The Resource enum encodes the storage mode in the high bits
+ // (MTLResourceStorageModeShift == 4), but the lower-level enum
+ // values overlap, so we compare directly with the resource enum
+ // values to remain robust on every SDK version.
+ if (opts == MTLResourceStorageModeShared) return
ffi::String("shared");
+ if (opts == MTLResourceStorageModeManaged) return
ffi::String("managed");
Review Comment:
`metal.GetStorageMode` compares against `MTLResourceStorageModeManaged`
unconditionally. On iOS targets, the Managed storage mode is unavailable and
this reference can fail compilation even though `GetMetalStorageOptions()`
already gates Managed behind `#if !TARGET_OS_IPHONE`. Consider guarding the
Managed comparison/return in this FFI helper as well (or mapping Managed to
"shared" on iOS).
--
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]
