Update on "[ET-VK] Adding push constant and ubo verison of select and slice ops to improve memory and performance."

Adding push constant and ubo verison of select and slice ops to improve memory and performance.
*   Updated `transfer_buffer.yaml` and `transfer_texture.yaml` to include `UBO_PARAMS` parameter and generate variants for `select` and `slice` ops with UBO parameters.
*   Updated `transfer.glsl` to generate ubo and push constant versions of `select` and `slice` ops with UBO parameters.

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

[ghstack-poisoned]

Author: trivedivivek

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-9b498d3bb28b8e3411ce464dd2755c5b96d92c8f
+7cda4017ddda554752e89069ae205be5e8388f59

.ci/scripts/check_c10_sync.sh

@@ -12,4 +12,4 @@ pushd pytorch
 git checkout "$pytorch_pin"
 popd
 "$(dirname "${BASH_SOURCE[0]}")"/compare_dirs.sh runtime/core/portable_type/c10/c10 pytorch/c10
-"$(dirname "${BASH_SOURCE[0]}")"/compare_dirs.sh runtime/core/portable_type/c10/torch/standalone pytorch/torch/standalone
+"$(dirname "${BASH_SOURCE[0]}")"/compare_dirs.sh runtime/core/portable_type/c10/torch/headeronly pytorch/torch/headeronly

.github/workflows/trunk.yml

@@ -240,11 +240,11 @@ jobs:
 
         cxx_flags="-fno-exceptions -fno-rtti -Wall -Werror -Wno-int-in-bool-context -DET_HAVE_PREAD=0"
         setup_script_args=""
-        if [[ ${{ matrix.os}} == "bare_metal" ]]; then 
+        if [[ ${{ matrix.os}} == "bare_metal" ]]; then
           toolchain_prefix=arm-none-eabi-
-          threshold="103268" # ~100KiB
+          threshold="104000" # should be ~103.7KB, set threshold to 104KB.
           toolchain_cmake=examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake
-        elif [[ ${{ matrix.os}} == "zephyr-preset" ]]; then 
+        elif [[ ${{ matrix.os}} == "zephyr-preset" ]]; then
           setup_script_args="--target-toolchain zephyr"
           toolchain_prefix=arm-zephyr-eabi-
           threshold="133120" # should be ~125KB, set threshold to 130KB

CMakeLists.txt

@@ -490,7 +490,7 @@ install(
   INCLUDES
   DESTINATION ${_common_include_directories}
 )
-install(FILES tools/cmake/executorch-config.cmake
+install(FILES tools/cmake/Utils.cmake tools/cmake/executorch-config.cmake
         DESTINATION lib/cmake/ExecuTorch
 )
 
@@ -732,4 +732,8 @@ if(EXECUTORCH_BUILD_VULKAN)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/vulkan)
 endif()
 
+if(EXECUTORCH_BUILD_ANDROID_JNI)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/android)
+endif()
+
 include(Test.cmake)

backends/cadence/aot/compiler.py

@@ -8,7 +8,7 @@
 
 import logging
 from pathlib import Path
-from typing import Callable, cast, Optional
+from typing import Optional
 
 import executorch.backends.cadence.aot.ops_registrations  # noqa
 import torch
@@ -32,7 +32,6 @@
     ExecutorchBackendConfig,
     ExecutorchProgramManager,
 )
-from executorch.exir.pass_base import PassResult
 from executorch.exir.passes import ToOutVarPass
 from executorch.exir.passes.sym_shape_eval_pass import HintBasedSymShapeEvalPass
 from executorch.exir.program._program import to_edge_with_preserved_ops
@@ -41,7 +40,7 @@
 from torch.export.exported_program import ExportedProgram
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
-from .passes import get_cadence_passes
+from .passes import apply_exir_ops_passes, apply_torch_ops_passes
 
 from .utils import print_ops_info
 
@@ -262,14 +261,20 @@ def export_to_edge(
     inputs: tuple[object, ...],
     dump_graphs: bool = False,
     constant_methods: Optional[dict[str, object]] = None,
+    core_aten_exceptions: Optional[list[torch._ops.OpOverload]] = None,
 ) -> EdgeProgramManager:
     assert isinstance(model, torch.nn.Module), "model should be an nn.Module"
 
     # Export the model into an ExportedProgram.
     expo_program = trace(model, inputs)
 
+    # Apply passes which transform the ExportedProgram before it gets lowered to edge.
+    expo_program = apply_torch_ops_passes(expo_program)
+
     # Lower the model to edge IR.
-    edge_prog_manager = _lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
+    edge_prog_manager = _lower_ep_to_edge(
+        expo_program, dump_graphs, constant_methods, core_aten_exceptions
+    )
 
     return edge_prog_manager
 
@@ -311,14 +316,7 @@ def _lower_ep_to_cadence(
     Lower an existing ExportedProgram to edge IR and apply frontend optimization passes.
     """
     edge_prog_manager = _lower_ep_to_edge(program, dump_graphs=dump_graphs)
-    cadence_passes = get_cadence_passes(opt_level)
-
-    # Run a couple required passes for quant/dequant ops
-    cadence_prog_manager = edge_prog_manager.transform(
-        cast(
-            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
-        )
-    )
+    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
     return cadence_prog_manager
 
 
@@ -329,14 +327,7 @@ def export_to_cadence(
     opt_level: int = 1,
 ) -> EdgeProgramManager:
     edge_prog_manager = export_to_edge(model, inputs, dump_graphs=dump_graphs)
-    cadence_passes = get_cadence_passes(opt_level)
-
-    # Run a couple required passes for quant/dequant ops
-    cadence_prog_manager = edge_prog_manager.transform(
-        cast(
-            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
-        )
-    )
+    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
     return cadence_prog_manager
 
 
@@ -373,15 +364,8 @@ def export_to_executorch_gen_etrecord(
     memory_config: Optional[MemoryConfig] = None,
     dump_graphs: bool = False,
 ) -> ExecutorchProgramManager:
-    cadence_passes = get_cadence_passes(opt_level)
     edge_prog_manager = export_to_edge(model, inputs, dump_graphs)
-
-    # Run a couple required passes for quant/dequant ops
-    cadence_prog_manager = edge_prog_manager.transform(
-        cast(
-            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
-        )
-    )
+    cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
 
     # Print some information to terminal
     print_ops_info(

backends/cadence/aot/fuse_ops.py

@@ -1127,6 +1127,7 @@ class CadenceFuseOpsInGraph:
         FuseCascadedTransposeOrPermuteOps,
         FuseCascadedViewOps,
         FuseQuantDequantToRequantizePass,
+        FuseMulTensorIntoQuantPass,
         FuseMulTensorIntoDequantPass,
         FuseMulScalarIntoDequantPass,
         FuseFullThenReshapePass,

backends/cadence/aot/passes.py

@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from typing import Any, List, Optional
+from typing import Any, Callable, cast, List, Optional
 
 import torch
 import torch.fx
@@ -28,13 +28,18 @@
     RemoveRedundantOps,
 )
 from executorch.backends.cadence.aot.reorder_ops import CadenceReorderOpsInGraph
-from executorch.backends.cadence.aot.replace_ops import CadenceReplaceOpsInGraph
+from executorch.backends.cadence.aot.replace_ops import (
+    CadenceReplaceOpsInGraph,
+    ReplaceMulTensorWithMulAndFullOpsPass,
+)
 from executorch.backends.cadence.aot.simplify_ops import CadenceSimplifyOpsInGraph
+from executorch.exir import EdgeProgramManager
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.pass_manager import PassManager, PassType
 from executorch.exir.passes import dead_code_elimination_pass
 from executorch.exir.passes.scalar_to_tensor_pass import ScalarToTensorPass
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
+from torch.export.exported_program import ExportedProgram
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=0))
@@ -89,14 +94,37 @@ def get_passes_in_default_order() -> List[ExportPass]:
     return pytree.tree_flatten(passes)[0]
 
 
-def get_cadence_passes(
+def apply_exir_ops_passes(
     opt_level: int,
-) -> List[Optional[PassResult]]:
+    edge_prog_manager: EdgeProgramManager,
+) -> EdgeProgramManager:
     passes = get_passes_in_default_order()
     pass_filter = create_cadence_pass_filter(opt_level)
-    filtered_passes = [
-        # pyre-ignore[20]: Expect argument graph_module
-        filtered_pass()
+    cadence_passes = [
+        (
+            lambda graph_module, filtered_pass=filtered_pass: filtered_pass()(
+                graph_module
+            )
+        )
         for filtered_pass in list(filter(pass_filter, passes))
     ]
-    return filtered_passes
+    cadence_prog_manager = edge_prog_manager.transform(
+        cast(
+            list[Callable[[torch.fx.GraphModule], Optional[PassResult]]], cadence_passes
+        )
+    )
+    return cadence_prog_manager
+
+
+def apply_torch_ops_passes(expo_program: ExportedProgram) -> ExportedProgram:
+    """
+    Applies compiler passes on torch.ops IR, including torch.ops.aten, torch.ops.cadence, etc.
+    expo_program is expected to be the output of the torch.export.export().
+    """
+
+    aten_passes: List[Callable[[torch.fx.GraphModule], Optional[PassResult]]] = [
+        ReplaceMulTensorWithMulAndFullOpsPass()
+    ]
+    # TODO(T230417247): Use PassResult which is currently ignored.
+    PassManager(aten_passes)(expo_program.graph_module)
+    return expo_program

backends/nxp/runtime/TARGETS

@@ -0,0 +1,5 @@
+load("targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()

backends/nxp/runtime/targets.bzl

@@ -0,0 +1,20 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+def define_common_targets():
+    runtime.cxx_library(
+        name = "nxp_backend",
+        srcs = ["NeutronBackend.cpp"],
+        headers = ["NeutronDriver.h", "NeutronErrors.h"],
+        compatible_with = ["ovr_config//cpu:arm32-embedded", "@fbsource//arvr/firmware/projects/smartglasses/config:embedded-mcu-rtos"],
+        # Neutron runtime needs to compile with executor as whole
+        # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole)
+        link_whole = True,
+        # Constructor needed for backend registration.
+        compiler_flags = ["-Wno-global-constructors", "-fno-rtti", "-DNO_HEAP_USAGE"],
+        visibility = ["@EXECUTORCH_CLIENTS"],
+        deps = [
+            "//executorch/runtime/backend:interface",
+            "//executorch/runtime/core:core",
+            "fbsource//arvr/third-party/toolchains/nxp-sdk/2.16.0/middleware/eiq/executorch/third-party/neutron/rt700:libNeutron",
+        ],
+    )

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -877,6 +877,8 @@ def compile_spirv(shader_paths_pair):
             src_file_name = shader_paths_pair[0]
             # path of template file used for codegen
             src_file_fullpath = shader_paths_pair[1][0]
+            # args used for codegen
+            codegen_params = shader_paths_pair[1][1]
 
             # Assume that generated files will have the same file extension as the
             # source template file.
@@ -917,6 +919,7 @@ def compile_spirv(shader_paths_pair):
                         shutil.copyfile(cached_spv_out_path, spv_out_path)
                         return (spv_out_path, gen_out_path)
 
+            vk_version = codegen_params.get("VK_VERSION", "1.1")
             # Only proceed if a GLSL compiler was specified
             if self.glslc_path is not None:
                 cmd_base = [
@@ -925,7 +928,7 @@ def compile_spirv(shader_paths_pair):
                     gen_out_path,
                     "-o",
                     spv_out_path,
-                    "--target-env=vulkan1.1",
+                    "--target-env=vulkan{}".format(vk_version),
                     "-Werror",
                 ] + [
                     arg