Update on "[Executorch] Refactor op_mul's broadcasting utils"

Summary:
Refactoring broadcast handling utils that were added for op_mul. This is
in prepartion use these utils to handle broadcast for other ops such as
add, sub, div.
Plus remove a redundant test

Test Plan:
optimized_kernels_test in CI

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

Author: kimishpatel

.buckconfig

@@ -33,3 +33,6 @@
       **/.git, \
       cmake-out, \
       pip-out
+
+[buck2]
+restarter=true

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-0a94bb432ed75cc2d950d81b2921363218a7e459
+27e35de6c288bffad1b4d18b393579c1d1a95547

.ci/docker/conda-env-ci.txt

@@ -1,4 +1,5 @@
 cmake=3.22.1
 ninja=1.10.2
 libuv
+llvm-openmp
 pkg-config

.ci/scripts/setup-linux.sh

@@ -23,4 +23,7 @@ fi
 # of nightly. This allows CI to test against latest commits from PyTorch
 install_executorch "use-pt-pinned-commit"
 build_executorch_runner "${BUILD_TOOL}"
-do_not_use_nightly_on_ci
+
+if [[ "${GITHUB_BASE_REF:-}" == *main* || "${GITHUB_BASE_REF:-}" == *gh* ]]; then
+  do_not_use_nightly_on_ci
+fi

.ci/scripts/setup-macos.sh

@@ -121,6 +121,7 @@ setup_macos_env_variables
 # NB: we need buck2 in all cases because cmake build also depends on calling
 # buck2 atm
 install_buck
+brew install libomp
 install_pip_dependencies
 
 # TODO(huydhn): Unlike our self-hosted runner, GitHub runner doesn't have access
@@ -136,4 +137,7 @@ install_pytorch_and_domains
 # the pinned commit from PyTorch
 install_executorch "use-pt-pinned-commit"
 build_executorch_runner "${BUILD_TOOL}"
-do_not_use_nightly_on_ci
+
+if [[ "${GITHUB_BASE_REF:-}" == *main* || "${GITHUB_BASE_REF:-}" == *gh* ]]; then
+  do_not_use_nightly_on_ci
+fi

.github/workflows/pull.yml

@@ -212,17 +212,14 @@ jobs:
       docker-image: executorch-ubuntu-22.04-clang12
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
+      timeout: 180
       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}"
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install Llava requirements
         bash examples/models/llama/install_requirements.sh
         bash examples/models/llava/install_requirements.sh
@@ -483,9 +480,6 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install phi-3-mini requirements
         bash examples/models/phi-3-mini/install_requirements.sh
 
@@ -513,9 +507,6 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install llama requirements
         bash examples/models/llama/install_requirements.sh
 
@@ -535,17 +526,14 @@ jobs:
       docker-image: executorch-ubuntu-22.04-clang12
       submodules: 'true'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
-      timeout: 90
+      timeout: 180
       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}"
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install llama requirements
         bash examples/models/llama/install_requirements.sh
 
@@ -573,9 +561,6 @@ jobs:
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
 
-        # install pybind
-        bash install_executorch.sh --pybind xnnpack --use-pt-pinned-commit
-
         # install llama requirements
         bash examples/models/llama/install_requirements.sh
 

.gitmodules

@@ -64,3 +64,6 @@
 [submodule "third-party/ao"]
 	path = third-party/ao
 	url = https://github.com/pytorch/ao.git
+[submodule "backends/cadence/utils/FACTO"]
+	path = backends/cadence/utils/FACTO
+	url = https://github.com/pytorch-labs/FACTO.git

CMakeLists.txt

@@ -240,6 +240,13 @@ cmake_dependent_option(
   "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF
 )
 
+
+if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
+  set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON)
+  set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON)
+  set(EXECUTORCH_BUILD_EXTENSION_MODULE ON)
+endif()
+
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
   set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON)
   set(EXECUTORCH_BUILD_KERNELS_CUSTOM ON)
@@ -791,6 +798,35 @@ if(EXECUTORCH_BUILD_PYBIND)
   install(TARGETS portable_lib
           LIBRARY DESTINATION executorch/extension/pybindings
   )
+
+  if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
+
+    set(_pybind_training_dep_libs
+        ${TORCH_PYTHON_LIBRARY}
+        etdump
+        executorch
+        util
+        torch
+        extension_training
+    )
+
+    if(EXECUTORCH_BUILD_XNNPACK)
+      # need to explicitly specify XNNPACK and microkernels-prod
+      # here otherwise uses XNNPACK and microkernel-prod symbols from libtorch_cpu
+      list(APPEND _pybind_training_dep_libs xnnpack_backend XNNPACK microkernels-prod)
+    endif()
+
+    # pybind training
+    pybind11_add_module(_training_lib SHARED extension/training/pybindings/_training_lib.cpp)
+
+    target_include_directories(_training_lib PRIVATE ${TORCH_INCLUDE_DIRS})
+    target_compile_options(_training_lib PUBLIC ${_pybind_compile_options})
+    target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs})
+
+    install(TARGETS _training_lib
+            LIBRARY DESTINATION executorch/extension/training/pybindings
+    )
+  endif()
 endif()
 
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM)

backends/arm/_passes/annotate_decomposed_matmul.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import itertools
 
 from typing import List

backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import copy
 
 from typing import cast, Dict, Set, Tuple

backends/arm/_passes/fuse_batchnorm2d_pass.py

@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import torch
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops

backends/arm/_passes/fuse_quantized_activation_pass.py

@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import torch
 from executorch.backends.arm.tosa_quant_utils import q_op
 from executorch.exir.dialects._ops import ops as exir_ops

backends/arm/_passes/insert_table_ops.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from typing import Callable, Dict
 
 import torch
@@ -31,7 +33,7 @@ class InsertTableOpsPass(ExportPass):
     """
     For ops in self.table_ops they need to be serialized as a TOSA TABLE. This pass replaces these
     edge ops with a tosa._table(input: Tensor, target_str: str) where target_str == str(node.target).
-    When loweringthe _table node target_str will be used to find the corresponding torch operator
+    When lowering the _table node target_str will be used to find the corresponding torch operator
     which will be used to produce the table values in operators/op_table.py.
     """
 
@@ -43,6 +45,7 @@ class InsertTableOpsPass(ExportPass):
         exir_ops.edge.aten.sigmoid.default: torch.sigmoid,
         exir_ops.edge.aten.tanh.default: torch.tanh,
         exir_ops.edge.aten.hardsigmoid.default: torch.nn.functional.hardsigmoid,
+        exir_ops.edge.aten.hardswish.default: torch.nn.functional.hardswish,
     }
 
     def __init__(self, exported_program: ExportedProgram) -> None:

backends/arm/_passes/mm_to_bmm_pass.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import torch
 from executorch.backends.arm._passes.arm_pass_utils import (
     create_node,

backends/arm/arm_partitioner.py

@@ -115,6 +115,7 @@ def ops_to_not_decompose(
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
         ops_to_not_decompose_if_quant_op = [
             torch.ops.aten.hardsigmoid.default,
+            torch.ops.aten.hardswish.default,
         ]
 
         def filter_fn(node: torch.fx.Node) -> bool:

backends/arm/operator_support/right_shift_support.py

@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 
 import logging
 

backends/arm/operator_support/tosa_supported_operators.py

@@ -81,6 +81,7 @@ def is_node_supported(self, submodules, node: fx.Node) -> bool:
             exir_ops.edge.aten.permute_copy.default,
             exir_ops.edge.aten.hardsigmoid.default,
             exir_ops.edge.aten.hardtanh.default,
+            exir_ops.edge.aten.hardswish.default,
             exir_ops.edge.aten.convolution.default,
             exir_ops.edge.aten.div.Tensor,
             exir_ops.edge.aten.eq.Tensor,

backends/arm/operators/op_clamp.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree
 
+# pyre-unsafe
+
 from typing import Any, List, Tuple
 
 import serializer.tosa_serializer as ts  # type: ignore

backends/arm/operators/op_rshift.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 from typing import List
 
 import serializer.tosa_serializer as ts  # type: ignore

backends/arm/quantizer/quantization_annotator.py

@@ -133,6 +133,7 @@ def _match_pattern(
     torch.ops.aten.tanh.default,
     torch.ops.aten.sum.dim_IntList,
     torch.ops.aten.hardsigmoid.default,
+    torch.ops.aten.hardswish.default,
 ]
 
 _one_to_one_shared_input_qspec = [

backends/arm/test/common.py

@@ -9,9 +9,17 @@
 
 import tempfile
 from datetime import datetime
+
 from pathlib import Path
+from typing import Any
 
+import pytest
 from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder
+from executorch.backends.arm.test.runner_utils import (
+    arm_executor_runner_exists,
+    corstone300_installed,
+    corstone320_installed,
+)
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 
@@ -145,3 +153,44 @@ def get_u85_compile_spec_unbuilt(
         .dump_intermediate_artifacts_to(artifact_path)
     )
     return compile_spec  # type: ignore[return-value]
+
+
+SkipIfNoCorstone300 = pytest.mark.skipif(
+    not corstone300_installed() or not arm_executor_runner_exists("corstone-300"),
+    reason="Did not find Corstone-300 FVP or executor_runner on path",
+)
+"""Skips a test if Corsone300 FVP is not installed, or if the executor runner is not built"""
+
+SkipIfNoCorstone320 = pytest.mark.skipif(
+    not corstone320_installed() or not arm_executor_runner_exists("corstone-320"),
+    reason="Did not find Corstone-320 FVP or executor_runner on path",
+)
+"""Skips a test if Corsone320 FVP is not installed, or if the executor runner is not built."""
+
+
+def parametrize(
+    arg_name: str, test_data: dict[str, Any], xfails: dict[str, str] = None
+):
+    """
+    Custom version of pytest.mark.parametrize with some syntatic sugar and added xfail functionality
+        - test_data is expected as a dict of (id, test_data) pairs
+        - alllows to specifiy a dict of (id, failure_reason) pairs to mark specific tests as xfail
+    """
+    if xfails is None:
+        xfails = {}
+
+    def decorator_func(func):
+        """Test data is transformed from a dict of (id, data) pairs to a list of pytest params to work with the native pytests parametrize function"""
+        pytest_testsuite = []
+        for id, test_parameters in test_data.items():
+            if id in xfails:
+                pytest_param = pytest.param(
+                    test_parameters, id=id, marks=pytest.mark.xfail(reason=xfails[id])
+                )
+            else:
+                pytest_param = pytest.param(test_parameters, id=id)
+            pytest_testsuite.append(pytest_param)
+
+        return pytest.mark.parametrize(arg_name, pytest_testsuite)(func)
+
+    return decorator_func