From 458be49c8c82f646fe444b769948966238cc0062 Mon Sep 17 00:00:00 2001 From: Egor Churaev Date: Mon, 31 May 2021 11:01:58 +0300 Subject: [PATCH] [METAL] Fix the rest memory leaks in Metal runtime When we throw exception from autoreleasepool, then the resources won't be released in proper way. In the documentation we can see that "When the block is exited with an exception, the pool is not drained.". Link on the documentation: https://clang.llvm.org/docs/AutomaticReferenceCounting.html#autoreleasepool Implemented a wrapper which handles all exceptions in autoreleasepool block and throw them after this block. --- src/runtime/metal/metal_common.h | 28 +++++++++++++++++++++ src/runtime/metal/metal_device_api.mm | 35 +++++++++++++-------------- src/runtime/metal/metal_module.mm | 26 ++++++++++++-------- 3 files changed, 61 insertions(+), 28 deletions(-) diff --git a/src/runtime/metal/metal_common.h b/src/runtime/metal/metal_common.h index 9ebe04efbe4c..1fd124f1e25c 100644 --- a/src/runtime/metal/metal_common.h +++ b/src/runtime/metal/metal_common.h @@ -42,9 +42,37 @@ #include "../workspace_pool.h" +#define AUTORELEASEPOOL tvm::runtime::metal::AutoReleasePoolWrapper::GetInstance() << [&]() + namespace tvm { namespace runtime { namespace metal { +class AutoReleasePoolWrapper { + public: + static AutoReleasePoolWrapper& GetInstance() { + static AutoReleasePoolWrapper instance; + return instance; + } + template + void operator<<(const T& f) { + std::exception_ptr eptr; + @autoreleasepool { + try { + f(); + } catch (...) { + eptr = std::current_exception(); + } + } + if (eptr) std::rethrow_exception(eptr); + } + + private: + AutoReleasePoolWrapper() = default; + ~AutoReleasePoolWrapper() = default; + AutoReleasePoolWrapper(const AutoReleasePoolWrapper&) = delete; + AutoReleasePoolWrapper& operator=(const AutoReleasePoolWrapper&) = delete; +}; + /*! * \brief Structure for error handling in queues */ diff --git a/src/runtime/metal/metal_device_api.mm b/src/runtime/metal/metal_device_api.mm index 193e4647733a..0cbbb0ec4c1f 100644 --- a/src/runtime/metal/metal_device_api.mm +++ b/src/runtime/metal/metal_device_api.mm @@ -30,16 +30,14 @@ namespace metal { MetalWorkspace* MetalWorkspace::Global() { - @autoreleasepool { - // NOTE: explicitly use new to avoid exit-time destruction of global state - // Global state will be recycled by OS as the process exits. - static MetalWorkspace* inst = new MetalWorkspace(); - return inst; - } + // NOTE: explicitly use new to avoid exit-time destruction of global state + // Global state will be recycled by OS as the process exits. + static MetalWorkspace* inst = new MetalWorkspace(); + return inst; } void MetalWorkspace::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) { - @autoreleasepool { + AUTORELEASEPOOL { this->Init(); size_t index = static_cast(dev.device_id); if (kind == kExist) { @@ -80,7 +78,7 @@ case kDriverVersion: return; } - } + }; } static const char* kDummyKernel = R"A0B0( @@ -161,7 +159,8 @@ int GetWarpSize(id dev) { void* MetalWorkspace::AllocDataSpace(Device device, size_t nbytes, size_t alignment, DLDataType type_hint) { - @autoreleasepool { + id buf; + AUTORELEASEPOOL { this->Init(); id dev = GetDevice(device); // GPU memory only @@ -173,20 +172,20 @@ int GetWarpSize(id dev) { storage_mode = MTLResourceStorageModeManaged; #endif */ - id buf = [dev newBufferWithLength:nbytes options:storage_mode]; + buf = [dev newBufferWithLength:nbytes options:storage_mode]; ICHECK(buf != nil); - return (void*)(buf); - } + }; + return (void*)(buf); } void MetalWorkspace::FreeDataSpace(Device dev, void* ptr) { - @autoreleasepool { + AUTORELEASEPOOL { // MTLBuffer PurgeableState should be set to empty before manual // release in order to prevent memory leak [(id)ptr setPurgeableState:MTLPurgeableStateEmpty]; // release the ptr. CFRelease(ptr); - } + }; } Stream* GetStream(TVMStreamHandle stream, int device_id) { @@ -199,7 +198,7 @@ int GetWarpSize(id dev) { void MetalWorkspace::CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset, size_t size, Device dev_from, Device dev_to, DLDataType type_hint, TVMStreamHandle stream) { - @autoreleasepool { + AUTORELEASEPOOL { this->Init(); Device dev = dev_from; Stream* s = GetStream(stream, dev.device_id); @@ -261,7 +260,7 @@ int GetWarpSize(id dev) { LOG(FATAL) << "Expect copy from/to Metal or between Metal" << ", from=" << from_dev_type << ", to=" << to_dev_type; } - } + }; } TVMStreamHandle MetalWorkspace::CreateStream(Device dev) { @@ -276,7 +275,7 @@ int GetWarpSize(id dev) { } void MetalWorkspace::StreamSync(Device dev, TVMStreamHandle stream) { - @autoreleasepool { + AUTORELEASEPOOL { Stream* s = GetStream(stream, dev.device_id); // commit an empty command buffer and wait until it completes. id cb = s->GetCommandBuffer(); @@ -285,7 +284,7 @@ int GetWarpSize(id dev) { if (s->HasErrorHappened()) { LOG(FATAL) << "Error! Some problems on GPU happaned!"; } - } + }; } void MetalWorkspace::SetStream(Device dev, TVMStreamHandle stream) { diff --git a/src/runtime/metal/metal_module.mm b/src/runtime/metal/metal_module.mm index 2920c60449d1..88501880557e 100644 --- a/src/runtime/metal/metal_module.mm +++ b/src/runtime/metal/metal_module.mm @@ -193,7 +193,7 @@ void Init(MetalModuleNode* m, ObjectPtr sptr, const std::string& func_na } // invoke the function with void arguments void operator()(TVMArgs args, TVMRetValue* rv, const ArgUnion64* pack_args) const { - @autoreleasepool { + AUTORELEASEPOOL { metal::MetalThreadEntry* t = metal::MetalThreadEntry::ThreadLocal(); int device_id = t->device.device_id; auto stream = static_cast(t->stream[device_id]); @@ -223,7 +223,7 @@ void operator()(TVMArgs args, TVMRetValue* rv, const ArgUnion64* pack_args) cons [encoder dispatchThreadgroups:dimGrid threadsPerThreadgroup:dimBlock]; [encoder endEncoding]; [cb commit]; - } + }; } private: @@ -248,27 +248,33 @@ void operator()(TVMArgs args, TVMRetValue* rv, const ArgUnion64* pack_args) cons PackedFunc MetalModuleNode::GetFunction(const std::string& name, const ObjectPtr& sptr_to_self) { - @autoreleasepool { + PackedFunc pf; + AUTORELEASEPOOL { ICHECK_EQ(sptr_to_self.get(), this); ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main"; auto it = fmap_.find(name); - if (it == fmap_.end()) return PackedFunc(); + if (it == fmap_.end()) { + pf = PackedFunc(); + return; + } const FunctionInfo& info = it->second; MetalWrappedFunc f; size_t num_buffer_args = NumBufferArgs(info.arg_types); f.Init(this, sptr_to_self, name, num_buffer_args, info.arg_types.size() - num_buffer_args, info.thread_axis_tags); - return PackFuncNonBufferArg(f, info.arg_types); - } + pf = PackFuncNonBufferArg(f, info.arg_types); + }; + return pf; } Module MetalModuleCreate(std::string data, std::string fmt, std::unordered_map fmap, std::string source) { - @autoreleasepool { + ObjectPtr n; + AUTORELEASEPOOL { metal::MetalWorkspace::Global()->Init(); - auto n = make_object(data, fmt, fmap, source); - return Module(n); - } + n = make_object(data, fmt, fmap, source); + }; + return Module(n); } // Load module from module.