From 87ab998bd6471f86974cdb2b01172ec1838d4e30 Mon Sep 17 00:00:00 2001 From: Raymond Kim Date: Thu, 6 Mar 2025 09:42:57 -0500 Subject: [PATCH] Revert "Extend the capabilities of the graph tooling to include extra information of the ttnn operations (#18380)" This reverts commit a33d64ad73db5b969ab3369cb8a3ec2377c8e9ce, which broke ttnn unit tests 1. --- tests/ttnn/unit_tests/gtests/CMakeLists.txt | 2 - .../test_graph_capture_arguments_morehdot.cpp | 113 ---------- ...test_graph_capture_arguments_transpose.cpp | 94 -------- .../unit_tests/gtests/ttnn_test_fixtures.hpp | 22 -- tests/ttnn/unit_tests/test_graph_capture.py | 203 ------------------ ttnn/CMakeLists.txt | 1 - .../ttnn/graph/graph_argument_serializer.cpp | 173 --------------- .../ttnn/graph/graph_argument_serializer.hpp | 46 ---- ttnn/cpp/ttnn/graph/graph_consts.hpp | 1 - ttnn/cpp/ttnn/graph/graph_processor.cpp | 38 ++-- ttnn/cpp/ttnn/graph/graph_processor.hpp | 1 - ttnn/cpp/ttnn/graph/graph_trace_utils.cpp | 17 -- ttnn/cpp/ttnn/graph/graph_trace_utils.hpp | 7 - ttnn/cpp/ttnn/tensor/types.cpp | 15 -- ttnn/cpp/ttnn/tensor/types.hpp | 2 - 15 files changed, 22 insertions(+), 713 deletions(-) delete mode 100644 tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_morehdot.cpp delete mode 100644 tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_transpose.cpp delete mode 100644 ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp delete mode 100644 ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp diff --git a/tests/ttnn/unit_tests/gtests/CMakeLists.txt b/tests/ttnn/unit_tests/gtests/CMakeLists.txt index d56ea2dfb08..4afd6a0cbf6 100644 --- a/tests/ttnn/unit_tests/gtests/CMakeLists.txt +++ b/tests/ttnn/unit_tests/gtests/CMakeLists.txt @@ -5,8 +5,6 @@ set(TTNN_UNIT_TESTS_SRC ${CMAKE_CURRENT_SOURCE_DIR}/test_async_runtime.cpp ${CMAKE_CURRENT_SOURCE_DIR}/test_multiprod_queue.cpp ${CMAKE_CURRENT_SOURCE_DIR}/test_multi_cq_multi_dev.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/test_graph_capture_arguments_morehdot.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/test_graph_capture_arguments_transpose.cpp ${CMAKE_CURRENT_SOURCE_DIR}/test_graph_query_op_constraints.cpp ${CMAKE_CURRENT_SOURCE_DIR}/test_graph_query_op_runtime.cpp ${CMAKE_CURRENT_SOURCE_DIR}/test_reflect.cpp diff --git a/tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_morehdot.cpp b/tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_morehdot.cpp deleted file mode 100644 index b862df915a2..00000000000 --- a/tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_morehdot.cpp +++ /dev/null @@ -1,113 +0,0 @@ -// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. -// -// SPDX-License-Identifier: Apache-2.0 - -#include "gtest/gtest.h" -#include "ttnn_test_fixtures.hpp" -#include "ttnn/device.hpp" -#include "ttnn/graph/graph_processor.hpp" -#include "ttnn/graph/graph_consts.hpp" -#include "ttnn/graph/graph_trace_utils.hpp" -#include "ttnn/operations/moreh/moreh_dot/moreh_dot.hpp" -#include -#include - -namespace ttnn::graph::arguments::test { - -class TestGraphCaptureArgumentsMorehDot : public TTNNFixtureWithTensor {}; - -TEST_P(TestGraphCaptureArgumentsMorehDot, MorehDot) { - auto tt_input1 = CreateTensor(); - auto tt_input2 = CreateTensor(); - ttnn::graph::GraphProcessor::begin_graph_capture(tt::tt_metal::IGraphProcessor::RunMode::NORMAL); - ttnn::moreh_dot(tt_input1, tt_input2, std::nullopt, DataType::BFLOAT16, std::nullopt, std::nullopt); - auto trace = ttnn::graph::GraphProcessor::end_graph_capture(); - auto operations = ttnn::graph::extract_arguments(trace); - - auto operation0 = operations[0]; - EXPECT_EQ(operation0.operation_name, "ttnn::moreh_dot"); - EXPECT_EQ(operation0.arguments.size(), 6); - EXPECT_EQ( - operation0.arguments[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, 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::L1,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))"); - EXPECT_EQ( - operation0.arguments[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, 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::L1,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))"); - EXPECT_EQ(operation0.arguments[2], "[ unsupported type , std::__1::reference_wrapper]"); - EXPECT_EQ(operation0.arguments[3], "BFLOAT16"); - EXPECT_EQ(operation0.arguments[4], "[ unsupported type , std::__1::reference_wrapper]"); - EXPECT_EQ(operation0.arguments[5], "[ unsupported type , std::__1::reference_wrapper]"); - - auto operation1 = operations[1]; - EXPECT_EQ(operation1.operation_name, "ttnn::prim::moreh_dot"); - EXPECT_EQ(operation1.arguments.size(), 6); - EXPECT_EQ( - operation1.arguments[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, 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::L1,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))"); - EXPECT_EQ( - operation1.arguments[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, 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::L1,shard_spec=std::nullopt),alignment=Alignment([32, 32]))))"); - EXPECT_EQ(operation1.arguments[2], "nullopt"); - EXPECT_EQ(operation1.arguments[3], "BFLOAT16"); - EXPECT_EQ(operation1.arguments[4], "nullopt"); - EXPECT_EQ( - operation1.arguments[5], - "[ unsupported type , " - "std::__1::reference_wrapper> const>]"); - - auto operation2 = operations[2]; - EXPECT_EQ(operation2.operation_name, "MorehDotOperation"); - EXPECT_EQ(operation2.arguments.size(), 2); - EXPECT_EQ( - operation2.arguments[0], - "[ unsupported type , " - "std::__1::reference_wrapper]"); - EXPECT_EQ( - operation2.arguments[1], - "[ unsupported type , " - "std::__1::reference_wrapper]"); - - auto operation3 = operations[3]; - EXPECT_EQ(operation3.operation_name, "tt::tt_metal::create_device_tensor"); - EXPECT_EQ(operation3.arguments.size(), 5); - EXPECT_EQ(operation3.arguments[0], "Shape([1, 1, 1, 1])"); - EXPECT_EQ(operation3.arguments[1], "BFLOAT16"); - EXPECT_EQ(operation3.arguments[2], "Tile"); - EXPECT_EQ(operation3.arguments[3], "[ unsupported type , std::__1::reference_wrapper]"); - EXPECT_EQ( - operation3.arguments[4], - "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::" - "nullopt)"); -} - -INSTANTIATE_TEST_SUITE_P( - TestGraphCaptureArgumentsMorehDot_MorehDot, - TestGraphCaptureArgumentsMorehDot, - ::testing::Values(CreateTensorParameters{ - .input_shape = ttnn::Shape({1, 1, 1, 32}), - .dtype = DataType::BFLOAT16, - .layout = TILE_LAYOUT, - .mem_cfg = L1_MEMORY_CONFIG})); -} // namespace ttnn::graph::arguments::test diff --git a/tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_transpose.cpp b/tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_transpose.cpp deleted file mode 100644 index ed1ab12ab98..00000000000 --- a/tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_transpose.cpp +++ /dev/null @@ -1,94 +0,0 @@ -// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc. -// -// SPDX-License-Identifier: Apache-2.0 - -#include "gtest/gtest.h" -#include "ttnn_test_fixtures.hpp" -#include "ttnn/device.hpp" -#include "ttnn/graph/graph_processor.hpp" -#include "ttnn/graph/graph_consts.hpp" -#include "ttnn/graph/graph_trace_utils.hpp" -#include "ttnn/operations/data_movement/transpose/transpose.hpp" -#include -#include - -namespace ttnn::graph::arguments::test { - -class TestGraphCaptureArgumentsTranspose : public TTNNFixtureWithTensor {}; - -TEST_P(TestGraphCaptureArgumentsTranspose, Transpose) { - auto tt_input = CreateTensor(); - tt_input.reshape(ttnn::Shape{1, 2048, 4, 128}); - ttnn::graph::GraphProcessor::begin_graph_capture(tt::tt_metal::IGraphProcessor::RunMode::NORMAL); - ttnn::transpose(tt_input, 1, 2); - auto trace = ttnn::graph::GraphProcessor::end_graph_capture(); - auto operations = ttnn::graph::extract_arguments(trace); - - auto operation0 = operations[0]; - EXPECT_EQ(operation0.operation_name, "ttnn::transpose"); - EXPECT_EQ(operation0.arguments.size(), 3); - EXPECT_EQ( - operation0.arguments[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, 1, 2048, " - "512]),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]))))"); - EXPECT_EQ(operation0.arguments[1], "1"); - EXPECT_EQ(operation0.arguments[2], "2"); - - auto operation1 = operations[1]; - EXPECT_EQ(operation1.operation_name, "ttnn::prim::permute"); - EXPECT_EQ(operation1.arguments.size(), 5); - EXPECT_EQ( - operation1.arguments[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, 1, 2048, " - "512]),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]))))"); - EXPECT_EQ(operation1.arguments[1], "SmallVector([0, 2, 1, 3])"); - EXPECT_EQ( - operation1.arguments[2], - "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::" - "nullopt)"); - EXPECT_EQ(operation1.arguments[3], "[ unsupported type , std::__1::reference_wrapper]"); - EXPECT_EQ(operation1.arguments[4], "0"); - - auto operation2 = operations[2]; - EXPECT_EQ(operation2.operation_name, "PermuteDeviceOperation"); - EXPECT_EQ(operation2.arguments.size(), 2); - EXPECT_EQ( - operation2.arguments[0], - "[ unsupported type , " - "std::__1::reference_wrapper]"); - EXPECT_EQ( - operation2.arguments[1], - "[ unsupported type , " - "std::__1::reference_wrapper]"); - - auto operation3 = operations[3]; - EXPECT_EQ(operation3.operation_name, "tt::tt_metal::create_device_tensor"); - EXPECT_EQ(operation3.arguments.size(), 5); - EXPECT_EQ(operation3.arguments[0], "Shape([1, 2048, 1, 512])"); - EXPECT_EQ(operation3.arguments[1], "BFLOAT16"); - EXPECT_EQ(operation3.arguments[2], "Row Major"); - EXPECT_EQ(operation3.arguments[3], "[ unsupported type , std::__1::reference_wrapper]"); - EXPECT_EQ( - operation3.arguments[4], - "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::" - "nullopt)"); -} - -INSTANTIATE_TEST_SUITE_P( - TestGraphCaptureArgumentsTranspose_Transpose, - TestGraphCaptureArgumentsTranspose, - ::testing::Values(CreateTensorParameters{ - .input_shape = ttnn::Shape({1, 1, 2048, 512}), - .dtype = DataType::BFLOAT16, - .layout = ROW_MAJOR_LAYOUT, - .mem_cfg = L1_MEMORY_CONFIG})); -} // namespace ttnn::graph::arguments::test diff --git a/tests/ttnn/unit_tests/gtests/ttnn_test_fixtures.hpp b/tests/ttnn/unit_tests/gtests/ttnn_test_fixtures.hpp index 2aff14b8074..c889afb6b1a 100644 --- a/tests/ttnn/unit_tests/gtests/ttnn_test_fixtures.hpp +++ b/tests/ttnn/unit_tests/gtests/ttnn_test_fixtures.hpp @@ -13,8 +13,6 @@ #include "ttnn/device.hpp" #include "ttnn/types.hpp" -#include "ttnn/tensor/tensor.hpp" -#include "ttnn/tensor/tensor_impl.hpp" #include "tests/tt_metal/test_utils/env_vars.hpp" #include #include "hostdevcommon/common_values.hpp" @@ -55,26 +53,6 @@ class TTNNFixtureWithDevice : public TTNNFixture { tt::tt_metal::IDevice& getDevice() { return *device_; } }; -struct CreateTensorParameters { - ttnn::Shape input_shape; - DataType dtype; - Layout layout; - MemoryConfig mem_cfg; -}; - -class TTNNFixtureWithTensor : public TTNNFixtureWithDevice, public testing::WithParamInterface { -protected: - [[nodiscard]] const Tensor CreateTensor() { - CreateTensorParameters params = GetParam(); - TensorSpec tensor_spec( - params.input_shape, TensorLayout(params.dtype, PageConfig(params.layout), params.mem_cfg)); - auto input_buffer = tt::tt_metal::tensor_impl::allocate_buffer_on_device(device_, tensor_spec); - auto input_storage = tt::tt_metal::DeviceStorage{input_buffer}; - Tensor input_tensor = Tensor(input_storage, params.input_shape, params.dtype, params.layout); - return std::move(input_tensor); - } -}; - } // namespace ttnn namespace ttnn::distributed::test { diff --git a/tests/ttnn/unit_tests/test_graph_capture.py b/tests/ttnn/unit_tests/test_graph_capture.py index b9ffb34f1d5..c3fd190f1d1 100644 --- a/tests/ttnn/unit_tests/test_graph_capture.py +++ b/tests/ttnn/unit_tests/test_graph_capture.py @@ -35,206 +35,3 @@ def test_graph_capture(tmp_path, device, scalar, size, mode): ttnn.graph.pretty_print(captured_graph) ttnn.graph.visualize(captured_graph, file_name=tmp_path / pathlib.Path("graph.svg")) - - -def test_graph_capture_with_all_parameters(device): - # Create input tensor - torch_input = torch.rand((1, 1, 2048, 512), dtype=torch.bfloat16) - - # TT operations - tt_input = ttnn.from_torch( - torch_input, - dtype=ttnn.DataType.BFLOAT16, - layout=ttnn.ROW_MAJOR_LAYOUT, - device=device, - memory_config=ttnn.L1_MEMORY_CONFIG, - ) - - tt_input = tt_input.reshape(1, 2048, 4, 128) - ttnn.graph.begin_graph_capture(ttnn.graph.RunMode.NORMAL) - 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 ( - 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 ( - 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 ( - 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 ( - node7[4] - == "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::L1,shard_spec=std::nullopt)" - ) - - -def test_graph_capture_without_memory_config(device): - # Create input tensor - input_shape = (1, 1, 1, 32) - torch_input = torch.rand(input_shape, dtype=torch.bfloat16) - torch_other = torch.rand(input_shape, dtype=torch.bfloat16) - - # TT operations - tt_input = ttnn.from_torch( - torch_input, - dtype=ttnn.DataType.BFLOAT16, - layout=ttnn.TILE_LAYOUT, - device=device, - ) - - tt_other = ttnn.from_torch( - torch_input, - dtype=ttnn.DataType.BFLOAT16, - layout=ttnn.TILE_LAYOUT, - device=device, - ) - - ttnn.graph.begin_graph_capture(ttnn.graph.RunMode.NORMAL) - tt_out = ttnn.operations.moreh.dot(tt_input, tt_other, dtype=ttnn.bfloat16, output=None) - captured_graph = ttnn.graph.end_graph_capture() - - # ttnn::moreh_dot - node1 = captured_graph[1]["arguments"] - assert ( - 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 ( - 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 ( - 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 ( - node10[4] - == "MemoryConfig(memory_layout=TensorMemoryLayout::INTERLEAVED,buffer_type=BufferType::DRAM,shard_spec=std::nullopt)" - ) - - -def test_graph_capture_without_dtype(device): - torch_input = torch.randint(0, 100, (32, 32), dtype=torch.int32) - tt_input = ttnn.from_torch(torch_input, layout=ttnn.TILE_LAYOUT, device=device) - ttnn.graph.begin_graph_capture(ttnn.graph.RunMode.NORMAL) - tt_output = ttnn.moreh_full_like(tt_input, 3) - captured_graph = ttnn.graph.end_graph_capture() - - # ttnn::moreh_full_like - node1 = captured_graph[1]["arguments"] - assert ( - 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 ( - 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 ( - 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 ( - 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 24d8384818e..47a3c3eb84d 100644 --- a/ttnn/CMakeLists.txt +++ b/ttnn/CMakeLists.txt @@ -202,7 +202,6 @@ 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 deleted file mode 100644 index 7aac2990c05..00000000000 --- a/ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp +++ /dev/null @@ -1,173 +0,0 @@ -// 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 tt::tt_metal::Tile& config) { - tt::stl::reflection::operator<<(os, config); - return os; -} - -std::ostream& operator<<(std::ostream& os, const tt::tt_metal::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 std::string_view name) { - int status = -4; - char* res = abi::__cxa_demangle(name.data(), NULL, NULL, &status); - const char* const demangled_name = (status == 0) ? res : name.data(); - std::string ret_val(demangled_name); - free(res); - return ret_val; -} - -GraphArgumentSerializer::GraphArgumentSerializer() { initialize(); } - -GraphArgumentSerializer& GraphArgumentSerializer::instance() { - static GraphArgumentSerializer new_instance; - return new_instance; -} - -std::unordered_map& GraphArgumentSerializer::registry() { - 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(); - } else if (value.type() == typeid(std::reference_wrapper)) { - auto referenced_value = std::any_cast>(value); - oss << referenced_value.get(); - } 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 deleted file mode 100644 index 3bc5f8280c8..00000000000 --- a/ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp +++ /dev/null @@ -1,46 +0,0 @@ -// 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 std::string_view name); - -class GraphArgumentSerializer { -public: - using ConvertionFunction = std::function; - std::vector to_list(const std::span& span); - - static GraphArgumentSerializer& instance(); - -private: - GraphArgumentSerializer(); - std::unordered_map& registry(); - - template - 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 - void register_special_type(); - - template - void register_type(); - - template - void register_optional_type(); - - void initialize(); - -private: - std::unordered_map map; -}; -} // namespace ttnn::graph diff --git a/ttnn/cpp/ttnn/graph/graph_consts.hpp b/ttnn/cpp/ttnn/graph/graph_consts.hpp index 28e64d3a8dd..7ede9bb7c23 100644 --- a/ttnn/cpp/ttnn/graph/graph_consts.hpp +++ b/ttnn/cpp/ttnn/graph/graph_consts.hpp @@ -10,7 +10,6 @@ constexpr auto kNodeType = "node_type"; constexpr auto kCounter = "counter"; constexpr auto kConnections = "connections"; constexpr auto kParams = "params"; -constexpr auto kArguments = "arguments"; // params keys constexpr auto kName = "name"; constexpr auto kInputs = "inputs"; diff --git a/ttnn/cpp/ttnn/graph/graph_processor.cpp b/ttnn/cpp/ttnn/graph/graph_processor.cpp index fd6852edfa4..77f415cb348 100644 --- a/ttnn/cpp/ttnn/graph/graph_processor.cpp +++ b/ttnn/cpp/ttnn/graph/graph_processor.cpp @@ -3,23 +3,36 @@ // 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 "ttnn/core.hpp" +#include +#include #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::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; +} + std::string tensorMemoryLayoutToString(TensorMemoryLayout layout) { switch (layout) { case TensorMemoryLayout::INTERLEAVED: return "INTERLEAVED"; @@ -41,7 +54,6 @@ nlohmann::json to_json(const ttnn::graph::GraphProcessor::Vertex& data) { j[ttnn::graph::kCounter] = data.counter; j[ttnn::graph::kNodeType] = data.node_type; j[ttnn::graph::kParams] = data.params; - j[ttnn::graph::kArguments] = data.arguments; j[ttnn::graph::kConnections] = data.connections; return j; } @@ -89,7 +101,6 @@ GraphProcessor::GraphProcessor(RunMode mode) : run_mode(mode) { ptr->end_function_process_tensor(val); }; } - void GraphProcessor::track_allocate(const tt::tt_metal::Buffer* buffer) { const std::lock_guard lock(mutex); auto buffer_id = add_buffer(buffer); @@ -187,17 +198,12 @@ void GraphProcessor::track_function_start(std::string_view function_name, std::s {kInputs, std::to_string(input_parameters.size())}, {kName, std::string(function_name)}, }; - - std::vector serialized_arguments; - serialized_arguments = GraphArgumentSerializer::instance().to_list(input_parameters); - auto counter = graph.size(); { - graph.push_back(Vertex{ + graph.push_back(Vertex{ .counter = counter, .node_type = kNodeFunctionStart, .params = params, - .arguments = serialized_arguments, .connections = {/*current_op_id.top()*/}}); if (last_finished_op_id != -1) { graph[last_finished_op_id].connections.push_back(counter); @@ -214,7 +220,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: {}", graph_demangle(any.type().name())); + tt::log_info("input any type name ignored: {}", demangle(any.type().name())); } } } @@ -249,7 +255,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: {}", graph_demangle(output_tensors.type().name())); + tt::log_info("output any type name ignored: {}", demangle(output_tensors.type().name())); } TT_ASSERT(current_op_id.size() > 0); // we should always have capture_start on top current_op_id.pop(); @@ -291,7 +297,7 @@ int GraphProcessor::add_tensor(const Tensor& t) { if (buffers.empty()) { tt::log_info( - "Tensor doesn't have buffer, but storage is {}", graph_demangle(get_type_in_var(t.get_storage()).name())); + "Tensor doesn't have buffer, but storage is {}", demangle(get_type_in_var(t.get_storage()).name())); } for (auto& buffer : buffers) { diff --git a/ttnn/cpp/ttnn/graph/graph_processor.hpp b/ttnn/cpp/ttnn/graph/graph_processor.hpp index 3c8f93457da..0736ff8a9b0 100644 --- a/ttnn/cpp/ttnn/graph/graph_processor.hpp +++ b/ttnn/cpp/ttnn/graph/graph_processor.hpp @@ -67,7 +67,6 @@ class GraphProcessor : public tt::tt_metal::IGraphProcessor { int counter = 0; std::string node_type; std::unordered_map params; - std::vector arguments; std::vector connections; }; using ProcessFunc = std::function; diff --git a/ttnn/cpp/ttnn/graph/graph_trace_utils.cpp b/ttnn/cpp/ttnn/graph/graph_trace_utils.cpp index 8d2c9656688..629ea0335a3 100644 --- a/ttnn/cpp/ttnn/graph/graph_trace_utils.cpp +++ b/ttnn/cpp/ttnn/graph/graph_trace_utils.cpp @@ -154,23 +154,6 @@ std::vector extract_calltrace(const nlohmann::json& trace) { return op_calls; } -std::vector extract_arguments(const nlohmann::json& trace) { - std::vector operations; - size_t i = 0; - while (i < trace.size()) { - const auto& v = trace[i]; - i++; - OperationInfo info; - if (v[kArguments].size() > 0) { - info.operation_name = v[kParams][kName]; - info.arguments = v[kArguments]; - operations.push_back(info); - } - } - - return operations; -} - std::unordered_set extract_output_tensors(const nlohmann::json& trace) { // Lambda to find the last 'function_end' node auto find_function_end_node = [](const auto& trace) -> const nlohmann::json& { diff --git a/ttnn/cpp/ttnn/graph/graph_trace_utils.hpp b/ttnn/cpp/ttnn/graph/graph_trace_utils.hpp index 67260062a67..408ae47acb5 100644 --- a/ttnn/cpp/ttnn/graph/graph_trace_utils.hpp +++ b/ttnn/cpp/ttnn/graph/graph_trace_utils.hpp @@ -11,11 +11,6 @@ namespace ttnn::graph { -struct OperationInfo { - std::string operation_name; - std::vector arguments; -}; - enum class ExecutionStatus { Success, Error }; uint32_t extract_peak_L1_memory_usage(const nlohmann::json& trace); @@ -29,8 +24,6 @@ std::pair count_intermediate_and_output_tensors(const nlohma std::vector extract_calltrace(const nlohmann::json& trace); -std::vector extract_arguments(const nlohmann::json& trace); - std::unordered_set extract_output_tensors(const nlohmann::json& trace); struct TensorInfo { diff --git a/ttnn/cpp/ttnn/tensor/types.cpp b/ttnn/cpp/ttnn/tensor/types.cpp index 852f0a69ece..32eb5daa1bb 100644 --- a/ttnn/cpp/ttnn/tensor/types.cpp +++ b/ttnn/cpp/ttnn/tensor/types.cpp @@ -8,21 +8,6 @@ namespace tt::tt_metal { -std::ostream& operator<<(std::ostream& os, const tt::tt_metal::DataType& data_type) { - switch (data_type) { - case DataType::BFLOAT16: return os << "BFLOAT16"; - case DataType::FLOAT32: return os << "FLOAT32"; - case DataType::UINT32: return os << "UINT32"; - case DataType::BFLOAT8_B: return os << "BFLOAT8_B"; - case DataType::BFLOAT4_B: return os << "BFLOAT4_B"; - case DataType::UINT8: return os << "UINT8"; - case DataType::UINT16: return os << "UINT16"; - case DataType::INT32: return os << "INT32"; - case DataType::INVALID: - default: return os << "Invalid"; - } -} - bool is_floating_point(DataType dtype) { switch (dtype) { case DataType::BFLOAT16: diff --git a/ttnn/cpp/ttnn/tensor/types.hpp b/ttnn/cpp/ttnn/tensor/types.hpp index 041c5b99995..09a2aeecf19 100644 --- a/ttnn/cpp/ttnn/tensor/types.hpp +++ b/ttnn/cpp/ttnn/tensor/types.hpp @@ -43,8 +43,6 @@ enum class DataType { INVALID = 8, }; -std::ostream& operator<<(std::ostream& os, const tt::tt_metal::DataType& data_type); - template consteval inline DataType convert_to_data_type() { if constexpr (std::is_same_v) {