diff --git a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py index 99e6dd2e4a7..0213505949e 100755 --- a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py +++ b/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 diff --git a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py b/backends/arm/cmsis_pack/test/all_ops/op_recipes.py index 0b3abe1ecfd..185ae5aae46 100644 --- a/backends/arm/cmsis_pack/test/all_ops/op_recipes.py +++ b/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",