Skip to content

Commit

Permalink
#15290: Fix UAF issue with Singleton destruction order
Browse files Browse the repository at this point in the history
- Decouple MeshDevice destruction from SystemMesh destruction.
- Refactor SystemMesh to encapsulate internals in pImpl class.
- Added a new test case for SystemMesh tear down with static mesh in the distributed tests.
  • Loading branch information
cfjchu committed Dec 4, 2024
1 parent 77e056e commit 69661ca
Show file tree
Hide file tree
Showing 4 changed files with 254 additions and 169 deletions.
6 changes: 5 additions & 1 deletion tests/ttnn/distributed/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
add_executable(test_distributed test_distributed.cpp)
add_executable(
test_distributed
test_distributed.cpp
test_distributed_atexit.cpp
)

# Set up properties for the target
setup_ttnn_test_target(test_distributed)
Expand Down
32 changes: 32 additions & 0 deletions tests/ttnn/distributed/test_distributed_atexit.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
//
// SPDX-License-Identifier: Apache-2.0

#include <gtest/gtest.h>

#include <cstddef>
#include <ttnn/core.hpp>
#include <ttnn/distributed/api.hpp>

namespace ttnn::distributed::test {

class DistributedTest : public ::testing::Test {
protected:
void SetUp() override {}
void TearDown() override {}
};

// Test that the mesh device is properly cleaned up when the program exits
static std::shared_ptr<ttnn::MeshDevice> mesh;

TEST_F(DistributedTest, TestSystemMeshTearDownWithoutClose_StaticMesh) {
auto& sys = tt::tt_metal::distributed::SystemMesh::instance();
mesh = ttnn::distributed::open_mesh_device(
{2, 4}, DEFAULT_L1_SMALL_SIZE, DEFAULT_TRACE_REGION_SIZE, 1, tt::tt_metal::DispatchCoreType::WORKER);

auto [rows, cols] = sys.get_shape();
EXPECT_GT(rows, 0);
EXPECT_GT(cols, 0);
}

} // namespace ttnn::distributed::test
Loading

0 comments on commit 69661ca

Please sign in to comment.