From 6a8371b57b935659c36372962f94e3c66d2d8258 Mon Sep 17 00:00:00 2001 From: Diego Gomez Date: Mon, 3 Mar 2025 17:40:54 +0000 Subject: [PATCH] Fixed format issues --- tests/ttnn/unit_tests/test_graph_capture.py | 135 ++++++++++--- ttnn/CMakeLists.txt | 1 + .../ttnn/graph/graph_argument_serializer.cpp | 169 +++++++++++++++++ .../ttnn/graph/graph_argument_serializer.hpp | 39 ++++ ttnn/cpp/ttnn/graph/graph_processor.cpp | 178 +++--------------- ttnn/cpp/ttnn/graph/graph_processor.hpp | 2 +- 6 files changed, 349 insertions(+), 175 deletions(-) create mode 100644 ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp create mode 100644 ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp diff --git a/tests/ttnn/unit_tests/test_graph_capture.py b/tests/ttnn/unit_tests/test_graph_capture.py index 2085aab3d12..b9ffb34f1d5 100644 --- a/tests/ttnn/unit_tests/test_graph_capture.py +++ b/tests/ttnn/unit_tests/test_graph_capture.py @@ -55,25 +55,52 @@ def test_graph_capture_with_all_parameters(device): ttnn.transpose(tt_input, 1, 2) captured_graph = ttnn.graph.end_graph_capture() + node1 = captured_graph[1]["arguments"] # ttnn:transpose + assert node1[0] == "\x00" assert ( - captured_graph[1]["arguments"] - == "[ \x00, std::__1::reference_wrapper>],[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 2048, 4, 128]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=RowMajorPageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt),alignment=Alignment([1])))), std::__1::reference_wrapper],[ 1, std::__1::reference_wrapper],[ 2, std::__1::reference_wrapper],[ nullopt, std::__1::reference_wrapper const>],[ 0, std::__1::reference_wrapper const>]" + node1[1] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 2048, 4, 128]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=RowMajorPageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt),alignment=Alignment([1]))))" ) + assert node1[2] == "1" + assert node1[3] == "2" + assert node1[4] == "nullopt" + assert node1[5] == "0" + # ttnn::prim::permute + node4 = captured_graph[4]["arguments"] + assert ( + node4[0] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 2048, 4, 128]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=RowMajorPageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt),alignment=Alignment([1]))))" + ) + assert node4[1] == "SmallVector([0, 2, 1, 3])" assert ( - captured_graph[4]["arguments"] - == "[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 2048, 4, 128]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=RowMajorPageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt),alignment=Alignment([1])))), std::__1::reference_wrapper],[ SmallVector([0, 2, 1, 3]), std::__1::reference_wrapper>],[ MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt), std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper],[ 0, std::__1::reference_wrapper const>]" + node4[2] + == "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)" ) + assert node4[3] == "[ unsupported type , std::__1::reference_wrapper]" + assert node4[4] == "0" + # PermuteDeviceOperation + node6 = captured_graph[6]["arguments"] assert ( - captured_graph[6]["arguments"] - == "[ unsupported type , std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper]" + node6[0] + == "[ unsupported type , std::__1::reference_wrapper]" ) + assert ( + node6[1] + == "[ unsupported type , std::__1::reference_wrapper]" + ) + # tt::tt_metal::create_device_tensor + node7 = captured_graph[7]["arguments"] + assert node7[0] == "Shape([1, 4, 2048, 128])" + assert node7[1] == "BFLOAT16" + assert node7[2] == "Row Major" + assert node7[3] == "[ unsupported type , std::__1::reference_wrapper]" assert ( - captured_graph[7]["arguments"] - == "[ Shape([1, 4, 2048, 128]), std::__1::reference_wrapper],[ BFLOAT16, std::__1::reference_wrapper],[ Row Major, std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper],[ MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt), std::__1::reference_wrapper]" + node7[4] + == "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)" ) @@ -103,24 +130,61 @@ def test_graph_capture_without_memory_config(device): captured_graph = ttnn.graph.end_graph_capture() # ttnn::moreh_dot + node1 = captured_graph[1]["arguments"] assert ( - captured_graph[1]["arguments"] - == "[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32])))), std::__1::reference_wrapper],[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32])))), std::__1::reference_wrapper],[ nullopt, std::__1::reference_wrapper const>],[ BFLOAT16, std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>],[ unsupported type , std::__1::reference_wrapper> const>]" + node1[0] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))" ) + assert ( + node1[1] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))" + ) + assert node1[2] == "nullopt" + assert node1[3] == "BFLOAT16" + assert node1[4] == "nullopt" + assert ( + node1[5] + == "[ unsupported type , std::__1::reference_wrapper> const>]" + ) + # ttnn::prim::moreh_dot + node6 = captured_graph[6]["arguments"] assert ( - captured_graph[6]["arguments"] - == "[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32])))), std::__1::reference_wrapper],[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32])))), std::__1::reference_wrapper],[ nullopt, std::__1::reference_wrapper const>],[ BFLOAT16, std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>],[ unsupported type , std::__1::reference_wrapper> const>]" + node6[0] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))" ) + assert ( + node6[1] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([1, 1, 1, 32]),tensor_layout=TensorLayout(dtype=BFLOAT16,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))" + ) + assert node6[2] == "nullopt" + assert node6[3] == "BFLOAT16" + assert node6[4] == "nullopt" + assert ( + node6[5] + == "[ unsupported type , std::__1::reference_wrapper> const>]" + ) + # MorehDotOperation + node9 = captured_graph[9]["arguments"] + assert ( + node9[0] + == "[ unsupported type , std::__1::reference_wrapper]" + ) assert ( - captured_graph[9]["arguments"] - == "[ unsupported type , std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper]" + node9[1] + == "[ unsupported type , std::__1::reference_wrapper]" ) + # tt::tt_metal::create_device_tensor + node10 = captured_graph[10]["arguments"] + assert node10[0] == "Shape([1, 1, 1, 1])" + assert node10[1] == "BFLOAT16" + assert node10[2] == "Tile" + assert node10[3] == "[ unsupported type , std::__1::reference_wrapper]" assert ( - captured_graph[10]["arguments"] - == "[ Shape([1, 1, 1, 1]), std::__1::reference_wrapper],[ BFLOAT16, std::__1::reference_wrapper],[ Tile, std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper],[ MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt), std::__1::reference_wrapper]" + node10[4] + == "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)" ) @@ -132,22 +196,45 @@ def test_graph_capture_without_dtype(device): captured_graph = ttnn.graph.end_graph_capture() # ttnn::moreh_full_like + node1 = captured_graph[1]["arguments"] assert ( - captured_graph[1]["arguments"] - == "[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([32, 32]),tensor_layout=TensorLayout(dtype=INT32,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32])))), std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper>],[ nullopt, std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>]" + node1[0] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([32, 32]),tensor_layout=TensorLayout(dtype=INT32,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))" ) + assert node1[1] == "[ unsupported type , std::__1::reference_wrapper>]" + assert node1[2] == "nullopt" + assert node1[3] == "nullopt" + assert node1[4] == "nullopt" + # ttnn::prim::moreh_full_like + node4 = captured_graph[4]["arguments"] assert ( - captured_graph[4]["arguments"] - == "[ Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([32, 32]),tensor_layout=TensorLayout(dtype=INT32,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32])))), std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>],[ nullopt, std::__1::reference_wrapper const>]" + node4[0] + == "Tensor(storage=DeviceStorage(memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)),tensor_spec=TensorSpec(logical_shape=Shape([32, 32]),tensor_layout=TensorLayout(dtype=INT32,page_config=PageConfig(config=TilePageConfig(tile=Tile(tile_shape={32, 32},face_shape={16, 16},num_faces=4))),memory_config=MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))" ) + assert node4[1] == "[ unsupported type , std::__1::reference_wrapper const>]" + assert node4[2] == "nullopt" + assert node4[3] == "nullopt" + assert node4[4] == "nullopt" + # FullLikeOperation + node6 = captured_graph[6]["arguments"] assert ( - captured_graph[6]["arguments"] - == "[ unsupported type , std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper]" + node6[0] + == "[ unsupported type , std::__1::reference_wrapper]" ) + assert ( + node6[1] + == "[ unsupported type , std::__1::reference_wrapper]" + ) + # tt::tt_metal::create_device_tensor + node7 = captured_graph[7]["arguments"] + assert node7[0] == "Shape([32, 32])" + assert node7[1] == "INT32" + assert node7[2] == "Tile" + assert node7[3] == "[ unsupported type , std::__1::reference_wrapper]" assert ( - captured_graph[7]["arguments"] - == "[ Shape([32, 32]), std::__1::reference_wrapper],[ INT32, std::__1::reference_wrapper],[ Tile, std::__1::reference_wrapper],[ unsupported type , std::__1::reference_wrapper],[ MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt), std::__1::reference_wrapper]" + node7[4] + == "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)" ) diff --git a/ttnn/CMakeLists.txt b/ttnn/CMakeLists.txt index dd9516750dd..7c57243fe1c 100644 --- a/ttnn/CMakeLists.txt +++ b/ttnn/CMakeLists.txt @@ -202,6 +202,7 @@ set(TTNN_BASE_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/distributed/distributed_tensor.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/distributed/distributed_tensor_config.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/graph/graph_processor.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/graph/graph_argument_serializer.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/graph/graph_trace_utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/creation.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cpp/ttnn/operations/sharding_utilities.cpp diff --git a/ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp b/ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp new file mode 100644 index 00000000000..af0d5ae6801 --- /dev/null +++ b/ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#include "graph_argument_serializer.hpp" +#include "ttnn/types.hpp" + +namespace ttnn::graph { +std::ostream& operator<<(std::ostream& os, const tt::tt_metal::Layout& layout) { + switch (layout) { + case Layout::ROW_MAJOR: return os << "Row Major"; + case Layout::TILE: return os << "Tile"; + case Layout::INVALID: return os << "Invalid"; + default: return os << "Unknown layout"; + } +} + +std::ostream& operator<<(std::ostream& os, const Tile& config) { + tt::stl::reflection::operator<<(os, config); + return os; +} + +std::ostream& operator<<(std::ostream& os, const Tensor& tensor) { + tt::stl::reflection::operator<<(os, tensor); + return os; +} + +std::ostream& operator<<(std::ostream& os, const tt::stl::StrongType& h) { + return os << *h; +} + +template +std::ostream& operator<<(std::ostream& os, const std::optional& optional_value) { + if (optional_value.has_value()) { + os << optional_value.value(); + } else { + os << "nullopt"; + } + return os; +} + +std::string graph_demangle(const char* name) { + int status = -4; + char* res = abi::__cxa_demangle(name, NULL, NULL, &status); + const char* const demangled_name = (status == 0) ? res : name; + std::string ret_val(demangled_name); + free(res); + return ret_val; +} + +std::unordered_map& GraphArgumentSerializer::registry() { + static std::unordered_map map; + return map; +} + +template +void GraphArgumentSerializer::register_small_vector() { + registry()[typeid(std::reference_wrapper>)] = [](const std::any& value) -> std::string { + std::ostringstream oss; + auto referenced_value = std::any_cast>>(value); + oss << referenced_value.get(); + return oss.str(); + }; +} + +template +void GraphArgumentSerializer::register_special_type() { + registry()[typeid(std::reference_wrapper)] = [](const std::any& value) -> std::string { + std::ostringstream oss; + auto referenced_value = std::any_cast>(value); + oss << referenced_value.get(); + return oss.str(); + }; +} + +template +void GraphArgumentSerializer::register_optional_type() { + registry()[typeid(std::reference_wrapper)] = [](const std::any& value) -> std::string { + std::ostringstream oss; + auto referenced_value = std::any_cast>(value); + auto& referenced_optional = referenced_value.get(); + if (referenced_optional.has_value()) { + oss << *referenced_optional; + } else { + oss << "nullopt"; + } + + return oss.str(); + }; +} + +template +void GraphArgumentSerializer::register_type() { + GraphArgumentSerializer::ConvertionFunction regular_function = [](const std::any& value) -> std::string { + std::ostringstream oss; + if (value.type() == typeid(std::reference_wrapper)) { + auto referenced_value = std::any_cast>(value); + oss << referenced_value.get(); + std::cout << "Parsed: " << typeid(referenced_value).name() << std::endl; + } else if (value.type() == typeid(std::reference_wrapper)) { + auto referenced_value = std::any_cast>(value); + oss << referenced_value.get(); + std::cout << "Parsed: " << typeid(referenced_value).name() << std::endl; + } else { + oss << "Unable to parse"; + } + + return oss.str(); + }; + + // regular cases + registry()[typeid(std::reference_wrapper)] = regular_function; + registry()[typeid(std::reference_wrapper)] = regular_function; + registry()[typeid(const std::reference_wrapper)] = regular_function; + + // Particular cases for optional + register_optional_type>(); + register_optional_type>(); + + // Handle complex types (feel free to add more in the future) + // Small vector + register_small_vector(); + register_small_vector(); + register_small_vector(); + register_small_vector(); +} + +std::vector GraphArgumentSerializer::to_list(const std::span& span) { + std::vector result; + for (const auto& element : span) { + if (!element.has_value()) { + result.push_back("[any, empty]"); + continue; + } + + auto it = registry().find(element.type()); + if (it != registry().end()) { + result.push_back(it->second(element)); + } else { + // for debugging reasons, I want to report the type that is not managed + std::ostringstream oss; + oss << "[ unsupported type" << " , "; + auto demangled_name = graph_demangle(element.type().name()); + oss << demangled_name; + oss << "]"; + result.push_back(oss.str()); + } + } + + return result; +} + +void GraphArgumentSerializer::initialize() { + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_type(); + GraphArgumentSerializer::register_special_type>(); +} +} // namespace ttnn::graph diff --git a/ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp b/ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp new file mode 100644 index 00000000000..6af5e8a1619 --- /dev/null +++ b/ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc. +// +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace ttnn::graph { +std::string graph_demangle(const char* name); + +struct GraphArgumentSerializer { + using ConvertionFunction = std::function; + + static std::unordered_map& registry(); + + template + static void register_small_vector(); + + // In case you don't care about all the variations of the type + // such as const T, const T&, T, T&, etc + template + static void register_special_type(); + + template + static void register_type(); + + template + static void register_optional_type(); + + static std::vector to_list(const std::span& span); + + static void initialize(); +}; +} // namespace ttnn::graph diff --git a/ttnn/cpp/ttnn/graph/graph_processor.cpp b/ttnn/cpp/ttnn/graph/graph_processor.cpp index 52ea9883a6b..67069d09847 100644 --- a/ttnn/cpp/ttnn/graph/graph_processor.cpp +++ b/ttnn/cpp/ttnn/graph/graph_processor.cpp @@ -3,118 +3,23 @@ // SPDX-License-Identifier: Apache-2.0 #include "graph_processor.hpp" +#include "graph_argument_serializer.hpp" #include "graph_consts.hpp" -#include #include "ttnn/types.hpp" -#include -#include +#include "ttnn/core.hpp" #include "ttnn/graph/graph_consts.hpp" #include #include #include +#include +#include +#include #include #include -#include "ttnn/core.hpp" using namespace tt::tt_metal; namespace { -std::ostream& operator<<(std::ostream& os, const tt::tt_metal::Layout& layout) { - switch (layout) { - case Layout::ROW_MAJOR: return os << "Row Major"; - case Layout::TILE: return os << "Tile"; - case Layout::INVALID: return os << "Invalid"; - default: return os << "Unknown layout"; - } -} - -std::ostream& operator<<(std::ostream& os, const Tile& config) { - tt::stl::reflection::operator<<(os, config); - return os; -} - -std::ostream& operator<<(std::ostream& os, const Tensor& tensor) { - tt::stl::reflection::operator<<(os, tensor); - return os; -} - -std::ostream& operator<<(std::ostream& os, const tt::stl::StrongType& h) { - return os << *h; -} - -template -std::ostream& operator<<(std::ostream& os, const std::optional& optional_value) { - if (optional_value.has_value()) { - os << optional_value.value(); - } else { - os << "nullopt"; - } - return os; -} - -std::string demangle(const char* name) { - int status = -4; - - char* res = abi::__cxa_demangle(name, NULL, NULL, &status); - - const char* const demangled_name = (status == 0) ? res : name; - - std::string ret_val(demangled_name); - - free(res); - - return ret_val; -} - -struct AnyToString { - using ConvertionFunction = std::function; - - static std::unordered_map& registry() { - static std::unordered_map map; - return map; - } - - template - static void register_type() { - registry()[typeid(T)] = [](const std::any& value) -> std::string { - std::ostringstream oss; - auto reference_value = std::any_cast(value); - oss << reference_value.get(); - std::string result = oss.str(); - return result; - }; - } - - static std::string to_string(const std::span& span) { - std::ostringstream oss; - for (const auto& element : span) { - if (!element.has_value()) { - oss << "[any, empty],"; - continue; - } - - auto it = registry().find(element.type()); - oss << "[ "; - if (it != registry().end()) { - oss << it->second(element) << ", "; - } else { - oss << "unsupported type" << " , "; - } - - oss << demangle(element.type().name()); - oss << "],"; - } - - std::string result = oss.str(); - if (!result.empty() && result.back() == ',') - { - result.pop_back(); // Remove last comma - } - - return result; - } -}; - std::string tensorMemoryLayoutToString(TensorMemoryLayout layout) { switch (layout) { case TensorMemoryLayout::INTERLEAVED: return "INTERLEAVED"; @@ -184,42 +89,8 @@ GraphProcessor::GraphProcessor(RunMode mode) : run_mode(mode) { ptr->end_function_process_tensor(val); }; - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type const>>(); - AnyToString::register_type const>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type const>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type const>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type const>>(); - AnyToString::register_type const>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type const>>(); - AnyToString::register_type>>(); - AnyToString::register_type>>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); - AnyToString::register_type>(); + // Add all the elements in the map to handle different datatypes + GraphArgumentSerializer::initialize(); } void GraphProcessor::track_allocate(const tt::tt_metal::Buffer* buffer) { @@ -238,7 +109,7 @@ void GraphProcessor::track_allocate(const tt::tt_metal::Buffer* buffer) { {kDeviceId, std::to_string(buffer->device()->id())}}; { graph.push_back( - Vertex{.counter = counter, .node_type = kNodeBufferAllocate, .params = params, .arguments = {}, .connections = {buffer_id}}); + Vertex{.counter = counter, .node_type = kNodeBufferAllocate, .params = params, .connections = {buffer_id}}); graph[current_op_id.top()].connections.push_back(counter); } } @@ -256,7 +127,7 @@ void GraphProcessor::track_deallocate(tt::tt_metal::Buffer* buffer) { {kDeviceId, std::to_string(buffer->device()->id())}}; { graph.push_back(Vertex{ - .counter = counter, .node_type = kNodeBufferDeallocate, .params = params, .arguments = {}, .connections = {buffer_id}}); + .counter = counter, .node_type = kNodeBufferDeallocate, .params = params, .connections = {buffer_id}}); graph[current_op_id.top()].connections.push_back(counter); } } @@ -291,7 +162,6 @@ void GraphProcessor::track_deallocate_cb(const tt::tt_metal::IDevice* device) { .counter = counter, .node_type = kNodeCBDeallocateAll, .params = {{kDeviceId, std::to_string(device->id())}}, - .arguments = {}, .connections = {current_op_id.top()}}); graph[current_op_id.top()].connections.push_back(counter); } @@ -321,12 +191,20 @@ void GraphProcessor::track_function_start(std::string_view function_name, std::s {kName, std::string(function_name)}, }; - auto serialized_arguments = AnyToString::to_string(input_parameters); - TT_ASSERT(serialized_arguments.size()); + std::vector serialized_arguments; + try{ + serialized_arguments = GraphArgumentSerializer::to_list(input_parameters); + for(auto item : serialized_arguments) + std::cout << item << std::endl; + } catch (const std::exception& e) { // ✅ Catches std::runtime_error, std::bad_any_cast, etc. + std::cerr << "Exception: " << e.what() << std::endl; + } catch (...) { // ✅ Catches anything else (segfaults, etc.) + std::cerr << "Unknown Exception occurred." << std::endl; + } auto counter = graph.size(); { - graph.push_back(Vertex{ + graph.push_back(Vertex{ .counter = counter, .node_type = kNodeFunctionStart, .params = params, @@ -347,7 +225,7 @@ void GraphProcessor::track_function_start(std::string_view function_name, std::s if (it != begin_function_any_map.end()) { it->second(any); } else { - tt::log_info("input any type name ignored: {}", demangle(any.type().name())); + tt::log_info("input any type name ignored: {}", graph_demangle(any.type().name())); } } } @@ -359,7 +237,7 @@ void GraphProcessor::track_function_end_impl() { auto counter = graph.size(); { graph.push_back( - Vertex{.counter = counter, .node_type = kNodeFunctionEnd, .params = {{kName, name}}, .arguments = {}, .connections = {}}); + Vertex{.counter = counter, .node_type = kNodeFunctionEnd, .params = {{kName, name}}, .connections = {}}); graph[current_op_id.top()].connections.push_back(counter); } last_finished_op_id = counter; @@ -382,7 +260,7 @@ void GraphProcessor::track_function_end(const std::any& output_tensors) { if (it != end_function_any_map.end()) { it->second(output_tensors); } else { - tt::log_info("output any type name ignored: {}", demangle(output_tensors.type().name())); + tt::log_info("output any type name ignored: {}", graph_demangle(output_tensors.type().name())); } TT_ASSERT(current_op_id.size() > 0); // we should always have capture_start on top current_op_id.pop(); @@ -418,13 +296,13 @@ int GraphProcessor::add_tensor(const Tensor& t) { if (tensor_id_to_counter.count(tensor_id) == 0) { graph.push_back( - Vertex{.counter = tensor_counter, .node_type = kNodeTensor, .params = params, .arguments = {}, .connections = {}}); + Vertex{.counter = tensor_counter, .node_type = kNodeTensor, .params = params, .connections = {}}); tensor_id_to_counter[tensor_id] = tensor_counter; } if (buffers.empty()) { tt::log_info( - "Tensor doesn't have buffer, but storage is {}", demangle(get_type_in_var(t.get_storage()).name())); + "Tensor doesn't have buffer, but storage is {}", graph_demangle(get_type_in_var(t.get_storage()).name())); } for (auto& buffer : buffers) { @@ -445,7 +323,7 @@ int GraphProcessor::add_buffer(const tt::tt_metal::Buffer* buffer) { {kLayout, tensorMemoryLayoutToString(buffer->buffer_layout())}, {kDeviceId, std::to_string(buffer->device()->id())}}; - graph.push_back(Vertex{.counter = counter, .node_type = kNodeBuffer, .params = params, .arguments = {}, .connections = {}}); + graph.push_back(Vertex{.counter = counter, .node_type = kNodeBuffer, .params = params, .connections = {}}); graph[current_op_id.top()].connections.push_back(counter); buffer_id_to_counter[buffer_id] = counter; return counter; @@ -554,7 +432,7 @@ void GraphProcessor::begin_capture(RunMode mode) { graph.clear(); buffer_id_to_counter.clear(); tensor_id_to_counter.clear(); - graph.push_back(Vertex{.counter = 0, .node_type = kNodeCaptureStart, .params = {}, .arguments = {}, .connections = {}}); + graph.push_back(Vertex{.counter = 0, .node_type = kNodeCaptureStart, .params = {}, .connections = {}}); if (!tt::tt_metal::GraphTracker::instance().get_hook()) { hook = std::make_shared(); @@ -566,7 +444,7 @@ void GraphProcessor::begin_capture(RunMode mode) { nlohmann::json GraphProcessor::end_capture() { const std::lock_guard lock(mutex); int counter = graph.size(); - graph.push_back(Vertex{.counter = counter, .node_type = kNodeCaptureEnd, .params = {}, .arguments = {}, .connections = {}}); + graph.push_back(Vertex{.counter = counter, .node_type = kNodeCaptureEnd, .params = {}, .connections = {}}); if (last_finished_op_id != -1) { graph[last_finished_op_id].connections.push_back(counter); } else { diff --git a/ttnn/cpp/ttnn/graph/graph_processor.hpp b/ttnn/cpp/ttnn/graph/graph_processor.hpp index 3cb469ce4f9..3c8f93457da 100644 --- a/ttnn/cpp/ttnn/graph/graph_processor.hpp +++ b/ttnn/cpp/ttnn/graph/graph_processor.hpp @@ -67,7 +67,7 @@ class GraphProcessor : public tt::tt_metal::IGraphProcessor { int counter = 0; std::string node_type; std::unordered_map params; - std::string arguments; + std::vector arguments; std::vector connections; }; using ProcessFunc = std::function;