Grim/explorer memory vis

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

@@ -93,18 +93,22 @@ 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)
 
+            #if memory_data:
+            #    graph = utils.add_to_dataclass(graph, "memory", memory_data)
+
             if overrides := self.model_runner.get_overrides(model_path):
                 graph = utils.add_to_dataclass(graph, "overrides", overrides)
         else:

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

@@ -10,6 +10,8 @@
 from ttmlir import ir, util
 from . import utils
 
+global memory_data
+global node_id_number
 
 def parse_loc_string(loc_str):
     """
@@ -394,6 +396,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 +465,32 @@ def parse_ttnn_ttnn_layout(attr):
             },
         )
     )
+
+    try:
+        result.append(
+            utils.add_to_dataclass(
+                graph_builder.KeyValue(
+                    key="dram_memory",
+                    value=str(memory_data[str(node_id_number)]["dram"]),
+                ),
+                'display_type', 
+                'memory'
+            )
+        )
+
+        result.append(
+            utils.add_to_dataclass(
+                graph_builder.KeyValue(
+                    key="l1_memory",
+                    value=str(memory_data[str(node_id_number)]["l1"]),
+                ),
+                'display_type', 
+                'memory'
+            )
+        )
+    except:
+        pass
+
     return result
 
 
@@ -528,6 +557,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(
@@ -544,6 +574,7 @@ def get_attributes(self):
                     )
             result.extend(output_attrs)
 
+
         # Add schedule as an attribute
         result.append(
             graph_builder.KeyValue(key="schedule", value=str(OpHandler.schedule))
@@ -580,7 +611,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 +629,19 @@ def build_graph(module, perf_trace=None):
             if loc:
                 loc_to_perf[loc] = row["DEVICE FW DURATION [ns]"]
 
+    global memory_data
+    global node_id_number
+    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
@@ -669,6 +710,7 @@ def build_graph(module, perf_trace=None):
                         )
 
                         output_attrs = []
+                        node_id_number = operand_index
                         if isinstance(operand.type, ir.RankedTensorType):
                             output_attrs = [
                                 graph_builder.KeyValue(
@@ -681,6 +723,7 @@ def build_graph(module, perf_trace=None):
                                     key="rank", value=str(operand.type.rank)
                                 ),
                             ]
+
                         if hasattr(operand.type, "encoding") and operand.type.encoding:
                             if "ttnn_layout" in str(operand.type.encoding):
                                 output_attrs.extend(
@@ -694,7 +737,7 @@ def build_graph(module, perf_trace=None):
                                     AttrHandler.parse_attr(
                                         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):
@@ -139,6 +140,19 @@ def get_perf_trace(self, model_path):
             raise FileNotFoundError(f"Performance file {op_perf_file} not found.")
 
         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. Try running command: ttrt run out.ttnn --memory --save-artifacts")
+
+        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)

tools/explorer/tt_adapter/src/tt_adapter/utils.py

@@ -115,3 +115,6 @@ def make_editable_kv(kv, editable):
     obj = asdict(kv)
     obj["editable"] = editable
     return to_dataclass(obj, "KeyValue")
+
+
+