From bae493e5e497784781380e11a1219f875b9b899c Mon Sep 17 00:00:00 2001 From: Joseph Chu Date: Sat, 22 Feb 2025 00:29:48 +0000 Subject: [PATCH] Merge disable async on SD from Stas --- ttnn/cpp/pybind11/device.cpp | 67 +++++++++++++++++++++------------ ttnn/cpp/ttnn/run_operation.cpp | 9 ++--- ttnn/cpp/ttnn/run_operation.hpp | 9 ++--- 3 files changed, 49 insertions(+), 36 deletions(-) diff --git a/ttnn/cpp/pybind11/device.cpp b/ttnn/cpp/pybind11/device.cpp index 92f93b1570c..074f851efde 100644 --- a/ttnn/cpp/pybind11/device.cpp +++ b/ttnn/cpp/pybind11/device.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include #include #include #include @@ -176,7 +178,24 @@ void device_module(py::module& m_device) { "num_program_cache_entries", &IDevice::num_program_cache_entries, "Number of entries in the program cache for this device") - .def("enable_async", &IDevice::enable_async) + .def( + "enable_async", + [](IDevice* device, bool enable) { + if (!enable) { + device->enable_async(false); + return; + } + + bool single_device = true; + if (auto mesh_device = dynamic_cast(device)) { + single_device = mesh_device->get_devices().size() == 1; + } + if (single_device) { + tt::log_warning("Async mode is disabled for single device, ignoring enable_async call"); + } else { + device->enable_async(true); + } + }) .def( "create_sub_device_manager", [](IDevice* device, @@ -260,13 +279,13 @@ void device_module(py::module& m_device) { Args: sub_device_manager_id (SubDeviceManagerId): The ID of the sub-device manager to remove. )doc") - .def( - "set_sub_device_stall_group", - [](IDevice* device, const std::vector& sub_device_ids) { - device->push_work([device, sub_device_ids] { device->set_sub_device_stall_group(sub_device_ids); }); - }, - py::arg("sub_device_ids"), - R"doc( + .def( + "set_sub_device_stall_group", + [](IDevice* device, const std::vector& sub_device_ids) { + device->push_work([device, sub_device_ids] { device->set_sub_device_stall_group(sub_device_ids); }); + }, + py::arg("sub_device_ids"), + R"doc( Set the SubDevice IDs that will be stalled on by default for Fast Dispatch commands such as reading, writing, synchronizing. Stalling here refers to the Fast Dispatch cores waiting for programs to complete execution on the specified SubDevices before proceeding with the specified instruction. The default SubDevice IDs to stall on are set to all SubDevice IDs, and whenever a new SubDevice Manager is loaded. @@ -274,25 +293,25 @@ void device_module(py::module& m_device) { Args: sub_device_ids (List[SubDeviceId]): The IDs of the SubDevices to stall on. )doc") - .def( - "reset_sub_device_stall_group", - [](IDevice* device) { device->push_work([device] { device->reset_sub_device_stall_group(); }); }, - R"doc( + .def( + "reset_sub_device_stall_group", + [](IDevice* device) { device->push_work([device] { device->reset_sub_device_stall_group(); }); }, + R"doc( Resets the sub_device_ids that will be stalled on by default for Fast Dispatch commands such as reading, writing, synchronizing back to all SubDevice IDs. )doc") - .def( - "sfpu_eps", - [](IDevice* device) { return tt::tt_metal::experimental::hal::get_eps(); }, - R"doc(Returns machine epsilon value for current architecture.)doc") - .def( - "sfpu_nan", - [](IDevice* device) { return tt::tt_metal::experimental::hal::get_nan(); }, - R"doc(Returns NaN value for current architecture.)doc") - .def( - "sfpu_inf", - [](IDevice* device) { return tt::tt_metal::experimental::hal::get_inf(); }, - R"doc(Returns Infinity value for current architecture.)doc"); + .def( + "sfpu_eps", + [](IDevice* device) { return tt::tt_metal::experimental::hal::get_eps(); }, + R"doc(Returns machine epsilon value for current architecture.)doc") + .def( + "sfpu_nan", + [](IDevice* device) { return tt::tt_metal::experimental::hal::get_nan(); }, + R"doc(Returns NaN value for current architecture.)doc") + .def( + "sfpu_inf", + [](IDevice* device) { return tt::tt_metal::experimental::hal::get_inf(); }, + R"doc(Returns Infinity value for current architecture.)doc"); auto pyDevice = static_cast>>(m_device.attr("Device")); pyDevice diff --git a/ttnn/cpp/ttnn/run_operation.cpp b/ttnn/cpp/ttnn/run_operation.cpp index 48ee11feb41..3c629ecef81 100644 --- a/ttnn/cpp/ttnn/run_operation.cpp +++ b/ttnn/cpp/ttnn/run_operation.cpp @@ -585,8 +585,7 @@ void launch_op_func( const Tensors input_tensors, OutputType& output_tensors, const OptionalConstTensors optional_input_tensors, - const OptionalTensors optional_output_tensors, - bool enable_autoformat_device) { + const OptionalTensors optional_output_tensors) { // Send host side op compile and run to the worker queue // Assert to ensure that worker threads are specified. ZoneScopedN("LaunchOp"); @@ -601,14 +600,12 @@ template void launch_op_func( const Tensors input_tensors, Tensors& output_tensors, const OptionalConstTensors optional_input_tensors, - const OptionalTensors optional_output_tensors, - bool enable_autoformat_device); + const OptionalTensors optional_output_tensors); template void launch_op_func( const std::function& op_func, const Tensors input_tensors, OptionalTensors& output_tensors, const OptionalConstTensors optional_input_tensors, - const OptionalTensors optional_output_tensors, - bool enable_autoformat_device); + const OptionalTensors optional_output_tensors); } // namespace tt::tt_metal::operation diff --git a/ttnn/cpp/ttnn/run_operation.hpp b/ttnn/cpp/ttnn/run_operation.hpp index f83319dd02f..132fe0e9b2a 100644 --- a/ttnn/cpp/ttnn/run_operation.hpp +++ b/ttnn/cpp/ttnn/run_operation.hpp @@ -126,8 +126,7 @@ __attribute__((noinline)) void launch_op_func( const Tensors input_tensors, OutputType& output_tensors, const OptionalConstTensors optional_input_tensors = {}, - const OptionalTensors optional_output_tensors = {}, - bool enable_autoformat_device = true); + const OptionalTensors optional_output_tensors = {}); /* */ @@ -137,16 +136,14 @@ void launch_op( const Tensors input_tensors, OutputType& output_tensors, const OptionalConstTensors optional_input_tensors = {}, - const OptionalTensors optional_output_tensors = {}, - bool enable_autoformat_device = true) { + const OptionalTensors optional_output_tensors = {}) { using FuncType = std::function; launch_op_func( FuncType(std::forward(op_func)), input_tensors, output_tensors, optional_input_tensors, - optional_output_tensors, - enable_autoformat_device); + optional_output_tensors); } void launch_with_autoformat(