Improve rope matching soundness

Author: GregoryComer

backends/xnnpack/_passes/convert_to_rope.py

@@ -88,6 +88,56 @@ def _trace_through_unsqueeze(node: torch.fx.Node) -> torch.fx.Node:
             return node.args[0]
         return node
 
+    @staticmethod
+    def _find_trig_source(node: torch.fx.Node) -> torch.fx.Node | None:
+        """Walk backwards through unsqueeze_copy ops to find cos/sin op."""
+        current = node
+        for _ in range(10):
+            if current.op != "call_function":
+                return None
+            if current.target in (
+                exir_ops.edge.aten.cos.default,
+                exir_ops.edge.aten.sin.default,
+            ):
+                return current
+            if current.target == exir_ops.edge.aten.unsqueeze_copy.default:
+                current = current.args[0]
+                continue
+            return None
+        return None
+
+    @classmethod
+    def _is_doubled_cat(cls, trig_node: torch.fx.Node) -> bool:
+        """Check that a cos/sin node's input is cat(x, x) with identical args."""
+        cat_node = trig_node.args[0]
+        if (
+            cat_node.op != "call_function"
+            or cat_node.target != exir_ops.edge.aten.cat.default
+        ):
+            return False
+        tensors = cat_node.args[0]
+        return len(tensors) == 2 and tensors[0] is tensors[1]
+
+    @classmethod
+    def _has_doubled_freqs(
+        cls,
+        cos_unsqueezed: torch.fx.Node,
+        sin_unsqueezed: torch.fx.Node,
+    ) -> bool:
+        """Verify that cos/sin frequencies are doubled (first half == second half).
+
+        Traces back through unsqueeze_copy ops to find the cos/sin producer,
+        then verifies its input is cat(x, x) where both args are the same
+        node — a structural proof that the first and second halves are identical.
+        """
+        cos_trig = cls._find_trig_source(cos_unsqueezed)
+        sin_trig = cls._find_trig_source(sin_unsqueezed)
+
+        if cos_trig is None or sin_trig is None:
+            return False
+
+        return cls._is_doubled_cat(cos_trig) and cls._is_doubled_cat(sin_trig)
+
     @staticmethod
     def _trace_through_permute(node: torch.fx.Node) -> torch.fx.Node | None:
         """If node is a permute_copy that swaps dims 1 and 2, return its input."""
@@ -148,6 +198,10 @@ def create_rope(
         cos_node = self._trace_through_unsqueeze(cos_unsqueezed)
         sin_node = self._trace_through_unsqueeze(sin_unsqueezed)
 
+        if not self._has_doubled_freqs(cos_unsqueezed, sin_unsqueezed):
+            logger.debug("Skipping RoPE fusion: cannot verify doubled frequencies")
+            return
+
         weights = self._build_weights(graph_module, cos_node, sin_node, output_node)
 
         with graph_module.graph.inserting_before(output_node):

backends/xnnpack/partition/config/generic_node_configs.py

@@ -216,7 +216,7 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         if not self.check_common_constraints(node, ep):
             return False
 
-        num_tensors = len(node.all_input_nodes)
+        num_tensors = len(node.args[0])
 
         if not (num_tensors >= 2):
             why(