diff --git a/backends/vulkan/_passes/fuse_patterns.py b/backends/vulkan/_passes/fuse_patterns.py index 1575dd6a4f6..b3c0c796236 100644 --- a/backends/vulkan/_passes/fuse_patterns.py +++ b/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 diff --git a/backends/vulkan/_passes/remove_asserts.py b/backends/vulkan/_passes/remove_asserts.py index 835f2ec1415..4dcf8fe1d8d 100644 --- a/backends/vulkan/_passes/remove_asserts.py +++ b/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, } diff --git a/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py b/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py index 25b28ce3117..bfc0ce899a6 100644 --- a/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py +++ b/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py @@ -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 @@ class SqueezeUnsqueezeInputs(ExportPass): 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 @@ def call_operator( 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 @@ def _squeezable(shape: List[int]) -> bool: return super().call_operator( exir_ops.edge.aten.view_copy.default, (linear_out, unsqueeze_shape), - kwargs, + {}, meta, ) diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index a074597466a..68399681699 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/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) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index a88df42ee7b..ec44b8dc8d4 100644 --- a/backends/vulkan/op_registry.py +++ b/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, diff --git a/backends/vulkan/partitioner/vulkan_partitioner.py b/backends/vulkan/partitioner/vulkan_partitioner.py index 298581ebef7..f500fd8e093 100644 --- a/backends/vulkan/partitioner/vulkan_partitioner.py +++ b/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, ) diff --git a/backends/vulkan/patterns/__init__.py b/backends/vulkan/patterns/__init__.py index 68df9905671..0f897e0e6e2 100644 --- a/backends/vulkan/patterns/__init__.py +++ b/backends/vulkan/patterns/__init__.py @@ -43,10 +43,15 @@ ) from executorch.backends.vulkan.patterns.rope import RotaryEmbeddingPattern -from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingPattern +from executorch.backends.vulkan.patterns.rope_hf import ( + HfRotaryEmbeddingPattern, + HfRotaryEmbeddingSinglePattern, +) from executorch.exir import ExportedProgram +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from torch._library.utils import is_impure from torch.fx.passes.utils.matcher_utils import SubgraphMatcher @@ -57,6 +62,7 @@ "CreateReplacementFn", "RotaryEmbeddingPattern", "HfRotaryEmbeddingPattern", + "HfRotaryEmbeddingSinglePattern", "fusable_patterns", "register_pattern_graph", "register_pattern_detector", @@ -64,6 +70,14 @@ ] +def _is_impure_node(node: torch.fx.Node) -> bool: + if node.is_impure(): + return True + if isinstance(node.target, EdgeOpOverload): + return is_impure(node.target._op, args=node.args, kwargs=node.kwargs) + return False + + def all_fusable_graph_patterns() -> List[torch.fx.GraphModule]: all_patterns = [] for entry in fusable_patterns.values(): @@ -112,7 +126,7 @@ def create_replacement_for_pattern( create_replacement_func(ep, graph_module, pattern) total_replaced += 1 # Remove dead code so they won't be matched again - graph_module.graph.eliminate_dead_code() + graph_module.graph.eliminate_dead_code(_is_impure_node) return total_replaced @@ -147,5 +161,5 @@ def replace_all_fusable_subgraphs( entry.create_replacement_fn(ep, graph_module, maybe_match) total_replaced += 1 - graph_module.graph.eliminate_dead_code() + graph_module.graph.eliminate_dead_code(_is_impure_node) return total_replaced diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py index 86a35298fa4..121d9c5141a 100644 --- a/backends/vulkan/patterns/quantized_linear.py +++ b/backends/vulkan/patterns/quantized_linear.py @@ -145,6 +145,14 @@ def __init__(self, mm_node: torch.fx.Node) -> None: # noqa: C901 ): self.quantize_input_node = input_to_dq_node self.pattern_input_node = self.quantize_input_node.args[0] + self.all_nodes.append(self.quantize_input_node) + if isinstance(self.input_scales_node, torch.fx.Node): + self.all_nodes.append(self.input_scales_node) + choose_qparams_node = self.input_scales_node.args[0] + if isinstance(choose_qparams_node, torch.fx.Node): + self.all_nodes.append(choose_qparams_node) + if isinstance(self.input_zeros_node, torch.fx.Node): + self.all_nodes.append(self.input_zeros_node) # The implementation has a limitation that input channels must be a # multiple of 4. This is to ensure that data loads are aligned well with @@ -595,6 +603,75 @@ def make_q8ta_linear_custom_op( match.quantize_output_node.replace_all_uses_with(qlinear_node) +def make_reused_linear_dq8ca_q4gsw_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: QuantizedLinearMatch, + packed_weight: torch.Tensor, + transposed_scales: torch.Tensor, +): + dequant_args = match.dequantize_weight_node.args + if ( + len(dequant_args) < 7 + or not isinstance(dequant_args[1], (list, tuple)) + or len(dequant_args[1]) != 2 + or dequant_args[1][0] != 1 + or dequant_args[4] != torch.int8 + or dequant_args[5:7] != (-8, 7) + ): + raise RuntimeError("shared 8da4w weight has an unexpected dequant ABI") + group_size = dequant_args[1][1] + if not isinstance(group_size, int) or group_size <= 0: + raise RuntimeError("shared 8da4w weight has an invalid group size") + if ( + packed_weight.dtype != torch.uint8 + or packed_weight.ndim != 2 + or transposed_scales.ndim != 2 + or transposed_scales.dtype != torch.float32 + or transposed_scales.shape[1] != packed_weight.shape[0] + or packed_weight.shape[1] * 2 != transposed_scales.shape[0] * group_size + ): + raise RuntimeError("shared 8da4w packed weight geometry mismatch") + + weight_name = utils.get_tensor_name(ep, match.weight_node) + sums_name = (weight_name + "_sums").replace(".", "_") + weight_sums_node = next( + ( + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.target == sums_name + ), + None, + ) + if weight_sums_node is None: + raise RuntimeError("shared 8da4w weight sums placeholder is missing") + weight_sums = get_param_tensor(ep, weight_sums_node) + if ( + weight_sums is None + or weight_sums.dtype != torch.int32 + or tuple(weight_sums.shape) != tuple(transposed_scales.shape) + ): + raise RuntimeError("shared 8da4w weight sums geometry mismatch") + + with graph_module.graph.inserting_before(match.output_node): + qlinear_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default, + args=( + match.pattern_input_node, + match.input_scales_node, + match.input_zeros_node, + match.weight_node, + weight_sums_node, + match.weight_scales_node, + group_size, + match.bias_node, + ), + ) + qlinear_node.meta["val"] = match.output_node.meta["val"] + match.output_node.replace_all_uses_with(qlinear_node) + + @register_pattern_replacement("quantized_linear") def replace_quantized_linear_patterns( ep: ExportedProgram, @@ -613,6 +690,27 @@ def replace_quantized_linear_patterns( weight_zeros_tensor = get_param_tensor(ep, match.weight_zeros_node) assert weight_zeros_tensor is not None + modification_tags = getattr(ep, "_et_vk_param_modification_tags", {}) + weight_tag = modification_tags.get(utils.get_tensor_name(ep, match.weight_node)) + scales_tag = modification_tags.get( + utils.get_tensor_name(ep, match.weight_scales_node) + ) + if weight_tag is not None or scales_tag is not None: + if ( + weight_tag != "4 bit linear weight" + or scales_tag != "4 bit linear scales" + or not match.is_input_dynamic_perchannel_quantized() + ): + raise RuntimeError("shared quantized linear mutation tags are inconsistent") + make_reused_linear_dq8ca_q4gsw_op( + ep, + graph_module, + match, + weight_tensor, + weight_scales_tensor, + ) + return + # Route to appropriate custom op. if ( match.is_input_static_per_tensor_quantized() diff --git a/backends/vulkan/patterns/rms_norm.py b/backends/vulkan/patterns/rms_norm.py index beb5e677ead..11904fdab8f 100644 --- a/backends/vulkan/patterns/rms_norm.py +++ b/backends/vulkan/patterns/rms_norm.py @@ -35,6 +35,16 @@ def _skip_casts(node: torch.fx.Node) -> torch.fx.Node: return node +def _is_rstd_node(node: torch.fx.Node) -> bool: + node = _skip_casts(node) + if node.target == exir_ops.edge.aten.rsqrt.default: + return True + if node.target != exir_ops.edge.aten.pow.Tensor_Scalar or len(node.args) < 2: + return False + exponent = node.args[1] + return isinstance(exponent, (int, float)) and exponent == -0.5 + + class RmsNormMatch(PatternMatch): """ Detects the decomposed RMSNorm pattern, including variants where dtype @@ -73,17 +83,18 @@ def __init__(self, final_mul_node: torch.fx.Node) -> None: # noqa: C901 self.all_nodes.append(norm_mul_node) # norm_mul: mul(x_f32, rstd_f32) - rsqrt_node, x_for_norm = self._identify_rsqrt_and_input(norm_mul_node) - if rsqrt_node is None: + rstd_node, x_for_norm = self._identify_rstd_and_input(norm_mul_node) + if rstd_node is None: return - self.all_nodes.append(rsqrt_node) + self.all_nodes.append(rstd_node) - # rsqrt -> add(mean_sq, eps) -> mean(x_sq, dim=-1, keepdim=True) - add_node = self._get_single_arg_node( - rsqrt_node, exir_ops.edge.aten.rsqrt.default - ) - if add_node is None or add_node.target != exir_ops.edge.aten.add.Tensor: + # rstd -> add(mean_sq, eps) -> mean(x_sq, dim=-1, keepdim=True) + add_node = rstd_node.args[0] if rstd_node.args else None + if ( + not isinstance(add_node, torch.fx.Node) + or add_node.target != exir_ops.edge.aten.add.Tensor + ): return self.all_nodes.append(add_node) @@ -182,49 +193,37 @@ def _identify_norm_mul_and_weight(self, final_mul_node): if ( isinstance(norm_candidate, torch.fx.Node) and norm_candidate.target == exir_ops.edge.aten.mul.Tensor - and self._has_rsqrt_ancestor(norm_candidate) + and self._has_rstd_ancestor(norm_candidate) ): return norm_candidate, weight_candidate_raw return None, None - def _has_rsqrt_ancestor(self, mul_node): - """Check if one of mul_node's args is an rsqrt node (possibly through casts).""" + def _has_rstd_ancestor(self, mul_node): + """Check if one arg computes reciprocal root mean square.""" for arg in mul_node.args[:2]: if not isinstance(arg, torch.fx.Node): continue - if _skip_casts(arg).target == exir_ops.edge.aten.rsqrt.default: + if _is_rstd_node(arg): return True return False - def _identify_rsqrt_and_input(self, norm_mul_node): - """From mul(x, rstd), find the rsqrt node and the input x. - The rsqrt may be wrapped in a cast node.""" + def _identify_rstd_and_input(self, norm_mul_node): + """From mul(x, rstd), find the rstd node and input x.""" if len(norm_mul_node.args) < 2: return None, None a, b = norm_mul_node.args[0], norm_mul_node.args[1] - for rsqrt_candidate_raw, input_candidate in [(a, b), (b, a)]: - if not isinstance(rsqrt_candidate_raw, torch.fx.Node): + for rstd_candidate_raw, input_candidate in [(a, b), (b, a)]: + if not isinstance(rstd_candidate_raw, torch.fx.Node): continue - rsqrt_candidate = _skip_casts(rsqrt_candidate_raw) - if ( - isinstance(rsqrt_candidate, torch.fx.Node) - and rsqrt_candidate.target == exir_ops.edge.aten.rsqrt.default - ): - return rsqrt_candidate, input_candidate + rstd_candidate = _skip_casts(rstd_candidate_raw) + if _is_rstd_node(rstd_candidate): + return rstd_candidate, input_candidate return None, None - def _get_single_arg_node(self, node, expected_target): - """Get the single input arg of a unary op node.""" - if node.target != expected_target: - return None - if len(node.args) < 1 or not isinstance(node.args[0], torch.fx.Node): - return None - return node.args[0] - @register_pattern_detector("rms_norm") def find_rms_norm_patterns( diff --git a/backends/vulkan/patterns/rope_hf.py b/backends/vulkan/patterns/rope_hf.py index 1514ab403b5..da1988610c5 100644 --- a/backends/vulkan/patterns/rope_hf.py +++ b/backends/vulkan/patterns/rope_hf.py @@ -63,6 +63,43 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) +class HfRotaryEmbeddingSinglePattern(torch.nn.Module): + """HuggingFace-style rotate-half RoPE for one BSHD tensor.""" + + def forward( + self, + x: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> torch.Tensor: + cos = freqs_cos.unsqueeze(1) + sin = freqs_sin.unsqueeze(1) + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return (x * cos) + (torch.cat((-x2, x1), dim=-1) * sin) + + +class HfRotaryEmbeddingFullTablePattern(torch.nn.Module): + """Rotate Q and K using full-width, possibly zero-padded tables.""" + + def forward( + self, + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + cos = freqs_cos.unsqueeze(1) + sin = freqs_sin.unsqueeze(1) + q1 = xq[..., : xq.shape[-1] // 2] + q2 = xq[..., xq.shape[-1] // 2 :] + k1 = xk[..., : xk.shape[-1] // 2] + k2 = xk[..., xk.shape[-1] // 2 :] + q = (xq * cos) + (torch.cat((-q2, q1), dim=-1) * sin) + k = (xk * cos) + (torch.cat((-k2, k1), dim=-1) * sin) + return q, k + + @lru_cache(maxsize=2) @register_pattern_graph("hf_rope") def get_hf_rope_graphs() -> List[torch.fx.GraphModule]: @@ -114,6 +151,46 @@ def get_hf_rope_graphs() -> List[torch.fx.GraphModule]: return graphs +@lru_cache(maxsize=1) +@register_pattern_graph("hf_rope_full_table") +def get_hf_rope_full_table_graphs() -> List[torch.fx.GraphModule]: + xq = torch.randn(1, 1, 4, 32, dtype=torch.float32) + xk = torch.randn(1, 1, 2, 32, dtype=torch.float32) + freqs_cos = torch.randn(1, 32, dtype=torch.float32) + freqs_sin = torch.randn(1, 32, dtype=torch.float32) + static_edge = to_edge( + export( + HfRotaryEmbeddingFullTablePattern(), + (xq, xk, freqs_cos, freqs_sin), + strict=True, + ), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + xq_dynamic = torch.randn(1, 2, 4, 32, dtype=torch.float32) + xk_dynamic = torch.randn(1, 2, 2, 32, dtype=torch.float32) + freqs_cos_dynamic = torch.randn(2, 32, dtype=torch.float32) + freqs_sin_dynamic = torch.randn(2, 32, dtype=torch.float32) + seq_dim = torch.export.Dim("hf_rope_full_table_seq", min=1, max=4) + dynamic_edge = to_edge( + export( + HfRotaryEmbeddingFullTablePattern(), + (xq_dynamic, xk_dynamic, freqs_cos_dynamic, freqs_sin_dynamic), + dynamic_shapes=( + {1: seq_dim}, + {1: seq_dim}, + {0: seq_dim}, + {0: seq_dim}, + ), + strict=True, + ), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + return [ + static_edge.exported_program().graph_module, + dynamic_edge.exported_program().graph_module, + ] + + def identify_hf_rotary_emb_io_nodes( ep: ExportedProgram, graph_module: torch.fx.GraphModule, @@ -134,6 +211,7 @@ def identify_hf_rotary_emb_io_nodes( return [xq, xk, freqs_cos, freqs_sin, xq_out, xk_out] +@register_pattern_replacement("hf_rope_full_table") @register_pattern_replacement("hf_rope") def create_hf_rotary_emb_custom_op( ep: ExportedProgram, @@ -152,6 +230,12 @@ def create_hf_rotary_emb_custom_op( freqs_cos.op == "call_function" and freqs_cos.target == exir_ops.edge.aten.slice_copy.Tensor ): + if ( + freqs_sin.op != "call_function" + or freqs_sin.target != exir_ops.edge.aten.slice_copy.Tensor + or freqs_cos.args[2:] != freqs_sin.args[2:] + ): + return full_freqs_cos = freqs_cos.args[0] start_pos = freqs_cos.args[2] full_freqs_sin = freqs_sin.args[0] @@ -186,3 +270,34 @@ def create_hf_rotary_emb_custom_op( xq_out.replace_all_uses_with(getitem_0) xk_out.replace_all_uses_with(getitem_1) + + +def create_hf_rotary_emb_single_custom_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: PatternMatch, +): + """Build the single-tensor replacement when a caller explicitly enables it.""" + if len(match.input_nodes) != 3 or len(match.output_nodes) != 1: + return + + x, freqs_cos, freqs_sin = match.input_nodes + if ( + freqs_cos.op != "call_function" + or freqs_cos.target != exir_ops.edge.aten.slice_copy.Tensor + or freqs_sin.op != "call_function" + or freqs_sin.target != exir_ops.edge.aten.slice_copy.Tensor + or freqs_cos.args[2:] != freqs_sin.args[2:] + ): + return + + x_out = match.output_nodes[0] + with graph_module.graph.inserting_before(x_out): + rotary_emb_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default, + args=(x, freqs_cos.args[0], freqs_sin.args[0], freqs_cos.args[2]), + ) + if "val" in x_out.meta: + rotary_emb_node.meta["val"] = x_out.meta["val"] + x_out.replace_all_uses_with(rotary_emb_node) diff --git a/backends/vulkan/patterns/sdpa.py b/backends/vulkan/patterns/sdpa.py index f67799f9b76..bc7dec9d177 100644 --- a/backends/vulkan/patterns/sdpa.py +++ b/backends/vulkan/patterns/sdpa.py @@ -39,7 +39,7 @@ def __init__(self, custom_sdpa_node: torch.fx.Node) -> None: # llama.custom_sdpa has signature: # custom_sdpa(query, key_cache, value_cache, start_pos, attn_mask, dropout_p, is_causal, scale) -> output - if len(custom_sdpa_node.args) < 4: + if len(custom_sdpa_node.args) < 7: return self.query_node = custom_sdpa_node.args[0] @@ -76,6 +76,13 @@ def __init__(self, custom_sdpa_node: torch.fx.Node) -> None: if self.update_value_cache_node is not None: self.value_projection_node = self.update_value_cache_node.args[0] + if ( + self.update_key_cache_node is None + or self.update_value_cache_node is None + or self.update_key_cache_node is self.update_value_cache_node + ): + return + # We have additional optional arguments but we don't need to capture them # since the new op doesn't use them diff --git a/backends/vulkan/patterns/select_as_symint.py b/backends/vulkan/patterns/select_as_symint.py index e7226b08188..402152d5932 100644 --- a/backends/vulkan/patterns/select_as_symint.py +++ b/backends/vulkan/patterns/select_as_symint.py @@ -94,7 +94,18 @@ def replace_select_local_scalar_dense_with_select_as_symint( ), ) - new_node.meta["val"] = match.anchor_node.meta["val"] + anchor_value = match.anchor_node.meta["val"] + new_node.meta["val"] = anchor_value + value_range = anchor_value.node.shape_env.bound_sympy(anchor_value.node.expr) + try: + lower_bound = int(value_range.lower) + except (AttributeError, TypeError, ValueError, OverflowError): + lower_bound = None + try: + upper_bound = int(value_range.upper) + except (AttributeError, TypeError, ValueError, OverflowError): + upper_bound = None + new_node.meta["et_vk_value_range"] = (lower_bound, upper_bound) match.anchor_node.replace_all_uses_with(new_node) # # Remove both the local_scalar_dense and select_copy nodes diff --git a/backends/vulkan/test/targets.bzl b/backends/vulkan/test/targets.bzl index fe13759243b..8d1a6784051 100644 --- a/backends/vulkan/test/targets.bzl +++ b/backends/vulkan/test/targets.bzl @@ -34,6 +34,10 @@ def define_common_targets(is_fbcode = False): srcs = [ "test_vulkan_passes.py", ], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], deps = [ "//caffe2:torch", "//executorch/backends/vulkan/_passes:vulkan_passes", @@ -89,6 +93,17 @@ def define_common_targets(is_fbcode = False): ], ) + python_unittest( + name = "test_vulkan_partitioner_extra_ops", + srcs = [ + "test_vulkan_partitioner_extra_ops.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + ], + ) + runtime.python_library( name = "tester", srcs = ["tester.py"], diff --git a/backends/vulkan/test/test_vulkan_partitioner_extra_ops.py b/backends/vulkan/test/test_vulkan_partitioner_extra_ops.py new file mode 100644 index 00000000000..d50eda2bda9 --- /dev/null +++ b/backends/vulkan/test/test_vulkan_partitioner_extra_ops.py @@ -0,0 +1,150 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import operator +import unittest +from unittest.mock import MagicMock + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.vulkan.partitioner.vulkan_partitioner import ( + CONTIGUOUS_BUFFER, + FP_T, + INT_T, + NO_STORAGE, + NONE_T, + OpFeatures, + VulkanPartitioner, + VulkanSupportedOperators, +) +from torch._subclasses.fake_tensor import FakeTensorMode + + +class TestScopedOperatorFeatures(unittest.TestCase): + def setUp(self) -> None: + self.op = torch.ops.aten.exp2.default + self.features = OpFeatures( + inputs_storage=CONTIGUOUS_BUFFER, + supports_resize=True, + ) + + def _node(self) -> torch.fx.Node: + node = MagicMock(spec=torch.fx.Node) + node.op = "call_function" + node.target = self.op + node.args = (MagicMock(),) + with FakeTensorMode() as mode: + node.meta = {"val": mode.from_tensor(torch.empty(1))} + return node + + def _support(self, **kwargs) -> VulkanSupportedOperators: + return VulkanSupportedOperators( + (16384, 16384, 2048), + 1 << 30, + **kwargs, + ) + + def test_public_feature_construction_surface(self) -> None: + features = OpFeatures( + inputs_dtypes=[FP_T, INT_T, NONE_T], + outputs_dtypes=[FP_T, INT_T], + inputs_storage=[ + CONTIGUOUS_BUFFER, + CONTIGUOUS_BUFFER, + NO_STORAGE, + ], + outputs_storage=[CONTIGUOUS_BUFFER, CONTIGUOUS_BUFFER], + supports_resize=True, + ) + self.assertTrue(features.supports_resize) + + def test_default_support_rejects_unregistered_operator(self) -> None: + self.assertFalse(self._support()._is_node_supported(self._node())) + + def test_explicit_support_accepts_only_the_scoped_operator(self) -> None: + support = self._support(extra_op_features={self.op: self.features}) + self.assertTrue(support._is_node_supported(self._node())) + self.assertFalse(self._support()._is_node_supported(self._node())) + + def test_allowlist_and_blocklist_run_before_scoped_features(self) -> None: + support = self._support( + operator_blocklist={self.op}, + extra_op_features={self.op: self.features}, + ) + self.assertFalse(support._is_node_supported(self._node())) + + def test_partitioner_copies_caller_mapping(self) -> None: + source = {self.op: self.features} + partitioner = VulkanPartitioner(extra_op_features=source) + source.clear() + self.assertEqual(partitioner.extra_op_features, {self.op: self.features}) + self.assertFalse(self._support()._is_node_supported(self._node())) + + def test_partitioner_rejects_global_registry_overlap(self) -> None: + with self.assertRaisesRegex(ValueError, "already registered"): + VulkanPartitioner( + extra_op_features={exir_ops.edge.aten.add.Tensor: self.features}, + ) + + +class TestGuardOnlySymbolicNodes(unittest.TestCase): + """`sym_min`/`sym_max` are non-tensor, so the guard-only check must run in + `node_is_compatible` before the `is_tensor_node` dispatch; routing it through + `op_node_is_compatible` made it unreachable.""" + + def _support(self) -> VulkanSupportedOperators: + return VulkanSupportedOperators((16384, 16384, 2048), 1 << 30) + + def _tensor_meta(self) -> dict: + with FakeTensorMode() as mode: + return {"val": mode.from_tensor(torch.empty(4))} + + def _sym_graph(self, sink: bool, tensor_user: bool): + graph = torch.fx.Graph() + a = graph.placeholder("a") + b = graph.placeholder("b") + sym = graph.call_function(torch.sym_max, (a, b)) + expr = graph.call_function(operator.add, (sym, 1)) + if sink: + graph.call_function(torch.ops.aten._assert_scalar.default, (expr, "ok")) + if tensor_user: + live = graph.call_function(torch.ops.aten.full.default, ([sym], 1.0)) + live.meta.update(self._tensor_meta()) + return sym + + def test_guard_only_chain_reaching_a_sink_is_accepted(self) -> None: + sym = self._sym_graph(sink=True, tensor_user=False) + ok, reason = self._support().node_is_compatible(sym) + self.assertTrue(ok, reason) + self.assertIsInstance(reason, str) + + def test_live_tensor_consumer_is_rejected(self) -> None: + sym = self._sym_graph(sink=True, tensor_user=True) + ok, reason = self._support().node_is_compatible(sym) + self.assertFalse(ok) + self.assertIsInstance(reason, str) + + def test_missing_sink_is_rejected(self) -> None: + sym = self._sym_graph(sink=False, tensor_user=False) + ok, reason = self._support().node_is_compatible(sym) + self.assertFalse(ok) + + def test_mixed_users_without_sink_are_rejected(self) -> None: + sym = self._sym_graph(sink=False, tensor_user=True) + ok, reason = self._support().node_is_compatible(sym) + self.assertFalse(ok) + + def test_contract_is_always_a_bool_str_pair(self) -> None: + for sink, tensor_user in ((True, False), (True, True), (False, False)): + with self.subTest(sink=sink, tensor_user=tensor_user): + result = self._support().node_is_compatible( + self._sym_graph(sink, tensor_user) + ) + self.assertIsInstance(result, tuple) + self.assertEqual(len(result), 2) + self.assertIsInstance(result[0], bool) + self.assertIsInstance(result[1], str) diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index f030b9268a1..69c4328d695 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -1,18 +1,45 @@ +import operator import unittest from typing import Optional, Tuple +from unittest.mock import MagicMock, patch import torch from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass +from executorch.backends.vulkan._passes.squeeze_unsqueeze_inputs import ( + SqueezeUnsqueezeInputs, +) +from executorch.backends.vulkan.patterns.quantized_linear import ( + find_quantized_linear_patterns, + replace_quantized_linear_patterns, +) +from executorch.backends.vulkan.patterns.sdpa import ( + CausalSDPAMatch, + is_custom_sdpa_node, + is_sdpa_with_kv_cache_node, + is_update_cache_node, +) from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge from executorch.exir.backend.canonical_partitioners.config_partitioner import ( format_target_name, ) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass +from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 +from torch.export._remove_auto_functionalized_pass import ( + unsafe_remove_auto_functionalized_pass, +) +from torchao.quantization.granularity import PerGroup from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from torchao.quantization.pt2e.quantizer import Quantizer +from torchao.quantization.quant_api import ( + Int8DynamicActivationIntxWeightConfig, + IntxWeightOnlyConfig, + quantize_, +) ################### ## Common Models ## @@ -34,6 +61,145 @@ def get_sample_inputs(self): return sample_inputs +class HfRopeFullTableModule(torch.nn.Module): + def __init__(self, broken_k_arm: bool = False, sin_start: int = 2) -> None: + super().__init__() + self.broken_k_arm = broken_k_arm + self.sin_start = sin_start + + def forward(self, xq, xk, cos_table, sin_table): + cos = cos_table[2:3].unsqueeze(1) + sin = sin_table[self.sin_start : self.sin_start + 1].unsqueeze(1) + q1 = xq[..., : xq.shape[-1] // 2] + q2 = xq[..., xq.shape[-1] // 2 :] + k1 = xk[..., : xk.shape[-1] // 2] + k2 = xk[..., xk.shape[-1] // 2 :] + q = (xq * cos) + (torch.cat((-q2, q1), dim=-1) * sin) + k_rotation = torch.cat((-k2, k1), dim=-1) * sin + if self.broken_k_arm: + return q, (xk * cos) - k_rotation + return q, (xk * cos) + k_rotation + + def get_sample_inputs(self): + return ( + torch.rand(1, 1, 4, 32), + torch.rand(1, 1, 2, 32), + torch.rand(8, 32), + torch.rand(8, 32), + ) + + +class TiedLinearPair(torch.nn.Module): + def __init__(self, in_features: int = 256, out_features: int = 128) -> None: + super().__init__() + self.in_features = in_features + self.linear1 = torch.nn.Linear(in_features, out_features) + self.linear2 = torch.nn.Linear(in_features, out_features) + self.linear2.weight = self.linear1.weight + + def forward(self, x): + return self.linear1(x) + self.linear2(x) + + def get_sample_inputs(self): + return (torch.rand(8, self.in_features),) + + +class RmsNormVariantModule(torch.nn.Module): + def __init__(self, exponent: Optional[float]) -> None: + super().__init__() + self.exponent = exponent + self.weight = torch.nn.Parameter(torch.ones(64)) + + def forward(self, x): + mean_sq = x.pow(2).mean(-1, keepdim=True) + 1e-6 + if self.exponent is None: + rstd = torch.rsqrt(mean_sq) + else: + rstd = torch.pow(mean_sq, self.exponent) + return (x * rstd) * self.weight + + def get_sample_inputs(self): + return (torch.rand(2, 8, 64),) + + +class SelectAsSymIntModule(torch.nn.Module): + def __init__(self, bounded: bool = True, use_select: bool = True) -> None: + super().__init__() + self.bounded = bounded + self.use_select = use_select + + def forward(self, index, x): + if self.use_select: + value = index.select(0, 0).item() + else: + value = index.sum().item() + if not self.bounded: + return x.sum() + value + torch._check(value >= 3) + torch._check(value <= 17) + return x[:, :value].sum() + + def get_sample_inputs(self): + return (torch.tensor([5], dtype=torch.int64), torch.rand(2, 32)) + + +class CausalSdpaModule(torch.nn.Module): + def __init__(self, update_key: bool, update_value: bool) -> None: + super().__init__() + self.update_key = update_key + self.update_value = update_value + self.register_buffer("key_cache", torch.zeros(1, 8, 2, 4)) + self.register_buffer("value_cache", torch.zeros(1, 8, 2, 4)) + + def forward(self, query, key, value): + if self.update_key: + torch.ops.llama.update_cache(key, self.key_cache, 0) + if self.update_value: + torch.ops.llama.update_cache(value, self.value_cache, 0) + return torch.ops.llama.custom_sdpa( + query, + self.key_cache, + self.value_cache, + 0, + None, + 0.0, + True, + None, + ) + + def get_sample_inputs(self): + return ( + torch.rand(1, 1, 2, 4), + torch.rand(1, 1, 2, 4), + torch.rand(1, 1, 2, 4), + ) + + +class SharedCacheCausalSdpaModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("cache", torch.zeros(1, 8, 2, 4)) + + def forward(self, query, projected_cache_value): + torch.ops.llama.update_cache(projected_cache_value, self.cache, 0) + return torch.ops.llama.custom_sdpa( + query, + self.cache, + self.cache, + 0, + None, + 0.0, + True, + None, + ) + + def get_sample_inputs(self): + return ( + torch.rand(1, 1, 2, 4), + torch.rand(1, 1, 2, 4), + ) + + ########### ## Tests ## ########### @@ -86,7 +252,88 @@ def op_node_count(graph_module: torch.fx.GraphModule, canonical_op_name: str) -> return count +def lower_module(model: torch.nn.Module): + program = torch.export.export(model, model.get_sample_inputs(), strict=True) + return to_edge( + program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + + +def run_fuse_patterns(exported_program): + fuse_pass = FusePatternsPass() + fuse_pass._exported_program = exported_program + return fuse_pass.call(exported_program.graph_module) + + +def nodes_named(graph_module: torch.fx.GraphModule, canonical_op_name: str): + return [ + node + for node in graph_module.graph.nodes + if get_target_canonical_name(node) == canonical_op_name + ] + + +def quantized_linear_matches(exported_program): + matches = [] + for node in exported_program.graph_module.graph.nodes: + match = find_quantized_linear_patterns(node) + if match is not None: + matches.append(match) + return matches + + +def quantize_tied_pair(model: TiedLinearPair, config) -> TiedLinearPair: + quantize_(model, config) + model.linear2.weight = model.linear1.weight + return model + + +def shared_8da4w_program( + in_features: int = 256, + out_features: int = 128, + group_size: int = 32, + weight_dtype=torch.int4, +): + model = TiedLinearPair(in_features, out_features).eval() + config = Int8DynamicActivationIntxWeightConfig( + weight_dtype=weight_dtype, + weight_granularity=PerGroup(group_size), + ) + return lower_module(quantize_tied_pair(model, config)) + + +def causal_sdpa_program(update_key: bool, update_value: bool): + exported_program = lower_module(CausalSdpaModule(update_key, update_value)) + return unsafe_remove_auto_functionalized_pass(exported_program) + + +def shared_cache_causal_sdpa_program(): + exported_program = lower_module(SharedCacheCausalSdpaModule()) + return unsafe_remove_auto_functionalized_pass(exported_program) + + class TestVulkanPasses(unittest.TestCase): + def test_squeeze_preserves_gelu_kwargs_only_on_gelu(self): + input_arg = MagicMock() + input_arg.node.meta = {"val": torch.empty(2, 1, 4)} + kwargs = {"approximate": "tanh"} + metadata = {"val": torch.empty(2, 1, 4)} + + with patch.object(ExportPass, "call_operator", autospec=True) as base_call: + base_call.side_effect = [MagicMock(), MagicMock(), MagicMock()] + SqueezeUnsqueezeInputs().call_operator( + exir_ops.edge.aten.gelu.default, + (input_arg,), + kwargs, + metadata, + ) + + calls = base_call.call_args_list + self.assertEqual(calls[0].args[3], {}) + self.assertEqual(calls[1].args[3], kwargs) + self.assertEqual(calls[2].args[3], {}) + def test_fuse_torchao_quantized_embedding(self): """A torchao-dialect 4-bit weight-only quantized embedding (torchao.dequantize_affine -> aten.embedding) should fuse into a single @@ -779,3 +1026,281 @@ def forward(self, x): gm = ep.graph_module self.assertEqual(op_node_count(gm, "q8ta_pixel_shuffle.default"), 0) + + +class GemmaMatcherWitnessTest(unittest.TestCase): + def test_full_table_hf_rope_preserves_tables_start_and_meta_shapes(self): + model = HfRopeFullTableModule() + exported_program = lower_module(model) + + result = run_fuse_patterns(exported_program) + fused_nodes = nodes_named(result.graph_module, "apply_rotary_emb_hf.default") + self.assertEqual(len(fused_nodes), 1) + fused = fused_nodes[0] + self.assertEqual(fused.args[2].target, "cos_table") + self.assertEqual(fused.args[3].target, "sin_table") + self.assertEqual(fused.args[4], 2) + + output_shapes = { + node.args[1]: tuple(node.meta["val"].shape) + for node in result.graph_module.graph.nodes + if node.op == "call_function" + and node.target is operator.getitem + and node.args[0] is fused + } + self.assertEqual(output_shapes, {0: (1, 1, 4, 32), 1: (1, 1, 2, 32)}) + + def test_full_table_hf_rope_rejects_changed_k_arithmetic(self): + result = run_fuse_patterns(lower_module(HfRopeFullTableModule(True))) + + self.assertFalse(result.modified) + self.assertEqual( + nodes_named(result.graph_module, "apply_rotary_emb_hf.default"), [] + ) + + def test_full_table_hf_rope_rejects_mismatched_frequency_slices(self): + result = run_fuse_patterns(lower_module(HfRopeFullTableModule(sin_start=3))) + + self.assertEqual( + nodes_named(result.graph_module, "apply_rotary_emb_hf.default"), [] + ) + + def test_shared_8da4w_reuses_only_weight_storage(self): + exported_program = shared_8da4w_program() + matches = quantized_linear_matches(exported_program) + self.assertEqual(len(matches), 2) + self.assertIs(matches[0].weight_node, matches[1].weight_node) + + result = run_fuse_patterns(exported_program) + fused = nodes_named(result.graph_module, "linear_dq8ca_q4gsw.default") + self.assertEqual(len(fused), 2) + for argument_index in (3, 4, 5): + self.assertIs(fused[0].args[argument_index], fused[1].args[argument_index]) + self.assertIsNot(fused[0].args[7], fused[1].args[7]) + self.assertEqual( + exported_program._et_vk_param_modification_tags, + { + "linear1.parametrizations.weight.original0": "4 bit linear weight", + "linear1.parametrizations.weight.original1": "4 bit linear scales", + }, + ) + + def test_shared_8da4w_rejects_incomplete_mutation_tags(self): + exported_program = shared_8da4w_program() + matches = quantized_linear_matches(exported_program) + self.assertEqual(len(matches), 2) + replace_quantized_linear_patterns( + exported_program, exported_program.graph_module, matches[0] + ) + tags = exported_program._et_vk_param_modification_tags + scales_name = next( + name for name, tag in tags.items() if tag == "4 bit linear scales" + ) + del tags[scales_name] + + with self.assertRaisesRegex( + RuntimeError, "shared quantized linear mutation tags are inconsistent" + ): + replace_quantized_linear_patterns( + exported_program, exported_program.graph_module, matches[1] + ) + + def test_shared_weight_only_4bit_linear_is_rejected_fail_closed(self): + model = TiedLinearPair().eval() + config = IntxWeightOnlyConfig( + weight_dtype=torch.int4, + granularity=PerGroup(32), + ) + exported_program = lower_module(quantize_tied_pair(model, config)) + + with self.assertRaisesRegex( + RuntimeError, "shared quantized linear mutation tags are inconsistent" + ): + run_fuse_patterns(exported_program) + + def test_shared_8da4w_rejects_packed_weight_geometry(self): + exported_program = shared_8da4w_program(in_features=252, group_size=4) + + with self.assertRaisesRegex( + RuntimeError, "shared 8da4w packed weight geometry mismatch" + ): + run_fuse_patterns(exported_program) + + def test_shared_8da4w_rejects_weight_sums_geometry(self): + for out_features, error in ( + (130, "shared 8da4w packed weight geometry mismatch"), + (132, "shared 8da4w weight sums geometry mismatch"), + ): + with self.subTest(out_features=out_features, error=error): + exported_program = shared_8da4w_program(out_features=out_features) + + with self.assertRaisesRegex(RuntimeError, error): + run_fuse_patterns(exported_program) + + def test_shared_8da4w_rejects_unexpected_dequant_abi(self): + exported_program = shared_8da4w_program(weight_dtype=torch.int2) + + with self.assertRaisesRegex( + RuntimeError, "shared 8da4w weight has an unexpected dequant ABI" + ): + run_fuse_patterns(exported_program) + + def test_rms_norm_accepts_rsqrt_and_pow_negative_half(self): + for exponent, expected_rsqrt, expected_pow in ( + (None, 1, 1), + (-0.5, 0, 2), + ): + with self.subTest(exponent=exponent): + exported_program = lower_module(RmsNormVariantModule(exponent)) + self.assertEqual( + op_node_count(exported_program.graph_module, "rsqrt.default"), + expected_rsqrt, + ) + self.assertEqual( + op_node_count(exported_program.graph_module, "pow.Tensor_Scalar"), + expected_pow, + ) + result = run_fuse_patterns(exported_program) + self.assertEqual( + op_node_count(result.graph_module, "rms_norm.default"), 1 + ) + + def test_rms_norm_rejects_other_reciprocal_exponents(self): + for exponent in (-1.0, -0.25): + with self.subTest(exponent=exponent): + exported_program = lower_module(RmsNormVariantModule(exponent)) + self.assertEqual( + op_node_count(exported_program.graph_module, "pow.Tensor_Scalar"), + 2, + ) + result = run_fuse_patterns(exported_program) + self.assertFalse(result.modified) + self.assertEqual( + op_node_count(result.graph_module, "rms_norm.default"), 0 + ) + + def test_select_as_symint_preserves_bounded_and_unbounded_ranges(self): + for bounded, expected_range in ((True, (3, 17)), (False, (None, None))): + with self.subTest(bounded=bounded): + result = run_fuse_patterns( + lower_module(SelectAsSymIntModule(bounded=bounded)) + ) + selected = nodes_named(result.graph_module, "select_as_symint.default") + self.assertEqual(len(selected), 1) + self.assertEqual(selected[0].meta["et_vk_value_range"], expected_range) + constraints = [ + node + for node in result.graph_module.graph.nodes + if "sym_constrain_range.default" in str(node.target) + ] + self.assertEqual(len(constraints), 1) + self.assertEqual( + constraints[0].kwargs, + {"min": expected_range[0], "max": expected_range[1]}, + ) + + def test_select_as_symint_rejects_non_select_scalar(self): + result = run_fuse_patterns(lower_module(SelectAsSymIntModule(use_select=False))) + + self.assertFalse(result.modified) + self.assertEqual( + nodes_named(result.graph_module, "select_as_symint.default"), [] + ) + + def test_select_as_symint_eager_requires_integral_input(self): + self.assertEqual( + torch.ops.et_vk.select_as_symint.default( + torch.tensor([5, 6], dtype=torch.int64), 0, 1 + ), + 6, + ) + with self.assertRaisesRegex(ValueError, "requires an integral input"): + torch.ops.et_vk.select_as_symint.default(torch.tensor([5.0, 6.0]), 0, 1) + + def test_causal_sdpa_requires_both_cache_updates(self): + for update_key, update_value, expected in ( + (True, True, (True, True, True)), + (True, False, (True, False, False)), + (False, True, (False, True, False)), + (False, False, (False, False, False)), + ): + with self.subTest(update_key=update_key, update_value=update_value): + exported_program = causal_sdpa_program(update_key, update_value) + sdpa_node = next( + node + for node in exported_program.graph_module.graph.nodes + if is_custom_sdpa_node(node) + ) + match = CausalSDPAMatch(sdpa_node) + actual = ( + match.update_key_cache_node is not None, + match.update_value_cache_node is not None, + match.match_found, + ) + self.assertEqual(actual, expected) + if match.match_found: + self.assertEqual(match.query_node.target, "query") + self.assertEqual(match.key_projection_node.target, "key") + self.assertEqual(match.value_projection_node.target, "value") + + def test_causal_sdpa_replaces_complete_cache_topology(self): + exported_program = causal_sdpa_program(True, True) + result = run_fuse_patterns(exported_program) + + self.assertTrue(result.modified) + self.assertEqual( + sum(is_update_cache_node(node) for node in result.graph_module.graph.nodes), + 0, + ) + self.assertEqual( + sum( + is_sdpa_with_kv_cache_node(node) + for node in result.graph_module.graph.nodes + ), + 1, + ) + + def test_causal_sdpa_rejects_shared_key_value_cache_update(self): + exported_program = shared_cache_causal_sdpa_program() + sdpa_node = next( + node + for node in exported_program.graph_module.graph.nodes + if is_custom_sdpa_node(node) + ) + match = CausalSDPAMatch(sdpa_node) + + self.assertIs(match.update_key_cache_node, match.update_value_cache_node) + self.assertFalse(match.match_found) + + result = run_fuse_patterns(exported_program) + self.assertFalse(result.modified) + self.assertEqual( + sum(is_update_cache_node(node) for node in result.graph_module.graph.nodes), + 1, + ) + self.assertEqual( + sum(is_custom_sdpa_node(node) for node in result.graph_module.graph.nodes), + 1, + ) + self.assertEqual( + sum( + is_sdpa_with_kv_cache_node(node) + for node in result.graph_module.graph.nodes + ), + 0, + ) + + def test_causal_sdpa_rejects_incomplete_optional_arguments(self): + for argument_count in (4, 5, 6): + with self.subTest(argument_count=argument_count): + graph = torch.fx.Graph() + query = graph.placeholder("query") + key = graph.placeholder("key") + value = graph.placeholder("value") + full_args = (query, key, value, 0, None, 0.0) + sdpa_node = graph.call_function( + torch.ops.llama.custom_sdpa.default, + full_args[:argument_count], + ) + + self.assertFalse(CausalSDPAMatch(sdpa_node).match_found) diff --git a/backends/vulkan/test/test_vulkan_tensor_repr.py b/backends/vulkan/test/test_vulkan_tensor_repr.py index 5a0fc664c17..e3796bc86da 100644 --- a/backends/vulkan/test/test_vulkan_tensor_repr.py +++ b/backends/vulkan/test/test_vulkan_tensor_repr.py @@ -22,6 +22,8 @@ CONTIGUOUS_ANY, CONTIGUOUS_BUFFER, DEFAULT_TEXTURE_LIMITS, + extents_are_valid, + filter_invalid_reprs, HEIGHT_PACKED_TEXTURE, make_tensor_repset, NO_STORAGE, @@ -31,6 +33,7 @@ PACKED_INT8_4W_BUFFER, PACKED_INT8_BUFFER, PACKED_INT8_CHANNELS_PACKED_BUFFER, + required_image_extents, TensorRepr, TensorReprList, TensorRepSet, @@ -38,6 +41,7 @@ WIDTH_PACKED_TEXTURE, ) from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv def _make_fake_tensor(shape, dtype=torch.float32): @@ -68,6 +72,94 @@ def _make_tensor_arg_node(shape, dtype=torch.float32): return node +class _AddOne(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +def _make_backed_extent(maximum: int) -> torch.SymInt: + extent = torch.export.Dim(f"extent_{maximum}", min=1, max=maximum) + program = torch.export.export( + _AddOne(), + (torch.randn(10),), + dynamic_shapes={"x": {0: extent}}, + ) + placeholder = next(node for node in program.graph.nodes if node.op == "placeholder") + return placeholder.meta["val"].shape[0] + + +class TestTextureExtentValidation(unittest.TestCase): + def test_concrete_texture_extents(self): + self.assertTrue(extents_are_valid((512, 1, 2048), DEFAULT_TEXTURE_LIMITS)) + self.assertFalse(extents_are_valid((512, 1, 2049), DEFAULT_TEXTURE_LIMITS)) + + def test_unbacked_texture_extent_is_invalid(self): + unbacked = ShapeEnv().create_unbacked_symint() + self.assertFalse( + extents_are_valid((512, 1, unbacked + 1), DEFAULT_TEXTURE_LIMITS) + ) + + def test_symbolic_texture_extent_uses_declared_range(self): + self.assertTrue( + extents_are_valid( + (512, 1, _make_backed_extent(1024)), DEFAULT_TEXTURE_LIMITS + ) + ) + self.assertFalse( + extents_are_valid( + (512, 1, _make_backed_extent(4096)), DEFAULT_TEXTURE_LIMITS + ) + ) + + def test_symbolic_texture_extents_after_packing(self): + width_within_limit = required_image_extents( + torch.Size((1, 1, 1, _make_backed_extent(65536))), + VkMemoryLayout.TENSOR_WIDTH_PACKED, + ) + width_crossing_limit = required_image_extents( + torch.Size((1, 1, 1, _make_backed_extent(65537))), + VkMemoryLayout.TENSOR_WIDTH_PACKED, + ) + channels_within_limit = required_image_extents( + torch.Size((2, _make_backed_extent(4096), 1, 1)), + VkMemoryLayout.TENSOR_CHANNELS_PACKED, + ) + channels_crossing_limit = required_image_extents( + torch.Size((2, _make_backed_extent(4097), 1, 1)), + VkMemoryLayout.TENSOR_CHANNELS_PACKED, + ) + + self.assertTrue(extents_are_valid(width_within_limit, DEFAULT_TEXTURE_LIMITS)) + self.assertFalse( + extents_are_valid(width_crossing_limit, DEFAULT_TEXTURE_LIMITS) + ) + self.assertTrue( + extents_are_valid(channels_within_limit, DEFAULT_TEXTURE_LIMITS) + ) + self.assertFalse( + extents_are_valid(channels_crossing_limit, DEFAULT_TEXTURE_LIMITS) + ) + + def test_unknown_texture_extent_preserves_only_buffer_layouts(self): + unbacked = ShapeEnv().create_unbacked_symint() + tensor_val = MagicMock() + tensor_val.shape = (1, unbacked + 1, 1, 512) + + result = filter_invalid_reprs( + tensor_val, ANY_STORAGE, DEFAULT_TEXTURE_LIMITS + ) + self.assertEqual( + result.valid_buffer_layouts, + ANY_STORAGE.valid_buffer_layouts, + ) + self.assertFalse(result.valid_texture_layouts) + + texture_only = filter_invalid_reprs( + tensor_val, ANY_TEXTURE, DEFAULT_TEXTURE_LIMITS + ) + self.assertTrue(texture_only.is_empty()) + + class TestTensorRepSet(unittest.TestCase): # -- Construction and emptiness -- diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 84b901b6b6e..aa5e09e466a 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -24,6 +24,7 @@ from torch.export import ExportedProgram from torch.export.exported_program import InputKind from torch.export.graph_signature import TensorArgument +from torch.fx.experimental.symbolic_shapes import statically_known_true TorchOpType = Union[EdgeOpOverload, torch._ops.OpOverload, str] @@ -840,7 +841,10 @@ def required_image_extents(sizes: torch.Size, layout: VkMemoryLayout) -> ImageEx def extents_are_valid(extents: ImageExtents, limits: ImageExtents) -> bool: - return all(extents[i] <= limits[i] for i in range(len(extents))) + return all( + statically_known_true(extents[index] <= limits[index]) + for index in range(len(extents)) + ) def valid_texture_memory_layouts(