From e2aff3cc0eedc673cdfea114d2621382d1020a7c Mon Sep 17 00:00:00 2001 From: Kyle Mabee Date: Thu, 29 Aug 2024 05:16:51 +0000 Subject: [PATCH] Temporary Workaround for tt-metal Segfaults during teardown (#408) - A hack, in createBufferFromTensorRef(), remove when proper bug fix is made in tt-metal and propagates here. - Cannot remove UNSUPPORTED on TTMetal/to_layout.mlir yet because CI mixes ttnn/ttmetal binaries fine, but closing device fails, need some kind of ttrt fix --- runtime/include/tt/runtime/detail/ttmetal.h | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/runtime/include/tt/runtime/detail/ttmetal.h b/runtime/include/tt/runtime/detail/ttmetal.h index f7838b39b0..e69fcad7b0 100644 --- a/runtime/include/tt/runtime/detail/ttmetal.h +++ b/runtime/include/tt/runtime/detail/ttmetal.h @@ -159,6 +159,19 @@ createBufferFromTensorRef(::tt::tt_metal::Device *device, std::shared_ptr<::tt::tt_metal::Buffer> buffer = ::tt::tt_metal::CreateBuffer(shardedBufferConfig); assert(tensorRef->address()); + + // Issue #408: Temporary Hack, remove when fix available. + // Update tt-metal BUFFER_MAP with updated address and remove + // entry for original alloc'd address. + auto &buffer_map = tt::tt_metal::detail::BUFFER_MAP; + auto map_copy = buffer_map.value(); + auto old_key = std::make_tuple(device->id(), buffer->address()); + if (auto it = map_copy.find(old_key); it != map_copy.end()) { + auto new_key = std::make_tuple(device->id(), tensorRef->address()); + buffer_map.insert(new_key, it->second); + buffer_map.erase(old_key); + } + buffer->set_address(tensorRef->address()); return buffer; }