Arm cmsis pack all ops test : change format of generated tests

backends/arm/cmsis_pack/test/all_ops/generate_test_models.py

@@ -32,6 +32,7 @@
 sys.path.insert(0, str(_SCRIPTS))
 
 import op_recipes  # type: ignore[import-not-found]  # noqa: E402
+
 from op_guards import (  # type: ignore[import-not-found]  # noqa: E402
     discover_components,
 )
@@ -44,47 +45,64 @@ class TensorSpec:
     shape: list
 
 
-def _save_tensor(t: torch.Tensor, path: Path) -> TensorSpec:
-    src = t.detach().cpu()
-    # The firmware copies these bytes directly into ExecuTorch tensor storage.
-    # For channels_last tensors, that storage is NHWC even though the logical
-    # tensor shape remains NCHW.
-    if (
-        src.dim() == 4
-        and src.is_contiguous(memory_format=torch.channels_last)
-        and not src.is_contiguous()
-    ):
-        storage = src.permute(0, 2, 3, 1).contiguous()
-    else:
-        storage = src.contiguous()
-    path.write_bytes(storage.numpy().tobytes())
-    return TensorSpec(
-        file=path.name, dtype=str(src.dtype).replace("torch.", ""), shape=list(src.shape)
-    )
 
+def _get_number_of_outputs(outputs) -> int:
+    if isinstance(outputs, torch.Tensor):
+        return 1
+    elif isinstance(outputs, (tuple, list)):
+        return len(outputs)
+    else:
+        raise TypeError(f"unsupported output type {type(outputs)}")
 
-def _flatten_outputs(out) -> list:
-    if isinstance(out, torch.Tensor):
-        return [out]
-    if isinstance(out, (tuple, list)):
-        flat = []
-        for o in out:
-            flat.extend(_flatten_outputs(o))
-        return flat
-    raise TypeError(f"unsupported output type {type(out)}")
 
+def _mk_metadata(inputs: tuple, outputs: tuple | torch.Tensor, atol: float, rtol: float) -> dict:
+    metadata = {}
+    metadata["nb_inputs"] = len(inputs)
+    metadata["nb_outputs"] = _get_number_of_outputs(outputs)
+    metadata["atol"] = atol
+    metadata["rtol"] = rtol
+
+    channel_last = False 
+    # We assume that when one input tensor is channel_last, all others too
+    # This assumption is true for the operators tested
+    for i, t in enumerate(inputs):
+        metadata[f"input_{i}"] = t
+        if t.is_contiguous(memory_format=torch.channels_last):
+            channel_last = True
+    if isinstance(outputs, torch.Tensor):
+        metadata["output_0"] = outputs
+    else:
+        for i,t in enumerate(outputs):
+           metadata[f"output_{i}"] = t
+
+    # Input / outputs are exported as channel_first.
+    # So we need to store the memory format information in the metadata so that
+    # we can change the input / output tensors to channel_last in the C++
+    # tests.
+    # It is a workaround for what looks like a bug in the to_edge_transform_and_lower 
+    # function, which does not preserve the memory format of the input / output tensors
+    # exported as constant_methods.
+    metadata["channel_last"] = channel_last
+
+    return metadata
 
-def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Recipe, display_quantized_values: bool = False) -> tuple[bytes, float, float]:
+def _export_portable(model: torch.nn.Module, inputs: tuple,recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> ExecutorchProgramManager:
     from executorch.exir import EdgeCompileConfig, to_edge
 
     model = model.eval()
     exported = torch.export.export(model, inputs, strict=True)
+
+    expected = model(*inputs)
+    metadata = _mk_metadata(inputs, expected,recipe.atol, recipe.rtol)
+    if display_metadata:
+        print(metadata)
     # The pack ships the full portable op set, including ops outside the Core
     # ATen opset (bitwise shifts, unfold, var_mean.correction, ...), so skip the
     # core-ATen IR validity gate; to_executorch still lowers to the .out kernels.
-    edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False))
+    edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False),
+                   constant_methods=metadata)
     program = edge.to_executorch()
-    return bytes(program.buffer), recipe.atol, recipe.rtol
+    return program
 
 def _compute_test_threshold(actual, expected):
     abs_err = (actual - expected).abs()
@@ -99,7 +117,7 @@ def _compute_test_threshold(actual, expected):
         rtol = 0.0
     return atol, rtol
 
-def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False) -> tuple[bytes, float, float]:
+def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.Recipe, display_quantized_values: bool = False, display_metadata: bool = False) -> ExecutorchProgramManager:
     from executorch.backends.cortex_m.passes.cortex_m_pass_manager import (
         CortexMPassManager,
     )
@@ -111,6 +129,7 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R
 
     model = model.eval()
     expected = model(*inputs)
+
     captured = torch.export.export(model, inputs, strict=True).module()
     prepared = prepare_pt2e(captured, CortexMQuantizer())
     prepared(*inputs)  # calibrate
@@ -119,10 +138,16 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R
     actual = quantized(*inputs)
     if display_quantized_values:
         print("=== quantized values ===")
-        print("Expected:", expected)
-        print("Actual:", actual)
+        print("Expected:", expected,expected.shape)
+        print("Actual:", actual,actual.shape)
     atol,rtol = _compute_test_threshold(actual, expected)
+
+    metadata = _mk_metadata(inputs, expected,atol, rtol)
+    if display_metadata:
+       print(metadata)
+
     exported = torch.export.export(quantized, inputs, strict=True)
+
     edge = to_edge_transform_and_lower(
         exported, compile_config=EdgeCompileConfig(
                       preserve_ops=[
@@ -134,12 +159,13 @@ def _export_cortex_m(model: torch.nn.Module, inputs: tuple, recipe: op_recipes.R
                       ],
                       _check_ir_validity=False,
                       _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default],
-                    )
+                    ),
+        constant_methods=metadata
     )
     edge._edge_programs["forward"] = CortexMPassManager(
         edge.exported_program()
     ).transform()
-    return bytes(edge.to_executorch().buffer),atol,rtol
+    return edge.to_executorch()
 
 
 def main() -> None:
@@ -165,6 +191,19 @@ def main() -> None:
         action="store_true",
         help="display quantized values during export",
     )
+    parser.add_argument(
+        "--display-metadata",
+        action="store_true",
+        help="display metadata during export",
+    )
+    # Model explorer (with pte extension) or Netron are not able to display
+    # all the .pte in the right way.
+    # So, another way to check is to print the final graph.
+    parser.add_argument(
+        "--display-graph",
+        action="store_true",
+        help="display graph during export",
+    )
     args = parser.parse_args()
 
     source_dir = Path(args.source_dir)
@@ -193,36 +232,34 @@ def main() -> None:
 
     for component in components:
         key = (component.category, component.name)
+        # For debugging. To select only the nodes under investigation
+        #if component.name != "convolution" and component.name != "embedding":
+        #    skipped.append((component.category, component.name, "For debug"))
+        #    continue
+
         if key in op_recipes.SKIPS:
             skipped.append((component.category, component.name, op_recipes.SKIPS[key]))
             continue
         recipe = op_recipes.RECIPES[key]
         cat_id = component.category.lower().replace("-", "_")
-        op_dir = out_dir / f"{cat_id}__{component.name}"
+        op_name = out_dir / f"{cat_id}__{component.name}"
+        op_pte = op_name.with_suffix(".pte")
         try:
-            if args.display_quantized_values:
+            if args.display_quantized_values or args.display_metadata:
                print(f"=== exporting {component.category}/{component.name} ===")
             model, inputs = recipe.make()
-            reference = model.eval()(*inputs)
-            pte,atol,rtol = exporters[component.category](model, inputs, recipe,display_quantized_values=args.display_quantized_values)
-            op_dir.mkdir(parents=True, exist_ok=True)
-            (op_dir / "model.pte").write_bytes(pte)
-            in_specs = [
-                _save_tensor(t, op_dir / f"input_{i}.bin") for i, t in enumerate(inputs)
-            ]
-            out_specs = [
-                _save_tensor(t, op_dir / f"expected_{i}.bin")
-                for i, t in enumerate(_flatten_outputs(reference))
-            ]
+
+            pte = exporters[component.category](model, inputs, recipe,display_quantized_values=args.display_quantized_values,display_metadata=args.display_metadata)
+            if args.display_graph:
+                _ = pte.exported_program("forward").graph_module.print_readable()
+
+            op_pte.write_bytes(bytes(pte.buffer))
+
             manifest.append(
                 {
                     "op": component.name,
                     "category": component.category,
-                    "dir": op_dir.name,
-                    "atol": atol,
-                    "rtol": rtol,
-                    "inputs": [s.__dict__ for s in in_specs],
-                    "outputs": [s.__dict__ for s in out_specs],
+                    "name": op_name.name
                 }
             )
         except Exception as exc:  # noqa: BLE001 - report, don't abort the sweep

backends/arm/cmsis_pack/test/all_ops/op_recipes.py

@@ -43,6 +43,60 @@ class Recipe:
 # Explicit skips keyed the same way: ops we deliberately do not execute.
 SKIPS: dict[tuple[str, str], str] = {}
 
+# torch.nn.Conv2d weights are initialized randomly 
+# To have reproducible results, we create a custom Conv2d module 
+# that initializes the weights to 1.0
+class FixedParametersConv2D(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(*args, **kwargs, bias=False)
+        self.conv.weight.data.fill_(1.0)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+# torch.nn.Conv2d weights are initialized randomly 
+# To have reproducible results, we create a custom Conv2d module 
+# that initializes the weights to 1.0
+class FixedParametersEmbedding(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.embedding = torch.nn.Embedding(*args, **kwargs)
+        self.embedding.weight.data.fill_(1.0)
+
+    def forward(self, x):
+        return self.embedding(x)
+
+# torch.nn.Linear weights are initialized randomly 
+# To have reproducible results, we create a custom Linear module 
+# that initializes the weights to 1.0
+class FixedParametersLinear(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.linear = torch.nn.Linear(*args, **kwargs)
+        self.linear.weight.data.fill_(1.0)
+        if self.linear.bias is not None:
+            self.linear.bias.data.fill_(1.0)
+
+    def forward(self, x):
+        return self.linear(x)
+
+# torch.nn.ConvTranspose2d weights are initialized randomly 
+# To have reproducible results, we create a custom ConvTranspose2d module 
+# that initializes the weights to 1.0
+class FixedParametersConvTranspose2d(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.conv_transpose = torch.nn.ConvTranspose2d(*args, **kwargs)
+        self.conv_transpose.weight.data.fill_(1.0)
+        if self.conv_transpose.bias is not None:
+            self.conv_transpose.bias.data.fill_(1.0)
+
+    def forward(self, x):
+        return self.conv_transpose(x)
+
+
 
 class _Fn(torch.nn.Module):
     """Wrap a plain callable as a module so torch.export can trace it."""
@@ -55,14 +109,11 @@ def forward(self, *args):
         return self.fn(*args)
 
 
-def _ramp(lo: float, hi: float, shape: tuple, channel_last=False) -> torch.Tensor:
+def _ramp(lo: float, hi: float, shape: tuple) -> torch.Tensor:
     n = 1
     for s in shape:
         n *= s
-    if channel_last:
-       return torch.linspace(lo, hi, n).reshape(shape).contiguous(memory_format=torch.channels_last)
-    else:
-       return torch.linspace(lo, hi, n).reshape(shape)
+    return torch.linspace(lo, hi, n).reshape(shape)
 
 
 def _reg(category: str, name: str, make: Callable, **kw) -> None:
@@ -289,7 +340,7 @@ def _red(name: str, fn: Callable, **kw):
 # Conv / pool / norm
 # --------------------------------------------------------------------------
 _portable(
-    "convolution", lambda: (torch.nn.Conv2d(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),))
+    "convolution", lambda: (FixedParametersConv2D(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),))
 )
 _portable("avg_pool2d", lambda: (torch.nn.AvgPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8)),)))
 _portable(
@@ -421,7 +472,7 @@ def _red(name: str, fn: Callable, **kw):
 )
 _portable(
     "embedding",
-    lambda: (torch.nn.Embedding(5, 3), (torch.tensor([[0, 2], [1, 4]]),)),
+    lambda: (FixedParametersEmbedding(5, 3), (torch.tensor([[0, 2], [1, 4]]),)),
 )
 _portable(
     "scatter_add",
@@ -525,6 +576,7 @@ def _red(name: str, fn: Callable, **kw):
 )
 
 
+
 # --------------------------------------------------------------------------
 # Cortex-M ops (quantized; exercised via the CortexM export flow). Inputs are
 # float; the CortexMQuantizer + passes lower them to cortex_m::* kernels.
@@ -551,7 +603,7 @@ def _cm(name: str, make: Callable, **kw):
 )
 _cm(
     "quantized_linear",
-    lambda: (torch.nn.Linear(8, 4, bias=False), (_ramp(-2, 2, (1, 8)),)),
+    lambda: (FixedParametersLinear(8, 4, bias=False), (_ramp(-2, 2, (1, 8)),)),
 )
 _cm(
     "softmax", lambda: (_Fn(lambda x: torch.softmax(x, dim=1)), (_ramp(-2, 2, (1, 8)),))
@@ -566,23 +618,23 @@ def _cm(name: str, make: Callable, **kw):
 )
 _cm(
     "quantized_conv2d",
-    lambda: (torch.nn.Conv2d(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),)),
+    lambda: (FixedParametersConv2D(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)),
 )
 _cm(
     "quantized_depthwise_conv2d",
-    lambda: (torch.nn.Conv2d(4, 4, 3, groups=4), (_ramp(-1, 1, (1, 4, 8, 8)),)),
+    lambda: (FixedParametersConv2D(4, 4, 3, groups=4), (_ramp(1, 5, (1, 4, 8, 8)).to(memory_format=torch.channels_last),)),
 )
 _cm(
     "quantized_transpose_conv2d",
-    lambda: (torch.nn.ConvTranspose2d(2, 3, 3), (_ramp(-1, 1, (1, 2, 8, 8)),)),
+    lambda: (FixedParametersConvTranspose2d(2, 4, 3), (_ramp(1, 5, (1, 2, 5, 5)).to(memory_format=torch.channels_last),)),
 )
 _cm(
     "quantized_avg_pool2d",
-    lambda: (torch.nn.AvgPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8), channel_last=True),)),
+    lambda: (torch.nn.AvgPool2d(kernel_size=2, stride=2), (_ramp(0, 15, (1, 1, 4, 4)),)),
 )
 _cm(
     "quantized_max_pool2d",
-    lambda: (torch.nn.MaxPool2d(2), (_ramp(-1, 1, (1, 2, 8, 8), channel_last=True),)),
+    lambda: (torch.nn.MaxPool2d(kernel_size=2, stride=2), (_ramp(-50, 50, (1, 1, 6, 6)),)),
 )
 _cm(
     "quantized_batch_matmul",