[ExecuTorch][Vulkan] Support scoped Gemma symbolic partitioning

backends/vulkan/_passes/fuse_patterns.py

@@ -11,6 +11,7 @@
 import torch
 
 from executorch.exir import ExportedProgram
+from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
 
@@ -27,6 +28,20 @@ def call(self, graph_module: torch.fx.GraphModule):
         )
 
         if total_replaced > 0:
+            for node in list(graph_module.graph.nodes):
+                if node.target != exir_ops.edge.et_vk.select_as_symint.default:
+                    continue
+                value_range = node.meta.get("et_vk_value_range")
+                if value_range is None:
+                    continue
+                lower_bound, upper_bound = value_range
+                with graph_module.graph.inserting_after(node):
+                    graph_module.graph.create_node(
+                        "call_function",
+                        exir_ops.edge.aten.sym_constrain_range.default,
+                        args=(node,),
+                        kwargs={"min": lower_bound, "max": upper_bound},
+                    )
             graph_module.recompile()
             # Re-trace the graph
             graph_module = super().call(graph_module).graph_module

backends/vulkan/_passes/remove_asserts.py

@@ -28,6 +28,7 @@ class RemoveAssertsTransform(ExportPass):
 
     assert_ops: Set[OpType] = {
         torch.ops.aten._assert_scalar.default,
+        torch.ops.aten.sym_constrain_range.default,
         torch.ops.aten.sym_constrain_range_for_size.default,
     }
 

backends/vulkan/_passes/squeeze_unsqueeze_inputs.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -16,6 +16,11 @@
 from torch._ops import OpOverload
 
 from torch.fx.node import Argument
+from torch.fx.experimental.symbolic_shapes import (
+    statically_known_false,
+    statically_known_true,
+    sym_and,
+)
 
 OpType = Union[str, OpOverload, EdgeOpOverload]
 
@@ -26,21 +31,47 @@
         exir_ops.edge.aten.gelu.default,
     }
 
+    @staticmethod
+    def _first_static_one(shape: List[int]) -> Union[int, None]:  # pyre-ignore
+        for index, dim in enumerate(shape):
+            if statically_known_true(dim == 1):
+                return index
+        return None
+
+    def _squeezed_shape(self, shape: List[int]) -> List[int]:  # pyre-ignore
+        squeezed_shape = list(shape)
+        while len(squeezed_shape) > 2:
+            index = self._first_static_one(squeezed_shape)
+            if index is None:
+                break
+            squeezed_shape.pop(index)
+        return squeezed_shape
+
     def should_squeeze(self, op, shape: List[int]) -> bool:  # pyre-ignore
         if len(shape) == 3:
-            return shape[1] == 1 and shape[0] > 1
+            return statically_known_true(sym_and(shape[1] == 1, shape[0] > 1))
         if len(shape) == 4:
-            # No need to squeeze if all dims are 1 except the width dim
-            if shape[0] == shape[1] == shape[2] == 1:
-                return False
-            # No need to squeeze if batch and channel dims are 1 and height and width are > 1
-            if shape[0] == shape[1] == 1 and shape[2] > 1 and shape[3] > 1:
-                return False
-            # No need to squeeze if batch dim is 1 and channel, height and width are > 1
-            if shape[0] == 1 and shape[1] > 1 and shape[2] > 1 and shape[3] > 1:
+            excluded_shapes = (
+                sym_and(shape[0] == 1, shape[1] == 1, shape[2] == 1),
+                sym_and(
+                    shape[0] == 1,
+                    shape[1] == 1,
+                    shape[2] > 1,
+                    shape[3] > 1,
+                ),
+                sym_and(
+                    shape[0] == 1,
+                    shape[1] > 1,
+                    shape[2] > 1,
+                    shape[3] > 1,
+                ),
+            )
+            if any(
+                not statically_known_false(excluded_shape)
+                for excluded_shape in excluded_shapes
+            ):
                 return False
-            # Otherwise, check for squeezable dim
-            return 1 in shape[:-1]
+            return self._first_static_one(shape[:-1]) is not None
 
         # Prefer not to introduce additional orchestration ops by default
         return False
@@ -61,18 +92,13 @@
         if not self.should_squeeze(op, input_shape):
             return super().call_operator(op, args, kwargs, meta)
 
-        def _squeezable(shape: List[int]) -> bool:
-            return len(shape) > 2 and 1 in shape
-
         # squeeze input tensor
-        squeeze_shape = list(input_shape)
-        while _squeezable(squeeze_shape):
-            squeeze_shape.remove(1)
+        squeeze_shape = self._squeezed_shape(input_shape)
 
         squeeze_out = super().call_operator(
             exir_ops.edge.aten.view_copy.default,
             (args[0], squeeze_shape),
-            kwargs,
+            {},
             meta,
         )
         # call linear on squeezed output
@@ -88,6 +114,6 @@
         return super().call_operator(
             exir_ops.edge.aten.view_copy.default,
             (linear_out, unsqueeze_shape),
-            kwargs,
+            {},
             meta,
         )

backends/vulkan/custom_ops_lib.py

@@ -882,11 +882,22 @@ def apply_rotary_emb_hf_impl(
     return pattern.forward(xq, xk, freqs_cos, freqs_sin)
 
 
+def apply_rotary_emb_hf_meta(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+    start_pos: int,
+):
+    return torch.empty_like(xq), torch.empty_like(xk)
+
+
 name = "apply_rotary_emb_hf"
 lib.define(
     f"{name}(Tensor xq, Tensor xk, Tensor freqs_cos, Tensor freqs_sin, SymInt start_pos) -> (Tensor, Tensor)"
 )
 lib.impl(name, apply_rotary_emb_hf_impl, "CompositeExplicitAutograd")
+lib.impl(name, apply_rotary_emb_hf_meta, "Meta")
 apply_rotary_emb_hf_op = getattr(getattr(torch.ops, namespace), name)
 
 ##################################
@@ -1074,7 +1085,7 @@ def embedding_q4gsw_impl(
     scales = (
         weight_scales.unsqueeze(-1)
         if weight_scales.dim() > 1
-        else weight_scales.reshape(1, 1, 1)
+        else weight_scales.reshape(weight.shape[0], 1, 1)
     )
     dequantized = unpacked_groups.float() * scales.float()
     dequantized = dequantized.reshape(weight.shape[0], -1)
@@ -1098,8 +1109,15 @@ def select_as_symint_impl(x: torch.Tensor, dim: int, index: int):
     return x.fake_mode.shape_env.create_unbacked_symint()
 
 
+def select_as_symint_eager_impl(x: torch.Tensor, dim: int, index: int):
+    if x.dtype not in {torch.int32, torch.int64}:
+        raise ValueError("select_as_symint requires an integral input")
+    return x.select(dim, index).item()
+
+
 name = "select_as_symint"
 lib.define(f"{name}(Tensor x, int dim, int index) -> SymInt")
+lib.impl(name, select_as_symint_eager_impl, "CompositeExplicitAutograd")
 lib.impl(name, select_as_symint_impl, "Meta")
 select_as_symint_op = getattr(getattr(torch.ops, namespace), name)
 

backends/vulkan/op_registry.py

@@ -152,6 +152,15 @@ def update_features_impl(op: OpKey):
 # =============================================================================
 
 
+@update_features(exir_ops.edge.et_vk.select_as_symint.default)
+def register_select_as_symint_op():
+    return OpFeatures(
+        inputs_dtypes=utils.INT_T,
+        inputs_storage=utils.ANY_STORAGE,
+        supports_resize=True,
+    )
+
+
 @update_features(
     [
         operator.getitem,
@@ -161,6 +170,7 @@ def update_features_impl(op: OpKey):
         operator.sub,
         operator.floordiv,
         operator.mul,
+        operator.and_,
         operator.lt,
         operator.gt,
         operator.ge,

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -7,6 +7,7 @@
 # pyre-strict
 
 import logging
+import operator
 from typing import Any, Callable, Dict, final, List, Mapping, Optional, Set, Tuple
 
 import executorch.backends.vulkan.patterns as vk_patterns
@@ -54,6 +55,58 @@
 logger: logging.Logger = logging.getLogger("")
 logger.setLevel(logging.INFO)
 
+FP_T = utils.FP_T
+INT_T = utils.INT_T
+NONE_T = utils.NONE_T
+CONTIGUOUS_BUFFER = utils.CONTIGUOUS_BUFFER
+NO_STORAGE = utils.NO_STORAGE
+
+
+_GUARD_ONLY_SYMBOLIC_OPS: Set[Any] = {torch.sym_min, torch.sym_max}
+_GUARD_ONLY_EXPRESSION_OPS: Set[Any] = _GUARD_ONLY_SYMBOLIC_OPS | {
+    operator.add,
+    operator.sub,
+    operator.floordiv,
+    operator.mul,
+    operator.and_,
+    operator.lt,
+    operator.gt,
+    operator.ge,
+    operator.le,
+    operator.eq,
+}
+_GUARD_ONLY_SINK_OPS: Set[Any] = {
+    torch.ops.aten._assert_scalar.default,
+    torch.ops.aten.sym_constrain_range_for_size.default,
+}
+
+
+def _is_guard_only_symbolic_node(node: torch.fx.Node) -> bool:
+    if node.target not in _GUARD_ONLY_SYMBOLIC_OPS or not node.users:
+        return False
+
+    pending = list(node.users)
+    visited: Set[torch.fx.Node] = set()
+    reached_sink = False
+    while pending:
+        user = pending.pop()
+        if user in visited:
+            continue
+        visited.add(user)
+
+        if user.op != "call_function" or utils.is_tensor_node(user):
+            return False
+        if user.target in _GUARD_ONLY_SINK_OPS:
+            if user.users:
+                return False
+            reached_sink = True
+            continue
+        if user.target not in _GUARD_ONLY_EXPRESSION_OPS or not user.users:
+            return False
+        pending.extend(user.users)
+
+    return reached_sink
+
 
 class VulkanSupportedOperators(OperatorSupportBase):
     def __init__(
@@ -67,6 +120,7 @@ def __init__(
         fusable_subgraphs: Optional[List[PatternMatch]] = None,
         nn_module_blocklist: Optional[Set[str]] = None,
         nn_module_allowlist: Optional[Set[str]] = None,
+        extra_op_features: Optional[Mapping[OpKey, OpFeatures]] = None,
     ) -> None:
         super().__init__()
         self.texture_limits: utils.ImageExtents = texture_limits
@@ -87,6 +141,7 @@ def __init__(
 
         self.nn_module_blocklist = nn_module_blocklist
         self.nn_module_allowlist = nn_module_allowlist
+        self.extra_op_features: Dict[OpKey, OpFeatures] = dict(extra_op_features or {})
 
     def op_node_is_compatible(  # noqa: C901: Function is too complex
         self, node: torch.fx.Node, features: Optional[OpFeatures] = None
@@ -147,11 +202,23 @@ def op_node_is_compatible(  # noqa: C901: Function is too complex
     def node_is_compatible(
         self, node: torch.fx.Node, features: Optional[OpFeatures] = None
     ) -> Tuple[bool, str]:
+        # Guard-only symbolic nodes (sym_min/sym_max feeding only asserts or
+        # range constraints) are non-tensor, so this must run before the
+        # is_tensor_node dispatch below or it is unreachable.
+        if getattr(node, "target", None) in _GUARD_ONLY_SYMBOLIC_OPS:
+            if _is_guard_only_symbolic_node(node):
+                return True, "guard-only symbolic node"
+            self.log_skip(node, "symbolic result has a live non-guard user")
+            return False, "symbolic result has a live non-guard user"
+
         if utils.is_tensor_node(node):
             return self.op_node_is_compatible(node, features=features)
         # For non-tensor nodes, just check if the op is registered
         elif hasattr(node, "target"):
-            return node.target in vulkan_supported_ops, "Op is compatible"
+            return (
+                features is not None or node.target in vulkan_supported_ops,
+                "Op is compatible",
+            )
 
         return False, f"Unsupported node type: {node.format_node()}"
 
@@ -244,17 +311,22 @@ def _is_node_supported(self, node: torch.fx.Node) -> bool:  # noqa: C901
             self.log_skip(node, "permute node of non compatible linear node")
             return False
 
-        features = None
-        if target not in vulkan_supported_ops:
+        features: Optional[OpFeatures] = None
+        if target in vulkan_supported_ops:
+            features = vulkan_supported_ops[target]
+        elif target in self.extra_op_features:
+            features = self.extra_op_features[target]
+        else:
             # For some ops, i.e. custom ops the name is registered instead of the
             # OpOverload object.
-            if hasattr(target, "name") and target.name() in vulkan_supported_ops:
-                features = vulkan_supported_ops[target.name()]
+            target_name = target.name() if hasattr(target, "name") else None
+            if target_name in vulkan_supported_ops:
+                features = vulkan_supported_ops[target_name]
+            elif target_name in self.extra_op_features:
+                features = self.extra_op_features[target_name]
             else:
                 self.log_skip(node, "no operator implementation")
                 return False
-        else:
-            features = vulkan_supported_ops[target]
 
         assert features is not None
 
@@ -341,6 +413,7 @@ def __init__(
         operator_allowlist: Optional[List[OpKey]] = None,
         nn_module_blocklist: Optional[List[str]] = None,
         nn_module_allowlist: Optional[List[str]] = None,
+        extra_op_features: Optional[Mapping[OpKey, OpFeatures]] = None,
     ) -> None:
         self.options: Dict[str, Any] = {}
         if compile_options is not None:
@@ -349,6 +422,14 @@ def __init__(
         compile_spec = parse_compile_options(self.options)
         self.delegation_spec = DelegationSpec(VulkanBackend.__name__, compile_spec)
 
+        self.extra_op_features: Dict[OpKey, OpFeatures] = dict(extra_op_features or {})
+        overlapping_ops = self.extra_op_features.keys() & vulkan_supported_ops.keys()
+        if overlapping_ops:
+            raise ValueError(
+                "extra_op_features contains operators already registered globally: "
+                f"{sorted(str(op) for op in overlapping_ops)}"
+            )
+
         self.operator_blocklist: Set[OpKey] = set()
         if operator_blocklist is not None:
             for entry in operator_blocklist or []:
@@ -415,6 +496,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 fusable_subgraphs=fusable_subgraphs,
                 nn_module_blocklist=self.nn_module_blocklist,
                 nn_module_allowlist=self.nn_module_allowlist,
+                extra_op_features=self.extra_op_features,
             ),
             allows_single_node_partition=True,
         )