Grim/explorer memory vis

tools/explorer/test/run_tests.py

@@ -211,6 +211,16 @@ def test_execute_and_check_perf_data_exists():
     assert "perf_data" in result["graphs"][0]
 
 
+def test_execute_and_check_memory_data_exists():
+    execute_command_and_wait(
+        MNIST_SHARDING_PATH,
+        {"optimizationPolicy": "DF Sharding"},
+        timeout=300,
+    )
+    result = convert_command_and_assert(MNIST_SHARDING_PATH)
+    assert "display_type" in str(result["graphs"])
+
+
 def test_execute_model_invalid_policy():
     with pytest.raises(AssertionError):
         execute_command_and_wait(

tools/explorer/tt_adapter/src/tt_adapter/main.py

@@ -93,14 +93,15 @@ def convert(
             model_path
         ):
             logging.info(f"Using optimized model: {optimized_model_path}")
-            # Get performance results.
+            # Get performance and memory results.
             perf_trace = self.model_runner.get_perf_trace(model_path)
+            memory_trace = self.model_runner.get_memory_usage(model_path)
 
             with open(optimized_model_path, "r") as model_file:
                 module = utils.parse_mlir_str(model_file.read())
 
             # Convert TTIR to Model Explorer Graphs and Display/Return
-            graph, perf_data = mlir.build_graph(module, perf_trace)
+            graph, perf_data = mlir.build_graph(module, perf_trace, memory_trace)
             if perf_data:
                 # TODO(odjuricic) We should replace the perf_data with overlays once this is fixed on FE.
                 graph = utils.add_to_dataclass(graph, "perf_data", perf_data.graphsData)

tools/explorer/tt_adapter/src/tt_adapter/mlir.py

@@ -394,6 +394,7 @@ def parse_ttnn_ttnn_layout(attr):
     result = []
     result.append(graph_builder.KeyValue(key="linear", value=str(layout.linear)))
     memory_layout = layout.tensor_memory_layout_as_int
+
     if memory_layout is not None:
         result.append(
             utils.make_editable_kv(
@@ -462,6 +463,7 @@ def parse_ttnn_ttnn_layout(attr):
             },
         )
     )
+
     return result
 
 
@@ -528,6 +530,7 @@ def get_attributes(self):
                         key="rank", value=str(output_tensor.type.rank)
                     ),
                 ]
+
             if hasattr(output_tensor.type, "encoding") and output_tensor.type.encoding:
                 if "ttnn_layout" in str(output_tensor.type.encoding):
                     output_attrs.extend(
@@ -549,15 +552,17 @@ def get_attributes(self):
             graph_builder.KeyValue(key="schedule", value=str(OpHandler.schedule))
         )
         OpHandler.schedule += 1
-
         return result
 
-    def make_graph_node(self):
+    def make_graph_node(self, extra_attrs=None):
+        attrs = self.get_attributes()
+        if extra_attrs is not None:
+            attrs.extend(extra_attrs)
         return graph_builder.GraphNode(
             id=self.id,
             label=str(self.op.name),
             namespace=self.get_namespace(),
-            attrs=self.get_attributes(),
+            attrs=attrs,
         )
 
     def make_constant_node(self, constant_name):
@@ -580,7 +585,7 @@ def make_constant_node(self, constant_name):
 ]
 
 
-def build_graph(module, perf_trace=None):
+def build_graph(module, perf_trace=None, memory_trace=None):
     output_connections = defaultdict(int)
     graph = graph_builder.Graph(id="tt-graph")
 
@@ -598,9 +603,23 @@ def build_graph(module, perf_trace=None):
             if loc:
                 loc_to_perf[loc] = row["DEVICE FW DURATION [ns]"]
 
+    memory_data = {}
+    if memory_trace is not None:
+        for node in memory_trace:
+            memory_data[node] = {}
+            memory_data[node]["dram"] = (
+                memory_trace[node]["dram"]["device_0"]["total_bytes_allocated_per_bank"]
+                / memory_trace[node]["dram"]["device_0"]["total_bytes_per_bank"]
+            )
+            memory_data[node]["l1"] = (
+                memory_trace[node]["l1"]["device_0"]["total_bytes_allocated_per_bank"]
+                / memory_trace[node]["l1"]["device_0"]["total_bytes_per_bank"]
+            )
+
     module_op = OpHandler(module.operation)
     module_attrs = module_op.get_attributes()
     module_attrs = dict((attr.key, attr.value) for attr in module_attrs)
+
     # Add module attributes to the graph as "namespace attributes"
     group_node_attrs = {}
     group_node_attrs[module_op.get_namespace()] = module_attrs
@@ -610,9 +629,38 @@ def build_graph(module, perf_trace=None):
         for region in op.regions:
             for block in region.blocks:
                 for op in block.operations:
+                    extra_attrs = []
+                    for operand_index, operand in enumerate(op.operands):
+                        if memory_data:
+                            extra_attrs.append(
+                                utils.add_to_dataclass(
+                                    graph_builder.KeyValue(
+                                        key="dram_memory",
+                                        value=str(
+                                            memory_data[str(operand_index)]["dram"]
+                                        ),
+                                    ),
+                                    "display_type",
+                                    "memory",
+                                )
+                            )
+                        if memory_data:
+                            extra_attrs.append(
+                                utils.add_to_dataclass(
+                                    graph_builder.KeyValue(
+                                        key="l1_memory",
+                                        value=str(
+                                            memory_data[str(operand_index)]["l1"]
+                                        ),
+                                    ),
+                                    "display_type",
+                                    "memory",
+                                )
+                            )
+
                     # Create all the nodes and constants in the first pass.
                     operation = OpHandler(op)
-                    graph_node = operation.make_graph_node()
+                    graph_node = operation.make_graph_node(extra_attrs)
 
                     if (
                         operation.named_location in loc_to_perf
@@ -681,6 +729,33 @@ def build_graph(module, perf_trace=None):
                                     key="rank", value=str(operand.type.rank)
                                 ),
                             ]
+
+                        if memory_data:
+                            output_attrs.append(
+                                utils.add_to_dataclass(
+                                    graph_builder.KeyValue(
+                                        key="dram_memory",
+                                        value=str(
+                                            memory_data[str(operand_index)]["dram"]
+                                        ),
+                                    ),
+                                    "display_type",
+                                    "memory",
+                                )
+                            )
+                            output_attrs.append(
+                                utils.add_to_dataclass(
+                                    graph_builder.KeyValue(
+                                        key="l1_memory",
+                                        value=str(
+                                            memory_data[str(operand_index)]["l1"]
+                                        ),
+                                    ),
+                                    "display_type",
+                                    "memory",
+                                )
+                            )
+
                         if hasattr(operand.type, "encoding") and operand.type.encoding:
                             if "ttnn_layout" in str(operand.type.encoding):
                                 output_attrs.extend(
@@ -695,6 +770,7 @@ def build_graph(module, perf_trace=None):
                                         operand.type.encoding.get_named("tt.layout")
                                     )
                                 )
+
                         source_node.outputsMetadata.append(
                             graph_builder.MetadataItem(
                                 id=str(output_connections[source_node.id]),

tools/explorer/tt_adapter/src/tt_adapter/runner.py

@@ -13,6 +13,7 @@
 import pandas as pd
 import threading
 import queue
+import json
 
 
 class ExplorerRunException(Exception):
@@ -140,6 +141,19 @@ def get_perf_trace(self, model_path):
 
         return pd.read_csv(op_perf_file)
 
+    def get_memory_usage(self, model_path):
+
+        mem_file = f"{self.model_state[model_path].model_output_dir}/run/program_0/memory_results.json"
+        if not os.path.exists(mem_file):
+            raise FileNotFoundError(
+                f"Memory file {mem_file} not found. Memory file may not have been created."
+            )
+
+        with open(mem_file, "r") as file:
+            memory_trace = json.load(file)
+
+        return memory_trace
+
     def run_in_subprocess(self, command):
         self.log(f"Running command:\n{' '.join(command)}\n")
 
@@ -299,6 +313,7 @@ def compile_and_run(self, model_path, overrides_string):
             "perf",
             flatbuffer_file,
             f"--artifact-dir={self._explorer_artifacts_dir}",
+            "--memory",
         ]
 
         ttrt_process = self.run_in_subprocess(ttrt_perf_command)