#265: Add optional save artifacts flag (#276)

tenstorrent · Aug 2, 2024 · 4ab5845 · 4ab5845
1 parent 65eea8b
commit 4ab5845
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 47 deletions.
diff --git a/docs/src/build.md b/docs/src/build.md
@@ -33,7 +33,7 @@ cmake --build build
 ```
 
 > - To enable the ttnn/metal runtime add `-DTTMLIR_ENABLE_RUNTIME=ON`
-> - To enable the ttnn/metal perf runtime add `-DTT_RUNTIME_ENABLE_PERF_TRACE=ON`
+> - To enable the ttnn/metal perf runtime add `-DTT_RUNTIME_ENABLE_PERF_TRACE=ON` and `export ENABLE_TRACY=1` to environment before building
 > - To accelerate the builds with ccache use `-DCMAKE_CXX_COMPILER_LAUNCHER=ccache`
 > - To accelerate builds further, if python bindings aren't needed, `-DTTMLIR_ENABLE_BINDINGS_PYTHON=OFF`. For some reason the python bindings link step is very slow.
 > - TTNN build is automatically integrated / handled by tt-mlir cmake build system.  For debugging and further information regarding the TTNN backend build step, please refer to [TTNN Documentation](https://tenstorrent.github.io/tt-metal/latest/ttnn/ttnn/installing.html).
@@ -46,11 +46,11 @@ cmake --build build
 >   For more information, please refer to
 >   [TT-NN and TT-Metailium installation documentation](https://tenstorrent.github.io/tt-metal/latest/ttnn/ttnn/installing.html#step-4-install-and-start-using-tt-nn-and-tt-metalium).
 
-| OS | Offline Compiler Only | Runtime Enabled Build |
-|----|-----------------------|-----------------------|
-| Ubuntu 22.04  | ✅ | ❌ |
-| Ubuntu 20.04  | ✅ | ✅ |
-| MacOS         | ✅ | ❌ |
+| OS | Offline Compiler Only | Runtime Enabled Build | Runtime + Perf Enabled Build |
+|----|-----------------------|-----------------------| -----------------------------|
+| Ubuntu 22.04  | ✅ | ❌ | ❌ |
+| Ubuntu 20.04  | ✅ | ✅ | ✅ |
+| MacOS         | ✅ | ❌ | ❌ |
 
 ## Test
 

diff --git a/docs/src/ttrt.md b/docs/src/ttrt.md
@@ -32,6 +32,7 @@ ttrt read --section inputs out.ttnn
 ttrt read --section outputs out.ttnn
 ttrt read --section all out.ttnn
 ttrt read --section all out.ttnn --clean-artifacts
+ttrt read --section all out.ttnn --save-artifacts
 ttrt read --section all /dir/of/flatbuffers
 ```
 
@@ -43,6 +44,7 @@ build `-DTTMLIR_ENABLE_RUNTIME=ON`.
 ttrt run --help
 ttrt run out.ttnn
 ttrt run out.ttnn --clean-artifacts
+ttrt run out.ttnn --save-artifacts
 ttrt run out.ttnn --loops 10
 ttrt run --program-index 0 out.ttnn
 ttrt run /dir/of/flatbuffers
@@ -58,8 +60,8 @@ ttrt query --help
 ttrt query --system-desc
 ttrt query --system-desc-as-json
 ttrt query --system-desc-as-dict
-ttrt query --save-system-desc n300.ttsys
-ttrt query --save-system-desc n300.ttsys --clean-artifacts
+ttrt query --save-artifacts
+ttrt query --clean-artifacts
 ```
 
 ### perf
@@ -70,9 +72,11 @@ build `-DTTMLIR_ENABLE_RUNTIME=ON`. Also need perf enabled build `-DTT_RUNTIME_E
 ttrt perf --help
 ttrt perf out.ttnn
 ttrt perf out.ttnn --clean-artifacts
+ttrt perf out.ttnn --save-artifacts
 ttrt perf out.ttnn --loops 10
 ttrt perf --program-index 0 out.ttnn
-ttrt perf --generate-params --perf-csv trace.csv # Generate params dict of all model attributes
+ttrt perf --device out.ttnn
+ttrt perf --generate-params --perf-csv trace.csv
 ttrt perf /dir/of/flatbuffers
 ttrt perf /dir/of/flatbuffers --loops 10
 ```

diff --git a/runtime/tools/python/ttrt/__init__.py b/runtime/tools/python/ttrt/__init__.py
@@ -48,6 +48,11 @@ def main():
         action="store_true",
         help="clean all artifacts from previous runs",
     )
+    read_parser.add_argument(
+        "--save-artifacts",
+        action="store_true",
+        help="save all artifacts during run",
+    )
     read_parser.add_argument("binary", help="flatbuffer binary file")
     read_parser.set_defaults(func=read)
 
@@ -70,6 +75,11 @@ def main():
         default=1,
         help="number of loops",
     )
+    run_parser.add_argument(
+        "--save-artifacts",
+        action="store_true",
+        help="save all artifacts during run",
+    )
     run_parser.add_argument("binary", help="flatbuffer binary file")
     run_parser.set_defaults(func=run)
 
@@ -94,17 +104,16 @@ def main():
         action="store_true",
         help="print the system desc as python dict",
     )
-    query_parser.add_argument(
-        "--save-system-desc",
-        nargs="?",
-        default="system_desc.ttsys",
-        help="serialize a system desc for the current system to a file",
-    )
     query_parser.add_argument(
         "--clean-artifacts",
         action="store_true",
         help="clean all artifacts from previous runs",
     )
+    query_parser.add_argument(
+        "--save-artifacts",
+        action="store_true",
+        help="save all artifacts during run",
+    )
     query_parser.set_defaults(func=query)
 
     """
@@ -143,6 +152,11 @@ def main():
         default=1,
         help="number of loops",
     )
+    perf_parser.add_argument(
+        "--save-artifacts",
+        action="store_true",
+        help="save all artifacts during run",
+    )
     perf_parser.add_argument("binary", help="flatbuffer binary file")
     perf_parser.set_defaults(func=perf)
 

diff --git a/runtime/tools/python/ttrt/common/api.py b/runtime/tools/python/ttrt/common/api.py
@@ -44,32 +44,41 @@ def read(args):
     # acquire parameters
     arg_binary = args.binary
     arg_clean_artifacts = args.clean_artifacts
+    arg_save_artifacts = args.save_artifacts
 
     # preprocessing
     if os.path.isdir(arg_binary):
+        print("provided directory of flatbuffers")
         binaries = find_ttnn_files(arg_binary)
     else:
         binaries.append(arg_binary)
 
     if arg_clean_artifacts:
+        print("cleaning artifacts")
         clean_artifacts()
 
-    setup_artifacts(binaries)
+    if arg_save_artifacts:
+        print("setting up artifact directories")
+        setup_artifacts(binaries)
 
     # constraint checking
+    print("executing constraint for all provided flatbuffers")
     for binary in binaries:
         check_file_exists(binary)
         fbb = ttrt.binary.load_from_path(binary)
         check_version(fbb.version)
         fbb_list.append(fbb)
 
     # execution
+    print("executing action for all provided flatbuffers")
     for fbb in fbb_list:
         read_actions[args.section](fbb)
 
     # save artifacts
-    for binary in binaries:
-        copy_ttnn_binary_into_artifact(binary)
+    if arg_save_artifacts:
+        print("saving artifacts")
+        for binary in binaries:
+            copy_ttnn_binary_into_artifact(binary)
 
 
 """
@@ -90,35 +99,42 @@ def run(args):
 
     # acquire parameters
     arg_binary = args.binary
-    arg_program_index = args.program_index
+    arg_program_index = int(args.program_index)
     arg_clean_artifacts = args.clean_artifacts
     arg_loops = int(args.loops)
+    arg_save_artifacts = args.save_artifacts
 
     # preprocessing
     if os.path.isdir(arg_binary):
+        print("provided directory of flatbuffers")
         binaries = find_ttnn_files(arg_binary)
     else:
         binaries.append(arg_binary)
 
     if arg_clean_artifacts:
+        print("cleaning artifacts")
         clean_artifacts()
 
-    setup_artifacts(binaries)
+    if arg_save_artifacts:
+        print("setting up artifact directories")
+        setup_artifacts(binaries)
 
     # constraint checking
+    print("executing constraint for all provided flatbuffers")
     for binary in binaries:
         check_file_exists(binary)
         fbb = ttrt.binary.load_binary_from_path(binary)
         check_version(fbb.version)
         assert fbb.file_identifier == "TTNN", "Only TTNN binaries are supported"
         fbb_dict = ttrt.binary.as_dict(fbb)
         fbb_list.append((os.path.splitext(os.path.basename(binary))[0], fbb, fbb_dict))
-        program_index = int(arg_program_index)
+        program_index = arg_program_index
         assert program_index <= len(
             fbb_dict["programs"]
         ), "args.program_index out of range"
 
     # execution
+    print("executing action for all provided flatbuffers")
     for (binary_name, fbb, fbb_dict) in fbb_list:
         torch_inputs[binary_name] = []
         torch_outputs[binary_name] = []
@@ -179,24 +195,28 @@ def run(args):
         print("outputs:\n", torch_outputs)
 
     # save artifacts
-    for binary in binaries:
-        copy_ttnn_binary_into_artifact(binary)
-        binary_name = os.path.splitext(os.path.basename(binary))[0]
-        torch_input_tensors = torch_inputs[binary_name]
-        torch_output_tensors = torch_outputs[binary_name]
+    if arg_save_artifacts:
+        print("saving artifacts")
+        for binary in binaries:
+            copy_ttnn_binary_into_artifact(binary)
+            binary_name = os.path.splitext(os.path.basename(binary))[0]
+            torch_input_tensors = torch_inputs[binary_name]
+            torch_output_tensors = torch_outputs[binary_name]
+
+            for i, input in enumerate(torch_input_tensors):
+                save_torch_tensor_into_ttrt_artifacts(
+                    input, f"{binary_name}/input_{i}.pt"
+                )
 
-        for i, input in enumerate(torch_input_tensors):
-            save_torch_tensor_into_ttrt_artifacts(input, f"{binary_name}/input_{i}.pt")
+            for i, output in enumerate(torch_output_tensors):
+                save_torch_tensor_into_ttrt_artifacts(
+                    output, f"{binary_name}/output_{i}.pt"
+                )
 
-        for i, output in enumerate(torch_output_tensors):
-            save_torch_tensor_into_ttrt_artifacts(
-                output, f"{binary_name}/output_{i}.pt"
+            save_system_desc_into_ttrt_artifacts(
+                system_desc, f"{binary_name}/system_desc.ttsys"
             )
 
-        save_system_desc_into_ttrt_artifacts(
-            system_desc, f"{binary_name}/system_desc.ttsys"
-        )
-
 
 """
 API: query
@@ -214,16 +234,20 @@ def query(args):
     arg_system_desc = args.system_desc
     arg_system_desc_as_json = args.system_desc_as_json
     arg_system_desc_as_dict = args.system_desc_as_dict
-    args_save_system_desc = args.save_system_desc
     arg_clean_artifacts = args.clean_artifacts
+    arg_save_artifacts = args.save_artifacts
 
     # preprocessing
     if arg_clean_artifacts:
+        print("cleaning artifacts")
         clean_artifacts()
 
-    setup_artifacts()
+    if arg_save_artifacts:
+        print("setting up artifact directories")
+        setup_artifacts()
 
     # execution
+    print("executing action to get system desc")
     system_desc = ttrt.runtime.get_current_system_desc()[0]
 
     if arg_system_desc or arg_system_desc_as_json:
@@ -232,7 +256,9 @@ def query(args):
         print(system_desc_as_dict(system_desc))
 
     # save artifacts
-    save_system_desc_into_ttrt_artifacts(system_desc, args_save_system_desc)
+    if arg_save_artifacts:
+        print("saving artifacts")
+        save_system_desc_into_ttrt_artifacts(system_desc, "system_desc.ttsys")
 
 
 """
@@ -247,34 +273,44 @@ def perf(args):
 
     # acquire parameters
     arg_binary = args.binary
-    arg_perf_csv = args.perf_csv
+    arg_program_index = int(args.program_index)
     arg_clean_artifacts = args.clean_artifacts
+    arg_perf_csv = args.perf_csv
     arg_loops = int(args.loops)
+    arg_save_artifacts = args.save_artifacts
+    arg_device = args.device
+    arg_generate_params = args.generate_params
 
     # preprocessing
     if os.path.isdir(arg_binary):
+        print("provided directory of flatbuffers")
         binaries = find_ttnn_files(arg_binary)
     else:
         binaries.append(arg_binary)
 
     if arg_clean_artifacts:
+        print("cleaning artifacts")
         clean_artifacts()
 
-    setup_artifacts(binaries)
+    if arg_save_artifacts:
+        print("setting up artifact directories")
+        setup_artifacts(binaries)
 
     # constraint checking
+    print("executing constraint for all provided flatbuffers")
     if not perf_trace:
         print("Perf mode is not enabled. Please rebuild tt-mlir with perf mode")
         sys.exit(1)
 
-    if args.generate_params:
+    if arg_generate_params:
         check_file_exists(arg_perf_csv)
 
     for binary in binaries:
         check_file_exists(binary)
 
     # execution
-    if args.generate_params:
+    print("executing action for all provided flatbuffers")
+    if arg_generate_params:
         generate_params_dict(arg_perf_csv)
     else:
         for binary in binaries:
@@ -291,17 +327,18 @@ def perf(args):
             envVars["TRACY_PORT"] = port
             envVars["TT_METAL_DEVICE_PROFILER_DISPATCH"] = "0"
 
-            if args.device:
+            if arg_device:
                 envVars["TT_METAL_DEVICE_PROFILER"] = "1"
-            else:
-                if "TT_METAL_DEVICE_PROFILER" in envVars.keys():
-                    del envVars["TT_METAL_DEVICE_PROFILER"]
 
             # run perf artifact setup
             captureProcess = run_perf_artifact_setup(port)
 
             # generate test command to execute
-            testCommand = f"python -m tracy -p {TTMLIR_VENV_DIR}/bin/ttrt run {binary} --loops {arg_loops}"
+            testCommandOptions = ""
+            if arg_save_artifacts:
+                testCommandOptions += f"--save-artifacts "
+
+            testCommand = f"python -m tracy -p {TTMLIR_VENV_DIR}/bin/ttrt run {binary} --loops {arg_loops} --program-index {arg_program_index} {testCommandOptions}"
             testProcess = subprocess.Popen(
                 [testCommand], shell=True, env=envVars, preexec_fn=os.setsid
             )

diff --git a/runtime/tools/python/ttrt/common/util.py b/runtime/tools/python/ttrt/common/util.py
@@ -196,6 +196,7 @@ def program_outputs(fbb):
 
 
 def find_ttnn_files(directory):
+    print("finding all ttnn files")
     ttnn_files = []
     for root, _, files in os.walk(directory):
         for file in files: