Update on "[executorch] Add coreml quant recipes"

Fixing tests for stack that got reverted: #13265

Adds coreml quant recipes after FP32/16 recipes added in #13121

Recipes added:

PT2E_INT8_STATIC
PT2E_INT8_WEIGHT_ONLY
INT4_WEIGHT_ONLY_PER_CHANNEL
INT4_WEIGHT_ONLY_PER_GROUP
INT8_WEIGHT_ONLY_PER_CHANNEL
INT8_WEIGHT_ONLY_PER_GROUP
CODEBOOK_WEIGHT_ONLY

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

[ghstack-poisoned]

Author: abhinaykukkadapu

.ci/scripts/zephyr-utils.sh

@@ -6,9 +6,9 @@
 # LICENSE file in the root directory of this source tree.
 
 download_arm_zephyr_sdk () {
-    wget https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.16.0/zephyr-sdk-0.16.0_linux-x86_64.tar.xz
-    tar -xf zephyr-sdk-0.16.0_linux-x86_64.tar.xz
-    rm -f zephyr-sdk-0.16.0_linux-x86_64.tar.xz
+    wget https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.2/zephyr-sdk-0.17.2_linux-x86_64.tar.xz
+    tar -xf zephyr-sdk-0.17.2_linux-x86_64.tar.xz
+    rm -f zephyr-sdk-0.17.2_linux-x86_64.tar.xz
 }
 
 setup_zephyr_et_module () {

.github/workflows/add-unanswered-to-project.yml

@@ -1,8 +1,8 @@
 name: Add Open External Contributor PRs and Issues to PyTorch Org Project 136
 
 on:
-  schedule:
-    - cron: '0 * * * *'
+ # schedule:
+  #  - cron: '0 * * * *'
   workflow_dispatch:
 
 jobs:
@@ -12,7 +12,7 @@ jobs:
       - name: Add open issues and open, non-draft PRs to org project (excluding certain authors)
         uses: actions/github-script@v7
         with:
-          github-token: ${{ secrets.PYTORCH_PROJECT_PAT }}
+          github-token: ${{ secrets.GITHUB_TOKEN }}
           script: |
             const projectId = "PVT_kwDOAUB9vs4A_PUL"; // PyTorch org project 136
             const owner = 'pytorch';

.github/workflows/trunk.yml

@@ -92,7 +92,7 @@ jobs:
 
         # TODO @Bujji: Should see if this can be moved into the docker image itself
         download_arm_zephyr_sdk
-        ./zephyr-sdk-0.16.0/setup.sh -c -t arm-zephyr-eabi
+        ./zephyr-sdk-0.17.2/setup.sh -c -t arm-zephyr-eabi
         cd $ZEPHYR_PROJ_ROOT
         setup_zephyr_et_module
 

CMakeLists.txt

@@ -50,7 +50,10 @@
 cmake_minimum_required(VERSION 3.29)
 project(executorch)
 
+set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR})
+
 include(${PROJECT_SOURCE_DIR}/tools/cmake/common/preset.cmake)
+include(${PROJECT_SOURCE_DIR}/tools/cmake/Codegen.cmake)
 include(${PROJECT_SOURCE_DIR}/tools/cmake/Utils.cmake)
 include(CMakeDependentOption)
 include(ExternalProject)
@@ -123,8 +126,6 @@ set(CMAKE_INSTALL_RPATH_USE_LINK_PATH ON)
 # Instead please use `find_package(executorch REQUIRED)` in the example
 # directory and add a new executable in the example `CMakeLists.txt`.
 
-set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR})
-
 if(NOT EXECUTORCH_ENABLE_LOGGING)
   # Avoid pulling in the logging strings, which can be large. Note that this
   # will set the compiler flag for all targets in this directory, and for all
@@ -320,6 +321,7 @@ if(NOT EXECUTORCH_SRCS_FILE)
   message(STATUS "executorch: Generating source lists")
   set(EXECUTORCH_SRCS_FILE "${CMAKE_CURRENT_BINARY_DIR}/executorch_srcs.cmake")
   extract_sources(${EXECUTORCH_SRCS_FILE})
+  executorch_validate_build_variables()
 endif()
 
 # This file defines the `_<target>__srcs` variables used below.

CMakePresets.json

@@ -6,6 +6,36 @@
       "hidden": true,
       "binaryDir": "${sourceDir}/cmake-out"
     },
+    {
+      "name": "android-arm64-v8a",
+      "displayName": "Build executorch core and JNI bindings on android arm64-v8a",
+      "inherits": ["common"],
+      "binaryDir": "${sourceDir}/cmake-out-android-arm64-v8a",
+      "cacheVariables": {
+        "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/android.cmake",
+        "ANDROID_ABI": "arm64-v8a"
+      },
+      "condition": {
+        "type": "inList",
+        "string": "${hostSystemName}",
+        "list": ["Darwin", "Linux", "Windows"]
+      }
+    },
+    {
+      "name": "android-x86_64",
+      "displayName": "Build executorch core and JNI bindings on android x86_64",
+      "inherits": ["common"],
+      "binaryDir": "${sourceDir}/cmake-out-android-x86_64",
+      "cacheVariables": {
+        "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/android.cmake",
+        "ANDROID_ABI": "x86_64"
+      },
+      "condition": {
+        "type": "inList",
+        "string": "${hostSystemName}",
+        "list": ["Darwin", "Linux", "Windows"]
+      }
+    },
     {
       "name": "macos",
       "displayName": "Build ExecuTorch for macOS",

backends/cadence/aot/functions.yaml

@@ -229,6 +229,16 @@
     - arg_meta: null
       kernel_name: impl::reference::quantized_linear_per_tensor_out
 
+- func: cadence::quantized_linear_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::reference::quantized_linear_asym8sxasym8s_asym8s_per_tensor_out
+
+- func: cadence::quantized_linear_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: impl::reference::quantized_linear_asym8uxasym8u_asym8u_per_tensor_out
+
 - func: cadence::im2row.out(Tensor input, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, Tensor in_zero_point, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null

backends/cadence/aot/functions_hifi.yaml

@@ -314,6 +314,16 @@
     - arg_meta: null
       kernel_name: cadence::impl::HiFi::quantized_linear_per_tensor_out
 
+- func: cadence::quantized_linear_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: cadence::impl::HiFi::quantized_linear_asym8sxasym8s_asym8s_per_tensor_out
+
+- func: cadence::quantized_linear_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: cadence::impl::HiFi::quantized_linear_asym8uxasym8u_asym8u_per_tensor_out
+
 - func: cadence::quantized_relu_per_tensor.out(Tensor X, Tensor X_zero_point, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null

backends/cadence/aot/ops_registrations.py

@@ -56,10 +56,26 @@
 lib.define(
     "quantized_linear.per_tensor_out(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)"
 )
+lib.define(
+    "quantized_linear_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, "
+    "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)"
+)
+lib.define(
+    "quantized_linear_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, "
+    "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)"
+)
 lib.define(
     "quantized_linear.per_tensor(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, "
     "SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset) -> Tensor"
 )
+lib.define(
+    "quantized_linear_asym8sxasym8s_asym8s.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, "
+    "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)"
+)
+lib.define(
+    "quantized_linear_asym8uxasym8u_asym8u.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, "
+    "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)"
+)
 
 lib.define(
     "quantized_relu(Tensor X, Tensor X_zero_point, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Y)"
@@ -446,6 +462,50 @@ def quantized_linear_per_tensor_meta(
     return src.new_empty(out_size, dtype=src.dtype)
 
 
+@register_fake("cadence::quantized_linear_asym8sxasym8s_asym8s.per_tensor")
+def quantized_linear_asym8sxasym8s_asym8s_per_tensor_meta(
+    src: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    in_zero_point: int,
+    weight_zero_point: int,
+    out_multiplier: int,
+    out_shift: int,
+    out_zero_point: int,
+    offset: Optional[torch.Tensor],
+) -> torch.Tensor:
+    # src comes in shape [leading_dims, in_dim]
+    # weight comes in shape [out_dim, in_dim]
+    # output comes in empty with shape [leading_dims, out_dim]
+    out_size = list(src.size())
+    weight_size = list(weight.size())
+    assert len(weight_size) == 2
+    out_size[-1] = weight_size[0]
+    return src.new_empty(out_size, dtype=src.dtype)
+
+
+@register_fake("cadence::quantized_linear_asym8uxasym8u_asym8u.per_tensor")
+def quantized_linear_asym8uxasym8u_asym8u_per_tensor_meta(
+    src: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    in_zero_point: int,
+    weight_zero_point: int,
+    out_multiplier: int,
+    out_shift: int,
+    out_zero_point: int,
+    offset: Optional[torch.Tensor],
+) -> torch.Tensor:
+    # src comes in shape [leading_dims, in_dim]
+    # weight comes in shape [out_dim, in_dim]
+    # output comes in empty with shape [leading_dims, out_dim]
+    out_size = list(src.size())
+    weight_size = list(weight.size())
+    assert len(weight_size) == 2
+    out_size[-1] = weight_size[0]
+    return src.new_empty(out_size, dtype=src.dtype)
+
+
 @register_fake("cadence::quantized_conv")
 def quantized_conv_meta(
     input: torch.Tensor,

backends/cadence/aot/tests/test_type_dispatch_passes.py

@@ -18,7 +18,7 @@
 
 
 class TestTypeDispatchPasses(unittest.TestCase):
-    def test_int8_dispatch(self) -> None:
+    def test_int8_dispatch_quantized_fully_connected(self) -> None:
         """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant"""
         x = torch.randint(-128, 127, (2, 3), dtype=torch.int8)
         w = torch.randint(-128, 127, (4, 3), dtype=torch.int8)
@@ -44,7 +44,7 @@ def test_int8_dispatch(self) -> None:
             1,
         )
 
-    def test_uint8_dispatch(self) -> None:
+    def test_uint8_dispatch_quantized_fully_connected(self) -> None:
         """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant"""
         x = torch.randint(0, 255, (2, 3), dtype=torch.uint8)
         w = torch.randint(0, 255, (4, 3), dtype=torch.uint8)
@@ -70,6 +70,58 @@ def test_uint8_dispatch(self) -> None:
             1,
         )
 
+    def test_int8_dispatch_quantized_linear(self) -> None:
+        """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant for quantized_linear"""
+        x = torch.randint(-128, 127, (2, 3), dtype=torch.int8)
+        w = torch.randint(-128, 127, (4, 3), dtype=torch.int8)
+        b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32)
+        gm = single_op_builder(
+            placeholders=(x, w, b),
+            op=exir_ops.edge.cadence.quantized_linear.per_tensor,
+            args=(x, w, b, 0, 0, 1, 0, 0, None),
+        )
+        p = CompileTimeTypeDispatchPass()
+        gm = cast(PassResult, p(gm)).graph_module
+        # Original op should be replaced
+        self.assertEqual(
+            count_node(gm, exir_ops.edge.cadence.quantized_linear.per_tensor),
+            0,
+        )
+        # Should be replaced with int8 specific variant
+        self.assertEqual(
+            count_node(
+                gm,
+                exir_ops.edge.cadence.quantized_linear_asym8sxasym8s_asym8s.per_tensor,
+            ),
+            1,
+        )
+
+    def test_uint8_quantized_linear_dispatch(self) -> None:
+        """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant for quantized_linear"""
+        x = torch.randint(0, 255, (2, 3), dtype=torch.uint8)
+        w = torch.randint(0, 255, (4, 3), dtype=torch.uint8)
+        b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32)
+        gm = single_op_builder(
+            placeholders=(x, w, b),
+            op=exir_ops.edge.cadence.quantized_linear.per_tensor,
+            args=(x, w, b, 0, 0, 1, 0, 0, None),
+        )
+        p = CompileTimeTypeDispatchPass()
+        gm = cast(PassResult, p(gm)).graph_module
+        # Original op should be replaced
+        self.assertEqual(
+            count_node(gm, exir_ops.edge.cadence.quantized_linear.per_tensor),
+            0,
+        )
+        # Should be replaced with uint8 specific variant
+        self.assertEqual(
+            count_node(
+                gm,
+                exir_ops.edge.cadence.quantized_linear_asym8uxasym8u_asym8u.per_tensor,
+            ),
+            1,
+        )
+
     def test_mixed_types_error(self) -> None:
         """Test mixed int8/uint8 inputs should raise RuntimeError"""
         x = torch.randint(-128, 127, (2, 3), dtype=torch.int8)

backends/cadence/aot/type_dispatch.py

@@ -23,40 +23,40 @@ class CompileTimeTypeDispatchPass(ExportPass):
     Replaces generic ops with ops that have explicit types.
     """
 
+    _TYPE_DISPATCH_MAP: dict[tuple[torch.dtype, torch.dtype], str] = {
+        (torch.int8, torch.int8): "asym8sxasym8s_asym8s",
+        (torch.uint8, torch.uint8): "asym8uxasym8u_asym8u",
+    }
+
+    _SUPPORTED_OPS: dict[OpOverload, str] = {
+        exir_ops.edge.cadence.quantized_fully_connected.per_tensor: "quantized_fully_connected",
+        exir_ops.edge.cadence.quantized_linear.per_tensor: "quantized_linear",
+    }
+
     def call_operator(
         self,
         op: OpOverload,
         args: tuple[Argument, ...],
         kwargs: dict[str, Argument],
         meta: NodeMetadata,
     ) -> ProxyValue:
-        if op not in {
-            exir_ops.edge.cadence.quantized_fully_connected.per_tensor,
-        }:
+        if op not in self._SUPPORTED_OPS:
             return super().call_operator(op, args, kwargs, meta)
 
-        if (
-            # pyre-ignore[16]: None has no attribute `to_tensor`.
-            args[0].to_tensor().dtype == torch.int8
-            and args[1].to_tensor().dtype == torch.int8
-        ):
-            return super().call_operator(
-                exir_ops.edge.cadence.quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor,
-                args,
-                kwargs,
-                meta,
-            )
-        elif (
-            args[0].to_tensor().dtype == torch.uint8
-            and args[1].to_tensor().dtype == torch.uint8
-        ):
-            return super().call_operator(
-                exir_ops.edge.cadence.quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor,
-                args,
-                kwargs,
-                meta,
-            )
-        else:
+        # pyre-ignore[16]: None has no attribute `to_tensor`.
+        input_dtype = args[0].to_tensor().dtype
+        weight_dtype = args[1].to_tensor().dtype
+        dtype_pair = (input_dtype, weight_dtype)
+
+        if dtype_pair not in self._TYPE_DISPATCH_MAP:
             raise RuntimeError(
-                f"Unsupported input types for {op}: {args[0].to_tensor().dtype} and {args[1].to_tensor().dtype}"
+                f"Unsupported input types for {op}: {input_dtype} and {weight_dtype}"
             )
+
+        base_op_name = self._SUPPORTED_OPS[op]
+        type_suffix = self._TYPE_DISPATCH_MAP[dtype_pair]
+
+        typed_op_name = f"{base_op_name}_{type_suffix}"
+        typed_op = getattr(exir_ops.edge.cadence, typed_op_name).per_tensor
+
+        return super().call_operator(typed_op, args, kwargs, meta)