From c9db3d0a2b03bbbe9e4e0e421e0b848d540f94ca Mon Sep 17 00:00:00 2001 From: Egor Churaev Date: Fri, 4 Jun 2021 22:49:13 +0300 Subject: [PATCH] [METAL] Fix the rest memory leaks in Metal runtime (#8175) * [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. * Apply comments * Add documentation comments to wrapper and macro --- src/runtime/metal/metal_common.h | 56 +++++++++++++++++++++++++++ src/runtime/metal/metal_device_api.mm | 40 ++++++++++--------- src/runtime/metal/metal_module.mm | 26 ++++++++----- 3 files changed, 94 insertions(+), 28 deletions(-) diff --git a/src/runtime/metal/metal_common.h b/src/runtime/metal/metal_common.h index 9ebe04efbe4c..7d2ef0c9367b 100644 --- a/src/runtime/metal/metal_common.h +++ b/src/runtime/metal/metal_common.h @@ -42,9 +42,65 @@ #include "../workspace_pool.h" +/* Macro for convenience in using AutoReleasePoolWrapper. + * With this macro we can add AutoReleasePoolWrapper to our ObjC code in more + * native way. + * + * For example, this is ObjC code with autoreleasepool: + * @autoreleasepool { + * // Some code + * } + * + * To avoid possible memory leaks when an exception will be generated, we + * should update this code: + * AUTORELEASEPOOL { // Replace @autoreleasepool -> AUTORELEASEPOOL + * // Some code + * }; // Add semicolon after close bracket + * + * In macro AUTORELEASEPOOL we get the instance of AutoReleasePoolWrapper and + * put a lambda function with code from autoreleasepool to the insertion + * operator of AutoReleasePoolWrapper class. + * + * Note: If you want to return a value from the autoreleasepool, you should + * declare the variable with result before AUTORELEASEPOOL macro. This variable + * will be captured by reference and you can use it in the code in autorelease + * pool. But you should write return statement after AUTORELEASEPOOL macro. + */ +#define AUTORELEASEPOOL tvm::runtime::metal::AutoReleasePoolWrapper::GetInstance() << [&]() + namespace tvm { namespace runtime { namespace metal { +/*! + * \brief Wrapper on autoreleasepool with exception handling + * + * \note In case when the exception was thrown from the autoreleasepool, the + * allocated resources won't be released in proper way. So, we handle exception + * in autoreleasepool and after the autoreleasepool we rethrow this exception. + */ +class AutoReleasePoolWrapper { + public: + static AutoReleasePoolWrapper& GetInstance(); + 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..1c5666dfc17f 100644 --- a/src/runtime/metal/metal_device_api.mm +++ b/src/runtime/metal/metal_device_api.mm @@ -29,17 +29,20 @@ namespace runtime { namespace metal { +AutoReleasePoolWrapper& AutoReleasePoolWrapper::GetInstance() { + static AutoReleasePoolWrapper instance; + return instance; +} + 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 +83,7 @@ case kDriverVersion: return; } - } + }; } static const char* kDummyKernel = R"A0B0( @@ -161,7 +164,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 +177,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 +203,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 +265,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 +280,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 +289,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.