From 5d8c602257c6af01a4838997566fd7b3a451f864 Mon Sep 17 00:00:00 2001
From: Vladimir Milosevic <157983820+vmilosevic@users.noreply.github.com>
Date: Thu, 16 Jan 2025 01:22:09 +0100
Subject: [PATCH] Uplift third_party/tt-metal to
 3a9637d52b003d5d5eda455cc72fbeb75a689f32 2025-01-15 (#1793)

This PR uplifts the third_party/tt-metal to the
3a9637d52b003d5d5eda455cc72fbeb75a689f32

- Use `MeshDeviceConfig` explicit constructor instead of passing in
`MeshShape` implicitly due to metal commit 2c4a641

---------

Co-authored-by: brataTT <achoudhury@tenstorrent.com>
---
 runtime/lib/common/system_desc.cpp | 6 ++++--
 runtime/lib/ttmetal/runtime.cpp    | 5 +++--
 runtime/lib/ttnn/runtime.cpp       | 5 +++--
 third_party/CMakeLists.txt         | 2 +-
 4 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/runtime/lib/common/system_desc.cpp b/runtime/lib/common/system_desc.cpp
index ad86edf47f..f11dcc09c8 100644
--- a/runtime/lib/common/system_desc.cpp
+++ b/runtime/lib/common/system_desc.cpp
@@ -262,7 +262,9 @@ std::pair<::tt::runtime::SystemDesc, DeviceIds> getCurrentSystemDesc() {
   ::tt::tt_metal::distributed::MeshShape meshShape = {1, numDevices};
   std::shared_ptr<::tt::tt_metal::distributed::MeshDevice> meshDevice =
       ::tt::tt_metal::distributed::MeshDevice::create(
-          meshShape, DEFAULT_L1_SMALL_SIZE, DEFAULT_TRACE_REGION_SIZE, 1,
+          ::tt::tt_metal::distributed::MeshDeviceConfig{.mesh_shape =
+                                                            meshShape},
+          DEFAULT_L1_SMALL_SIZE, DEFAULT_TRACE_REGION_SIZE, 1,
           ::tt::tt_metal::DispatchCoreType::WORKER);
   std::exception_ptr eptr = nullptr;
   std::unique_ptr<::tt::runtime::SystemDesc> desc;
@@ -271,7 +273,7 @@ std::pair<::tt::runtime::SystemDesc, DeviceIds> getCurrentSystemDesc() {
   } catch (...) {
     eptr = std::current_exception();
   }
-  meshDevice->close_devices();
+  meshDevice->close();
   if (eptr) {
     std::rethrow_exception(eptr);
   }
diff --git a/runtime/lib/ttmetal/runtime.cpp b/runtime/lib/ttmetal/runtime.cpp
index 0e84c17cef..f2702b52a3 100644
--- a/runtime/lib/ttmetal/runtime.cpp
+++ b/runtime/lib/ttmetal/runtime.cpp
@@ -74,7 +74,8 @@ Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs,
   size_t l1SmallSizeValue = l1SmallSize.value_or(DEFAULT_L1_SMALL_SIZE);
   std::shared_ptr<::tt::tt_metal::distributed::MeshDevice> meshDevice =
       ::tt::tt_metal::distributed::MeshDevice::create(
-          grid, l1SmallSizeValue, DEFAULT_TRACE_REGION_SIZE, numHWCQs,
+          ::tt::tt_metal::distributed::MeshDeviceConfig{.mesh_shape = grid},
+          l1SmallSizeValue, DEFAULT_TRACE_REGION_SIZE, numHWCQs,
           ::tt::tt_metal::DispatchCoreType::WORKER);
 
   return Device(std::static_pointer_cast<void>(meshDevice),
@@ -92,7 +93,7 @@ void closeDevice(Device device) {
     ::tt::tt_metal::detail::DumpDeviceProfileResults(ttmetalDevice);
   }
 #endif
-  ttmetalMeshDevice.close_devices();
+  ttmetalMeshDevice.close();
 }
 
 void deallocateBuffers(Device deviceHandle) {
diff --git a/runtime/lib/ttnn/runtime.cpp b/runtime/lib/ttnn/runtime.cpp
index 58b694482e..72378f881c 100644
--- a/runtime/lib/ttnn/runtime.cpp
+++ b/runtime/lib/ttnn/runtime.cpp
@@ -184,7 +184,8 @@ Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs,
   ::tt::tt_metal::distributed::MeshShape grid = {1, deviceIds.size()};
   size_t l1SmallSizeValue = l1SmallSize.value_or(kL1SmallSize);
   std::shared_ptr<::ttnn::MeshDevice> meshDevice = ::ttnn::MeshDevice::create(
-      grid, l1SmallSizeValue, DEFAULT_TRACE_REGION_SIZE, numHWCQs,
+      ::tt::tt_metal::distributed::MeshDeviceConfig{.mesh_shape = grid},
+      l1SmallSizeValue, DEFAULT_TRACE_REGION_SIZE, numHWCQs,
       ::tt::tt_metal::DispatchCoreType::WORKER);
 
   bool enableAsync = debug::Env::get().enableAsyncTTNN;
@@ -205,7 +206,7 @@ void closeDevice(Device device) {
   }
 #endif
 
-  ttnnMeshDevice.close_devices();
+  ttnnMeshDevice.close();
 }
 
 void deallocateBuffers(Device deviceHandle) {
diff --git a/third_party/CMakeLists.txt b/third_party/CMakeLists.txt
index b1bf234e6e..42ebb830b9 100644
--- a/third_party/CMakeLists.txt
+++ b/third_party/CMakeLists.txt
@@ -1,6 +1,6 @@
 include(ExternalProject)
 
-set(TT_METAL_VERSION "5d4c047dacf2606dd56c7b4d51d5049bf2c6846a")
+set(TT_METAL_VERSION "3a9637d52b003d5d5eda455cc72fbeb75a689f32")
 
 if ("$ENV{ARCH_NAME}" STREQUAL "grayskull")
   set(ARCH_NAME "grayskull")