Update base for Update on "[ET-VK] Add custom VkInt4WeightOnlyQuantizer for vulkan"

## Context

This diff adds the `VkInt4WeightOnlyQuantizer` class which enables 4-bit quantization of linear layers via source transformation. This quantizer class is copied from `torchao.quantization.GPTQ.WeightOnlyInt4Linear` with some minor changes as annotated in the implementation.

Note that the pt2e quantization flow does not yet support groupwise quantization, so source transformation is the only way to perform groupwise quantization at the moment.

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

[ghstack-poisoned]

Author: SS-JIA

.ci/docker/build.sh

@@ -41,7 +41,7 @@ case "${IMAGE_NAME}" in
     LINTRUNNER=""
     CLANG_VERSION=12
     # From https://developer.android.com/ndk/downloads
-    ANDROID_NDK_VERSION=r26c
+    ANDROID_NDK_VERSION=r27b
     ;;
   *)
     echo "Invalid image name ${IMAGE_NAME}"

.github/workflows/android-release-artifacts.yml

@@ -13,8 +13,24 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  check-if-aar-exists:
+    name: check-if-aar-exists
+    runs-on: ubuntu-22.04
+    timeout-minutes: 10
+    steps:
+      - name: Check if this RC version is already in S3
+        shell: bash
+        run: |
+          VERSION="${{ inputs.version }}"
+          if curl -I "https://ossci-android.s3.amazonaws.com/executorch/release/${VERSION}/executorch.aar" | grep "200 OK"; then
+            echo "AAR already exists at https://ossci-android.s3.amazonaws.com/executorch/release/${VERSION}/executorch.aar"
+            echo "Will skip build/upload"
+            exit 1
+          fi
+
   build-aar:
     name: build-aar
+    needs: check-if-aar-exists
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     with:
       runner: linux.2xlarge

.github/workflows/android.yml

@@ -53,7 +53,7 @@ jobs:
     # NB: Use metal install for KVM support to run the emulator faster
     runs-on: linux.24xl.spr-metal
     env:
-      ANDROID_NDK_VERSION: r26c
+      ANDROID_NDK_VERSION: r27b
       API_LEVEL: 34
     steps:
       - name: Setup SSH (Click me for login details)

backends/arm/_passes/arm_pass_manager.py

@@ -19,6 +19,9 @@
     ConvertSplitToSlicePass,
 )
 from executorch.backends.arm._passes.decompose_div_pass import DecomposeDivPass
+from executorch.backends.arm._passes.insert_squeeze_after_sum_pass import (
+    InsertSqueezeAfterSumPass,
+)
 from executorch.backends.arm._passes.meandim_to_averagepool_pass import (
     ConvertMeanDimToAveragePool,
 )
@@ -47,6 +50,7 @@ def transform_to_backend_pipeline(
         self.add_pass(ConvertExpandCopyToRepeatPass())
         self.add_pass(ConvertMeanDimToAveragePool())
         self.add_pass(DecomposeDivPass())
+        self.add_pass(InsertSqueezeAfterSumPass())
         self.add_pass(ConvertSplitToSlicePass())
         for spec in compile_spec:
             if spec.key == "permute_memory_format":

backends/arm/_passes/insert_squeeze_after_sum_pass.py

@@ -0,0 +1,69 @@
+# Copyright 2024 Arm Limited and/or its 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.
+
+from typing import cast
+
+import torch
+import torch.fx
+from executorch.backends.arm._passes.arm_pass_utils import create_node, insert_q_dq_pair
+
+from executorch.backends.arm.tosa_quant_utils import get_quant_node_args, is_quant_node
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+class InsertSqueezeAfterSumPass(ExportPass):
+    """
+    In Pytorch, the default behaviour of Tensor.sum is to squeeze
+    the dimension that is summed (keep_dim = False).
+    However, in TOSA, REDUCE_SUM always preserves the
+    rank of the input (keep_dim = True).
+    To get a 1-1 mapping in the sum lowering, normalize the
+    keep_dim = False case to keep_dim = True and add squeeze ops.
+
+    Original:
+        sum(dims, keep_dim = False)
+    After pass:
+        sum(dims, keep_dim = True)
+        (q)
+        (dq)
+        squeeze(dim = dims)
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target != exir_ops.edge.aten.sum.dim_IntList:
+                continue
+            sum_node = cast(torch.fx.Node, node)
+            keep_dim = cast(bool, sum_node.args[2] if len(sum_node.args) > 2 else False)
+            if keep_dim:
+                continue
+
+            dim_list = cast(list[int], sum_node.args[1])
+            quantized = is_quant_node(sum_node)
+            if quantized:
+                qparams = get_quant_node_args(sum_node.all_input_nodes[0])
+                qparams = qparams + (torch.int8,)
+            else:
+                qparams = None
+
+            # Add keep_dim = True arg to sum node.
+            sum_node.args = sum_node.args[0:2] + (True,)
+
+            with graph_module.graph.inserting_after(sum_node):
+                squeeze_node = create_node(
+                    graph_module.graph, exir_ops.edge.aten.squeeze_copy.dims, ()
+                )
+                sum_node.replace_all_uses_with(squeeze_node)
+                squeeze_node.args = (sum_node, dim_list)
+                if quantized:
+                    sum_node = insert_q_dq_pair(graph_module.graph, sum_node, qparams)
+        graph_module.graph.eliminate_dead_code()
+        graph_module.recompile()
+        graph_module = super().call(graph_module).graph_module
+        return PassResult(graph_module, True)

backends/arm/arm_partitioner.py

@@ -63,6 +63,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten._softmax.default,
             exir_ops.edge.aten.slice_copy.Tensor,
             exir_ops.edge.aten.sub.Tensor,
+            exir_ops.edge.aten.sum.dim_IntList,
             exir_ops.edge.aten.view_copy.default,
             exir_ops.edge.aten.clone.default,
             exir_ops.edge.aten.mean.dim,

backends/arm/operators/__init__.py

@@ -34,6 +34,7 @@
     op_softmax,
     op_squeeze,
     op_sub,
+    op_sum,
     op_unsqueeze,
     op_view,
 )

backends/arm/operators/op_sum.py

@@ -0,0 +1,96 @@
+# Copyright 2023-2024 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 typing import cast, List
+
+import executorch.backends.arm.tosa_quant_utils as tqutils
+import executorch.backends.arm.tosa_utils as tutils
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+from torch.fx import Node
+
+
+@register_node_visitor
+class AddVisitor(NodeVisitor):
+    target = "aten.sum.dim_IntList"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+        input_node = inputs[0]
+        input_shape = list(input_node.shape)
+        dim_list = cast(list[int], inputs[1].special)
+        dim_list = [dim % len(input_node.shape) for dim in dim_list]
+        keep_dim = cast(bool, inputs[2].number if len(inputs) > 2 else False)
+        assert keep_dim, "This case should be handled by InsertSqueezeAfterSumPass"
+
+        if is_quant_node:
+
+            # Rescale input to 32 bit
+            rescaled_inputs, scale = tqutils.rescale_nodes_to_int32(
+                [node.all_input_nodes[0]], tosa_graph
+            )
+
+            prev_node = rescaled_inputs[0]
+            reduced_shape = input_shape
+
+            # Reduce all dims in dim_list one-by-one.
+            for dim in dim_list:
+                # When reduced, the size of the dim becomes 1.
+                reduced_shape[dim] = 1
+
+                attr = ts.TosaSerializerAttribute()
+                attr.AxisAttribute(input_node.dim_order.index(dim))
+
+                next_node = tosa_graph.addIntermediate(
+                    tutils.tosa_shape(reduced_shape, input_node.dim_order),
+                    dtype=ts.DType.INT32,
+                )
+
+                tosa_graph.addOperator(
+                    TosaOp.Op().REDUCE_SUM, [prev_node.name], [next_node.name], attr
+                )
+
+                prev_node = next_node
+            tqutils.rescale_node_back_to_int8(node, prev_node, scale, tosa_graph)
+        else:
+            input_name = input_node.name
+            reduced_shape = input_shape
+
+            # Reduce all dims in dim_list one-by-one.
+            for dim in dim_list:
+                # When reduced, the size of the dim becomes 1
+                reduced_shape[dim] = 1
+
+                attr = ts.TosaSerializerAttribute()
+                attr.AxisAttribute(input_node.dim_order.index(dim))
+
+                if dim == dim_list[-1]:
+                    output_name = output.name
+                else:
+                    output_name = tosa_graph.addIntermediate(
+                        tutils.tosa_shape(reduced_shape, input_node.dim_order),
+                        dtype=ts.DType.FP32,
+                    ).name
+
+                tosa_graph.addOperator(
+                    TosaOp.Op().REDUCE_SUM, [input_name], [output_name], attr
+                )
+
+                input_name = output_name

backends/arm/quantizer/arm_quantizer.py

@@ -267,11 +267,11 @@ class ArmQuantizer(Quantizer):
         "add",
         "sub",
         "mul",
-        "sigmoid",
         "mm",
         "cat",
         "one_to_one",
         "generic",
+        "sum",
     ]
 
     def __init__(self) -> None:

backends/arm/quantizer/quantization_annotation/__init__.py

@@ -59,6 +59,6 @@ def decorator(annotator: AnnotatorType):
     mm_annotator,
     mul_annotator,
     one_to_one_annotator,
-    sigmoid_annotator,
     sub_annotator,
+    sum_annotator,
 )