Arm backend: Pass location as string instead of TosaOpLocation (pytorch#15994)

* Fixes issue with serialization
* Update TOSA schema

Author: tom-arm

backends/arm/operators/node_visitor.py

@@ -87,7 +87,7 @@ def _serialize_operator(
             None: Mutates ``tosa_graph`` in place.
 
         """
-        op_location = ts.TosaOpLocation()
+        op_location = None
         if self.debug_hook:
             debug_info = self.debug_hook.add(
                 node,
@@ -96,7 +96,7 @@ def _serialize_operator(
             )
 
             if self.debug_hook.mode == ArmCompileSpec.DebugMode.TOSA:
-                op_location.text = json.dumps(debug_info.to_dict())
+                op_location = json.dumps(debug_info.to_dict())
 
         tosa_graph.addOperator(
             tosa_op,

backends/arm/test/misc/test_debug_feats.py

@@ -16,6 +16,7 @@
 import torch
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.test import common
+from executorch.backends.arm.test.runner_utils import dbg_tosa_fb_to_json
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
     TosaPipelineFP,
@@ -238,25 +239,47 @@ def test_dump_tosa_debug_json(test_data: input_t1):
 
 @common.parametrize("test_data", Linear.inputs)
 def test_dump_tosa_debug_tosa(test_data: input_t1):
-    with tempfile.TemporaryDirectory() as tmpdir:
-        aten_ops: list[str] = []
-        exir_ops: list[str] = []
-        pipeline = TosaPipelineINT[input_t1](
-            module=Linear(),
-            test_data=test_data,
-            aten_op=aten_ops,
-            exir_op=exir_ops,
-            custom_path=tmpdir,
-            tosa_debug_mode=ArmCompileSpec.DebugMode.TOSA,
-        )
+    output_dir = "test_dump_tosa_debug"
 
-        pipeline.pop_stage("run_method_and_compare_outputs")
-        pipeline.run()
+    aten_ops: list[str] = []
+    exir_ops: list[str] = []
+    pipeline = TosaPipelineFP[input_t1](
+        module=Linear(),
+        test_data=test_data,
+        use_to_edge_transform_and_lower=True,
+        aten_op=aten_ops,
+        exir_op=exir_ops,
+        custom_path=output_dir,
+        tosa_debug_mode=ArmCompileSpec.DebugMode.TOSA,
+    )
 
-        json_output_path = Path(tmpdir) / "debug.json"
+    pipeline.pop_stage("run_method_and_compare_outputs")
+    pipeline.run()
+
+    output_path = Path(output_dir)
+    json_output_path = output_path / "debug.json"
+
+    # A JSON file should not be created when TOSA mode used
+    assert not json_output_path.exists()
+
+    # At least one TOSA file should exist
+    tosa_files = list(output_path.glob("*.tosa"))
+    assert len(tosa_files) > 0
+
+    tosa_file = tosa_files[0]
+    with tosa_file.open("rb") as f:
+        tosa_json = dbg_tosa_fb_to_json(f.read())
+
+    # Check all non-empty JSON strings are valid
+    ops = tosa_json["regions"][0]["blocks"][0]["operators"]
+    for op in ops:
+        if op["location"]["text"]:
+            try:
+                json.loads(op["location"]["text"])
+            except json.JSONDecodeError:
+                pytest.fail("Failed to load debug JSON string")
 
-        # A JSON file should not be created when TOSA mode used
-        assert not json_output_path.exists()
+    shutil.rmtree(output_dir, ignore_errors=True)
 
 
 @common.parametrize("test_data", Linear.inputs)

backends/arm/tosa/schemas/tosa_1.0.fbs

@@ -510,13 +510,31 @@ table TosaTensor {
   variable: bool;                   // is this a variable tensor
   is_unranked: bool;                // whether this is an unranked tensor
   variable_name:string;             // name for variable attribute
+
+  // In a model that is larger than 2GB, then tensors instead uses the following
+  // attributes to find stored data, which is outside of flatbuffers
+  // the offset is calculated relative to the beginning of the file and is only
+  // valid if > 1.
+  offset: ulong;
+  size: ulong;
+}
+
+table TosaShape {
+  name: string;                     // name of the shape
+  rank: uint32;                     // rank of the shape
+  data: [ubyte] (force_align: 8);   // raw data array if it's a constant shape
+}
+
+table OpLocation {
+  text: string;       // Opaque string, interpretted by user
 }
 
 table TosaOperator {
   op:Op;                    // operator enum
   attribute:Attribute;      // union structure. operator attribute
-  inputs:[string];          // list of input tensor names
-  outputs:[string];         // list of output tensor names
+  inputs:[string];          // list of input tensor or shape names
+  outputs:[string];         // list of output tensor or shape names
+  location: OpLocation;     // location of this Op in mlir
 }
 
 table TosaBasicBlock {
@@ -525,6 +543,7 @@ table TosaBasicBlock {
   tensors:[TosaTensor];     // tensors array
   inputs:[string];          // name of graph inputs
   outputs:[string];         // name of graph outputs
+  shapes:[TosaShape];       // shapes array
 }
 
 table TosaRegion {
@@ -537,4 +556,4 @@ table TosaGraph {
   regions:[TosaRegion];       // regions array
 }
 
-root_type TosaGraph;
+root_type TosaGraph;