From 66abe6e82c63881d0a09c93a031ee8b76ad98dc0 Mon Sep 17 00:00:00 2001
From: Brett Grady <bgrady@tenstorrent.com>
Date: Wed, 5 Mar 2025 18:24:35 +0000
Subject: [PATCH] more qol improvements

---
 tt_npe/scripts/analyze_noc_trace_dir.py | 14 +++++++++-----
 tt_npe/scripts/npe_trace_and_analyze.py | 17 +++++++++++------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/tt_npe/scripts/analyze_noc_trace_dir.py b/tt_npe/scripts/analyze_noc_trace_dir.py
index 2eba7fb..5c502fa 100755
--- a/tt_npe/scripts/analyze_noc_trace_dir.py
+++ b/tt_npe/scripts/analyze_noc_trace_dir.py
@@ -130,7 +130,7 @@ def run_npe(opname, workload_file, output_dir, emit_stats_as_json):
     cfg = npe.Config()
     cfg.workload_json_filepath = workload_file
     #cfg.congestion_model_name = "fast"
-    #cfg.cycles_per_timestep = 64
+    cfg.cycles_per_timestep = 32
     cfg.set_verbosity_level(0)
     if emit_stats_as_json:
         cfg.emit_stats_as_json = True
@@ -169,7 +169,7 @@ def analyze_noc_traces_in_dir(noc_trace_dir, emit_stats_as_json):
         raise FileNotFoundError(f"The directory {noc_trace_dir} does not exist")
 
     output_dir = os.path.join(
-        os.path.basename(os.path.normpath(noc_trace_dir)), "npe_stats"
+        os.path.dirname(os.path.normpath(noc_trace_dir)), "npe_stats"
     )
     if emit_stats_as_json:
         Path(output_dir).mkdir(parents=True, exist_ok=True)
@@ -202,6 +202,7 @@ def analyze_noc_traces_in_dir(noc_trace_dir, emit_stats_as_json):
     # Print header
     BOLD = '\033[1m'
     RESET = '\033[0m'
+    GREEN = '\033[32m'
     print(
             f"{BOLD}{'opname':42} {'op_id':>5}, {'AVG LINK UTIL':>14}, {'DRAM_BW_UTIL':>14}, {'CONG IMPACT':>14}{RESET}"
     )
@@ -215,9 +216,12 @@ def analyze_noc_traces_in_dir(noc_trace_dir, emit_stats_as_json):
     #print(f"error percentiles : ")
     #for k, v in stats.getErrorPercentiles().items():
     #    print(f"  {k:15} : {v:4.1f}%")
-    print(f"average link util                : {stats.getAvgLinkUtil():.2f} ")
-    print(f"cycle-weighted overall link util : {stats.getWeightedAvgLinkUtil():.2f} ")
-    print(f"cycle-weighted dram bw util      : {stats.getWeightedAvgDramBWUtil():.2f} ")
+    print(f"average link util                : {stats.getAvgLinkUtil():.1f}% ")
+    print(f"cycle-weighted overall link util : {stats.getWeightedAvgLinkUtil():.1f}% ")
+    print(f"cycle-weighted dram bw util      : {stats.getWeightedAvgDramBWUtil():.1f}% ")
+
+    if emit_stats_as_json:
+        print(f"\n👉 {BOLD}{GREEN}ttnn-visualizer files located in: '{output_dir}'{RESET}")
 
 def main():
     args = get_cli_args()
diff --git a/tt_npe/scripts/npe_trace_and_analyze.py b/tt_npe/scripts/npe_trace_and_analyze.py
index f0424c7..d28950a 100644
--- a/tt_npe/scripts/npe_trace_and_analyze.py
+++ b/tt_npe/scripts/npe_trace_and_analyze.py
@@ -30,10 +30,10 @@ def get_cli_args():
     )
 
     parser.add_argument(
-        "-c", "--command", type=str, help="command to run", default=None
+        "-c", "--command", help="command to run", default=None
     )
     parser.add_argument(
-        "-o", "--output_dir", type=str, help="Output result directory", default=None
+        "-o", "--output_dir", help="Output result directory", default=None
     )
 
     return parser.parse_args()
@@ -62,15 +62,20 @@ def run_command_with_noc_tracing(command, output_dir):
 def main():
     args = get_cli_args()
 
+    if args.output_dir is None:
+        print(f"Error: Must specify an output subdirectory with -o <path> !")
+        sys.exit(1)
+
     output_path = Path(args.output_dir)
     if output_path.exists() and not output_path.is_dir():
         print(f"Error: requested output dir {args.output_dir} exists but is not a directory!")
         sys.exit(1)
+    elif output_path.is_dir():
+        print(f"\nOutput directory '{output_path}' already exists!")
+        if prompt_yn("Delete old data before continuing? "):
+        	shutil.rmtree(output_path)
 
     trace_output_subdir = os.path.join(args.output_dir, "noc_traces")
-    if os.path.isdir(trace_output_subdir):
-        if prompt_yn("Output directory already exists! Delete old data before continuing? "):
-        	shutil.rmtree(trace_output_subdir)
     Path(trace_output_subdir).mkdir(parents=True, exist_ok=True)
 
     run_command_with_noc_tracing(args.command, trace_output_subdir)
@@ -80,7 +85,7 @@ def main():
     print("-------------------------------------------")
     stats_output_subdir = os.path.join(args.output_dir, "npe_stats")
     Path(stats_output_subdir).mkdir(parents=True, exist_ok=True)
-    analyze_noc_traces_in_dir(trace_output_subdir, emit_stats_as_json=True)
+    analyze_noc_traces_in_dir(trace_output_subdir, True)
 
 if __name__ == "__main__":
     main()