add flag for dynamic shapes to filter out static ops (#3733)

Summary:

Since we have updated XNNPACK to support almost 100% of dynamic shape ops, we can now create static_op lists which do not have any dynamic shape support and filter them out instead

Differential Revision: D57787384

Author: mcr229

backends/xnnpack/partition/configs.py

@@ -144,14 +144,12 @@
 
 SUPPORTED_DYN_QUANT_MODULES = SUPPORTED_DYN_QUANT_LINEAR_MODULES
 
-# TODO delete this once we catch up to 100% of the supported op with dynamic shape support.
-# This is tobe used only during the transition when we may not want to partition all the
-# nodes for a dynamic model.
-_SUPPORTED_OPS_WITH_DYNAMIC_SHAPE = [
-    exir_ops.edge.aten.add.Tensor,
-    exir_ops.edge.aten.mul.Tensor,
-]
-_SUPPORTED_MODULES_WITH_DYNAMIC_SHAPE = [
-    torch.nn.Conv1d,
-    torch.nn.Conv2d,
+# XNNPACK supports majority of shape dynamism, however some ops are
+# explicitly static, so we maintain a set here to exclude them from
+# dynamic shape support.
+STATIC_OPS = [
+    exir_ops.edge.aten.cat.default,
+    exir_ops.edge.aten.slice_copy.Tensor,
 ]
+
+STATIC_MODULES = []

backends/xnnpack/partition/xnnpack_partitioner.py

@@ -12,8 +12,8 @@
 import torch
 
 from executorch.backends.xnnpack.partition.configs import (
-    _SUPPORTED_MODULES_WITH_DYNAMIC_SHAPE,
-    _SUPPORTED_OPS_WITH_DYNAMIC_SHAPE,
+    STATIC_MODULES,
+    STATIC_OPS,
     SUPPORTED_DYN_QUANT_LINEAR_MODULES,
     SUPPORTED_DYN_QUANT_MODULES,
     SUPPORTED_MODULES,
@@ -94,6 +94,7 @@ def __init__(
         ] = _OP_SUPPORT_CONSTRAINTS,
         supported_ops: Optional[List] = None,
         unsupported_modules: Optional[List] = None,
+        dynamic_shapes=False,
     ):
         """
         @Arg constraints_dict: Dict mapping each node to a lambda function that
@@ -111,6 +112,7 @@ def __init__(
             exir_ops.edge.aten.mm.default,
             exir_ops.edge.aten.bmm.default,
         }
+        self.dynamic_shapes = dynamic_shapes
         assert len(self.constraints)
 
     def _check_inputs_are_valid_dtypes(self, node, valid_dtypes):
@@ -838,7 +840,7 @@ def __init__(
         supported_quant_modules: List[Callable] = SUPPORTED_QUANT_MODULES,
         supported_quant_ops: Optional[List[Callable]] = SUPPORTED_QUANT_OPS,
         quant: Optional[bool] = None,
-        _only_ops_with_dynamic_shape_support: Optional[bool] = False,
+        dynamic_shape: bool = False,
         _lower_recomposed_sdpa: Optional[bool] = True,
     ):
         super().__init__()
@@ -851,44 +853,16 @@ def __init__(
 
         self.quant = quant
 
-        if _only_ops_with_dynamic_shape_support is True:
-            self._update_op_lists_for_dynamic_shapes()
-
         # TODO(T174256335) - remove this once we have a better way to handle >2d Mask
         self._lower_recomposed_sdpa: bool = _lower_recomposed_sdpa or True
 
         self.delegation_spec = DelegationSpec(XnnpackBackend.__name__, [])
         self.partition_tags: Dict[str, DelegationSpec] = {}
 
-    def _update_op_lists_for_dynamic_shapes(self):
-        # Not ready for quants yet
-        assert (
-            self.quant is not True
-        ), "Dynamic shape only supported for valid FP32 ops, no quants support yet."
-        self.supported_quant_ops = set()
-        self.supported_quant_modules = set()
-
-        # for supported ops
-        self.supported_ops_with_dynamic_shape = set(_SUPPORTED_OPS_WITH_DYNAMIC_SHAPE)
-        assert self.supported_ops_with_dynamic_shape.issubset(
-            self.supported_ops
-        ), "All ops with dynamic shape support must be in SUPPORTED_OPS"
-        self.supported_ops = self.supported_ops_with_dynamic_shape
-        log.info(
-            f"Xnnpack Partitioner updated supported op for dynamic shapes: {self.supported_ops}"
-        )
-
-        # for supported modules
-        self.supported_modules_with_dynamic_shape = set(
-            _SUPPORTED_MODULES_WITH_DYNAMIC_SHAPE
-        )
-        assert self.supported_modules_with_dynamic_shape.issubset(
-            self.supported_modules
-        ), "All modules with dynamic shape support must be in SUPPORTED_MODULES"
-        self.supported_modules = self.supported_modules_with_dynamic_shape
-        log.info(
-            f"Xnnpack Partitioner updated supported modules with dynamic shapes: {self.supported_modules}"
-        )
+        self.dynamic_shape = dynamic_shape
+        if dynamic_shape:
+            self.supported_ops = self.supported_ops - set(STATIC_OPS)
+            self.supported_modules = self.supported_modules - set(STATIC_MODULES)
 
     def get_supported_modules(self, quant: bool) -> Set[Callable]:
         """