rebase

[ghstack-poisoned]

Author: swolchok

.github/scripts/run_nm.py

@@ -0,0 +1,171 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import re
+import subprocess
+import sys
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+
+@dataclass
+class Symbol:
+    name: str
+    addr: int
+    size: int
+    symbol_type: str
+
+
+class Parser:
+    def __init__(self, elf: str, toolchain_prefix: str = "", filter=None):
+        self.elf = elf
+        self.toolchain_prefix = toolchain_prefix
+        self.symbols: Dict[str, Symbol] = self._get_nm_output()
+        self.filter = filter
+
+    @staticmethod
+    def run_nm(
+        elf_file_path: str, args: Optional[List[str]] = None, nm: str = "nm"
+    ) -> str:
+        """
+        Run the nm command on the specified ELF file.
+        """
+        args = [] if args is None else args
+        cmd = [nm] + args + [elf_file_path]
+        try:
+            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+            return result.stdout
+        except FileNotFoundError:
+            print(f"Error: 'nm' command not found. Please ensure it's installed.")
+            sys.exit(1)
+        except subprocess.CalledProcessError as e:
+            print(f"Error running nm on {elf_file_path}: {e}")
+            print(f"stderr: {e.stderr}")
+            sys.exit(1)
+
+    def _get_nm_output(self) -> Dict[str, Symbol]:
+        args = [
+            "--print-size",
+            "--size-sort",
+            "--reverse-sort",
+            "--demangle",
+            "--format=bsd",
+        ]
+        output = Parser.run_nm(
+            self.elf,
+            args,
+            nm=self.toolchain_prefix + "nm" if self.toolchain_prefix else "nm",
+        )
+        lines = output.splitlines()
+        symbols = []
+        symbol_pattern = re.compile(
+            r"(?P<addr>[0-9a-fA-F]+)\s+(?P<size>[0-9a-fA-F]+)\s+(?P<type>\w)\s+(?P<name>.+)"
+        )
+
+        def parse_line(line: str) -> Optional[Symbol]:
+
+            match = symbol_pattern.match(line)
+            if match:
+                addr = int(match.group("addr"), 16)
+                size = int(match.group("size"), 16)
+                type_ = match.group("type").strip().strip("\n")
+                name = match.group("name").strip().strip("\n")
+                return Symbol(name=name, addr=addr, size=size, symbol_type=type_)
+            return None
+
+        for line in lines:
+            symbol = parse_line(line)
+            if symbol:
+                symbols.append(symbol)
+
+        assert len(symbols) > 0, "No symbols found in nm output"
+        if len(symbols) != len(lines):
+            print(
+                "** Warning: Not all lines were parsed, check the output of nm. Parsed {len(symbols)} lines, given {len(lines)}"
+            )
+        if any(symbol.size == 0 for symbol in symbols):
+            print("** Warning: Some symbols have zero size, check the output of nm.")
+
+        # TODO: Populate the section and module fields from the linker map if available (-Wl,-Map=linker.map)
+        return {symbol.name: symbol for symbol in symbols}
+
+    def print(self):
+        print(f"Elf: {self.elf}")
+
+        def print_table(filter=None, filter_name=None):
+            print("\nAddress\t\tSize\tType\tName")
+            # Apply filter and sort symbols
+            symbols_to_print = {
+                name: sym
+                for name, sym in self.symbols.items()
+                if not filter or filter(sym)
+            }
+            sorted_symbols = sorted(
+                symbols_to_print.items(), key=lambda x: x[1].size, reverse=True
+            )
+
+            # Print symbols and calculate total size
+            size_total = 0
+            for name, sym in sorted_symbols:
+                print(f"{hex(sym.addr)}\t\t{sym.size}\t{sym.symbol_type}\t{sym.name}")
+                size_total += sym.size
+
+            # Print summary
+            symbol_percent = len(symbols_to_print) / len(self.symbols) * 100
+            print("-----")
+            print(f"> Total bytes: {size_total}")
+            print(
+                f"Counted: {len(symbols_to_print)}/{len(self.symbols)}, {symbol_percent:0.2f}% (filter: '{filter_name}')"
+            )
+            print("=====\n")
+
+        # Print tables with different filters
+        def is_executorch_symbol(s):
+            return "executorch" in s.name or s.name.startswith("et")
+
+        FILTER_NAME_TO_FILTER_AND_LABEL = {
+            "all": (None, "All"),
+            "executorch": (is_executorch_symbol, "ExecuTorch"),
+            "executorch_text": (
+                lambda s: is_executorch_symbol(s) and s.symbol_type.lower() == "t",
+                "ExecuTorch .text",
+            ),
+        }
+
+        filter_func, label = FILTER_NAME_TO_FILTER_AND_LABEL.get(
+            self.filter, FILTER_NAME_TO_FILTER_AND_LABEL["all"]
+        )
+        print_table(filter_func, label)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Process ELF file and linker map file."
+    )
+    parser.add_argument(
+        "-e", "--elf-file-path", required=True, help="Path to the ELF file"
+    )
+    parser.add_argument(
+        "-f",
+        "--filter",
+        required=False,
+        default="all",
+        help="Filter symbols by pre-defined filters",
+        choices=["all", "executorch", "executorch_text"],
+    )
+    parser.add_argument(
+        "-p",
+        "--toolchain-prefix",
+        required=False,
+        default="",
+        help="Optional toolchain prefix for nm",
+    )
+
+    args = parser.parse_args()
+    p = Parser(args.elf_file_path, args.toolchain_prefix, filter=args.filter)
+    p.print()

.github/workflows/trunk.yml

@@ -231,6 +231,60 @@ jobs:
         # Run arm unit tests using the simulator
         backends/arm/test/test_arm_baremetal.sh test_pytest_ethosu_fvp
 
+  test-arm-cortex-m-size-test:
+    name: test-arm-cortex-m-size-test
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-arm-sdk
+      submodules: 'true'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 90
+      script: |
+        # The generic Linux job chooses to use base env, not the one setup by the image
+        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+        conda activate "${CONDA_ENV}"
+
+        source .ci/scripts/utils.sh
+        install_executorch "--use-pt-pinned-commit"
+        .ci/scripts/setup-arm-baremetal-tools.sh
+        source examples/arm/ethos-u-scratch/setup_path.sh
+
+        # User baremetal toolchain
+        arm-none-eabi-c++ --version
+        toolchain_cmake=examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake
+        toolchain_cmake=$(realpath ${toolchain_cmake})
+
+        # Build and test size test
+        bash test/build_size_test.sh "-DCMAKE_TOOLCHAIN_FILE=${toolchain_cmake} -DEXECUTORCH_BUILD_ARM_BAREMETAL=ON"
+        elf="cmake-out/test/size_test"
+
+        # Dump basic info
+        ls -al ${elf}
+        arm-none-eabi-size ${elf}
+
+        # Dump symbols
+        python .github/scripts/run_nm.py -e ${elf}
+        python .github/scripts/run_nm.py -e ${elf} -f "executorch" -p "arm-none-eabi-"
+        python .github/scripts/run_nm.py -e ${elf} -f "executorch_text" -p "arm-none-eabi-"
+
+        # Add basic guard - TODO: refine this!
+        arm-none-eabi-strip ${elf}
+        output=$(ls -la ${elf})
+        arr=($output)
+        size=${arr[4]}
+        threshold="102400" # 100KiB
+        echo "size: $size, threshold: $threshold"
+        if [[ "$size" -le "$threshold" ]]; then
+          echo "Success $size <= $threshold"
+        else
+          echo "Fail $size > $threshold"
+          exit 1
+        fi
+
   test-coreml-delegate:
     name: test-coreml-delegate
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main

CMakeLists.txt

@@ -645,13 +645,18 @@ target_link_options_shared_lib(executorch)
 # Real integrations should supply their own YAML file that only lists the
 # operators necessary for the models that will run.
 #
+if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
+  # find pytorch lib here to make it available to all
+  # sub-directories. Find it before including portable so that
+  # optimized_portable_kernels can use it.
+  find_package_torch_headers()
+endif()
+
 if(BUILD_EXECUTORCH_PORTABLE_OPS)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels/portable)
 endif()
 
 if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
-  # find pytorch lib here to make it available to all sub-directories
-  find_package_torch_headers()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels/optimized)
 endif()
 

backends/arm/_passes/__init__.py

@@ -39,6 +39,7 @@
 from .insert_table_ops import InsertTableOpsPass  # noqa
 from .keep_dims_false_to_squeeze_pass import KeepDimsFalseToSqueezePass  # noqa
 from .match_arg_ranks_pass import MatchArgRanksPass  # noqa
+from .match_where_self_arg_dtype_pass import MatchWhereSelfDtypePass  # noqa
 from .meandim_to_averagepool_pass import ConvertMeanDimToAveragePoolPass  # noqa
 from .mm_to_bmm_pass import ConvertMmToBmmPass  # noqa
 from .remove_clone_pass import RemoveClonePass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -40,6 +40,7 @@
     InsertTableOpsPass,
     KeepDimsFalseToSqueezePass,
     MatchArgRanksPass,
+    MatchWhereSelfDtypePass,
     QuantizeOperatorArguments,
     RemoveClonePass,
     ReplaceScalarWithTensorArgPassTOSABI,
@@ -80,6 +81,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(ConvertToClampPass())
         self.add_pass(ConvertMinMaxPass())
         self.add_pass(ConvertAnyDefaultDimDimsPass())
+        self.add_pass(MatchWhereSelfDtypePass())
         if isinstance(self.tosa_spec, Tosa_0_80) and self.tosa_spec.is_U55_subset:
             self.add_pass(CastToInt32Pass())
 
@@ -130,6 +132,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(ConvertToClampPass())
         self.add_pass(ConvertMinMaxPass())
         self.add_pass(ConvertAnyDefaultDimDimsPass())
+        self.add_pass(MatchWhereSelfDtypePass())
 
         self.add_pass(AnnotateDecomposedMatmulPass())
         self.add_pass(QuantizeOperatorArguments())

backends/arm/_passes/arm_pass_utils.py

@@ -13,6 +13,7 @@
 
 import torch
 import torch.fx
+from executorch.backends.arm.tosa_utils import get_node_debug_info
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 
@@ -169,9 +170,13 @@ def get_first_fake_tensor(node: torch.fx.Node) -> FakeTensor:
     else:
         fake_tensor = node.meta["val"]
 
-    assert isinstance(
-        fake_tensor, FakeTensor
-    ), f'Found {fake_tensor} in meta["val"] of {node}, expected to find FakeTensor.'
+    if not isinstance(fake_tensor, FakeTensor):
+        raise TypeError(
+            f'Expected a FakeTensor in meta["val"] of node {node}, but got '
+            f"{type(fake_tensor).__name__}\n"
+            f"{get_node_debug_info(node)}"
+        )
+
     return fake_tensor
 
 

backends/arm/_passes/match_arg_ranks_pass.py

@@ -49,6 +49,7 @@ def __init__(self, exported_program):
         exir_ops.edge.aten.bitwise_left_shift.Tensor,
         exir_ops.edge.aten.eq.Tensor,
         exir_ops.edge.aten.pow.Tensor_Tensor,
+        exir_ops.edge.aten.where.self,
     ]
 
     def _match_op_rank(self, graph_module, node, arg, max_rank):