Extend the capabilities of the graph tooling to include extra information of the ttnn operations #18380

dgomezTT · 2025-02-26T22:05:38Z

Ticket

Problem description

The org will be creating new operations and they need some tooling to identify what is actually being sent via TTNN, so we extended the graph solution to include extra information, in this case, the arguments of the ttnn calls

What's changed

Extended the capabilities of the graph tool to add extra information for those vertex that include arguments, so whoever wants to use them, will have more information available

Checklist

All post commit CI passes
Blackhole Post commit CI passes (if applicable)
Model regression CI passes (if applicable)
Device performance regression CI passes (if applicable)
(For models and ops writers) Full new models tests CI passes (if applicable)
New/Existing tests provide coverage for changes

tests/ttnn/unit_tests/test_graph_capture.py

ttnn/cpp/ttnn/graph/graph_processor.hpp

ayerofieiev-tt · 2025-02-26T23:53:40Z

ttnn/cpp/ttnn/graph/graph_processor.cpp

@@ -33,6 +66,55 @@ std::string demangle(const char* name) {
    return ret_val;
 }

+struct AnyToString {


lets move to a different file
and lets move specific serializations to a different file too please

ttnn/cpp/ttnn/graph/graph_processor.cpp

ayerofieiev-tt · 2025-03-03T18:46:38Z

ttnn/cpp/ttnn/graph/graph_processor.hpp

@@ -67,6 +67,7 @@ class GraphProcessor : public tt::tt_metal::IGraphProcessor {
        int counter = 0;
        std::string node_type;
        std::unordered_map<std::string, std::string> params;
+        std::vector<std::string> arguments;


I think its good for this PR
But I am now think if this should store the copy of the any value itself and stringify only when we try to dump the trace 🤔

Without context, it is not clear what the difference is between "args" and "params". Consider documenting, since this is a public API.

I will definitely document it because it took me a good while to get the difference

ayerofieiev-tt · 2025-03-03T18:49:06Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp

+namespace ttnn::graph {
+std::string graph_demangle(const char* name);
+
+struct GraphArgumentSerializer {


You can make this a class, make constructor private, move init inside constructor.
GraphArgumentSerializer::instance() or GraphArgumentSerializer::registry() will create an instance and the instance will get itself initialized.

+1. If we use a static registry, right now it isn't clear what exactly is static here.

Also use tt::stl::Indestructible for static data.

ayerofieiev-tt · 2025-03-03T18:56:24Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp

+    register_small_vector<T, 2>();
+    register_small_vector<T, 4>();
+    register_small_vector<T, 8>();
+    register_small_vector<T, 16>();


I only saw a single usage for Shape related things I think.
I'd reduce to a single default value for now.

CC @sminakov-tt

ayerofieiev-tt

Can you please also add a sample C++ test?
We have some C++ consumers interested in this.

ayerofieiev-tt

Great!
I think it is a good first step as is.
There are a couple things which I commented on, which are good to discuss separately, but they should not be done in this PR. Many things going to evolve over a couple iterations.

omilyutin-tt · 2025-03-03T22:49:51Z

ttnn/cpp/ttnn/tensor/types.cpp

+        case DataType::UINT16: return os << "UINT16";
+        case DataType::INT32: return os << "INT32";
+        case DataType::INVALID: return os << "Invalid";
+        default: return os << "Unknown data type";


Remove the default - not having it triggers the compiler error when new enums are added - this is useful for future changes.

omilyutin-tt · 2025-03-03T22:57:10Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp

+#include <unordered_map>
+
+namespace ttnn::graph {
+std::string graph_demangle(const char* name);


can name be std::string_view?

omilyutin-tt · 2025-03-03T22:57:38Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp

+    static std::unordered_map<std::type_index, ConvertionFunction>& registry();
+
+    template <typename T, std::size_t N>
+    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 <typename T>
+    static void register_special_type();
+
+    template <typename T>
+    static void register_type();
+
+    template <typename OptionalT>
+    static void register_optional_type();
+
+    static std::vector<std::string> to_list(const std::span<std::any>& span);
+
+    static void initialize();


Does this need to be public?

omilyutin-tt · 2025-03-03T22:57:57Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp

+        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";


Ditto for default

ayerofieiev-tt · 2025-03-04T17:37:45Z

ttnn/cpp/ttnn/graph/graph_trace_utils.cpp

@@ -154,6 +154,22 @@ std::vector<std::string> extract_calltrace(const nlohmann::json& trace) {
    return op_calls;
 }

+std::vector<std::vector<std::string>> extract_arguments(const nlohmann::json& trace) {


If you want it in this way, it does not belong to common utils and need to be moved to tests.
If you want a useful common util, wdyt if the return type is slightly different? Something like

using OperationName = std::string; // tttn.op name using Arguments = std::vector<std::string>; // vector of serialized args using Returns = std::vector<std::string>; // vector of serialized TensorSpec for returned tensors vector<OperationName, Arguments, Return>

sounds like a great idea! Right on it!

ayerofieiev-tt · 2025-03-04T17:40:35Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp

+    using ConvertionFunction = std::function<std::string(const std::any&)>;
+    std::vector<std::string> to_list(const std::span<std::any>& span);
+
+    static GraphArgumentSerializer* instance();


Suggested change

static GraphArgumentSerializer* instance();

static GraphArgumentSerializer& instance();

ayerofieiev-tt · 2025-03-04T17:41:39Z

ttnn/cpp/ttnn/graph/graph_processor.cpp

@@ -100,7 +88,11 @@ GraphProcessor::GraphProcessor(RunMode mode) : run_mode(mode) {
    end_function_any_map[typeid(std::reference_wrapper<Tensor>)] = [ptr = this](const std::any& val) mutable {
        ptr->end_function_process_tensor(val);
    };
+
+    // Add all the elements in the map to handle different datatypes
+    TT_ASSERT(GraphArgumentSerializer::instance());


this is not needed really

ayerofieiev-tt · 2025-03-04T17:42:14Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp

+GraphArgumentSerializer::GraphArgumentSerializer() { initialize(); }
+
+GraphArgumentSerializer* GraphArgumentSerializer::instance() {
+    static GraphArgumentSerializer* new_instance = new GraphArgumentSerializer();


this can as well be

Suggested change

static GraphArgumentSerializer* new_instance = new GraphArgumentSerializer();

static GraphArgumentSerializer new_instance;

ayerofieiev-tt · 2025-03-04T17:42:24Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp

+
+GraphArgumentSerializer::GraphArgumentSerializer() { initialize(); }
+
+GraphArgumentSerializer* GraphArgumentSerializer::instance() {


Suggested change

GraphArgumentSerializer* GraphArgumentSerializer::instance() {

GraphArgumentSerializer& GraphArgumentSerializer::instance() {

ayerofieiev-tt · 2025-03-04T19:16:11Z

tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_morehdot.cpp

+    auto tt_input1 = CreateTensor();
+    auto tt_input2 = CreateTensor();
+    ttnn::graph::GraphProcessor::begin_graph_capture(Mode);
+    ttnn::operations::moreh::moreh_dot::MorehDot::invoke(


ttnn::moreh_dot(...) it is a registered operation, do not call invoke directly please because it ends up being an undecorated call. Meaning you do not capture the actual call of ttnn.op but only underlying device operations or parts of the composite call

omilyutin-tt · 2025-03-04T19:38:32Z

tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_morehdot.cpp

+class TestGraphCaptureArgumentsMorehDot : public TTNNFixtureWithTensor {
+protected:
+    tt::tt_metal::IGraphProcessor::RunMode Mode = tt::tt_metal::IGraphProcessor::RunMode::NORMAL;
+};


mode_ for data members

But if it is just the constant, I think we can inline in the one place it is used?

omilyutin-tt · 2025-03-04T19:39:07Z

tests/ttnn/unit_tests/gtests/test_graph_capture_arguments_transpose.cpp

+    TestGraphCaptureArgumentsTranspose_Transpose,
+    TestGraphCaptureArgumentsTranspose,
+    ::testing::Values(CreateTensorParameters{
+        .input_shape = ttnn::Shape(tt::tt_metal::Array4D{1, 1, 2048, 512}),


Can instantiate a shape directly with {1, 1, 2048, 512}

omilyutin-tt · 2025-03-04T19:41:18Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.hpp

+private:
+    GraphArgumentSerializer();
+    std::unordered_map<std::type_index, ConvertionFunction>& registry();
+
+    template <typename T, std::size_t N>
+    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 <typename T>
+    void register_special_type();
+
+    template <typename T>
+    void register_type();
+
+    template <typename OptionalT>
+    void register_optional_type();
+
+    void initialize();
+
+private:
+    std::unordered_map<std::type_index, GraphArgumentSerializer::ConvertionFunction> map;


Yep, but this can be further moved to .cpp, so that we don't change the header and the users don't need to re-compile it?

@omilyutin-tt how could I use this in the graph_processor class if moved to a cpp? or do you mean an specific part? For instance, keeping the map the constructor and registry, while the others are functions in the cpp?

Hmm, I was thinking have the map as a function local static variable in .cpp, along with all of "register" functions:

// Anonymous namespace to enforce internal linkage namespace { std::unordered_map<std::type_index, GraphArgumentSerializer::ConvertionFunction>& registry() { static tt::stl::Indestructible<...> registry; return registry.get(); template <typename T, std::size_t N> 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 <typename T> void register_special_type() {...} template <typename T> void register_type() {...} template <typename OptionalT> void register_optional_type() {...} } // namespace GraphArgumentSerializer::initialize() { // As before... }

Unless I am missing something:)

omilyutin-tt · 2025-03-05T15:40:13Z

ttnn/cpp/ttnn/graph/graph_argument_serializer.cpp

+    return os;
+}
+
+std::string graph_demangle(const std::string_view& name) {


Pass string_view by value

great catch!

…mers

… information of the ttnn operations (#18380)" This reverts commit a33d64a, which broke ttnn unit tests 1.

…tion of the ttnn operations (#18380)"" bring back commit 87ab998.

dgomezTT requested review from ayerofieiev-tt, dmakoviichuk-tt, rfurko-tt, cfjchu, TT-BrianLiu, razorback3, dongjin-na, bbradelTT and omilyutin-tt as code owners February 26, 2025 22:05

ayerofieiev-tt reviewed Feb 26, 2025

View reviewed changes

tests/ttnn/unit_tests/test_graph_capture.py Outdated Show resolved Hide resolved

ayerofieiev-tt reviewed Feb 26, 2025

View reviewed changes

tests/ttnn/unit_tests/test_graph_capture.py Outdated Show resolved Hide resolved

ayerofieiev-tt reviewed Feb 26, 2025

View reviewed changes

ttnn/cpp/ttnn/graph/graph_processor.hpp Outdated Show resolved Hide resolved

ayerofieiev-tt reviewed Feb 26, 2025

View reviewed changes

ttnn/cpp/ttnn/graph/graph_processor.cpp Outdated Show resolved Hide resolved

ayerofieiev-tt reviewed Feb 27, 2025

View reviewed changes

ttnn/cpp/ttnn/graph/graph_processor.cpp Outdated Show resolved Hide resolved

ayerofieiev-tt reviewed Feb 27, 2025

View reviewed changes

ttnn/cpp/ttnn/graph/graph_processor.cpp Outdated Show resolved Hide resolved

dgomezTT force-pushed the dgomeztt/18097 branch from bac4b25 to 6a8371b Compare March 3, 2025 17:42

dgomezTT requested a review from sminakov-tt as a code owner March 3, 2025 17:42

ayerofieiev-tt reviewed Mar 3, 2025

View reviewed changes

ayerofieiev-tt approved these changes Mar 3, 2025

View reviewed changes

bbradelTT approved these changes Mar 3, 2025

View reviewed changes

dmakoviichuk-tt approved these changes Mar 3, 2025

View reviewed changes

omilyutin-tt reviewed Mar 3, 2025

View reviewed changes

ayerofieiev-tt reviewed Mar 4, 2025

View reviewed changes

omilyutin-tt reviewed Mar 4, 2025

View reviewed changes

omilyutin-tt reviewed Mar 5, 2025

View reviewed changes

dgomezTT force-pushed the dgomeztt/18097 branch 2 times, most recently from 2bdf296 to 95f740c Compare March 5, 2025 22:13

dgomezTT added 10 commits March 6, 2025 04:25

Added argument information to the graph

09cf7c0

Removed pragma optimize off that I accidentally left in the file

3118ce0

Fixed format issues

ca68cbb

removing some debug information I used to test the change

06aefea

added c++ tests, addressed general feedback

ee4727f

removing some debug info

0987580

improved the trace utils method for arguments so is clearer for custo…

df87a67

…mers

Changed to use the ttnn operation instead of invoke

5f875ca

fixing compilation issues:

4fd4adf

fix compilation issue

26b0652

dgomezTT force-pushed the dgomeztt/18097 branch from 95f740c to 26b0652 Compare March 6, 2025 04:26

dgomezTT merged commit a33d64a into main Mar 6, 2025
220 of 222 checks passed

dgomezTT deleted the dgomeztt/18097 branch March 6, 2025 13:02

tt-rkim pushed a commit that referenced this pull request Mar 6, 2025

Revert "Extend the capabilities of the graph tooling to include extra…

87ab998

… information of the ttnn operations (#18380)" This reverts commit a33d64a, which broke ttnn unit tests 1.

dgomezTT added a commit that referenced this pull request Mar 6, 2025

Extend the capabilities of the graph tooling to include extra informa…

1bbbff3

…tion of the ttnn operations (#18380)"" bring back commit 87ab998.

dgomezTT added a commit that referenced this pull request Mar 6, 2025

Extend the capabilities of the graph tooling to include extra informa…

28017ea

…tion of the ttnn operations (#18380)"" bring back commit 87ab998.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Extend the capabilities of the graph tooling to include extra information of the ttnn operations #18380

Extend the capabilities of the graph tooling to include extra information of the ttnn operations #18380

dgomezTT commented Feb 26, 2025

ayerofieiev-tt Feb 26, 2025

dgomezTT Mar 3, 2025

ayerofieiev-tt Mar 3, 2025

omilyutin-tt Mar 3, 2025

dgomezTT Mar 3, 2025

ayerofieiev-tt Mar 3, 2025

omilyutin-tt Mar 3, 2025

ayerofieiev-tt Mar 3, 2025

ayerofieiev-tt left a comment

ayerofieiev-tt left a comment

omilyutin-tt Mar 3, 2025

omilyutin-tt Mar 3, 2025

omilyutin-tt Mar 3, 2025

omilyutin-tt Mar 3, 2025

ayerofieiev-tt Mar 4, 2025

dgomezTT Mar 4, 2025

ayerofieiev-tt Mar 4, 2025

ayerofieiev-tt Mar 4, 2025

ayerofieiev-tt Mar 4, 2025

ayerofieiev-tt Mar 4, 2025

ayerofieiev-tt Mar 4, 2025 •

edited

Loading

omilyutin-tt Mar 4, 2025

omilyutin-tt Mar 4, 2025

omilyutin-tt Mar 4, 2025

dgomezTT Mar 5, 2025 •

edited

Loading

omilyutin-tt Mar 5, 2025

omilyutin-tt Mar 5, 2025

dgomezTT Mar 5, 2025

	static GraphArgumentSerializer* instance();
	static GraphArgumentSerializer& instance();

	static GraphArgumentSerializer* new_instance = new GraphArgumentSerializer();
	static GraphArgumentSerializer new_instance;


		GraphArgumentSerializer::GraphArgumentSerializer() { initialize(); }

		GraphArgumentSerializer* GraphArgumentSerializer::instance() {

	GraphArgumentSerializer* GraphArgumentSerializer::instance() {
	GraphArgumentSerializer& GraphArgumentSerializer::instance() {

Extend the capabilities of the graph tooling to include extra information of the ttnn operations #18380

Extend the capabilities of the graph tooling to include extra information of the ttnn operations #18380

Conversation

dgomezTT commented Feb 26, 2025

Ticket

Problem description

What's changed

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ayerofieiev-tt left a comment

Choose a reason for hiding this comment

ayerofieiev-tt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ayerofieiev-tt Mar 4, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dgomezTT Mar 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ayerofieiev-tt Mar 4, 2025 •

edited

Loading

dgomezTT Mar 5, 2025 •

edited

Loading