From df46c1505b1fa8e5cf66cc139ad12b64ec338426 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Tue, 25 Feb 2025 18:48:36 +0000 Subject: [PATCH 01/17] 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 02/17] 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 03/17] 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 04/17] 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 05/17] 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 06/17] 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 07/17] 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") - - - From 4ee508ae440a97e3d6ab805f017f4a0b0e03413f Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Mon, 3 Mar 2025 16:10:03 -0600 Subject: [PATCH 08/17] Update mlir.py --- .../tt_adapter/src/tt_adapter/mlir.py | 99 +++++++++---------- 1 file changed, 46 insertions(+), 53 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 19f407a130..1e3dd843bf 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -606,15 +606,22 @@ 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] = {} - memory_data[node]["dram"] = ( + loc = parse_loc_string(memory_trace[node]["loc"]) + memory_data[loc] = {} + memory_data[loc]["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_data[loc]["l1"] = ( memory_trace[node]["l1"]["device_0"]["total_bytes_allocated_per_bank"] / memory_trace[node]["l1"]["device_0"]["total_bytes_per_bank"] ) + memory_data[loc]["l1_small"] = ( + memory_trace[node]["l1_small"]["device_0"][ + "total_bytes_allocated_per_bank" + ] + / memory_trace[node]["l1_small"]["device_0"]["total_bytes_per_bank"] + ) module_op = OpHandler(module.operation) module_attrs = module_op.get_attributes() @@ -630,36 +637,48 @@ def build_graph(module, perf_trace=None, memory_trace=None): 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"] - ), + operation = OpHandler(op) + if memory_data and operation.named_location in memory_data: + extra_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="dram_memory", + value=str( + memory_data[operation.named_location]["dram"] ), - "display_type", - "memory", - ) + ), + "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"] - ), + ) + extra_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="l1_memory", + value=str( + memory_data[operation.named_location]["l1"] ), - "display_type", - "memory", - ) + ), + "display_type", + "memory", ) + ) + extra_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="l1_small_memory", + value=str( + memory_data[operation.named_location][ + "l1_small" + ] + ), + ), + "display_type", + "memory", + ) + ) # Create all the nodes and constants in the first pass. - operation = OpHandler(op) graph_node = operation.make_graph_node(extra_attrs) if ( @@ -730,32 +749,6 @@ def build_graph(module, perf_trace=None, memory_trace=None): ), ] - 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( From d13b53c581495f69ece024b94be7d5fd721fc6e3 Mon Sep 17 00:00:00 2001 From: Julia Grim Date: Mon, 3 Mar 2025 16:11:11 -0600 Subject: [PATCH 09/17] Update run_tests.py --- tools/explorer/test/run_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/explorer/test/run_tests.py b/tools/explorer/test/run_tests.py index 4e1ad20699..9b12f70a11 100644 --- a/tools/explorer/test/run_tests.py +++ b/tools/explorer/test/run_tests.py @@ -214,7 +214,7 @@ def test_execute_and_check_perf_data_exists(): def test_execute_and_check_memory_data_exists(): execute_command_and_wait( MNIST_SHARDING_PATH, - {"optimizationPolicy": "DF Sharding"}, + {"optimizationPolicy": "Optimizer Disabled"}, timeout=300, ) result = convert_command_and_assert(MNIST_SHARDING_PATH) From 8b929a7573f5b8bf5103d728d99662f6b6612ed8 Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Tue, 4 Mar 2025 00:33:14 +0000 Subject: [PATCH 10/17] Modified mlir.py to accomodate recent changes --- .../tt_adapter/src/tt_adapter/mlir.py | 451 ++++++++++++------ 1 file changed, 294 insertions(+), 157 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 1e3dd843bf..c4c4034788 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -599,9 +599,11 @@ def build_graph(module, perf_trace=None, memory_trace=None): if perf_trace is not None: for _, row in perf_trace.iterrows(): loc = parse_loc_string(row["LOC"]) - assert loc not in loc_to_perf - if loc: - loc_to_perf[loc] = row["DEVICE FW DURATION [ns]"] + if not loc: + continue + if loc not in loc_to_perf: + loc_to_perf[loc] = 0 + loc_to_perf[loc] += row["DEVICE FW DURATION [ns]"] memory_data = {} if memory_trace is not None: @@ -623,159 +625,17 @@ def build_graph(module, perf_trace=None, memory_trace=None): / memory_trace[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) - - # Add module attributes to the graph as "namespace attributes" - group_node_attrs = {} - group_node_attrs[module_op.get_namespace()] = module_attrs - - for op in module.body.operations: - append_later = [] - for region in op.regions: - for block in region.blocks: - for op in block.operations: - extra_attrs = [] - operation = OpHandler(op) - if memory_data and operation.named_location in memory_data: - extra_attrs.append( - utils.add_to_dataclass( - graph_builder.KeyValue( - key="dram_memory", - value=str( - memory_data[operation.named_location]["dram"] - ), - ), - "display_type", - "memory", - ) - ) - extra_attrs.append( - utils.add_to_dataclass( - graph_builder.KeyValue( - key="l1_memory", - value=str( - memory_data[operation.named_location]["l1"] - ), - ), - "display_type", - "memory", - ) - ) - extra_attrs.append( - utils.add_to_dataclass( - graph_builder.KeyValue( - key="l1_small_memory", - value=str( - memory_data[operation.named_location][ - "l1_small" - ] - ), - ), - "display_type", - "memory", - ) - ) - - # Create all the nodes and constants in the first pass. - graph_node = operation.make_graph_node(extra_attrs) - - if ( - operation.named_location in loc_to_perf - and operation.op.name not in EMPTY_OPS - ): - perf_node_data[operation.id] = node_data_builder.NodeDataResult( - loc_to_perf[operation.named_location] - ) - - if op.name not in FILTERED_OPS and op.name in EMPTY_OPS: - append_later.append(graph_node) - elif op.name not in FILTERED_OPS: - graph.nodes.append(graph_node) - - op_to_graph_node[op] = graph_node - - for operand in op.operands: - if isinstance(operand, ir.Value) and not isinstance( - operand.owner, ir.Operation - ): - # If the owner is not an op, then it is a constant provided from the toplevel FuncOp. - - if operand not in operands_in_graph: - # This is a constant and we need to create a node for it. - operand_node = operation.make_constant_node( - operand.get_name() - ) - graph.nodes.append(operand_node) - op_to_graph_node[operand] = operand_node - operands_in_graph.add(operand) - - # This puts the node at the far right when viewing which is a bit more consistant with it being the last operand. - for node in append_later: - graph.nodes.append(node) - - for op in block.operations: - # Create all the edges in the second pass. - for operand_index, operand in enumerate(op.operands): - if operand.owner == block: - source_node = op_to_graph_node[operand] - else: - source_node = op_to_graph_node[operand.owner] - - target_node = op_to_graph_node[op] - - target_node.incomingEdges.append( - graph_builder.IncomingEdge( - sourceNodeId=source_node.id, - sourceNodeOutputId=str( - output_connections[source_node.id] - ), - targetNodeInputId=str(operand_index), - ) - ) - - output_attrs = [] - if isinstance(operand.type, ir.RankedTensorType): - output_attrs = [ - graph_builder.KeyValue( - key="shape", value=str(operand.type.shape) - ), - graph_builder.KeyValue( - key="dtype", value=str(operand.type.element_type) - ), - graph_builder.KeyValue( - 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( - AttrHandler.parse_attr( - operand.type.encoding.get_named("ttnn_layout") - ) - ) - else: - # Parse as a standard layout - output_attrs.extend( - AttrHandler.parse_attr( - operand.type.encoding.get_named("tt.layout") - ) - ) - - source_node.outputsMetadata.append( - graph_builder.MetadataItem( - id=str(output_connections[source_node.id]), - attrs=[ - graph_builder.KeyValue( - key="__tensor_tag", value=str(target_node.label) - ), - ] - + output_attrs, - ) - ) - output_connections[source_node.id] += 1 + # Process the module hierarchy recursively + process_module( + module, + graph, + op_to_graph_node, + operands_in_graph, + output_connections, + loc_to_perf, + perf_node_data, + memory_data, + ) # Add performance data to the graph color overlay, if it exists overlay_data = None @@ -791,6 +651,283 @@ def build_graph(module, perf_trace=None, memory_trace=None): graphsData={"tt-graph": graph_node_data} ) - graph.groupNodeAttributes = group_node_attrs OpHandler.schedule = 0 return graph, overlay_data + + +def process_module( + module, + graph, + op_to_graph_node, + operands_in_graph, + output_connections, + loc_to_perf, + perf_node_data, + memory_data, +): + """ + Process a module's operations. Only works on top-level module, any nested modules won't have a body so they need to directly call process_operations instead. + + Args: + module: The module to process + graph: The graph being built + op_to_graph_node: Mapping from operations to graph nodes + operands_in_graph: Set of operands already added to graph + output_connections: Tracking of output connections + loc_to_perf: Mapping from locations to performance data + perf_node_data: Performance data for nodes + """ + 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" + if not graph.groupNodeAttributes: + graph.groupNodeAttributes = {} + + # Add this module's namespace attributes + namespace = module_op.get_namespace() + if namespace not in graph.groupNodeAttributes: + graph.groupNodeAttributes[namespace] = module_attrs + else: + # Merge with existing attributes if namespace already exists + graph.groupNodeAttributes[namespace].update(module_attrs) + + # Process operations in this module + process_operations( + module.body.operations, + graph, + op_to_graph_node, + operands_in_graph, + output_connections, + loc_to_perf, + perf_node_data, + memory_data, + ) + + +def process_operations( + operations, + graph, + op_to_graph_node, + operands_in_graph, + output_connections, + loc_to_perf, + perf_node_data, + memory_data, +): + """ + Recursively process a list of operations, including handling nested modules. + + Args: + operations: List of operations to process + graph: The graph being built + op_to_graph_node: Mapping from operations to graph nodes + operands_in_graph: Set of operands already added to graph + output_connections: Tracking of output connections + loc_to_perf: Mapping from locations to performance data + perf_node_data: Performance data for nodes + """ + append_later = [] + + # First pass: create all nodes and constants + for op in operations: + # Check if this operation is a nested module + if is_module_op(op): + # Process the nested module's ops recursively + process_operations( + op.regions[0].blocks[0], + graph, + op_to_graph_node, + operands_in_graph, + output_connections, + loc_to_perf, + perf_node_data, + memory_data, + ) + continue + + # Process regions in the operation + for region in op.regions: + for block in region.blocks: + # Recursively process operations in this block + process_operations( + block.operations, + graph, + op_to_graph_node, + operands_in_graph, + output_connections, + loc_to_perf, + perf_node_data, + memory_data, + ) + + # Create graph node for this operation + operation = OpHandler(op) + + if ( + operation.named_location in loc_to_perf + and operation.op.name not in EMPTY_OPS + ): + perf_node_data[operation.id] = node_data_builder.NodeDataResult( + loc_to_perf[operation.named_location] + ) + extra_attrs = [] + if memory_data and operation.named_location in memory_data: + extra_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="dram_memory", + value=str( + memory_data[operation.named_location]["dram"] + ), + ), + "display_type", + "memory", + ) + ) + extra_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="l1_memory", + value=str( + memory_data[operation.named_location]["l1"] + ), + ), + "display_type", + "memory", + ) + ) + extra_attrs.append( + utils.add_to_dataclass( + graph_builder.KeyValue( + key="l1_small_memory", + value=str( + memory_data[operation.named_location][ + "l1_small" + ] + ), + ), + "display_type", + "memory", + ) + ) + if not op.name == "func.func": + graph_node = operation.make_graph_node(extra_attrs) + + if op.name not in FILTERED_OPS and op.name in EMPTY_OPS: + append_later.append(graph_node) + elif op.name not in FILTERED_OPS: + graph.nodes.append(graph_node) + + op_to_graph_node[op] = graph_node + + # Process operands + for operand in op.operands: + if isinstance(operand, ir.Value) and not isinstance( + operand.owner, ir.Operation + ): + # If the owner is not an op, then it is a constant provided from the toplevel FuncOp. + if operand not in operands_in_graph: + # This is a constant and we need to create a node for it. + operand_node = operation.make_constant_node(operand.get_name()) + graph.nodes.append(operand_node) + op_to_graph_node[operand] = operand_node + operands_in_graph.add(operand) + + # Add the nodes that should be appended later + for node in append_later: + graph.nodes.append(node) + + # Second pass: create all edges + for op in operations: + # Skip module + func operations as they've been processed recursively + if is_module_op(op) or op.name == "func.func": + continue + + # Process regions in the operation + for region in op.regions: + for block in region.blocks: + create_edges_for_block(block, op_to_graph_node, output_connections) + + +def create_edges_for_block(block, op_to_graph_node, output_connections): + """ + Create edges between nodes for operations in a block. + + Args: + block: The block containing operations + op_to_graph_node: Mapping from operations to graph nodes + output_connections: Tracking of output connections + """ + for op in block.operations: + # Skip module operations as they've been processed recursively + if is_module_op(op): + continue + + # Create edges for this operation + for operand_index, operand in enumerate(op.operands): + if operand.owner == block: + source_node = op_to_graph_node[operand] + else: + source_node = op_to_graph_node[operand.owner] + + target_node = op_to_graph_node[op] + + target_node.incomingEdges.append( + graph_builder.IncomingEdge( + sourceNodeId=source_node.id, + sourceNodeOutputId=str(output_connections[source_node.id]), + targetNodeInputId=str(operand_index), + ) + ) + + output_attrs = [] + if isinstance(operand.type, ir.RankedTensorType): + output_attrs = [ + graph_builder.KeyValue(key="shape", value=str(operand.type.shape)), + graph_builder.KeyValue( + key="dtype", value=str(operand.type.element_type) + ), + graph_builder.KeyValue(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( + AttrHandler.parse_attr( + operand.type.encoding.get_named("ttnn_layout") + ) + ) + else: + # Parse as a standard layout + output_attrs.extend( + AttrHandler.parse_attr( + operand.type.encoding.get_named("tt.layout") + ) + ) + source_node.outputsMetadata.append( + graph_builder.MetadataItem( + id=str(output_connections[source_node.id]), + attrs=[ + graph_builder.KeyValue( + key="__tensor_tag", value=str(target_node.label) + ), + ] + + output_attrs, + ) + ) + output_connections[source_node.id] += 1 + + +def is_module_op(op): + """ + Check if an operation represents a module. + + Args: + op: The operation to check + + Returns: + bool: True if the operation is a module, False otherwise + """ + # Check for tt.device_module or builtin.module operations + return op.name == "tt.device_module" or op.name == "builtin.module" From 34b18f53ff3b27166dad41c30a79da9cc3362783 Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Tue, 4 Mar 2025 18:20:44 +0000 Subject: [PATCH 11/17] Updated MODEL_EXPLORER_VERSION --- tools/explorer/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/explorer/CMakeLists.txt b/tools/explorer/CMakeLists.txt index 280f2d8bb4..50afec7b7c 100644 --- a/tools/explorer/CMakeLists.txt +++ b/tools/explorer/CMakeLists.txt @@ -3,7 +3,7 @@ include(ExternalProject) set(TT_EXPLORER_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/run.py) set(TTMLIR_BUILD_BIN_DIR ${TTMLIR_BINARY_DIR}/bin) -set(MODEL_EXPLORER_VERSION "5be0dda742f81a7dcc8aeec90afd40198cc686a2") +set(MODEL_EXPLORER_VERSION "77c7bda7d56fc2a1c80aa21eeb44e089d27c1330") ExternalProject_Add( model-explorer PREFIX ${CMAKE_CURRENT_SOURCE_DIR}/model-explorer From 3c5465d5edde8e0806a71565c3a8ffbc37cbc7b6 Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Tue, 4 Mar 2025 20:33:42 +0000 Subject: [PATCH 12/17] update --- tools/explorer/tt_adapter/src/tt_adapter/mlir.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index c4c4034788..13f980cb00 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -778,9 +778,7 @@ def process_operations( utils.add_to_dataclass( graph_builder.KeyValue( key="dram_memory", - value=str( - memory_data[operation.named_location]["dram"] - ), + value=str(memory_data[operation.named_location]["dram"]), ), "display_type", "memory", @@ -790,9 +788,7 @@ def process_operations( utils.add_to_dataclass( graph_builder.KeyValue( key="l1_memory", - value=str( - memory_data[operation.named_location]["l1"] - ), + value=str(memory_data[operation.named_location]["l1"]), ), "display_type", "memory", @@ -802,11 +798,7 @@ def process_operations( utils.add_to_dataclass( graph_builder.KeyValue( key="l1_small_memory", - value=str( - memory_data[operation.named_location][ - "l1_small" - ] - ), + value=str(memory_data[operation.named_location]["l1_small"]), ), "display_type", "memory", From 2bf17f183224f0932dfe2885668a786311729004 Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Thu, 6 Mar 2025 14:59:07 +0000 Subject: [PATCH 13/17] Rounded memory data, changed run test --- .../tt_adapter/src/tt_adapter/mlir.py | 15 +- .../tt_adapter/src/tt_adapter/runner.py | 555 +++++++----------- 2 files changed, 222 insertions(+), 348 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 13f980cb00..913a1ccc09 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -610,19 +610,22 @@ def build_graph(module, perf_trace=None, memory_trace=None): for node in memory_trace: loc = parse_loc_string(memory_trace[node]["loc"]) memory_data[loc] = {} - memory_data[loc]["dram"] = ( + memory_data[loc]["dram"] = round( memory_trace[node]["dram"]["device_0"]["total_bytes_allocated_per_bank"] - / memory_trace[node]["dram"]["device_0"]["total_bytes_per_bank"] + / memory_trace[node]["dram"]["device_0"]["total_bytes_per_bank"], + 4, ) - memory_data[loc]["l1"] = ( + memory_data[loc]["l1"] = round( memory_trace[node]["l1"]["device_0"]["total_bytes_allocated_per_bank"] - / memory_trace[node]["l1"]["device_0"]["total_bytes_per_bank"] + / memory_trace[node]["l1"]["device_0"]["total_bytes_per_bank"], + 4, ) - memory_data[loc]["l1_small"] = ( + memory_data[loc]["l1_small"] = round( memory_trace[node]["l1_small"]["device_0"][ "total_bytes_allocated_per_bank" ] - / memory_trace[node]["l1_small"]["device_0"]["total_bytes_per_bank"] + / memory_trace[node]["l1_small"]["device_0"]["total_bytes_per_bank"], + 4, ) # Process the module hierarchy recursively diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index 31d2a3bdd7..9dafc63056 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -1,360 +1,231 @@ # SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC # # SPDX-License-Identifier: Apache-2.0 -import subprocess -import os -import logging - -# TODO(odjuricic) Cleaner to implement ttrt --quiet flag. -# os.environ["TTRT_LOGGER_LEVEL"] = "ERROR" -from ttrt import API as ttrt -from ttmlir import passes -from . import utils, mlir -import pandas as pd -import threading -import queue -import json - - -class ExplorerRunException(Exception): - pass - - -class ModelState: - """ - After a model is compiled and executed we keep track of all additional data that was created. - """ - - # Path to the compiled TTNN IR file. - optimized_model_path = None - # Path to the output directory where ttrt dumps all model files (perf trace, memory state, etc) - model_output_dir = None - # Overrides, changes that the user made to op configurations. - overrides = None - - -class ModelRunner: - """ - ModelRunner is a singleton class used for compilation and running of models. Ensuring only one can be run at a time. - This is necessary because the adaptor class is reinitialized on every request from the frontend, so it cannot keep state. - """ - - # Global static runner state. Initialized once. - _instance = None - _explorer_artifacts_dir = None - _build_dir = None - - # Singleton runner state. Initialized on every run. - runner_thread = None - runner_error = None - log_queue = queue.Queue() - # progress should be a number between 0 and 100. - progress = 0 - - # State for models that have been executed. - # Contains a mapping from model path to ModelState. - model_state = dict() - - def __new__(cls, *args, **kwargs): - if not cls._instance: - logging.info("Creating a new ModelRunner instance.") - cls._instance = super(ModelRunner, cls).__new__(cls, *args, **kwargs) - cls._instance.initialize() - return cls._instance - - def initialize(self): - # Initialize machine to generate SystemDesc and load up functionality to begin - logging.info("Running ttrt initialization.") - ttrt.initialize_apis() - - if "TT_MLIR_HOME" not in os.environ: - raise RuntimeError("TT_MLIR_HOME not set. Did you run source env/activate?") - - # TODO(odjuricic, #1200) ttrt perf breaks if artifacts dir is changed from default. - # self._explorer_artifacts_dir = os.environ['TT_MLIR_HOME'] + '/explorer-artifacts' - self._explorer_artifacts_dir = os.environ["TT_MLIR_HOME"] + "/ttrt-artifacts" - self._build_dir = os.environ["TT_MLIR_HOME"] + "/build" - os.makedirs(self._explorer_artifacts_dir, exist_ok=True) - - # Save the system descriptor. - ttrt.Query( - args={ - "--save-artifacts": True, - "--artifact-dir": self._explorer_artifacts_dir, - "--quiet": True, - } - )() - - logging.info("ModelRunner initialized.") - - def get_optimized_model_path(self, model_path): - if model_path in self.model_state: - return self.model_state[model_path].optimized_model_path - return None - - def get_output_dir(self, model_path): - return self.model_state[model_path].model_output_dir - - def get_overrides(self, model_path): - if model_path in self.model_state: - return self.model_state[model_path].overrides - return None - - def get_error(self): - return self.runner_error - - def get_progress(self): - return self.progress - - def get_artifacts_dir(self): - return self._explorer_artifacts_dir - - def is_busy(self): - return self.runner_thread and self.runner_thread.is_alive() - - def get_logs(self): - logs = [] - while not self.log_queue.empty(): - logs.append(self.log_queue.get()) - return "\n".join(logs) - - def reset_state(self, model_path): - assert not self.is_busy() - self.runner_thread = None - self.runner_error = None - self.progress = 0 - self.log_queue.queue.clear() - - if model_path in self.model_state: - del self.model_state[model_path] - - def log(self, message, severity=logging.info): - severity(message) - self.log_queue.put(message) - - def get_perf_trace(self, model_path): - op_perf_file = ( - f"{self.model_state[model_path].model_output_dir}/perf/ops_perf_results.csv" - ) - if not os.path.exists(op_perf_file): - 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." - ) +import model_explorer - 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") - - process = subprocess.Popen( - command, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - - for line in process.stdout: - self.log(line.strip()) - - process.stdout.close() - process.wait() - - return process +import requests +import time +import multiprocessing +import pytest +import glob +import os - def compile_and_run_wrapper(self, model_path, overrides_string): +HOST = "localhost" +PORT = 8002 +COMMAND_URL = "http://" + HOST + ":" + str(PORT) + "/apipost/v1/send_command" +TEST_LOAD_MODEL_PATHS = [ + "test/ttmlir/Dialect/TTNN/optimizer/mnist_sharding.mlir", + "test/ttmlir/Explorer/**/*.mlir", + "test/ttmlir/Silicon/TTNN/**/*.mlir", +] +MNIST_SHARDING_PATH = "test/ttmlir/Silicon/TTNN/n150/optimizer/mnist_sharding.mlir" +TEST_EXECUTE_MODEL_PATHS = [ + MNIST_SHARDING_PATH, +] + +if "TT_EXPLORER_GENERATED_TEST_DIR" in os.environ: + TEST_LOAD_MODEL_PATHS.append( + os.environ["TT_EXPLORER_GENERATED_TEST_DIR"] + "/**/*.mlir" + ) + + +def get_test_files(paths): + files = [] + for path in paths: + files.extend(glob.glob(path, recursive=True)) + return files + + +def execute_command(model_path, settings): + cmd = { + "extensionId": "tt_adapter", + "cmdId": "execute", + "modelPath": model_path, + "deleteAfterConversion": False, + "settings": settings, + } + + result = requests.post(COMMAND_URL, json=cmd) + assert result.ok + if "error" in result.json(): + print(result.json()) + assert False + + +@pytest.fixture(scope="function", autouse=True) +def start_server(request): + """Start the model explorer server before running tests and stop it after.""" + server_thread = multiprocessing.Process( + target=model_explorer.visualize_from_config, + kwargs={ + "extensions": ["tt_adapter"], + "host": HOST, + "port": PORT, + "no_open_in_browser": True, + }, + ) + server_thread.start() + + # Wait for the server to start + for _ in range(200): # Try for up to 20 seconds try: - self.compile_and_run(model_path, overrides_string) - except ExplorerRunException as e: - self.runner_error = str(e) - raise e - except Exception as e: - self.runner_error = "An unexpected error occurred: " + str(e) - self.log(self.runner_error, severity=logging.error) - raise e + response = requests.get(f"http://{HOST}:{PORT}/check_health", timeout=1) + if response.status_code == 200: + print("Explorer server started") + break + except requests.ConnectionError: + pass finally: - self.progress = 100 + time.sleep(0.1) + else: + raise RuntimeError("Server did not start within the expected time") - def compile_and_run(self, model_path, overrides_string): - FLATBUFFER = False - if model_path.endswith(".ttnn"): - # This is being run from a Flatbuffer. Need To Render TTIR from Flatbuffer - FLATBUFFER = True - # Write the TTIR from this file into a temporary file to run through the compiler - ttir_module_str = utils.parse_flatbuffer_file( - model_path, at_pass="PRE-PIPELINE" - ) - ttir_module_path = f"{model_path}_ttir.mlir" - with open(ttir_module_path, "w+") as temp_module: - temp_module.write(ttir_module_str) + # Terminate the server and wait for it to finish. + def server_shutdown(): + server_thread.terminate() + server_thread.join() - model_name = os.path.basename(model_path) - flatbuffer_file = model_name + ".ttnn" - state = self.model_state[model_path] + request.addfinalizer(server_shutdown) - state.model_output_dir = self._explorer_artifacts_dir + "/" + flatbuffer_file - if os.path.exists(state.model_output_dir): - self.log("Removing artifacts of previous run.") - os.system(f"rm -rf {state.model_output_dir}") +def get_test_files(paths): + files = [] + for path in paths: + files.extend(glob.glob(path)) + return files - os.makedirs(state.model_output_dir) - # Copy the model to the run directory. - os.system(f"cp {model_path} {state.model_output_dir}") - self.progress = 10 +def send_command(command, model_path, settings={}): + cmd = { + "extensionId": "tt_adapter", + "cmdId": command, + "modelPath": model_path, + "deleteAfterConversion": False, + "settings": settings, + } - ############################### Compile ################################## + return requests.post(COMMAND_URL, json=cmd, timeout=10) - ttnn_ir_file = ( - f"{state.model_output_dir}/{model_name.replace('.mlir', '_ttnn.mlir')}" - ) - if FLATBUFFER: - ttnn_ir_file = f"{state.model_output_dir}/{model_name}.mlir" - - compile_command = [ - f"{self._build_dir}/bin/ttmlir-opt", - f"--ttir-to-ttnn-backend-pipeline={overrides_string}", - model_path if not FLATBUFFER else ttir_module_path, - "-o", - ttnn_ir_file, - "--mlir-print-debuginfo", - ] - - self.log("Running compile TTIR to TTNN Backend Pipeline") - self.log("With options: " + overrides_string) - - compile_process = self.run_in_subprocess(compile_command) - if compile_process.returncode != 0: - error = "Error running compile TTIR to TTNN Backend Pipeline" - self.log(error, severity=logging.error) - raise ExplorerRunException(error) - self.progress = 20 - - ############################## Translate ################################# - - # Need this flatbuffer file to inherit the golden data - if FLATBUFFER: - golden_map = utils.golden_map_from_flatbuffer(model_path) - # need to parse this golden_map - kept_alive_data_arrs = [] - rendered_golden_map = {} - - for entry in golden_map: - data = entry["value"] - # Turn this into a Torch Tensor to easily format it for the GoldenMap - # data is a uint8_t buffer type that contains the data in the format of dtype - # We will need to render this data as a buffer reference for the GoldenTensor constructor - import array - - # B is unsigned char in the array library - # This will parse the data as a 1D Buffer of uint8_t, exactly the pointer type expected - data_arr = array.array("B", data["data"]) - - rendered_golden_map[entry["key"]] = passes.GoldenTensor( - data["name"], - data["shape"], - data["stride"], - passes.lookup_dtype(data["dtype"]), - data_arr.buffer_info()[0], - data_arr.buffer_info()[1], - ) - - # Get module from file - with open(ttnn_ir_file, "r") as f: - ttnn_module = utils.parse_mlir_str(f.read()) - - self.log("Running TTNN to Flatbuffer File") - - # Run through pybound translation so we can pass golden_map - try: - if golden_map: - passes.ttnn_to_flatbuffer_file( - ttnn_module, flatbuffer_file, rendered_golden_map - ) - else: - passes.ttnn_to_flatbuffer_file(ttnn_module, flatbuffer_file) - except: - self.log("Error while running TTNN to Flatbuffer File") - raise ExplorerRunException() - else: - # Translate to Flatbuffer normally. - to_flatbuffer_command = [ - f"{self._build_dir}/bin/ttmlir-translate", - "--ttnn-to-flatbuffer", - ttnn_ir_file, - "-o", - flatbuffer_file, - ] - - translate_process = self.run_in_subprocess(to_flatbuffer_command) - if translate_process.returncode != 0: - error = "Error while running TTNN to Flatbuffer File" - self.log(error, severtity=logging.error) - raise ExplorerRunException(error) - - self.progress = 30 - - ############################## TTRT Perf ################################# - - ttrt_perf_command = [ - "ttrt", - "perf", - flatbuffer_file, - f"--artifact-dir={self._explorer_artifacts_dir}", - "--memory", - ] - - ttrt_process = self.run_in_subprocess(ttrt_perf_command) - - if ttrt_process.returncode != 0: - error = "Error while running TTRT perf" - self.log(error, severity=logging.error) - raise ExplorerRunException(error) - - perf = self.get_perf_trace(model_path) - columns = [ - "GLOBAL CALL COUNT", - "OP CODE", - "DEVICE FW DURATION [ns]", - "CORE COUNT", - "OUTPUT_0_MEMORY", - "LOC", - ] - perf = perf[columns] - logging.info(perf) - - logging.info( - "Total device duration: ", perf["DEVICE FW DURATION [ns]"].sum(), "ns" - ) +def execute_command(model_path, settings): + result = send_command("execute", model_path, settings) + assert result.ok + if "error" in result.json(): + print(result.json()) + assert False - # TTNN_IR_FILE from flatbuffer is still relevant since model_path is the FB with golden data and it will rented optimized_model_path instead - state.optimized_model_path = ttnn_ir_file - self.progress = 100 - - def run(self, model_path, compile_options, overrides): - # Check if a run is already in progress - if self.is_busy(): - raise RuntimeError( - "A model is already being processed. Please wait for it to finish." - ) - self.reset_state(model_path) - self.model_state[model_path] = ModelState() - self.model_state[model_path].overrides = overrides - - # Start compile and run in a new thread - self.runner_thread = threading.Thread( - target=self.compile_and_run_wrapper, args=(model_path, compile_options) + +def wait_for_execution_to_finish(timeout): + for _ in range(timeout): + try: + response = send_command("status_check", "") + if response.status_code == 200 and response.json().get("graphs")[0].get( + "isDone" + ): + return response.json() + except requests.RequestException as e: + print(f"Request failed: {e}") + raise Exception("Status check request failed") + time.sleep(1) + raise RuntimeError(f"Execution did not finish within {timeout} seconds") + + +def execute_command_and_wait(model_path, settings, timeout): + execute_command(model_path, settings) + adapter_response = wait_for_execution_to_finish(timeout) + assert "graphs" in adapter_response + assert len(adapter_response["graphs"]) == 1 + response = adapter_response["graphs"][0] + assert response["isDone"] + assert response["error"] is None + + +def convert_command_and_assert(model_path): + result = send_command("convert", model_path) + assert result.ok + if "error" in result.json(): + print(result.json()) + assert False + return result.json() + + +@pytest.mark.parametrize("model_path", get_test_files(TEST_LOAD_MODEL_PATHS)) +def test_load_model(model_path): + convert_command_and_assert(model_path) + + +@pytest.mark.parametrize("model_path", get_test_files(TEST_EXECUTE_MODEL_PATHS)) +def test_execute_model(model_path): + execute_command_and_wait( + model_path, {"optimizationPolicy": "DF Sharding"}, timeout=300 + ) + convert_command_and_assert(model_path) + + +def test_execute_mnist_l1_interleaved(): + execute_command_and_wait( + MNIST_SHARDING_PATH, + {"optimizationPolicy": "Greedy L1 Interleaved"}, + timeout=300, + ) + convert_command_and_assert(MNIST_SHARDING_PATH) + + +def test_execute_mnist_optimizer_disabled(): + execute_command_and_wait( + MNIST_SHARDING_PATH, + {"optimizationPolicy": "Optimizer Disabled"}, + timeout=300, + ) + convert_command_and_assert(MNIST_SHARDING_PATH) + + +def test_execute_mnist_with_overrides(): + overrides = { + 'loc("matmul_1"("MNISTLinear":4294967295:10))__17': { + "named_location": "matmul_1", + "attributes": [ + {"key": "data_type", "value": "f32"}, + {"key": "memory_layout", "value": "tile"}, + {"key": "buffer_type", "value": "dram"}, + {"key": "tensor_memory_layout", "value": "interleaved"}, + {"key": "grid_shape", "value": "[8,8]"}, + ], + } + } + execute_command_and_wait( + MNIST_SHARDING_PATH, + {"optimizationPolicy": "DF Sharding", "overrides": overrides}, + timeout=300, + ) + convert_command_and_assert(MNIST_SHARDING_PATH) + + +def test_execute_and_check_perf_data_exists(): + execute_command_and_wait( + MNIST_SHARDING_PATH, + {"optimizationPolicy": "DF Sharding"}, + timeout=300, + ) + result = convert_command_and_assert(MNIST_SHARDING_PATH) + assert "perf_data" in result["graphs"][0] + + +def test_execute_and_check_memory_data_exists(): + execute_command_and_wait( + MNIST_SHARDING_PATH, + {"optimizationPolicy": "Optimizer Disabled"}, + timeout=300, + ) + result = convert_command_and_assert(MNIST_SHARDING_PATH) + assertion = "display_type" in str(result["graphs"][0]["nodes"]) + assert assertion, str(result["graphs"][0]["nodes"]) + + +def test_execute_model_invalid_policy(): + with pytest.raises(AssertionError): + execute_command_and_wait( + TEST_EXECUTE_MODEL_PATHS[0], + {"optimizationPolicy": "Invalid Policy"}, + timeout=300, ) - self.runner_thread.start() From e789059a73ae05bd027e66746a48cea3970aaa10 Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Thu, 6 Mar 2025 15:02:26 +0000 Subject: [PATCH 14/17] redoing the last commit --- tools/explorer/test/run_tests.py | 3 +- .../tt_adapter/src/tt_adapter/runner.py | 555 +++++++++++------- 2 files changed, 344 insertions(+), 214 deletions(-) diff --git a/tools/explorer/test/run_tests.py b/tools/explorer/test/run_tests.py index 9b12f70a11..9dafc63056 100644 --- a/tools/explorer/test/run_tests.py +++ b/tools/explorer/test/run_tests.py @@ -218,7 +218,8 @@ def test_execute_and_check_memory_data_exists(): timeout=300, ) result = convert_command_and_assert(MNIST_SHARDING_PATH) - assert "display_type" in str(result["graphs"]) + assertion = "display_type" in str(result["graphs"][0]["nodes"]) + assert assertion, str(result["graphs"][0]["nodes"]) def test_execute_model_invalid_policy(): diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index 9dafc63056..31d2a3bdd7 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -1,231 +1,360 @@ # SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC # # SPDX-License-Identifier: Apache-2.0 -import model_explorer - -import requests -import time -import multiprocessing -import pytest -import glob +import subprocess import os +import logging + +# TODO(odjuricic) Cleaner to implement ttrt --quiet flag. +# os.environ["TTRT_LOGGER_LEVEL"] = "ERROR" +from ttrt import API as ttrt +from ttmlir import passes +from . import utils, mlir +import pandas as pd +import threading +import queue +import json + + +class ExplorerRunException(Exception): + pass + + +class ModelState: + """ + After a model is compiled and executed we keep track of all additional data that was created. + """ + + # Path to the compiled TTNN IR file. + optimized_model_path = None + # Path to the output directory where ttrt dumps all model files (perf trace, memory state, etc) + model_output_dir = None + # Overrides, changes that the user made to op configurations. + overrides = None + + +class ModelRunner: + """ + ModelRunner is a singleton class used for compilation and running of models. Ensuring only one can be run at a time. + This is necessary because the adaptor class is reinitialized on every request from the frontend, so it cannot keep state. + """ + + # Global static runner state. Initialized once. + _instance = None + _explorer_artifacts_dir = None + _build_dir = None + + # Singleton runner state. Initialized on every run. + runner_thread = None + runner_error = None + log_queue = queue.Queue() + # progress should be a number between 0 and 100. + progress = 0 + + # State for models that have been executed. + # Contains a mapping from model path to ModelState. + model_state = dict() + + def __new__(cls, *args, **kwargs): + if not cls._instance: + logging.info("Creating a new ModelRunner instance.") + cls._instance = super(ModelRunner, cls).__new__(cls, *args, **kwargs) + cls._instance.initialize() + return cls._instance + + def initialize(self): + # Initialize machine to generate SystemDesc and load up functionality to begin + logging.info("Running ttrt initialization.") + ttrt.initialize_apis() + + if "TT_MLIR_HOME" not in os.environ: + raise RuntimeError("TT_MLIR_HOME not set. Did you run source env/activate?") + + # TODO(odjuricic, #1200) ttrt perf breaks if artifacts dir is changed from default. + # self._explorer_artifacts_dir = os.environ['TT_MLIR_HOME'] + '/explorer-artifacts' + self._explorer_artifacts_dir = os.environ["TT_MLIR_HOME"] + "/ttrt-artifacts" + self._build_dir = os.environ["TT_MLIR_HOME"] + "/build" + os.makedirs(self._explorer_artifacts_dir, exist_ok=True) + + # Save the system descriptor. + ttrt.Query( + args={ + "--save-artifacts": True, + "--artifact-dir": self._explorer_artifacts_dir, + "--quiet": True, + } + )() + + logging.info("ModelRunner initialized.") + + def get_optimized_model_path(self, model_path): + if model_path in self.model_state: + return self.model_state[model_path].optimized_model_path + return None + + def get_output_dir(self, model_path): + return self.model_state[model_path].model_output_dir + + def get_overrides(self, model_path): + if model_path in self.model_state: + return self.model_state[model_path].overrides + return None + + def get_error(self): + return self.runner_error + + def get_progress(self): + return self.progress + + def get_artifacts_dir(self): + return self._explorer_artifacts_dir + + def is_busy(self): + return self.runner_thread and self.runner_thread.is_alive() + + def get_logs(self): + logs = [] + while not self.log_queue.empty(): + logs.append(self.log_queue.get()) + return "\n".join(logs) + + def reset_state(self, model_path): + assert not self.is_busy() + self.runner_thread = None + self.runner_error = None + self.progress = 0 + self.log_queue.queue.clear() + + if model_path in self.model_state: + del self.model_state[model_path] + + def log(self, message, severity=logging.info): + severity(message) + self.log_queue.put(message) + + def get_perf_trace(self, model_path): + op_perf_file = ( + f"{self.model_state[model_path].model_output_dir}/perf/ops_perf_results.csv" + ) + if not os.path.exists(op_perf_file): + 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." + ) -HOST = "localhost" -PORT = 8002 -COMMAND_URL = "http://" + HOST + ":" + str(PORT) + "/apipost/v1/send_command" -TEST_LOAD_MODEL_PATHS = [ - "test/ttmlir/Dialect/TTNN/optimizer/mnist_sharding.mlir", - "test/ttmlir/Explorer/**/*.mlir", - "test/ttmlir/Silicon/TTNN/**/*.mlir", -] -MNIST_SHARDING_PATH = "test/ttmlir/Silicon/TTNN/n150/optimizer/mnist_sharding.mlir" -TEST_EXECUTE_MODEL_PATHS = [ - MNIST_SHARDING_PATH, -] - -if "TT_EXPLORER_GENERATED_TEST_DIR" in os.environ: - TEST_LOAD_MODEL_PATHS.append( - os.environ["TT_EXPLORER_GENERATED_TEST_DIR"] + "/**/*.mlir" - ) - - -def get_test_files(paths): - files = [] - for path in paths: - files.extend(glob.glob(path, recursive=True)) - return files - - -def execute_command(model_path, settings): - cmd = { - "extensionId": "tt_adapter", - "cmdId": "execute", - "modelPath": model_path, - "deleteAfterConversion": False, - "settings": settings, - } - - result = requests.post(COMMAND_URL, json=cmd) - assert result.ok - if "error" in result.json(): - print(result.json()) - assert False - - -@pytest.fixture(scope="function", autouse=True) -def start_server(request): - """Start the model explorer server before running tests and stop it after.""" - server_thread = multiprocessing.Process( - target=model_explorer.visualize_from_config, - kwargs={ - "extensions": ["tt_adapter"], - "host": HOST, - "port": PORT, - "no_open_in_browser": True, - }, - ) - server_thread.start() - - # Wait for the server to start - for _ in range(200): # Try for up to 20 seconds + 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") + + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + for line in process.stdout: + self.log(line.strip()) + + process.stdout.close() + process.wait() + + return process + + def compile_and_run_wrapper(self, model_path, overrides_string): try: - response = requests.get(f"http://{HOST}:{PORT}/check_health", timeout=1) - if response.status_code == 200: - print("Explorer server started") - break - except requests.ConnectionError: - pass + self.compile_and_run(model_path, overrides_string) + except ExplorerRunException as e: + self.runner_error = str(e) + raise e + except Exception as e: + self.runner_error = "An unexpected error occurred: " + str(e) + self.log(self.runner_error, severity=logging.error) + raise e finally: - time.sleep(0.1) - else: - raise RuntimeError("Server did not start within the expected time") + self.progress = 100 - # Terminate the server and wait for it to finish. - def server_shutdown(): - server_thread.terminate() - server_thread.join() + def compile_and_run(self, model_path, overrides_string): + FLATBUFFER = False + if model_path.endswith(".ttnn"): + # This is being run from a Flatbuffer. Need To Render TTIR from Flatbuffer + FLATBUFFER = True + # Write the TTIR from this file into a temporary file to run through the compiler + ttir_module_str = utils.parse_flatbuffer_file( + model_path, at_pass="PRE-PIPELINE" + ) + ttir_module_path = f"{model_path}_ttir.mlir" + with open(ttir_module_path, "w+") as temp_module: + temp_module.write(ttir_module_str) - request.addfinalizer(server_shutdown) + model_name = os.path.basename(model_path) + flatbuffer_file = model_name + ".ttnn" + state = self.model_state[model_path] + state.model_output_dir = self._explorer_artifacts_dir + "/" + flatbuffer_file -def get_test_files(paths): - files = [] - for path in paths: - files.extend(glob.glob(path)) - return files + if os.path.exists(state.model_output_dir): + self.log("Removing artifacts of previous run.") + os.system(f"rm -rf {state.model_output_dir}") + os.makedirs(state.model_output_dir) + # Copy the model to the run directory. + os.system(f"cp {model_path} {state.model_output_dir}") -def send_command(command, model_path, settings={}): - cmd = { - "extensionId": "tt_adapter", - "cmdId": command, - "modelPath": model_path, - "deleteAfterConversion": False, - "settings": settings, - } + self.progress = 10 - return requests.post(COMMAND_URL, json=cmd, timeout=10) + ############################### Compile ################################## + ttnn_ir_file = ( + f"{state.model_output_dir}/{model_name.replace('.mlir', '_ttnn.mlir')}" + ) -def execute_command(model_path, settings): - result = send_command("execute", model_path, settings) - assert result.ok - if "error" in result.json(): - print(result.json()) - assert False - + if FLATBUFFER: + ttnn_ir_file = f"{state.model_output_dir}/{model_name}.mlir" + + compile_command = [ + f"{self._build_dir}/bin/ttmlir-opt", + f"--ttir-to-ttnn-backend-pipeline={overrides_string}", + model_path if not FLATBUFFER else ttir_module_path, + "-o", + ttnn_ir_file, + "--mlir-print-debuginfo", + ] + + self.log("Running compile TTIR to TTNN Backend Pipeline") + self.log("With options: " + overrides_string) + + compile_process = self.run_in_subprocess(compile_command) + if compile_process.returncode != 0: + error = "Error running compile TTIR to TTNN Backend Pipeline" + self.log(error, severity=logging.error) + raise ExplorerRunException(error) + self.progress = 20 + + ############################## Translate ################################# + + # Need this flatbuffer file to inherit the golden data + if FLATBUFFER: + golden_map = utils.golden_map_from_flatbuffer(model_path) + # need to parse this golden_map + kept_alive_data_arrs = [] + rendered_golden_map = {} + + for entry in golden_map: + data = entry["value"] + # Turn this into a Torch Tensor to easily format it for the GoldenMap + # data is a uint8_t buffer type that contains the data in the format of dtype + # We will need to render this data as a buffer reference for the GoldenTensor constructor + import array + + # B is unsigned char in the array library + # This will parse the data as a 1D Buffer of uint8_t, exactly the pointer type expected + data_arr = array.array("B", data["data"]) + + rendered_golden_map[entry["key"]] = passes.GoldenTensor( + data["name"], + data["shape"], + data["stride"], + passes.lookup_dtype(data["dtype"]), + data_arr.buffer_info()[0], + data_arr.buffer_info()[1], + ) + + # Get module from file + with open(ttnn_ir_file, "r") as f: + ttnn_module = utils.parse_mlir_str(f.read()) + + self.log("Running TTNN to Flatbuffer File") + + # Run through pybound translation so we can pass golden_map + try: + if golden_map: + passes.ttnn_to_flatbuffer_file( + ttnn_module, flatbuffer_file, rendered_golden_map + ) + else: + passes.ttnn_to_flatbuffer_file(ttnn_module, flatbuffer_file) + except: + self.log("Error while running TTNN to Flatbuffer File") + raise ExplorerRunException() + else: + # Translate to Flatbuffer normally. + to_flatbuffer_command = [ + f"{self._build_dir}/bin/ttmlir-translate", + "--ttnn-to-flatbuffer", + ttnn_ir_file, + "-o", + flatbuffer_file, + ] + + translate_process = self.run_in_subprocess(to_flatbuffer_command) + if translate_process.returncode != 0: + error = "Error while running TTNN to Flatbuffer File" + self.log(error, severtity=logging.error) + raise ExplorerRunException(error) + + self.progress = 30 + + ############################## TTRT Perf ################################# + + ttrt_perf_command = [ + "ttrt", + "perf", + flatbuffer_file, + f"--artifact-dir={self._explorer_artifacts_dir}", + "--memory", + ] + + ttrt_process = self.run_in_subprocess(ttrt_perf_command) + + if ttrt_process.returncode != 0: + error = "Error while running TTRT perf" + self.log(error, severity=logging.error) + raise ExplorerRunException(error) + + perf = self.get_perf_trace(model_path) + columns = [ + "GLOBAL CALL COUNT", + "OP CODE", + "DEVICE FW DURATION [ns]", + "CORE COUNT", + "OUTPUT_0_MEMORY", + "LOC", + ] + perf = perf[columns] + logging.info(perf) + + logging.info( + "Total device duration: ", perf["DEVICE FW DURATION [ns]"].sum(), "ns" + ) -def wait_for_execution_to_finish(timeout): - for _ in range(timeout): - try: - response = send_command("status_check", "") - if response.status_code == 200 and response.json().get("graphs")[0].get( - "isDone" - ): - return response.json() - except requests.RequestException as e: - print(f"Request failed: {e}") - raise Exception("Status check request failed") - time.sleep(1) - raise RuntimeError(f"Execution did not finish within {timeout} seconds") - - -def execute_command_and_wait(model_path, settings, timeout): - execute_command(model_path, settings) - adapter_response = wait_for_execution_to_finish(timeout) - assert "graphs" in adapter_response - assert len(adapter_response["graphs"]) == 1 - response = adapter_response["graphs"][0] - assert response["isDone"] - assert response["error"] is None - - -def convert_command_and_assert(model_path): - result = send_command("convert", model_path) - assert result.ok - if "error" in result.json(): - print(result.json()) - assert False - return result.json() - - -@pytest.mark.parametrize("model_path", get_test_files(TEST_LOAD_MODEL_PATHS)) -def test_load_model(model_path): - convert_command_and_assert(model_path) - - -@pytest.mark.parametrize("model_path", get_test_files(TEST_EXECUTE_MODEL_PATHS)) -def test_execute_model(model_path): - execute_command_and_wait( - model_path, {"optimizationPolicy": "DF Sharding"}, timeout=300 - ) - convert_command_and_assert(model_path) - - -def test_execute_mnist_l1_interleaved(): - execute_command_and_wait( - MNIST_SHARDING_PATH, - {"optimizationPolicy": "Greedy L1 Interleaved"}, - timeout=300, - ) - convert_command_and_assert(MNIST_SHARDING_PATH) - - -def test_execute_mnist_optimizer_disabled(): - execute_command_and_wait( - MNIST_SHARDING_PATH, - {"optimizationPolicy": "Optimizer Disabled"}, - timeout=300, - ) - convert_command_and_assert(MNIST_SHARDING_PATH) - - -def test_execute_mnist_with_overrides(): - overrides = { - 'loc("matmul_1"("MNISTLinear":4294967295:10))__17': { - "named_location": "matmul_1", - "attributes": [ - {"key": "data_type", "value": "f32"}, - {"key": "memory_layout", "value": "tile"}, - {"key": "buffer_type", "value": "dram"}, - {"key": "tensor_memory_layout", "value": "interleaved"}, - {"key": "grid_shape", "value": "[8,8]"}, - ], - } - } - execute_command_and_wait( - MNIST_SHARDING_PATH, - {"optimizationPolicy": "DF Sharding", "overrides": overrides}, - timeout=300, - ) - convert_command_and_assert(MNIST_SHARDING_PATH) - - -def test_execute_and_check_perf_data_exists(): - execute_command_and_wait( - MNIST_SHARDING_PATH, - {"optimizationPolicy": "DF Sharding"}, - timeout=300, - ) - result = convert_command_and_assert(MNIST_SHARDING_PATH) - assert "perf_data" in result["graphs"][0] - - -def test_execute_and_check_memory_data_exists(): - execute_command_and_wait( - MNIST_SHARDING_PATH, - {"optimizationPolicy": "Optimizer Disabled"}, - timeout=300, - ) - result = convert_command_and_assert(MNIST_SHARDING_PATH) - assertion = "display_type" in str(result["graphs"][0]["nodes"]) - assert assertion, str(result["graphs"][0]["nodes"]) - - -def test_execute_model_invalid_policy(): - with pytest.raises(AssertionError): - execute_command_and_wait( - TEST_EXECUTE_MODEL_PATHS[0], - {"optimizationPolicy": "Invalid Policy"}, - timeout=300, + # TTNN_IR_FILE from flatbuffer is still relevant since model_path is the FB with golden data and it will rented optimized_model_path instead + state.optimized_model_path = ttnn_ir_file + self.progress = 100 + + def run(self, model_path, compile_options, overrides): + # Check if a run is already in progress + if self.is_busy(): + raise RuntimeError( + "A model is already being processed. Please wait for it to finish." + ) + self.reset_state(model_path) + self.model_state[model_path] = ModelState() + self.model_state[model_path].overrides = overrides + + # Start compile and run in a new thread + self.runner_thread = threading.Thread( + target=self.compile_and_run_wrapper, args=(model_path, compile_options) ) + self.runner_thread.start() From 7386f9183e84f6abd8f8e7229de09cb436f1db57 Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Thu, 6 Mar 2025 16:40:04 +0000 Subject: [PATCH 15/17] Updated pre-commit --- tools/explorer/tt_adapter/src/tt_adapter/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index 2495733249..f8534487a0 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -153,7 +153,7 @@ def get_memory_usage(self, model_path): memory_trace = json.load(file) return memory_trace - + def get_golden_results(self, model_path): accuracy_res = f"{self.model_state[model_path].model_output_dir}/run/program_0/golden_results.json" From 02f79cd49e00bd727c269f866a36aeae9e9352fe Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Fri, 7 Mar 2025 18:51:51 +0000 Subject: [PATCH 16/17] Updating memory visualization for recent explorer changes --- tools/explorer/test/run_tests.py | 21 ++++++++-------- .../tt_adapter/src/tt_adapter/main.py | 10 +++----- .../tt_adapter/src/tt_adapter/mlir.py | 25 ++++++++++--------- 3 files changed, 27 insertions(+), 29 deletions(-) diff --git a/tools/explorer/test/run_tests.py b/tools/explorer/test/run_tests.py index a1ec1a615c..76d64dbfc8 100644 --- a/tools/explorer/test/run_tests.py +++ b/tools/explorer/test/run_tests.py @@ -221,17 +221,6 @@ def test_execute_and_check_perf_data_exists(): assert "perf_data" in result["graphs"][0]["overlays"] -def test_execute_and_check_memory_data_exists(): - execute_command_and_wait( - MNIST_SHARDING_PATH, - {"optimizationPolicy": "Optimizer Disabled"}, - timeout=300, - ) - result = convert_command_and_assert(MNIST_SHARDING_PATH) - assertion = "display_type" in str(result["graphs"][0]["nodes"]) - assert assertion, str(result["graphs"][0]["nodes"]) - - def test_execute_model_invalid_policy(): with pytest.raises(AssertionError): execute_command_and_wait( @@ -241,6 +230,16 @@ def test_execute_model_invalid_policy(): ) +def test_execute_and_check_memory_data_exists(): + execute_command_and_wait( + MNIST_SHARDING_PATH, + {"optimizationPolicy": "Optimizer Disabled"}, + timeout=300, + ) + result = convert_command_and_assert(MNIST_SHARDING_PATH) + assert "display_type" in str(result) + + def test_execute_and_check_accuracy_data_exists(): # Get the test_mnist path test_mnist_path = GET_TTNN_TEST() diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py index d932ef177f..501504d87f 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/main.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/main.py @@ -93,7 +93,7 @@ def convert( model_path ): logging.info(f"Using optimized model: {optimized_model_path}") - # Get performance and memory results. + # Get performance results. perf_trace = self.model_runner.get_perf_trace(model_path) memory_trace = self.model_runner.get_memory_usage(model_path) golden_results = self.model_runner.get_golden_results(model_path) @@ -102,11 +102,9 @@ def convert( 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, 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) - graph, overlays = mlir.build_graph(module, perf_trace, golden_results) + graph, overlays = mlir.build_graph( + module, perf_trace, memory_trace, golden_results + ) if overlays: graph = utils.add_to_dataclass(graph, "overlays", overlays) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py index 1346910003..07fd9593a5 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/mlir.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/mlir.py @@ -396,7 +396,6 @@ 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( @@ -465,7 +464,6 @@ def parse_ttnn_ttnn_layout(attr): }, ) ) - return result @@ -532,7 +530,6 @@ 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( @@ -554,6 +551,7 @@ def get_attributes(self): graph_builder.KeyValue(key="schedule", value=str(OpHandler.schedule)) ) OpHandler.schedule += 1 + return result def make_graph_node(self, extra_attrs=None): @@ -629,6 +627,7 @@ def build_graph(module, perf_trace=None, memory_trace=None, golden_results=None) / memory_trace[node]["l1_small"]["device_0"]["total_bytes_per_bank"], 4, ) + accuracy_node_data = {} loc_to_accuracy = {} if golden_results is not None: @@ -711,6 +710,7 @@ def process_module( loc_to_perf: Mapping from locations to performance data loc_to_accuracy: Locs to Golden Results perf_node_data: Performance data for nodes + memory_data: Memory usage for nodes accuracy_node_data: Acccuracy Node Data """ module_op = OpHandler(module.operation) @@ -818,6 +818,16 @@ def process_operations( perf_node_data[operation.id] = node_data_builder.NodeDataResult( loc_to_perf[operation.named_location] ) + + if ( + operation.named_location in loc_to_accuracy + and operation.op.name not in EMPTY_OPS + ): + accuracy_node_data[operation.id] = node_data_builder.NodeDataResult( + loc_to_accuracy[operation.named_location]["actual_pcc"] + - loc_to_accuracy[operation.named_location]["expected_pcc"] + ) + extra_attrs = [] if memory_data and operation.named_location in memory_data: extra_attrs.append( @@ -851,15 +861,6 @@ def process_operations( ) ) - if ( - operation.named_location in loc_to_accuracy - and operation.op.name not in EMPTY_OPS - ): - accuracy_node_data[operation.id] = node_data_builder.NodeDataResult( - loc_to_accuracy[operation.named_location]["actual_pcc"] - - loc_to_accuracy[operation.named_location]["expected_pcc"] - ) - if not op.name == "func.func": graph_node = operation.make_graph_node(extra_attrs) From 11fe71b1b6e6cbc93c71d4d777a55a3727b01e6c Mon Sep 17 00:00:00 2001 From: jgrimTT Date: Sat, 8 Mar 2025 23:47:27 +0000 Subject: [PATCH 17/17] Trying agin to pass those tests --- .../explorer/tt_adapter/src/tt_adapter/runner.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tools/explorer/tt_adapter/src/tt_adapter/runner.py b/tools/explorer/tt_adapter/src/tt_adapter/runner.py index f8534487a0..19bf19797a 100644 --- a/tools/explorer/tt_adapter/src/tt_adapter/runner.py +++ b/tools/explorer/tt_adapter/src/tt_adapter/runner.py @@ -142,16 +142,18 @@ def get_perf_trace(self, model_path): return pd.read_csv(op_perf_file) def get_memory_usage(self, model_path): - + memory_trace = {} 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." + error = ( + "Error while fetching memory file,", + mem_file, + "may not have been created ... Continuing Explorer Execution", ) - - with open(mem_file, "r") as file: - memory_trace = json.load(file) - + self.log(error, severity=logging.error) + else: + with open(mem_file, "r") as file: + memory_trace = json.load(file) return memory_trace def get_golden_results(self, model_path):