update dynamic shape detection

Summary: Updating Dynamic Shape Detection re: #5794

Differential Revision: D68036835

Author: mcr229

backends/xnnpack/operators/op_squeeze.py

@@ -16,6 +16,10 @@
     XNNStaticReshape,
     XNode,
 )
+from torch.fx.experimental.symbolic_shapes import (
+    free_symbols,
+)
+
 from executorch.backends.xnnpack.utils.utils import check_or_raise, get_input_node
 
 
@@ -57,7 +61,7 @@ def define_node(
 
         num_dynamic_dims = 0
         for dim in dynamic_shape:
-            if isinstance(dim, torch.SymInt):
+            if free_symbols(dim):
                 num_dynamic_dims += 1
                 new_shape.append(0)
             else:
@@ -119,7 +123,7 @@ def define_node(
 
         num_dynamic_dims = 0
         for dim in dynamic_shape:
-            if isinstance(dim, torch.SymInt):
+            if free_symbols(dim):
                 num_dynamic_dims += 1
                 new_shape.append(0)
             else:

exir/backend/utils.py

@@ -23,6 +23,7 @@
 
 from executorch.exir.lowered_backend_module import create_submodule_from_nodes
 from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
+from torch.fx.experimental.symbolic_shapes import has_free_symbols
 from torch.fx.node import Node
 from torch.fx.passes.utils.source_matcher_utils import SourcePartition
 
@@ -424,10 +425,7 @@ def is_shape_dynamic(node: torch.fx.Node) -> bool:
     Check if the node shape is dynamic.
     """
 
-    # Shape is dynamic if any of the dimensions don't evaluate to a static value
-    return "val" in node.meta and any(
-        isinstance(d, torch.SymInt) for d in node.meta["val"].shape
-    )
+    return has_free_symbols(node.meta["val"].shape)
 
 
 # TODO - style: use templated types