Arm Backend: Expose Vela's Debug Database (#14511)

Follow on from #14401
Enables dumping of Vela's debug database to a specified directory . This
gives us generic information on operators in our model, and can be
combined with the trace output to provide more detailed profiling
analysis.

Co-authored-by: Zingo Andersen <[email protected]>

Author: jmahbs

backends/arm/arm_vela.py

@@ -49,7 +49,12 @@ def vela_bin_pack_io(prefix, data):
 # Output via Vela to binary stream for ArmBackendEthosU
 # WARNING: Do not change this without changing VelaBinStream.cpp as that
 #          function consumes this format and the two need to align.
-def vela_compile(tosa_flatbuffer: bytes, args: List[str], verbose: bool = False):
+def vela_compile(
+    tosa_flatbuffer: bytes,
+    args: List[str],
+    verbose: bool = False,
+    intermediate_path: str | None = None,
+):
     """
     Compile a TOSA graph to a binary stream for ArmBackendEthosU using Vela.
     """
@@ -58,14 +63,14 @@ def vela_compile(tosa_flatbuffer: bytes, args: List[str], verbose: bool = False)
             "ethos-u-vela pip package couldn't be imported. Make sure it's installed!"
         )
 
-    with tempfile.TemporaryDirectory() as tmpdir:
+    def run(dir: str) -> bytes:
         tosaname = "out.tosa"
-        tosa_path = os.path.join(tmpdir, tosaname)
+        tosa_path = os.path.join(dir, tosaname)
         with open(tosa_path, "wb") as f:
             f.write(tosa_flatbuffer)
 
         # invoke vela
-        output_dir = os.path.join(tmpdir, "output")
+        output_dir = os.path.join(dir, "output")
         args.append(f"--output-dir={output_dir}")
         args.append(tosa_path)
         if verbose:
@@ -75,9 +80,9 @@ def vela_compile(tosa_flatbuffer: bytes, args: List[str], verbose: bool = False)
         if any("ethos-u85" in arg for arg in args) or any(
             "debug-force-regor" in arg for arg in args
         ):
-            np_path = os.path.join(tmpdir, "output", "out_vela.npz")
+            np_path = os.path.join(dir, "output", "out_vela.npz")
         else:
-            np_path = os.path.join(tmpdir, "output", "out_sg0_vela.npz")
+            np_path = os.path.join(dir, "output", "out_sg0_vela.npz")
 
         blocks = b""
         with np.load(np_path, allow_pickle=False) as data:
@@ -125,3 +130,9 @@ def vela_compile(tosa_flatbuffer: bytes, args: List[str], verbose: bool = False)
                 blocks = blocks + block
 
         return blocks
+
+    if intermediate_path is not None:
+        return run(intermediate_path)
+    else:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            return run(tmpdir)

backends/arm/ethosu/backend.py

@@ -56,6 +56,7 @@ def _compile_tosa_flatbuffer(
             tosa_flatbuffer,
             compile_flags,
             verbose=logger.getEffectiveLevel() == logging.INFO,
+            intermediate_path=compile_spec.get_intermediate_path(),
         )
         return binary
 

examples/arm/aot_arm_compiler.py

@@ -323,11 +323,14 @@ def get_compile_spec(
             tosa_spec = TosaSpecification.create_from_string("TOSA-1.0+INT")
         compile_spec = TosaCompileSpec(tosa_spec)
     elif "ethos-u" in target:
+        extra_flags = ["--verbose-operators", "--verbose-cycle-estimate"]
+        if debug_mode is not None:
+            extra_flags.append("--enable-debug-db")
         compile_spec = EthosUCompileSpec(
             target,
             system_config=system_config,
             memory_mode=memory_mode,
-            extra_flags=["--verbose-operators", "--verbose-cycle-estimate"],
+            extra_flags=extra_flags,
             config_ini=config,
         )
     elif "vgf" in target:
@@ -473,7 +476,7 @@ def get_args():
         "--config",
         required=False,
         default="Arm/vela.ini",
-        help="Specify custom vela configuration file (vela.ini)",
+        help="Specify custom vela configuration file (vela.ini) for Ethos-U targets.",
     )
     parser.add_argument(
         "--non_strict_export",
@@ -491,7 +494,7 @@ def get_args():
         "--enable_debug_mode",
         required=False,
         choices=["json", "tosa"],
-        help="Flag to enable ATen-to-TOSA debug mode.",
+        help="Flag to enable ATen-to-TOSA debug mode and dumping of Vela's debug database.",
     )
     args = parser.parse_args()