Support per-channel weights end to end

backends/cadence/aot/BUCK

@@ -434,6 +434,7 @@ fbcode_target(_kind = python_unittest,
         "//executorch/backends/cadence/aot:compiler",
         "//executorch/backends/test:graph_builder",
         "//executorch/backends/cadence/aot:pass_utils",
+        "//executorch/backends/cadence/aot/quantizer:utils",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/passes:lib",
@@ -645,6 +646,7 @@ fbcode_target(_kind = python_unittest,
     deps = [
         ":typing_stubs",
         "//executorch/backends/cadence/aot:ops_registrations",
+        "//executorch/backends/cadence/aot/quantizer:utils",
         "//caffe2:torch",
     ]
 )
@@ -664,6 +666,21 @@ fbcode_target(_kind = python_unittest,
     ],
 )
 
+fbcode_target(_kind = python_unittest,
+    name = "test_pattern_utils",
+    srcs = [
+        "tests/test_pattern_utils.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/test:graph_builder",
+        "//executorch/backends/test:program_builder",
+        "//executorch/backends/cadence/aot:compiler_funcs",
+        "//executorch/backends/cadence/aot:ops_registrations",
+        "//executorch/backends/cadence/aot/quantizer:quantizer",
+    ],
+)
+
 fbcode_target(_kind = python_unittest,
     name = "test_to_out_var_pass",
     srcs = [

backends/cadence/aot/compiler.py

@@ -170,7 +170,7 @@ def apply_pre_edge_transform_passes(
     PassManager(
         [
             FuseQATConvBN(converted_program),
-            QuantFusionPass(patterns),
+            QuantFusionPass(patterns, converted_program),
         ]
     )(converted_program.graph_module)
 

backends/cadence/aot/compiler_funcs.py

@@ -13,8 +13,12 @@
 
 import torch
 
+from executorch.backends.cadence.aot.quantizer.pattern_utils import (
+    EXPORTED_PROGRAM_META_KEY,
+)
 from executorch.backends.transforms.permute_pass_utils import get_arg
 from torch._inductor.decomposition import remove_decompositions
+from torch.export.exported_program import ExportedProgram
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, prepare_qat_pt2e
@@ -714,26 +718,41 @@ class QuantFusionPass(PassBase):
     """
     Iterates patterns, finds anchor ops in the converted graph, and calls
     pattern.fuse() to replace dq-op-q subgraphs with fused ops.
+
+    ``exported_program`` is optional but required for per-channel weights:
+    fusion has to read the weight scale vector and materialize the derived
+    qparam tensors, and by this point both live in the program's constants
+    rather than in the graph.
     """
 
-    def __init__(self, patterns: Sequence[object]) -> None:
+    def __init__(
+        self,
+        patterns: Sequence[object],
+        exported_program: Optional[ExportedProgram] = None,
+    ) -> None:
         super().__init__()
         self.patterns = patterns
+        self.exported_program = exported_program
 
     def call(self, graph_module: GraphModule) -> Optional[PassResult]:
         changed = False
-        for pattern in self.patterns:
-            pattern_changed = False
-            for target in pattern.anchor_ops():  # pyre-ignore[16]
-                for node in graph_module.graph.find_nodes(
-                    op="call_function", target=target
-                ):
-                    result = pattern.fuse(graph_module, node)  # pyre-ignore[16]
-                    if result is not None:
-                        changed = True
-                        pattern_changed = True
-            if pattern_changed:
-                graph_module.graph.eliminate_dead_code()
+        if self.exported_program is not None:
+            graph_module.meta[EXPORTED_PROGRAM_META_KEY] = self.exported_program
+        try:
+            for pattern in self.patterns:
+                pattern_changed = False
+                for target in pattern.anchor_ops():  # pyre-ignore[16]
+                    for node in graph_module.graph.find_nodes(
+                        op="call_function", target=target
+                    ):
+                        result = pattern.fuse(graph_module, node)  # pyre-ignore[16]
+                        if result is not None:
+                            changed = True
+                            pattern_changed = True
+                if pattern_changed:
+                    graph_module.graph.eliminate_dead_code()
+        finally:
+            graph_module.meta.pop(EXPORTED_PROGRAM_META_KEY, None)
         if changed:
             graph_module.recompile()
         return PassResult(graph_module, changed)

backends/cadence/aot/ops_registrations.py

@@ -268,12 +268,24 @@ def register_fake(
 lib.define(
     "quantized_conv1d_nlc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, Tensor? offset=None, *, Tensor(a!) out) -> Tensor(a!)"
 )
+lib.define(
+    "quantized_depthwise_conv1d_ncl(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)"
+)
+lib.define(
+    "quantized_depthwise_conv1d_ncl.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)"
+)
 lib.define(
     "quantized_depthwise_conv1d_ncl.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)"
 )
 lib.define(
     "quantized_depthwise_conv1d_ncl.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)"
 )
+lib.define(
+    "quantized_depthwise_conv1d_nlc(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)"
+)
+lib.define(
+    "quantized_depthwise_conv1d_nlc.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)"
+)
 lib.define(
     "quantized_depthwise_conv1d_nlc.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)"
 )
@@ -1396,6 +1408,76 @@ def quantized_depthwise_conv1d_nlc_per_tensor_meta(
     return input.new_empty(output_size, dtype=input.dtype)
 
 
+@register_fake("cadence::quantized_depthwise_conv1d_ncl")
+def quantized_depthwise_conv1d_ncl_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    stride: Tuple[int],
+    padding: Tuple[int],
+    dilation: Tuple[int],
+    groups: int,
+    in_zero_point: int,
+    weight_zero_point: torch.Tensor,
+    bias_scale: torch.Tensor,
+    output_scale: float,
+    output_zero_point: int,
+    out_multiplier: torch.Tensor,
+    out_shift: torch.Tensor,
+) -> torch.Tensor:
+    return quantized_depthwise_conv1d_ncl_per_tensor_meta(
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        groups,
+        in_zero_point,
+        0,
+        1.0,
+        output_scale,
+        output_zero_point,
+        1,
+        0,
+    )
+
+
+@register_fake("cadence::quantized_depthwise_conv1d_nlc")
+def quantized_depthwise_conv1d_nlc_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    stride: Tuple[int],
+    padding: Tuple[int],
+    dilation: Tuple[int],
+    groups: int,
+    in_zero_point: int,
+    weight_zero_point: torch.Tensor,
+    bias_scale: torch.Tensor,
+    output_scale: float,
+    output_zero_point: int,
+    out_multiplier: torch.Tensor,
+    out_shift: torch.Tensor,
+) -> torch.Tensor:
+    return quantized_depthwise_conv1d_nlc_per_tensor_meta(
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        groups,
+        in_zero_point,
+        0,
+        1.0,
+        output_scale,
+        output_zero_point,
+        1,
+        0,
+    )
+
+
 @register_fake("cadence::quantized_conv2d_nchw")
 def quantized_conv2d_nchw_meta(
     input: torch.Tensor,