From df46c1505b1fa8e5cf66cc139ad12b64ec338426 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Tue, 25 Feb 2025 18:48:36 +0000 Subject: [PATCH 1/7] First draft explorer memory visualization, incomplete --- tools/explorer/tt_adapter/src/tt_adapter/main.py | 8 ++++++-- tools/explorer/tt_adapter/src/tt_adapter/mlir.py | 13 +++++++++++-- tools/explorer/tt_adapter/src/tt_adapter/runner.py | 13 +++++++++++++ 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py index 3c5371ebbf..ad16e1dab8 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/main.py +++ b/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, memory_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.graphsData) + if overrides := self.model_runner.get_overrides(model_path): graph = utils.add_to_dataclass(graph, "overrides", overrides) else: diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 7050ab94a4..0a31218be5 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -580,7 +580,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 +598,18 @@ 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]["dram"] = memory_data[node]["dram"]["device_0"]["total_bytes_allocated_per_bank"]/memory_data[node]["dram"]["device_0"]["total_bytes_per_bank"] + memory_data[node]["l1"] = memory_data[node]["l1"]["device_0"]["total_bytes_allocated_per_bank"]/memory_data[node]["l1"]["device_0"]["total_bytes_per_bank"] + #memory_data[node]["l1_small"] = memory_data[node]["l1_small"]["device_0"]["total_bytes_allocated_per_bank"]/memory_data[node]["l1_small"]["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) + #Is this where ^ to add the 2 or three new attributes? + # Add module attributes to the graph as "namespace attributes" group_node_attrs = {} group_node_attrs[module_op.get_namespace()] = module_attrs @@ -724,4 +733,4 @@ def build_graph(module, perf_trace=None): graph.groupNodeAttributes = group_node_attrs OpHandler.schedule = 0 - return graph, overlay_data + return graph, overlay_data, memory_data diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index e388c820ed..1afe0fb80b 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/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,18 @@ 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") From 8b8975bddaec8c146742ba1bd47af042ac53e3d5 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Wed, 26 Feb 2025 18:07:17 +0000 Subject: [PATCH 2/7] origin/grim/explorer_memory_vis --- .../tt_adapter/src/tt_adapter/main.py | 11 +++- .../tt_adapter/src/tt_adapter/mlir.py | 57 ++++++++++++++++++- .../tt_adapter/src/tt_adapter/runner.py | 4 +- .../tt_adapter/src/tt_adapter/utils.py | 4 ++ 4 files changed, 69 insertions(+), 7 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py index ad16e1dab8..37f8446151 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/main.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/main.py @@ -95,8 +95,9 @@ def convert( logging.info(f"Using optimized model: {optimized_model_path}") # Get performance and memory results. perf_trace = self.model_runner.get_perf_trace(model_path) + print("Main, convert, running get_memory_usage") memory_trace = self.model_runner.get_memory_usage(model_path) - + #print("Main, convert, running get_memory_usage success: ", memory_trace) with open(optimized_model_path, "r") as model_file: module = utils.parse_mlir_str(model_file.read()) @@ -105,9 +106,13 @@ def convert( 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) + print("perf_data type: ", type(perf_data)) + print("perf_data.graphsData type: ", type(perf_data.graphsData)) if memory_data: - graph = utils.add_to_dataclass(graph, "memory", memory_data.graphsData) + print("Main, convert, running add_to_dataclass") + graph = utils.add_to_dataclass(graph, "memory", memory_data) + print("Main, convert, ran add_to_dataclass") if overrides := self.model_runner.get_overrides(model_path): graph = utils.add_to_dataclass(graph, "overrides", overrides) @@ -127,7 +132,7 @@ def convert( module = utils.parse_mlir_str(model_file.read()) # Convert TTIR to Model Explorer Graphs and Display/Return - graph, _ = mlir.build_graph(module) + graph, _, _ = mlir.build_graph(module) return {"graphs": [graph]} diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 0a31218be5..59e2b83031 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -462,6 +462,25 @@ def parse_ttnn_ttnn_layout(attr): }, ) ) + temp_var = utils.add_to_dataclass( + graph_builder.KeyValue( + key="dram_memory", + value=str(.4), #whatever the % is I think + ), + "dram_memory", + str(.4) + ) + #print("trying to add to dataclass", type(temp_var), temp_var) + #result.append( + # utils.add_to_dataclass( + # graph_builder.KeyValue( + # key="dram_memory", + # value=str(.4), #whatever the % is I think + # ), + # "dram_memory", + # str(.4) + # ) + #) return result @@ -601,18 +620,21 @@ def build_graph(module, perf_trace=None, memory_trace=None): memory_data = {} if memory_trace is not None: for node in memory_trace: - memory_data[node]["dram"] = memory_data[node]["dram"]["device_0"]["total_bytes_allocated_per_bank"]/memory_data[node]["dram"]["device_0"]["total_bytes_per_bank"] - memory_data[node]["l1"] = memory_data[node]["l1"]["device_0"]["total_bytes_allocated_per_bank"]/memory_data[node]["l1"]["device_0"]["total_bytes_per_bank"] - #memory_data[node]["l1_small"] = memory_data[node]["l1_small"]["device_0"]["total_bytes_allocated_per_bank"]/memory_data[node]["l1_small"]["device_0"]["total_bytes_per_bank"] + 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) + #print("mlir: module_attrs type and self: ", type(module_attrs), module_attrs) + #module_attrs[] #Is this where ^ to add the 2 or three new attributes? # Add module attributes to the graph as "namespace attributes" group_node_attrs = {} group_node_attrs[module_op.get_namespace()] = module_attrs + #print("mlir: group_node_attrs: ", group_node_attrs) for op in module.body.operations: append_later = [] @@ -690,6 +712,11 @@ def build_graph(module, perf_trace=None, memory_trace=None): key="rank", value=str(operand.type.rank) ), ] + #print("0: output_attrs: ", output_attrs) + #print(type(operand), operand) + #output_attrs.append(graph_builder.KeyValue(key="dram_memory", value=str(.3), display_type = 'memory')) + #output_attrs.append(graph_builder.KeyValue(key="l1_memory", value=str(operand.index), display_type = 'memory')) + #print("1: output_attrs: ", output_attrs) if hasattr(operand.type, "encoding") and operand.type.encoding: if "ttnn_layout" in str(operand.type.encoding): output_attrs.extend( @@ -704,6 +731,30 @@ def build_graph(module, perf_trace=None, memory_trace=None): operand.type.encoding.get_named("tt.layout") ) ) + #Going to try to figure out how to add memory here, then later maybe move it + 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' + ) + ) + print(operand_index, memory_data) + output_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="l1_memory", + value=str(memory_data[str(operand_index)]["l1"]), + ), + 'display_type', + 'memory' + ) + ) + print("output_attrs: ", output_attrs) source_node.outputsMetadata.append( graph_builder.MetadataItem( id=str(output_connections[source_node.id]), diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index 1afe0fb80b..e1cfa98d7f 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -142,6 +142,7 @@ def get_perf_trace(self, model_path): return pd.read_csv(op_perf_file) def get_memory_usage(self, model_path): + print("Opening get_memory_usage") mem_file = ( f"{self.model_state[model_path].model_output_dir}/run/program_0/memory_results.json" ) @@ -150,7 +151,7 @@ def get_memory_usage(self, model_path): with open(mem_file, "r") as file: memory_trace = json.load(file) - + #print("memory_trace: ", memory_trace) return memory_trace def run_in_subprocess(self, command): @@ -312,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) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/utils.py b/tools/explorer/tt_adapter/src/tt_adapter/utils.py index 69fab1c623..727d732bf3 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/utils.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/utils.py @@ -114,4 +114,8 @@ def to_adapter_format(*objs): def make_editable_kv(kv, editable): obj = asdict(kv) obj["editable"] = editable + #print("in utils.make_editable_kv: obj type and self: ", type(obj), obj) return to_dataclass(obj, "KeyValue") + + + From ff692402d7816a7840914eada9e8ee6bb9c2b715 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Wed, 26 Feb 2025 19:15:24 +0000 Subject: [PATCH 3/7] Cleaned up code, functional but needs to be tested --- .../tt_adapter/src/tt_adapter/main.py | 7 +----- .../tt_adapter/src/tt_adapter/mlir.py | 24 +------------------ .../tt_adapter/src/tt_adapter/runner.py | 4 ++-- .../tt_adapter/src/tt_adapter/utils.py | 1 - 4 files changed, 4 insertions(+), 32 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py index 37f8446151..b832af8e32 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/main.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/main.py @@ -95,9 +95,8 @@ def convert( logging.info(f"Using optimized model: {optimized_model_path}") # Get performance and memory results. perf_trace = self.model_runner.get_perf_trace(model_path) - print("Main, convert, running get_memory_usage") memory_trace = self.model_runner.get_memory_usage(model_path) - #print("Main, convert, running get_memory_usage success: ", memory_trace) + with open(optimized_model_path, "r") as model_file: module = utils.parse_mlir_str(model_file.read()) @@ -106,13 +105,9 @@ def convert( 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) - print("perf_data type: ", type(perf_data)) - print("perf_data.graphsData type: ", type(perf_data.graphsData)) if memory_data: - print("Main, convert, running add_to_dataclass") graph = utils.add_to_dataclass(graph, "memory", memory_data) - print("Main, convert, ran add_to_dataclass") if overrides := self.model_runner.get_overrides(model_path): graph = utils.add_to_dataclass(graph, "overrides", overrides) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 59e2b83031..a142209515 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -470,17 +470,7 @@ def parse_ttnn_ttnn_layout(attr): "dram_memory", str(.4) ) - #print("trying to add to dataclass", type(temp_var), temp_var) - #result.append( - # utils.add_to_dataclass( - # graph_builder.KeyValue( - # key="dram_memory", - # value=str(.4), #whatever the % is I think - # ), - # "dram_memory", - # str(.4) - # ) - #) + return result @@ -627,14 +617,10 @@ def build_graph(module, perf_trace=None, memory_trace=None): module_op = OpHandler(module.operation) module_attrs = module_op.get_attributes() module_attrs = dict((attr.key, attr.value) for attr in module_attrs) - #print("mlir: module_attrs type and self: ", type(module_attrs), module_attrs) - #module_attrs[] - #Is this where ^ to add the 2 or three new attributes? # Add module attributes to the graph as "namespace attributes" group_node_attrs = {} group_node_attrs[module_op.get_namespace()] = module_attrs - #print("mlir: group_node_attrs: ", group_node_attrs) for op in module.body.operations: append_later = [] @@ -712,11 +698,6 @@ def build_graph(module, perf_trace=None, memory_trace=None): key="rank", value=str(operand.type.rank) ), ] - #print("0: output_attrs: ", output_attrs) - #print(type(operand), operand) - #output_attrs.append(graph_builder.KeyValue(key="dram_memory", value=str(.3), display_type = 'memory')) - #output_attrs.append(graph_builder.KeyValue(key="l1_memory", value=str(operand.index), display_type = 'memory')) - #print("1: output_attrs: ", output_attrs) if hasattr(operand.type, "encoding") and operand.type.encoding: if "ttnn_layout" in str(operand.type.encoding): output_attrs.extend( @@ -731,7 +712,6 @@ def build_graph(module, perf_trace=None, memory_trace=None): operand.type.encoding.get_named("tt.layout") ) ) - #Going to try to figure out how to add memory here, then later maybe move it if memory_data: output_attrs.append( utils.add_to_dataclass( @@ -743,7 +723,6 @@ def build_graph(module, perf_trace=None, memory_trace=None): 'memory' ) ) - print(operand_index, memory_data) output_attrs.append( utils.add_to_dataclass( graph_builder.KeyValue( @@ -754,7 +733,6 @@ def build_graph(module, perf_trace=None, memory_trace=None): 'memory' ) ) - print("output_attrs: ", output_attrs) source_node.outputsMetadata.append( graph_builder.MetadataItem( id=str(output_connections[source_node.id]), diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index e1cfa98d7f..d83d18c0a4 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -142,7 +142,7 @@ def get_perf_trace(self, model_path): return pd.read_csv(op_perf_file) def get_memory_usage(self, model_path): - print("Opening get_memory_usage") + mem_file = ( f"{self.model_state[model_path].model_output_dir}/run/program_0/memory_results.json" ) @@ -151,7 +151,7 @@ def get_memory_usage(self, model_path): with open(mem_file, "r") as file: memory_trace = json.load(file) - #print("memory_trace: ", memory_trace) + return memory_trace def run_in_subprocess(self, command): diff --git a/tools/explorer/tt_adapter/src/tt_adapter/utils.py b/tools/explorer/tt_adapter/src/tt_adapter/utils.py index 727d732bf3..3e6c774ff1 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/utils.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/utils.py @@ -114,7 +114,6 @@ def to_adapter_format(*objs): def make_editable_kv(kv, editable): obj = asdict(kv) obj["editable"] = editable - #print("in utils.make_editable_kv: obj type and self: ", type(obj), obj) return to_dataclass(obj, "KeyValue") From f06677628124d16530258782f2f8fd7fd63cfa01 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Thu, 27 Feb 2025 16:21:09 +0000 Subject: [PATCH 4/7] Memory Visualization Successful --- .../tt_adapter/src/tt_adapter/main.py | 8 +-- .../tt_adapter/src/tt_adapter/mlir.py | 65 ++++++++++--------- 2 files changed, 39 insertions(+), 34 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py index b832af8e32..5b299164c9 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/main.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/main.py @@ -101,13 +101,13 @@ def convert( module = utils.parse_mlir_str(model_file.read()) # Convert TTIR to Model Explorer Graphs and Display/Return - graph, perf_data, memory_data = mlir.build_graph(module, perf_trace, memory_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 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) @@ -127,7 +127,7 @@ def convert( module = utils.parse_mlir_str(model_file.read()) # Convert TTIR to Model Explorer Graphs and Display/Return - graph, _, _ = mlir.build_graph(module) + graph, _ = mlir.build_graph(module) return {"graphs": [graph]} diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index a142209515..31c58130f9 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/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,15 +465,32 @@ def parse_ttnn_ttnn_layout(attr): }, ) ) - temp_var = utils.add_to_dataclass( - graph_builder.KeyValue( - key="dram_memory", - value=str(.4), #whatever the % is I think - ), - "dram_memory", - str(.4) + + 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 @@ -537,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( @@ -553,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)) @@ -607,6 +629,8 @@ def build_graph(module, perf_trace=None, memory_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: @@ -686,6 +710,7 @@ def build_graph(module, perf_trace=None, memory_trace=None): ) output_attrs = [] + node_id_number = operand_index if isinstance(operand.type, ir.RankedTensorType): output_attrs = [ graph_builder.KeyValue( @@ -698,6 +723,7 @@ def build_graph(module, perf_trace=None, memory_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( @@ -711,28 +737,7 @@ def build_graph(module, perf_trace=None, memory_trace=None): AttrHandler.parse_attr( operand.type.encoding.get_named("tt.layout") ) - ) - 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' - ) - ) + ) source_node.outputsMetadata.append( graph_builder.MetadataItem( id=str(output_connections[source_node.id]), @@ -762,4 +767,4 @@ def build_graph(module, perf_trace=None, memory_trace=None): graph.groupNodeAttributes = group_node_attrs OpHandler.schedule = 0 - return graph, overlay_data, memory_data + return graph, overlay_data From 2552d47ed0ae6a3c70206a3b444d2d1cbfab6795 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Fri, 28 Feb 2025 15:44:58 +0000 Subject: [PATCH 5/7] Explorer Memory Visualization --- .../tt_adapter/src/tt_adapter/main.py | 3 - .../tt_adapter/src/tt_adapter/mlir.py | 96 +++++++++++-------- .../tt_adapter/src/tt_adapter/runner.py | 2 +- 3 files changed, 58 insertions(+), 43 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py index 5b299164c9..4e465f2a11 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/main.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/main.py @@ -106,9 +106,6 @@ def convert( # 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: diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 31c58130f9..3c1af96f12 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -10,8 +10,6 @@ from ttmlir import ir, util from . import utils -global memory_data -global node_id_number def parse_loc_string(loc_str): """ @@ -62,10 +60,8 @@ def parse_attr_name(attr: ir.Attribute) -> List[graph_builder.KeyValue]: def decorator(handler): AttrHandler.ATTR_HANDLERS[attr_name] = handler return handler - return decorator - @AttrHandler.register_handler("tt.device") def parse_tt_device(attr): device = tt.ir.DeviceAttr.maybe_downcast(attr) @@ -466,31 +462,6 @@ 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 @@ -557,7 +528,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( @@ -580,15 +551,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): @@ -629,8 +602,6 @@ def build_graph(module, perf_trace=None, memory_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: @@ -651,9 +622,34 @@ def build_graph(module, perf_trace=None, memory_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 @@ -710,7 +706,6 @@ def build_graph(module, perf_trace=None, memory_trace=None): ) output_attrs = [] - node_id_number = operand_index if isinstance(operand.type, ir.RankedTensorType): output_attrs = [ graph_builder.KeyValue( @@ -723,7 +718,29 @@ def build_graph(module, perf_trace=None, memory_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( @@ -737,7 +754,8 @@ def build_graph(module, perf_trace=None, memory_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]), diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index d83d18c0a4..862d88b695 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -147,7 +147,7 @@ def get_memory_usage(self, model_path): 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") + 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) From aed612ef2065d018145f2b94c611df775cc93db3 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Fri, 28 Feb 2025 15:56:00 +0000 Subject: [PATCH 6/7] Added a test for the explorer memory visualization --- tools/explorer/test/run_tests.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tools/explorer/test/run_tests.py b/tools/explorer/test/run_tests.py index 9fea9ee9cd..6966353cc9 100644 --- a/tools/explorer/test/run_tests.py +++ b/tools/explorer/test/run_tests.py @@ -114,7 +114,6 @@ def execute_command(model_path, settings): print(result.json()) assert False - def wait_for_execution_to_finish(timeout): for _ in range(timeout): try: @@ -211,6 +210,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( @@ -218,3 +227,4 @@ def test_execute_model_invalid_policy(): {"optimizationPolicy": "Invalid Policy"}, timeout=300, ) + From 5bcab1927530e9c7d2ea75b9f43fcc8640b94185 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Fri, 28 Feb 2025 16:00:45 +0000 Subject: [PATCH 7/7] Updated pre-commit --- tools/explorer/test/run_tests.py | 2 +- .../tt_adapter/src/tt_adapter/mlir.py | 49 ++++++++++++------- .../tt_adapter/src/tt_adapter/runner.py | 12 ++--- .../tt_adapter/src/tt_adapter/utils.py | 3 -- 4 files changed, 39 insertions(+), 27 deletions(-) diff --git a/tools/explorer/test/run_tests.py b/tools/explorer/test/run_tests.py index 6966353cc9..4e1ad20699 100644 --- a/tools/explorer/test/run_tests.py +++ b/tools/explorer/test/run_tests.py @@ -114,6 +114,7 @@ def execute_command(model_path, settings): print(result.json()) assert False + def wait_for_execution_to_finish(timeout): for _ in range(timeout): try: @@ -227,4 +228,3 @@ def test_execute_model_invalid_policy(): {"optimizationPolicy": "Invalid Policy"}, timeout=300, ) - diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 3c1af96f12..19f407a130 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -60,8 +60,10 @@ def parse_attr_name(attr: ir.Attribute) -> List[graph_builder.KeyValue]: def decorator(handler): AttrHandler.ATTR_HANDLERS[attr_name] = handler return handler + return decorator + @AttrHandler.register_handler("tt.device") def parse_tt_device(attr): device = tt.ir.DeviceAttr.maybe_downcast(attr) @@ -528,7 +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( @@ -545,7 +547,6 @@ def get_attributes(self): ) result.extend(output_attrs) - # Add schedule as an attribute result.append( graph_builder.KeyValue(key="schedule", value=str(OpHandler.schedule)) @@ -606,8 +607,14 @@ def build_graph(module, perf_trace=None, memory_trace=None): 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"] + 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() @@ -629,10 +636,12 @@ def build_graph(module, perf_trace=None, memory_trace=None): utils.add_to_dataclass( graph_builder.KeyValue( key="dram_memory", - value=str(memory_data[str(operand_index)]["dram"]), + value=str( + memory_data[str(operand_index)]["dram"] + ), ), - 'display_type', - 'memory' + "display_type", + "memory", ) ) if memory_data: @@ -640,10 +649,12 @@ def build_graph(module, perf_trace=None, memory_trace=None): utils.add_to_dataclass( graph_builder.KeyValue( key="l1_memory", - value=str(memory_data[str(operand_index)]["l1"]), + value=str( + memory_data[str(operand_index)]["l1"] + ), ), - 'display_type', - 'memory' + "display_type", + "memory", ) ) @@ -718,26 +729,30 @@ def build_graph(module, perf_trace=None, memory_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"]), + value=str( + memory_data[str(operand_index)]["dram"] + ), ), - 'display_type', - 'memory' + "display_type", + "memory", ) ) output_attrs.append( utils.add_to_dataclass( graph_builder.KeyValue( key="l1_memory", - value=str(memory_data[str(operand_index)]["l1"]), + value=str( + memory_data[str(operand_index)]["l1"] + ), ), - 'display_type', - 'memory' + "display_type", + "memory", ) ) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index 862d88b695..31d2a3bdd7 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -140,14 +140,14 @@ 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" - ) + 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.") + 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) @@ -313,7 +313,7 @@ def compile_and_run(self, model_path, overrides_string): "perf", flatbuffer_file, f"--artifact-dir={self._explorer_artifacts_dir}", - "--memory" + "--memory", ] ttrt_process = self.run_in_subprocess(ttrt_perf_command) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/utils.py b/tools/explorer/tt_adapter/src/tt_adapter/utils.py index 3e6c774ff1..69fab1c623 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/utils.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/utils.py @@ -115,6 +115,3 @@ def make_editable_kv(kv, editable): obj = asdict(kv) obj["editable"] = editable return to_dataclass(obj, "KeyValue") - - -