Update base for Update on "[ET] correcting cpu ref quantize_per_channel logic to align with ATen"

# Context

The quantize_per_channel was not perfectly aligned with the ATen implementation, and demonstrated errors when specifying different axis. This bug wasn't distinctly acknowledged given that the test cases only has one test for the whole operator. In order to align more closely with ATen this change simply does a single loop imlpementation with direct channel index calculation over the old `apply_over_dim_list` approach.

# Changes

We change the core logic for quantize_per_channel to more properly align with ATen's implementation, and we also change it from `apply_over_dim_list` approach to a single loop implementation with direct channel index calculation. This also adds more comprehensive testing for quantize_per_channel so that a bug isn't missed again.

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

[ghstack-poisoned]

Author: morelos

.gitignore

@@ -26,6 +26,7 @@ pip-out/
 *.model
 tokenizer.json
 *.pte
+*.ptd
 !test_bpe_tokenizer.bin
 !test_tiktoken_tokenizer.model
 

backends/vulkan/serialization/vulkan_graph_builder.py

@@ -45,7 +45,7 @@ def __init__(
         self,
         program: ExportedProgram,
         delegate_mapping_builder: DelegateMappingBuilder,
-        downcast_64_bit: bool = False,
+        downcast_64_bit: bool = True,
     ) -> None:
         self.program = program
         self.delegate_mapping_builder = delegate_mapping_builder

backends/xnnpack/_passes/__init__.py

@@ -21,9 +21,7 @@
 )
 from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
-from executorch.backends.xnnpack._passes.fuse_batch_norm_with_conv import (
-    FuseBatchNormWithConvPass,
-)
+from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 
@@ -60,7 +58,7 @@ def __init__(
                 ConvertToLinearPass,
                 ConvertToSDPAPass,
                 ConstPropPass,
-                FuseBatchNormWithConvPass,
+                FuseBatchNormPass,
                 FuseActivationPass,
                 DecomposeConcatenate,
                 RemoveGetItemPass,

backends/xnnpack/_passes/fuse_batch_norm.py

@@ -0,0 +1,247 @@
+# 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 operator
+
+import torch
+from executorch.backends.transforms.utils import (
+    create_constant_placeholder,
+    delete_constant_placeholder,
+)
+
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
+
+from executorch.backends.xnnpack.utils.utils import (
+    get_param_tensor,
+    get_tensor_name,
+    is_param_node,
+)
+from executorch.exir import ExportedProgram
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import PassResult
+from torch.export.graph_signature import InputKind
+
+from torch.nn.utils.fusion import fuse_conv_bn_weights, fuse_linear_bn_weights
+
+
+class FuseBatchNormPass(XNNPACKPass):
+    """
+    BatchNorm can be implemented using 1x1 Depthwise Convolution. However, doing so will increase
+    memory usage since we serialize new weights to represent the convolution. In most cases,
+    BatchNorm is used after convolution or linear. The 1x1 depthwise convolution can then be fused
+    with the previous convolution. For linear cases, BatchNorm can be folded into the previous linear layer.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        graph = graph_module.graph
+        constant_placeholders_to_delete = set()
+        for input_node in graph.nodes:
+            # We want to discover a chain of conv -> batch_norm or linear -> batch_norm.
+            # Only proceed if the current node is a conv or linear, and has a single user/successor.
+            is_conv = input_node.target == exir_ops.edge.aten.convolution.default
+            is_linear = input_node.target == exir_ops.edge.aten.linear.default
+
+            if not (is_conv or is_linear) or len(input_node.users) != 1:
+                continue
+
+            # The single user of the conv or linear node must be batch_norm. If not, bail.
+            bn = list(input_node.users.keys())[0]
+            if (
+                bn.target != exir_ops.edge.aten.native_batch_norm.default
+                and bn.target
+                != exir_ops.edge.aten._native_batch_norm_legit_no_training.default
+            ):
+                continue
+
+            if not self.can_fuse(input_node, bn, self.exported_program):
+                continue
+
+            self._fuse_ops(
+                graph_module,
+                graph,
+                input_node,
+                bn,
+                is_conv,
+                constant_placeholders_to_delete,
+            )
+
+        if len(constant_placeholders_to_delete) > 0:
+            graph_module.graph.eliminate_dead_code()
+            for node in constant_placeholders_to_delete:
+                if (node is not None) and (len(node.users) == 0):
+                    delete_constant_placeholder(self.exported_program, node)
+
+        graph_module.recompile()
+        # To regenerate metadata and shape information, retrace module.
+        graph_module = super().call(graph_module).graph_module
+
+        return PassResult(graph_module, True)
+
+    @staticmethod
+    def can_fuse(
+        input_node: torch.fx.Node, bn: torch.fx.Node, program: ExportedProgram
+    ) -> bool:
+        """
+        Determine whether a BatchNorm node can be fused with the preceding convolution or linear node.
+        """
+
+        # All users of the batch_norm node must be getitem ops.
+        # batch_norm returns a 3-element tuple.
+        # Each user must only access the first element of the tuple.
+        if [
+            (user.target == operator.getitem and user.args[1] == 0) for user in bn.users
+        ].count(False):
+            return False
+
+        input_node_weights = input_node.args[1]
+        bn_weights = bn.args[1]
+
+        # Check that the weights for conv or linear and batch_norm are both params.
+        if not isinstance(input_node_weights, torch.fx.Node) or not isinstance(
+            bn_weights, torch.fx.Node
+        ):
+            return False
+
+        if [
+            is_param_node(program, node) for node in {input_node_weights, bn_weights}
+        ].count(False):
+            return False
+
+        return True
+
+    def _fuse_ops(
+        self,
+        graph_module: torch.fx.GraphModule,
+        graph: torch.fx.Graph,
+        input_node: torch.fx.Node,
+        bn: torch.fx.Node,
+        is_conv: bool,
+        constant_placeholders_to_delete: set,
+    ) -> None:
+        """
+        Fuse a BatchNorm node into the preceding convolution or linear node.
+        Update the fused node's weight and bias, rewire users of the BatchNorm output,
+        and remove the BatchNorm node.
+        """
+
+        if is_conv:
+            assert len(input_node.args) == 9
+            has_bias_arg = True
+        else:
+            # Otherwise, this is a linear node.
+            # Linear has 2 or 3 args depending on whether bias is used: (input, weight, bias).
+            assert len(input_node.args) in (2, 3)
+            has_bias_arg = len(input_node.args) == 3
+
+        # Get the weight and bias parameters from the conv or linear op.
+        input_node_weight = get_param_tensor(self.exported_program, input_node.args[1])
+        input_node_weight_name = get_tensor_name(
+            self.exported_program, input_node.args[1]
+        )
+        assert input_node_weight is not None
+
+        if has_bias_arg:
+            input_node_bias = get_param_tensor(
+                self.exported_program, input_node.args[2]
+            )
+            input_node_bias_name = get_tensor_name(
+                self.exported_program, input_node.args[2]
+            )
+        else:
+            input_node_bias = None
+            input_node_bias_name = ""
+
+        # Get the parameters from the batch_norm op.
+        assert (
+            bn.target == exir_ops.edge.aten.native_batch_norm.default
+            and len(bn.args) == 8
+        ) or (
+            bn.target == exir_ops.edge.aten._native_batch_norm_legit_no_training.default
+            and len(bn.args) == 7
+        )
+        bn_weight = get_param_tensor(self.exported_program, bn.args[1])
+        bn_bias = get_param_tensor(self.exported_program, bn.args[2])
+
+        running_mean = get_param_tensor(self.exported_program, bn.args[3])
+        assert running_mean is not None
+
+        running_var = get_param_tensor(self.exported_program, bn.args[4])
+        assert running_var is not None
+
+        # args[7] for native_batch_norm, but args[6] for
+        # _native_batch_norm_legit_no_training (which doesn't have training
+        # as an arg).
+        eps = bn.args[-1]
+
+        # Compute the updated weight and bias after fusing the conv or linear op with the batch_norm op.
+        fuse_args = (
+            input_node_weight,
+            input_node_bias,
+            running_mean,
+            running_var,
+            eps,
+            bn_weight,
+            bn_bias,
+        )
+
+        if is_conv:
+            is_transpose = input_node.args[6]
+            fused_weight, fused_bias = fuse_conv_bn_weights(*fuse_args, is_transpose)
+        else:
+            # Otherwise, this is a linear node.
+            fused_weight, fused_bias = fuse_linear_bn_weights(*fuse_args)
+
+        fused_weight_name = (input_node_weight_name + "_fused_bn").replace(".", "_")
+        if input_node_bias_name == "":
+            fused_bias_name = (input_node_weight_name + "_bias_fused_bn").replace(
+                ".", "_"
+            )
+        else:
+            fused_bias_name = (input_node_bias_name + "_fused_bn").replace(".", "_")
+
+        # Modify the graph by updating the weight and bias of the conv or linear op
+        # with the fused weight and bias params, and replacing all the users
+        # of getitem(batch_norm) with the conv or linear op.
+        with graph.inserting_before(input_node.args[1]):
+            fused_op_weight_node = create_constant_placeholder(
+                exp_program=self.exported_program,
+                graph=graph_module.graph,
+                kind=InputKind.PARAMETER,
+                name=fused_weight_name,
+                data=fused_weight,
+            )
+            if fused_bias is not None:
+                fused_op_bias_node = create_constant_placeholder(
+                    exp_program=self.exported_program,
+                    graph=graph_module.graph,
+                    kind=InputKind.PARAMETER,
+                    name=fused_bias_name,
+                    data=fused_bias,
+                )
+            else:
+                fused_op_bias_node = None
+
+            # Replace the original weight and bias with the fused batch_norm values.
+            args = list(input_node.args)
+            args[1] = fused_op_weight_node
+
+            if has_bias_arg:
+                # Overwrite original bias with the fused bias.
+                args[2] = fused_op_bias_node
+            elif fused_op_bias_node is not None:
+                # Add the fused bias as a new argument if no bias had originally existed in the input_node.
+                args.append(fused_op_bias_node)
+
+            input_node.args = tuple(args)
+
+            # Remove any use of batch_norm from the graph.
+            for user in bn.users.copy():
+                assert user.target == operator.getitem
+                user.replace_all_uses_with(input_node)
+                graph.erase_node(user)
+
+            graph.erase_node(bn)
+            constant_placeholders_to_delete.update(input_node.args[1:3] + bn.args[1:5])