ttrt deallocate_buffers (#588)

Plumb through deallocate_buffers API to clear metal allocator between
programs.  Previously we were running OOM because of leaked buffers.
Leaked buffers should largely be cleaned up after the allocate=false
support is in #408.

It might still be needed even after #408 lands, because metal internally
uses its allocator for some fast dispatch things, but we can revisit
what policy to adopt.

runtime/include/tt/runtime/detail/ttmetal.h

@@ -68,6 +68,8 @@ Device openDevice(std::vector<int> const &deviceIds = {0},
 
 void closeDevice(Device device);
 
+void deallocateBuffers(Device device);
+
 Event submit(Device device, Binary executable, std::uint32_t programIndex,
              std::vector<Tensor> const &inputs,
              std::vector<Tensor> const &outputs);

runtime/include/tt/runtime/detail/ttnn.h

@@ -83,6 +83,8 @@ Device openDevice(std::vector<int> const &deviceIds = {0},
 
 void closeDevice(Device device);
 
+void deallocateBuffers(Device device);
+
 Event submit(Device device, Binary executable, std::uint32_t programIndex,
              std::vector<Tensor> const &inputs,
              std::vector<Tensor> const &outputs);

runtime/include/tt/runtime/runtime.h

@@ -17,6 +17,10 @@ namespace system_desc {
 std::pair<SystemDesc, DeviceIds> getCurrentSystemDesc();
 } // namespace system_desc
 
+namespace detail {
+void deallocateBuffers(Device device);
+}
+
 DeviceRuntime getCurrentRuntime();
 
 std::vector<DeviceRuntime> getAvailableRuntimes();

runtime/lib/runtime.cpp

@@ -26,6 +26,21 @@ DeviceRuntime globalCurrentRuntime = DeviceRuntime::TTMetal;
 DeviceRuntime globalCurrentRuntime = DeviceRuntime::Disabled;
 #endif
 
+void deallocateBuffers(Device device) {
+#if defined(TT_RUNTIME_ENABLE_TTNN)
+  if (getCurrentRuntime() == DeviceRuntime::TTNN) {
+    return ::tt::runtime::ttnn::deallocateBuffers(device);
+  }
+#endif
+
+#if defined(TT_RUNTIME_ENABLE_TTMETAL)
+  if (getCurrentRuntime() == DeviceRuntime::TTMetal) {
+    return ::tt::runtime::ttmetal::deallocateBuffers(device);
+  }
+#endif
+  throw std::runtime_error("runtime is not enabled");
+}
+
 } // namespace detail
 
 DeviceRuntime getCurrentRuntime() {

runtime/lib/ttmetal/runtime.cpp

@@ -78,6 +78,13 @@ void closeDevice(Device device) {
   }
 }
 
+void deallocateBuffers(Device deviceHandle) {
+  DeviceMesh &deviceMesh = deviceHandle.as<DeviceMesh>(DeviceRuntime::TTMetal);
+  for (::tt::tt_metal::Device *device : deviceMesh) {
+    device->deallocate_buffers();
+  }
+}
+
 static std::pair<std::shared_ptr<::tt::tt_metal::Buffer>,
                  std::shared_ptr<::tt::tt_metal::Event>>
 prepareInput(::tt::tt_metal::Device *device, MetalTensor const &metalTensor,

runtime/lib/ttnn/runtime.cpp

@@ -69,6 +69,10 @@ void closeDevice(Device device) {
   ::ttnn::close_device(ttnn_device);
 }
 
+void deallocateBuffers(Device deviceHandle) {
+  deviceHandle.as<::ttnn::Device>(DeviceRuntime::TTNN).deallocate_buffers();
+}
+
 static ::tt::target::ttnn::TTNNBinary const *getBinary(Flatbuffer binary) {
   bool isTTNN = ::tt::target::ttnn::SizePrefixedTTNNBinaryBufferHasIdentifier(
       binary.handle.get());

runtime/tools/python/ttrt/common/api.py

@@ -900,6 +900,8 @@ def _execute(binaries):
                             )
                             for tensor in program.output_tensors:
                                 self.logging.debug(f"{tensor}\n")
+
+                            device.deallocate_buffers()
                 finally:
                     ttrt.runtime.close_device(device)
 

runtime/tools/python/ttrt/runtime/module.cpp

@@ -13,7 +13,8 @@ PYBIND11_MODULE(_C, m) {
   m.doc() = "ttrt.runtime python extension for interacting with the "
             "Tenstorrent devies";
 
-  py::class_<tt::runtime::Device>(m, "Device");
+  py::class_<tt::runtime::Device>(m, "Device")
+      .def("deallocate_buffers", &tt::runtime::detail::deallocateBuffers);
   py::class_<tt::runtime::Event>(m, "Event");
   py::class_<tt::runtime::Tensor>(m, "Tensor");
   py::enum_<::tt::target::DataType>(m, "DataType")