Add async ttnn parameter in open device

.github/workflows/build-and-test.yml

@@ -164,12 +164,14 @@ jobs:
       fail-fast: false
       matrix:
         build: [
-          {runs-on: n150, enable_perf: OFF, name: "run", ttrt_flags: "--non-zero"},
-          {runs-on: n150, enable_perf: ON, name: "perf"},
-          {runs-on: n300, enable_perf: OFF, name: "run", ttrt_flags: "--non-zero"},
-          {runs-on: n300, enable_perf: ON, name: "perf"},
+          {runs-on: n150, enable_perf: OFF, enable_async: OFF, name: "run", ttrt_flags: "--non-zero"},
+          {runs-on: n150, enable_perf: OFF, enable_async: ON, name: "run", ttrt_flags: "--non-zero --enable-async-ttnn"},
+          {runs-on: n150, enable_perf: ON, enable_async: OFF, name: "perf"},
+          {runs-on: n300, enable_perf: OFF, enable_async: OFF, name: "run", ttrt_flags: "--non-zero"},
+          {runs-on: n300, enable_perf: OFF, enable_async: ON, name: "run", ttrt_flags: "--non-zero --enable-async-ttnn"},
+          {runs-on: n300, enable_perf: ON, enable_async: OFF, name: "perf"},
         ]
-    name: "run-tests (${{ matrix.build.runs-on }}, ${{ matrix.build.enable_perf }}, ${{ matrix.build.name }})"
+    name: "run-tests (${{ matrix.build.runs-on }}, ${{ matrix.build.enable_perf }}, ${{ matrix.build.enable_async }}, ${{ matrix.build.name }})"
 
     runs-on:
       - in-service
@@ -194,7 +196,7 @@ jobs:
       id: strings
       shell: bash
       env:
-        job-name: "run-tests (${{ matrix.build.runs-on }}, ${{ matrix.build.enable_perf }}, ${{ matrix.build.name }})"
+        job-name: "run-tests (${{ matrix.build.runs-on }}, ${{ matrix.build.enable_perf }}, ${{ matrix.build.enable_async }}, ${{ matrix.build.name }})"
       run: |
         echo "work-dir=$(pwd)" >> "$GITHUB_OUTPUT"
         echo "build-output-dir=$(pwd)/build" >> "$GITHUB_OUTPUT"
@@ -260,11 +262,18 @@ jobs:
 
     - name: Run functional tests
       shell: bash
-      if: matrix.build.enable_perf == 'OFF'
+      if: matrix.build.enable_perf == 'OFF' && matrix.build.enable_async == 'OFF'
       run: |
         source env/activate
         ttrt ${{ matrix.build.name }} ${{ matrix.build.ttrt_flags }} ${{ steps.strings.outputs.build-output-dir }}/test/ttmlir/Silicon
 
+    - name: Run async TTNN tests
+      shell: bash
+      if: matrix.build.enable_async == 'ON'
+      run: |
+        source env/activate
+        ttrt ${{ matrix.build.name }} ${{ matrix.build.ttrt_flags }} ${{ steps.strings.outputs.build-output-dir }}/test/ttmlir/Silicon/TTNN
+
     - name: Run perf tests
       shell: bash
       if: matrix.build.enable_perf == 'ON'
@@ -277,7 +286,7 @@ jobs:
       if: always()
       uses: actions/upload-artifact@v4
       with:
-        name: ${{ matrix.build.runs-on }}_${{ matrix.build.name }}_results.json
+        name: ${{ matrix.build.runs-on }}_${{ matrix.build.name }}_async_${{ matrix.build.enable_async }}_results.json
         path: ${{ matrix.build.name }}_results.json
 
     - name: Upload Test Report xml

runtime/include/tt/runtime/detail/debug.h

@@ -19,28 +19,25 @@ struct Env {
 #else
   constexpr static Env
 #endif
-  get(bool loadKernelsFromDisk = false, bool enableAsyncTTNN = false)
+  get(bool loadKernelsFromDisk = false)
 #if defined(TT_RUNTIME_DEBUG) && TT_RUNTIME_DEBUG == 1
       ;
 #else
   {
-    return Env(false, false);
+    return Env(false);
   }
 #endif
 
   bool loadKernelsFromDisk;
-  bool enableAsyncTTNN;
 
 private:
-  constexpr Env(bool loadKernelsFromDisk, bool enableAsyncTTNN)
-      : loadKernelsFromDisk(loadKernelsFromDisk),
-        enableAsyncTTNN(enableAsyncTTNN) {}
+  constexpr Env(bool loadKernelsFromDisk)
+      : loadKernelsFromDisk(loadKernelsFromDisk) {}
 };
 
 inline std::ostream &operator<<(std::ostream &os, Env const &env) {
   os << "debug::Env{\n"
-     << "\t" << "loadKernelsFromDisk: " << env.loadKernelsFromDisk << ",\n"
-     << "\t" << "enableAsyncTTNN: " << env.enableAsyncTTNN << "\n"
+     << "\t" << "loadKernelsFromDisk: " << env.loadKernelsFromDisk << "\n"
      << "}";
   return os;
 }

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

@@ -35,8 +35,10 @@ tt::target::DataType getTensorDataType(Tensor tensor);
 
 size_t getNumAvailableDevices();
 
-Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs = 1,
-                  std::optional<size_t> l1SmallSize = std::nullopt);
+Device
+openDevice(DeviceIds const &deviceIds, size_t numHWCQs = 1,
+           std::optional<size_t> l1SmallSize = std::nullopt,
+           [[maybe_unused]] std::optional<bool> enableAsyncTTNN = std::nullopt);
 
 void closeDevice(Device device);
 

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

@@ -87,7 +87,8 @@ tt::target::DataType getTensorDataType(Tensor tensor);
 size_t getNumAvailableDevices();
 
 Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs = 1,
-                  std::optional<size_t> l1SmallSize = std::nullopt);
+                  std::optional<size_t> l1SmallSize = std::nullopt,
+                  std::optional<bool> enableAsyncTTNN = std::nullopt);
 
 void closeDevice(Device device);
 

runtime/include/tt/runtime/runtime.h

@@ -72,7 +72,8 @@ tt::target::DataType getTensorDataType(Tensor tensor);
 size_t getNumAvailableDevices();
 
 Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs = 1,
-                  std::optional<size_t> l1SmallSize = std::nullopt);
+                  std::optional<size_t> l1SmallSize = std::nullopt,
+                  std::optional<bool> enableAsyncTTNN = std::nullopt);
 
 void closeDevice(Device device);
 

runtime/lib/common/debug.cpp

@@ -8,8 +8,8 @@
 
 namespace tt::runtime::debug {
 
-Env const &Env::get(bool loadKernelsFromDisk, bool enableAsyncTTNN) {
-  static Env config(loadKernelsFromDisk, enableAsyncTTNN);
+Env const &Env::get(bool loadKernelsFromDisk) {
+  static Env config(loadKernelsFromDisk);
   return config;
 }
 

runtime/lib/runtime.cpp

@@ -217,16 +217,19 @@ size_t getNumAvailableDevices() {
 }
 
 Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs,
-                  std::optional<size_t> l1SmallSize) {
+                  std::optional<size_t> l1SmallSize,
+                  std::optional<bool> enableAsyncTTNN) {
 #if defined(TT_RUNTIME_ENABLE_TTNN)
   if (getCurrentRuntime() == DeviceRuntime::TTNN) {
-    return ::tt::runtime::ttnn::openDevice(deviceIds, numHWCQs, l1SmallSize);
+    return ::tt::runtime::ttnn::openDevice(deviceIds, numHWCQs, l1SmallSize,
+                                           enableAsyncTTNN);
   }
 #endif
 
 #if defined(TT_RUNTIME_ENABLE_TTMETAL)
   if (getCurrentRuntime() == DeviceRuntime::TTMetal) {
-    return ::tt::runtime::ttmetal::openDevice(deviceIds, numHWCQs, l1SmallSize);
+    return ::tt::runtime::ttmetal::openDevice(deviceIds, numHWCQs, l1SmallSize,
+                                              enableAsyncTTNN);
   }
 #endif
   LOG_FATAL("runtime is not enabled");

runtime/lib/ttmetal/runtime.cpp

@@ -67,7 +67,8 @@ size_t getNumAvailableDevices() {
 }
 
 Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs,
-                  std::optional<size_t> l1SmallSize) {
+                  std::optional<size_t> l1SmallSize,
+                  [[maybe_unused]] std::optional<bool> enableAsyncTTNN) {
   LOG_ASSERT(deviceIds.size(), "No devices specified");
 
   ::tt::tt_metal::distributed::MeshShape grid = {1, deviceIds.size()};

runtime/lib/ttnn/runtime.cpp

@@ -179,17 +179,18 @@ size_t getNumAvailableDevices() {
 }
 
 Device openDevice(DeviceIds const &deviceIds, size_t numHWCQs,
-                  std::optional<size_t> l1SmallSize) {
+                  std::optional<size_t> l1SmallSize,
+                  std::optional<bool> enableAsyncTTNN) {
   LOG_ASSERT(deviceIds.size(), "No devices specified");
   ::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::DispatchCoreType::WORKER);
 
-  bool enableAsync = debug::Env::get().enableAsyncTTNN;
+  bool enableAsyncValue = enableAsyncTTNN.value_or(false);
   for (::ttnn::IDevice *device : meshDevice->get_devices()) {
-    device->enable_async(enableAsync);
+    device->enable_async(enableAsyncValue);
   }
 
   return Device(std::static_pointer_cast<void>(meshDevice),

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

@@ -378,9 +378,7 @@ def _execute(binaries):
                 self.logging.warning(f"no binaries found to run - returning early")
                 return
 
-            debug_env = ttrt.runtime.DebugEnv.get(
-                self["--load-kernels-from-disk"], self["--enable-async-ttnn"]
-            )
+            debug_env = ttrt.runtime.DebugEnv.get(self["--load-kernels-from-disk"])
             self.logging.debug(f"setting tt runtime debug env={debug_env}")
             workaround_env = ttrt.runtime.WorkaroundEnv.get(
                 not self["--disable-maxpool2d-preshard"],
@@ -394,7 +392,9 @@ def _execute(binaries):
             ttrt.runtime.set_compatible_runtime(binaries[0].fbb)
             current_runtime = ttrt.runtime.get_current_runtime()
             self.logging.debug(f"opening devices={self.query.device_ids}")
-            device = ttrt.runtime.open_device(self.query.device_ids)
+            device = ttrt.runtime.open_device(
+                self.query.device_ids, enable_async_ttnn=self["--enable-async-ttnn"]
+            )
 
             callback_runtime_config = CallbackRuntimeConfig(
                 device,

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

@@ -100,6 +100,7 @@ PYBIND11_MODULE(_C, m) {
   m.def("open_device", &tt::runtime::openDevice, py::arg("device_ids"),
         py::arg("num_hw_cqs") = size_t{1},
         py::arg("l1_small_size") = py::none(),
+        py::arg("enable_async_ttnn") = py::none(),
         "Open a mesh of devices for execution");
   m.def("close_device", &tt::runtime::closeDevice, "Close a mesh device");
   m.def("to_host", &tt::runtime::toHost, py::arg("tensor"),