Move depthwise conv check to helper function in utils

Author: keyprocedure

backends/xnnpack/partition/config/gemm_configs.py

@@ -28,6 +28,7 @@
 )
 from executorch.backends.xnnpack.utils.utils import (
     get_input_node,
+    is_depthwise_conv,
     is_getitem,
     is_node,
     is_param_node,
@@ -365,21 +366,10 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         groups = cast(int, node.args[8])
         is_transpose = node.args[6]
 
-        if is_transpose:
-            group_input_channels = int(kernel_shape[0] / groups)
-            group_output_channels = kernel_shape[1]
-        else:
-            group_input_channels = kernel_shape[1]
-            group_output_channels = int(kernel_shape[0] / groups)
-
-        is_depthwise = (
-            group_input_channels == 1
-            and group_output_channels % group_input_channels == 0
-        )
-
         # XNNPACK does not support dynamic quantization convs that are not 2D or are depthwise
         if self._detect_precision(node) == ConfigPrecisionType.DYNAMIC_QUANT and (
-            len(conv_stride) != 2 or is_depthwise
+            len(conv_stride) != 2
+            or is_depthwise_conv(kernel_shape, groups, is_transpose)
         ):
             why(
                 node,

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

@@ -6,6 +6,7 @@
 
 import torch
 import torch.nn.functional as F
+from executorch.backends.xnnpack.utils.utils import is_depthwise_conv
 from torch._subclasses import FakeTensor
 from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix
 from torch.ao.quantization.pt2e.export_utils import _WrapperModule
@@ -29,7 +30,6 @@
 )
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
 
-
 __all__ = [
     "OperatorConfig",
     "OperatorPatternType",
@@ -336,20 +336,8 @@ def _do_annotate_conv(
             if weight_shape is not None and len(weight_shape) != 4:
                 continue
 
-            # Default to 1 since groups is not available in the node
-            groups = 1
-            if is_conv_transpose:
-                group_input_channels = int(weight_shape[0] / groups)
-                group_output_channels = weight_shape[1]
-            else:
-                group_input_channels = weight_shape[1]
-                group_output_channels = int(weight_shape[0] / groups)
-
-            # Skip if depthwise
-            if (
-                group_input_channels == 1
-                and group_output_channels % group_input_channels == 0
-            ):
+            # Skip if depthwise (default to groups=1 since it's not an arg)
+            if is_depthwise_conv(weight_shape, 1, is_conv_transpose):
                 continue
 
         # adding weight node to the partition as well

backends/xnnpack/utils/utils.py

@@ -158,3 +158,33 @@ def get_source_fn(node: torch.fx.Node) -> Optional[torch.fx.Node]:
         return None
     source_fn = source_fn_st[-1]
     return source_fn[1]
+
+
+def is_depthwise_conv(
+    kernel_shape: Tuple[int, ...], groups: int = 1, is_transpose: bool = False
+) -> bool:
+    """
+    A convolution is depthwise if:
+        1) groups = input_channels (i.e. group_input_channels = 1)
+        2) output_channels is a positive integer multiple of input channels
+
+    For standard convolutions:
+        weight shape = (out_channels, in_channels_per_group, height, width)
+    For transposed convolutions:
+        weight shape = (in_channels, out_channels_per_group, height, width)
+
+    Returns True if the convolution is depthwise
+    """
+    if len(kernel_shape) < 2 or groups < 1:
+        return False
+
+    if is_transpose:
+        group_input_channels = int(kernel_shape[0] / groups)
+        group_output_channels = kernel_shape[1]
+    else:
+        group_input_channels = kernel_shape[1]
+        group_output_channels = int(kernel_shape[0] / groups)
+
+    return (
+        group_input_channels == 1 and group_output_channels % group_input_channels == 0
+    )