[ExecuTorch][XNNPACK] Don't partition 3d and transposed convs (#4796)

Little bug got past partitioning, because 3d and transposed convolutions were now being partitioned.

We expand the scope of the ConvConfig's check_constraint to also fail when the convolutions are either transposed or 3d

Co-authored-by: Max Ren <[email protected]>

Pull Request resolved: #4763

Author: kirklandsign

backends/xnnpack/partition/config/gemm_configs.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from itertools import chain
-from typing import List, Optional, Tuple
+from typing import cast, List, Optional, Tuple
 
 import torch
 from executorch.backends.xnnpack.partition.config.xnnpack_config import (
@@ -267,6 +267,23 @@ def __init__(self):
             fused_acts=["relu.default", "hardtanh.default"],
         )
 
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        """
+        Currently we have no support for convolution 3d and transposed convolution
+        """
+        if not super().check_constraints(node, ep):
+            return False
+
+        conv_stride = cast(List[int], node.args[3])
+        if len(conv_stride) > 2:
+            return False  # Only support 1D + 2D Conv
+
+        transposed = cast(bool, node.args[6])
+        if transposed:
+            return False  # Currently don't support transposed conv
+
+        return True
+
     def supported_precision_types(self):
         return [
             ConfigPrecisionType.FP32,