Update base for Update on "[ET-VK] 5/n Split dispatches between multiple command buffers. Track previously submitted command buffers in context and add function to execute all previous command buffers."

The diff adds changes to store command buffers submitted with final_use set to false.

Storing these buffers is necessary for `execute()` function. Since, `encode_execute()` function is typically called once but `execute()` can be called multiple times, `submit_all_non_final_cmds` function is added so all recorded command buffers with `final_use = False` can be called multiple times in `execute()`.

#### Key Changes
*   Added a flag `execute_pending_first_submission` to the `ComputeGraph` class to track whether execute nodes have been freshly encoded and need to be submitted first.
*   Added a new function `submit_all_non_final_cmds` to the `Context` class, which submits all non-final command buffers to the GPU.
*   Modified the `submit_cmd_to_gpu` function to add the submitted command buffer to the `non_final_cmds_` list if it's not marked as final use.
*   Updated the `execute` function in `ComputeGraph` to submit all non-final command buffers before executing the graph.

Differential Revision: [D78360038](https://our.internmc.facebook.com/intern/diff/D78360038/)

[ghstack-poisoned]

Author: trivedivivek

.github/workflows/trunk.yml

@@ -302,6 +302,37 @@ jobs:
           exit 1
         fi
 
+  test-arm-ootb-linux:
+    name: test-arm-ootb-linux
+    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: 'recursive'
+      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}"
+
+        # Follow the steps required before running the notebooks
+        # Try to mirror these as closely as possible
+        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
+
+        # Install requirements for converting notebooks
+        pip install notebook
+
+        # Run OOTB tests
+        backends/arm/test/test_arm_ootb.sh
+
   test-coreml-delegate:
     name: test-coreml-delegate
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main

CMakeLists.txt

@@ -48,8 +48,6 @@
 cmake_minimum_required(VERSION 3.24)
 project(executorch)
 
-# MARK: - Start EXECUTORCH_H12025_BUILD_MIGRATION
-
 include(${PROJECT_SOURCE_DIR}/tools/cmake/common/preset.cmake)
 include(${PROJECT_SOURCE_DIR}/tools/cmake/Utils.cmake)
 include(CMakeDependentOption)
@@ -82,6 +80,7 @@ announce_configured_options(BUCK2)
 
 announce_configured_options(CMAKE_CXX_COMPILER_ID)
 announce_configured_options(CMAKE_TOOLCHAIN_FILE)
+announce_configured_options(BUILD_TESTING)
 
 load_build_preset()
 include(${PROJECT_SOURCE_DIR}/tools/cmake/preset/default.cmake)
@@ -97,11 +96,6 @@ else()
 endif()
 announce_configured_options(CCACHE_PROGRAM)
 
-# Print all the configs that were called with announce_configured_options.
-print_configured_options()
-
-# MARK: - End EXECUTORCH_H12025_BUILD_MIGRATION
-
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # Setup RPATH. See
@@ -750,3 +744,6 @@ if(EXECUTORCH_BUILD_ANDROID_JNI)
 endif()
 
 include(Test.cmake)
+
+# Print all the configs that were called with announce_configured_options.
+print_configured_options()

CMakePresets.json

@@ -8,7 +8,7 @@
     },
     {
       "name": "macos",
-      "displayName": "Build everything buildable on macOS",
+      "displayName": "Build ExecuTorch for macOS",
       "inherits": ["common"],
       "generator": "Xcode",
       "cacheVariables": {
@@ -25,7 +25,7 @@
     },
     {
       "name": "ios",
-      "displayName": "Build everything buildable on iOS",
+      "displayName": "Build ExecuTorch for iOS",
       "inherits": ["common"],
       "generator": "Xcode",
       "cacheVariables": {
@@ -42,7 +42,7 @@
     },
     {
       "name": "ios-simulator",
-      "displayName": "Build everything buildable on iOS simulator",
+      "displayName": "Build ExecuTorch for iOS Simulator",
       "inherits": ["common"],
       "generator": "Xcode",
       "cacheVariables": {
@@ -59,7 +59,7 @@
     },
     {
       "name": "linux",
-      "displayName": "Build everything buildable on Linux",
+      "displayName": "Build ExecuTorch for Linux",
       "inherits": ["common"],
       "cacheVariables": {
         "CMAKE_SYSTEM_NAME": "Linux",
@@ -88,29 +88,21 @@
     {
         "name": "llm",
         "displayName": "Build LLM libraries",
-        "inherits": [
-            "common"
-        ],
+        "inherits": ["common"],
         "cacheVariables": {
             "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/llm.cmake",
             "CMAKE_OSX_DEPLOYMENT_TARGET": "12.0"
         },
         "condition": {
             "type": "inList",
             "string": "${hostSystemName}",
-            "list": [
-                "Darwin",
-                "Linux",
-                "Windows"
-            ]
+            "list": ["Darwin", "Linux", "Windows"]
         }
     },
     {
         "name": "zephyr",
-        "displayName": "Build everything buildable on Zephyr RTOS",
-        "inherits": [
-            "common"
-        ],
+        "displayName": "Build ExecuTorch for Zephyr RTOS",
+        "inherits": ["common"],
         "cacheVariables": {
             "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/zephyr.cmake",
             "CMAKE_TOOLCHAIN_FILE": "${sourceDir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake"

backends/arm/_passes/__init__.py

@@ -26,6 +26,7 @@
 from .decompose_adaptive_avg_pool2d_pass import DecomposeAdaptiveAvgPool2dPass  # noqa
 from .decompose_asin_pass import DecomposeAsinPass  # noqa
 from .decompose_atan_pass import DecomposeAtanPass  # noqa
+from .decompose_atanh_pass import DecomposeAtanhPass  # noqa
 from .decompose_avg_pool2d import DecomposeAvgPool2d  # noqa
 from .decompose_batch_norm_no_stats import DecomposeBatchNormNoStatsPass  # noqa
 from .decompose_cosine_similarity_pass import DecomposeCosineSimilarityPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -30,6 +30,7 @@
     DecomposeAcoshPass,
     DecomposeAdaptiveAvgPool2dPass,
     DecomposeAsinPass,
+    DecomposeAtanhPass,
     DecomposeAtanPass,
     DecomposeAvgPool2d,
     DecomposeBatchNormNoStatsPass,
@@ -163,6 +164,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(DecomposeAsinPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeAtanPass())
+        self.add_pass(DecomposeAtanhPass())
         self.add_pass(ConvertIntPowToMuls())
         self.add_pass(CastBoolToInt8Pass())
         self.add_pass(DecomposeSinhPass())

backends/arm/_passes/cast_int64_pass.py

@@ -47,7 +47,7 @@ def _to_int32(self, graph_module: torch.fx.GraphModule):
                 buffer_name = self.exported_program.graph_signature.inputs_to_buffers[
                     node.name
                 ]
-                buffer = self.exported_program.state_dict[node.name]
+                buffer = self.exported_program.state_dict[buffer_name]
                 self._assert_within_int32(buffer, node)
                 logger.warning(
                     f"Casting buffer {node.name} from torch.int64 to torch.int32"

backends/arm/_passes/decompose_asin_pass.py

@@ -85,12 +85,11 @@ def _build_polynomial(
         return result
 
     def call_operator(self, op, args, kwargs, meta):
+        if op not in edge_asin_op:
+            return super().call_operator(op, args, kwargs, meta)
         logging.info(
             f"Approximating asin. This may introduce small numerical errors. For details, see {__file__}."
         )
-        if op not in edge_asin_op:
-            return super().call_operator(op, args, kwargs, meta)
-
         x = args[0]
         half = 0.5
         one = 1.0

backends/arm/_passes/decompose_atanh_pass.py

@@ -0,0 +1,62 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from executorch.backends.arm._passes import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
+
+
+edge_atanh = exir_ops.edge.aten.atanh.default  # MI case
+
+
+def _get_atanh_ops(op):
+    """Return the primitive ops required.."""
+    if op is not edge_atanh:
+        raise RuntimeError(f"Can't decompose atanh for op {op}")
+    return (
+        exir_ops.edge.aten.mul.Tensor,
+        exir_ops.edge.aten.mul.Scalar,
+        exir_ops.edge.aten.add.Scalar,
+        exir_ops.edge.aten.reciprocal.default,
+        exir_ops.edge.aten.log.default,
+        exir_ops.edge.aten.neg.default,
+    )
+
+
+class DecomposeAtanhPass(ArmPass):
+    """
+    Decomposes the atanh operator into primitive ops.
+    atanh(x) = 0.5 * log((1 + x) / (1 - x))
+    """
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op is not edge_atanh:
+            return super().call_operator(op, args, kwargs, meta, updated=False)
+
+        ops = _get_atanh_ops(op)
+        (
+            op_mul_tensor,
+            op_mul_scalar,
+            op_add_scalar,
+            op_reciprocal,
+            op_log,
+            op_neg,
+        ) = ops
+
+        x = args[0]
+
+        nom = super().call_operator(op_add_scalar, (x, 1.0), {}, meta, updated=True)
+
+        neg_x = super().call_operator(op_neg, (x,), {}, meta, updated=True)
+        denom = super().call_operator(
+            op_add_scalar, (neg_x, 1.0), {}, meta, updated=True
+        )
+        recip = super().call_operator(op_reciprocal, (denom,), {}, meta, updated=True)
+
+        log_input = super().call_operator(
+            op_mul_tensor, (nom, recip), {}, meta, updated=True
+        )
+        log = super().call_operator(op_log, (log_input,), {}, meta, updated=True)
+
+        return super().call_operator(op_mul_scalar, (log, 0.5), {}, meta, updated=True)

backends/arm/_passes/insert_table_ops.py

@@ -52,6 +52,7 @@ class TableOps:
         exir_ops.edge.aten.sin.default: torch.sin,
         exir_ops.edge.aten.tanh.default: torch.tanh,
         exir_ops.edge.aten.atan.default: torch.atan,
+        exir_ops.edge.aten.atanh.default: torch.atanh,
         exir_ops.edge.aten.hardsigmoid.default: torch.nn.functional.hardsigmoid,
         exir_ops.edge.aten.hardswish.default: torch.nn.functional.hardswish,
         exir_ops.edge.aten.sinh.default: torch.sinh,

backends/arm/operator_support/tosa_supported_operators.py

@@ -252,6 +252,7 @@ def is_node_supported(
             exir_ops.edge.aten._adaptive_avg_pool2d.default,
             exir_ops.edge.aten.sign.default,
             exir_ops.edge.aten.asin.default,
+            exir_ops.edge.aten.atanh.default,
         ]
 
         return supported