[ET-VK][AOT] Enable exporting Q8 Quantized Linear + Convolution

As title. Introduce fusion patterns to enable fusing quantized convolution and linear graph patterns into a custom op.

## Changes

Introduce the concept of using custom pattern detection functions to detect graph patterns rather than solely relying on SubgraphMatcher. The issue with SubgraphMatcher is that a large number of graph patterns may need to be exported to obtain variants for different combinations of decompositions/quantization workflows. Having a custom detection function improves maintainability.

Implement detection + replacement functions for quantized linear and quantized conv2d.

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

[ghstack-poisoned]

Author: ssjia

.github/workflows/pull.yml

@@ -933,6 +933,10 @@ jobs:
         ./cmake-out/backends/vulkan/test/custom_ops/quantized_linear
         ./cmake-out/backends/vulkan/test/custom_ops/quantized_conv2d
 
+        # Run e2e testing for selected operators. More operators will be tested via this
+        # route in the future.
+        python -m unittest backends/vulkan/test/test_vulkan_delegate.py -k "*pt2e*"
+
   nxp-build-test:
     name: nxp-build-test
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main

backends/vulkan/_passes/TARGETS

@@ -118,6 +118,19 @@ runtime.python_library(
     ],
 )
 
+runtime.python_library(
+    name = "fold_qdq",
+    srcs = ["fold_qdq.py"],
+    visibility = [
+        "//executorch/backends/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/vulkan:utils_lib",
+        "//executorch/exir:pass_base",
+    ],
+)
+
 runtime.python_library(
     name = "fuse_patterns",
     srcs = ["fuse_patterns.py"],
@@ -144,6 +157,7 @@ runtime.python_library(
         "//executorch/examples/...",
     ],
     deps = [
+        ":fold_qdq",
         ":fuse_patterns",
         ":fuse_quantized_ops",
         ":insert_prepack_nodes",

backends/vulkan/_passes/__init__.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+from executorch.backends.vulkan._passes.fold_qdq import FoldQDQPass
 from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass
 from executorch.backends.vulkan._passes.fuse_quantized_ops import (
     FuseQuantizedOpsTransform,
@@ -30,6 +31,7 @@
 from executorch.backends.vulkan._passes.tag_memory_meta_pass import TagMemoryMetaPass
 
 __all__ = [
+    "FoldQDQPass",
     "FusePatternsPass",
     "FuseQuantizedOpsTransform",
     "insert_prepack_nodes",

backends/vulkan/_passes/fold_qdq.py

@@ -0,0 +1,48 @@
+# Copyright (c) Meta Platforms, Inc. and 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.
+
+import executorch.backends.vulkan.utils as utils
+import torch
+
+from executorch.exir.pass_base import ExportPass, PassResult
+from executorch.exir.passes import dead_code_elimination_pass
+
+
+class FoldQDQPass(ExportPass):
+    """
+    Erase Q/DQ chain introduced by PT2E quantization workflow. It is assumed that all
+    valid quant op patterns have already been fused before this pass.
+    """
+
+    def __init__(self, edge_program: torch.export.ExportedProgram):
+        super(FoldQDQPass, self).__init__()
+        self.edge_program = edge_program
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        for node in graph_module.graph.nodes:
+            # Criteria for a foldable Q/DQ node:
+            # - only one user (dequantize)
+            if utils.is_quant_node(node):
+                if len(node.users) > 1:
+                    continue
+
+                dq_node = None
+                for user in node.users:
+                    if utils.is_dequant_node(user):
+                        dq_node = user
+
+                if dq_node is None:
+                    continue
+
+                original_node = node.args[0]
+                assert isinstance(original_node, torch.fx.Node)
+                dq_node.replace_all_uses_with(original_node)
+
+        graph_module.recompile()
+        dead_code_elimination_pass(graph_module)
+        # Re-trace to validate everything is ok
+        graph_module = super().call(graph_module).graph_module
+        return PassResult(graph_module, True)

backends/vulkan/custom_ops_lib.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.
 
+from typing import Optional
+
 import executorch.backends.vulkan.patterns as vk_patterns
 import torch.library
 
@@ -321,6 +323,134 @@ def linear_qta8a_qga4w(
 lib.impl(name, linear_qta8a_qga4w, "CompositeExplicitAutograd")
 linear_qta8a_qga4w_op = getattr(getattr(torch.ops, namespace), name)
 
+#################
+## qaqw_linear ##
+#################
+
+
+def linear_q8ta_q8csw(
+    x: torch.Tensor,
+    input_scale: float,
+    input_zero_point: int,
+    qweights: torch.Tensor,
+    weight_sums: torch.Tensor,
+    weight_scales: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+):
+    weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32)
+    qweights = qweights.transpose(0, 1)
+    weights = torch.ops.quantized_decomposed.dequantize_per_channel(
+        qweights,
+        weight_scales,
+        weight_zeros,
+        0,
+        -127,
+        127,
+        torch.int8,
+    )
+
+    # Perform linear operation
+    out = torch.nn.functional.linear(x, weights)
+    if bias is not None:
+        out = out + bias
+
+    return out
+
+
+name = "linear_q8ta_q8csw"
+lib.define(
+    f"""
+    {name}(
+        Tensor x,
+        float input_scale,
+        int input_zero_point,
+        Tensor qweight,
+        Tensor weight_sums,
+        Tensor weight_scales,
+        Tensor? bias = None) -> Tensor
+    """
+)
+lib.impl(name, linear_q8ta_q8csw, "CompositeExplicitAutograd")
+qa_q8csw_linear = getattr(getattr(torch.ops, namespace), name)
+
+##################
+## conv2d_q8ta_q8csw ##
+##################
+
+
+def conv2d_q8ta_q8csw(
+    x: torch.Tensor,
+    input_scale: float,
+    input_zero_point: int,
+    qweights: torch.Tensor,
+    weight_sums: torch.Tensor,
+    weight_scales: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    kernel_size: list,
+    stride: list,
+    padding: list,
+    dilation: list,
+    groups: int,
+):
+    weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32)
+
+    # Restore weight tensor from 2D format (IC * H * W, OC) back to 4D format (OC, IC, H, W)
+    # First transpose to get (OC, IC * H * W)
+    qweights_transposed = qweights.transpose(0, 1)
+
+    # Extract kernel dimensions from the provided kernel_size
+    H, W = kernel_size[0], kernel_size[1]
+
+    # Calculate dimensions
+    OC = qweights_transposed.shape[0]
+    IC_H_W = qweights_transposed.shape[1]
+    IC = IC_H_W // (H * W)
+
+    # Reshape to original 4D format (OC, IC, H, W)
+    qweights_4d = qweights_transposed.view(OC, IC, H, W)
+    print(qweights_4d.shape)
+
+    # Dequantize weights
+    weights = torch.ops.quantized_decomposed.dequantize_per_channel(
+        qweights_4d,
+        weight_scales,
+        weight_zeros,
+        0,  # axis=0 for output channel quantization
+        -127,
+        127,
+        torch.int8,
+    )
+    print(weights.shape)
+    print(x.shape)
+
+    # Perform convolution
+    out = torch.nn.functional.conv2d(
+        x, weights, bias, stride, padding, dilation, groups
+    )
+
+    return out
+
+
+name = "conv2d_q8ta_q8csw"
+lib.define(
+    f"""
+    {name}(
+        Tensor x,
+        float input_scale,
+        int input_zero_point,
+        Tensor qweight,
+        Tensor weight_sums,
+        Tensor weight_scales,
+        Tensor? bias,
+        SymInt[] kernel_size,
+        SymInt[] stride,
+        SymInt[] padding,
+        SymInt[] dilation,
+        SymInt groups) -> Tensor
+    """
+)
+lib.impl(name, conv2d_q8ta_q8csw, "CompositeExplicitAutograd")
+conv2d_q8ta_q8csw_op = getattr(getattr(torch.ops, namespace), name)
 ######################
 ## apply_rotary_emb ##
 ######################

backends/vulkan/op_registry.py

@@ -318,6 +318,19 @@ def register_int8_mm_op():
     )
 
 
+@update_features(
+    [
+        exir_ops.edge.et_vk.linear_q8ta_q8csw.default,
+    ]
+)
+def register_qa_qw_linear():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        supports_prepacking=True,
+        supports_resize=False,
+    )
+
+
 @update_features(
     [
         exir_ops.edge.et_vk.linear_weight_int4.default,
@@ -457,6 +470,32 @@ def register_convolution_op():
     )
 
 
+@update_features(
+    [
+        exir_ops.edge.et_vk.conv2d_q8ta_q8csw.default,
+    ]
+)
+def register_quantized_conv_op():
+    return OpFeatures(
+        inputs_storage=[
+            utils.CHANNELS_PACKED_TEXTURE,  # input
+            utils.NO_STORAGE,  # input_scale (non tensor)
+            utils.NO_STORAGE,  # input_zero_point (non tensor)
+            utils.NO_STORAGE,  # weight (prepacked)
+            utils.NO_STORAGE,  # weight_sums (prepacked)
+            utils.NO_STORAGE,  # weight_scales (prepacked)
+            utils.NO_STORAGE,  # bias (prepacked)
+            utils.NO_STORAGE,  # kernel_size (non tensor)
+            utils.NO_STORAGE,  # stride (non tensor)
+            utils.NO_STORAGE,  # padding (non tensor)
+            utils.NO_STORAGE,  # dilation (non tensor)
+            utils.NO_STORAGE,  # groups (non tensor)
+        ],
+        supports_resize=False,
+        supports_prepacking=True,
+    )
+
+
 @update_features("llama::sdpa_with_kv_cache")
 def register_sdpa_with_kv_cache_op():
     return OpFeatures(

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -22,6 +22,8 @@
     vulkan_supported_ops,
 )
 
+from executorch.backends.vulkan.patterns import PatternMatch
+
 from executorch.backends.vulkan.serialization.vulkan_graph_schema import (
     VkMemoryLayout,
     VkStorageType,
@@ -41,7 +43,6 @@
 
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupportBase
-from torch.fx.passes.utils.matcher_utils import InternalMatch
 
 # pyre-ignore
 ops_not_to_decompose = [
@@ -60,7 +61,7 @@ def __init__(
         require_dynamic_shape: bool = False,
         operator_blocklist: Optional[Set[OpKey]] = None,
         operator_allowlist: Optional[Set[OpKey]] = None,
-        fusable_subgraphs: Optional[List[InternalMatch]] = None,
+        fusable_subgraphs: Optional[List[PatternMatch]] = None,
         nn_module_blocklist: Optional[Set[str]] = None,
         nn_module_allowlist: Optional[Set[str]] = None,
     ) -> None:
@@ -72,13 +73,13 @@ def __init__(
             operator_blocklist if operator_blocklist is not None else set()
         )
         self.operator_allowlist = operator_allowlist
-        self.fusable_subgraphs: List[InternalMatch] = (
+        self.fusable_subgraphs: List[PatternMatch] = (
             fusable_subgraphs if fusable_subgraphs is not None else []
         )
         # Create a set of all nodes that are part of fusable subgraphs for quick lookup
         self.fusable_nodes: Set[torch.fx.Node] = set()
         for match in self.fusable_subgraphs:
-            self.fusable_nodes.update(match.nodes_map.values())
+            self.fusable_nodes.update(match.all_nodes)
 
         self.nn_module_blocklist = nn_module_blocklist
         self.nn_module_allowlist = nn_module_allowlist

backends/vulkan/patterns/TARGETS

@@ -10,6 +10,7 @@ runtime.python_library(
         "pattern_registry.py",
         "rope.py",
         "quantized_linear.py",
+        "quantized_convolution.py",
     ],
     visibility = [
         "//executorch/backends/...",