Improve layout handling for RoPE fusion

Author: GregoryComer

backends/xnnpack/_passes/convert_to_rope.py

@@ -66,6 +66,18 @@ def _build_weights(
 
         return weights
 
+    @staticmethod
+    def _get_unsqueeze_dim(node: torch.fx.Node) -> int:
+        """Return the unsqueeze dim if node is an unsqueeze_copy, else -1."""
+        if (
+            node.op == "call_function"
+            and node.target == exir_ops.edge.aten.unsqueeze_copy.default
+        ):
+            return node.args[1]
+        return -1
+
+    _BHSD_TO_BSHD_PERM = [0, 2, 1, 3]
+
     @staticmethod
     def _trace_through_unsqueeze(node: torch.fx.Node) -> torch.fx.Node:
         """If node is an unsqueeze_copy, return its input. Otherwise return node as-is."""
@@ -76,6 +88,33 @@ def _trace_through_unsqueeze(node: torch.fx.Node) -> torch.fx.Node:
             return node.args[0]
         return node
 
+    @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."""
+        if (
+            node.op == "call_function"
+            and node.target == exir_ops.edge.aten.permute_copy.default
+            and list(node.args[1]) == [0, 2, 1, 3]
+        ):
+            return node.args[0]
+        return None
+
+    def _get_layout(self, cos_unsqueezed: torch.fx.Node) -> str | None:
+        """Determine the tensor layout from the cos unsqueeze dimension.
+
+        Returns "BSHD", "BHSD", or None if the layout cannot be determined.
+        """
+        unsqueeze_dim = self._get_unsqueeze_dim(cos_unsqueezed)
+        if unsqueeze_dim == -1:
+            return None
+        ndim = len(cos_unsqueezed.meta["val"].shape)
+        normalized = unsqueeze_dim if unsqueeze_dim >= 0 else unsqueeze_dim + ndim
+        if normalized == 2:
+            return "BSHD"
+        if normalized == 1:
+            return "BHSD"
+        return None
+
     def create_rope(
         self,
         graph_module: torch.fx.GraphModule,
@@ -90,9 +129,22 @@ def create_rope(
         sin_unsqueezed = match.placeholder_nodes[2]
         output_node = match.returning_nodes[0]
 
-        # Trace back through unsqueeze to get raw cos/sin for weight construction.
-        # The pattern excludes unsqueeze ops (they're shared between q/k RoPE),
-        # so the matched placeholders are the unsqueeze outputs.
+        # xnn_define_rope expects NTHC (batch, tokens, heads, channels) input.
+        # BSHD (unsqueeze_dim=2) maps directly to NTHC.
+        # BHSD (unsqueeze_dim=1) requires tracing through the BSHD→BHSD permute
+        # to recover the BSHD input, then re-permuting the output back to BHSD.
+        layout = self._get_layout(cos_unsqueezed)
+        if layout == "BSHD":
+            rope_input = x_node
+        elif layout == "BHSD":
+            rope_input = self._trace_through_permute(x_node)
+            if rope_input is None:
+                logger.debug("Skipping RoPE fusion: BHSD but x is not a permute_copy")
+                return
+        else:
+            logger.debug("Skipping RoPE fusion: unrecognized layout")
+            return
+
         cos_node = self._trace_through_unsqueeze(cos_unsqueezed)
         sin_node = self._trace_through_unsqueeze(sin_unsqueezed)
 
@@ -102,12 +154,23 @@ def create_rope(
             rope_node = graph_module.graph.create_node(
                 "call_function",
                 torch.ops.xnnpack.rope.default,
-                args=(x_node, weights),
+                args=(rope_input, weights),
             )
-
-        rope_node.meta["val"] = torch.empty_like(x_node.meta["val"])
-
-        output_node.replace_all_uses_with(rope_node)
+            rope_node.meta["val"] = torch.empty_like(rope_input.meta["val"])
+
+            if layout == "BHSD":
+                permute_node = graph_module.graph.call_function(
+                    exir_ops.edge.aten.permute_copy.default,
+                    args=(rope_node, self._BHSD_TO_BSHD_PERM),
+                )
+                permute_node.meta["val"] = rope_node.meta["val"].permute(
+                    self._BHSD_TO_BSHD_PERM
+                )
+                result_node = permute_node
+            else:
+                result_node = rope_node
+
+        output_node.replace_all_uses_with(result_node)
         graph_module.graph.eliminate_dead_code()
 
     # override

backends/xnnpack/test/ops/test_rope.py

@@ -68,3 +68,51 @@ def test_fp32_rope_dynamic_seq_len(self):
             {0: None, 1: seq, 2: None},
         )
         self._test_rope(inputs, dynamic_shapes=dynamic_shapes)
+
+    class HFRopeBHSD(torch.nn.Module):
+        """HF-style RoPE with BHSD layout (transpose before/after RoPE)."""
+
+        def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+            x = x.transpose(1, 2)
+            cos = cos.unsqueeze(1)
+            sin = sin.unsqueeze(1)
+            x1 = x[..., : x.shape[-1] // 2]
+            x2 = x[..., x.shape[-1] // 2 :]
+            rot = torch.cat((-x2, x1), dim=-1)
+            out = (x * cos) + (rot * sin)
+            return out.transpose(1, 2)
+
+    def _test_rope_bhsd(self, inputs, dynamic_shapes=None):
+        (
+            Tester(self.HFRopeBHSD(), inputs, dynamic_shapes=dynamic_shapes)
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs(inputs=inputs)
+        )
+
+    def test_fp32_rope_bhsd(self):
+        batch, seq_len, n_heads, head_dim = 1, 8, 4, 32
+        cos, sin = _hf_freqs(seq_len, head_dim)
+        inputs = (torch.randn(batch, seq_len, n_heads, head_dim), cos, sin)
+        self._test_rope_bhsd(inputs)
+
+    def test_fp32_rope_bhsd_large_head_dim(self):
+        batch, seq_len, n_heads, head_dim = 1, 16, 8, 128
+        cos, sin = _hf_freqs(seq_len, head_dim)
+        inputs = (torch.randn(batch, seq_len, n_heads, head_dim), cos, sin)
+        self._test_rope_bhsd(inputs)
+
+    def test_fp32_rope_bhsd_dynamic_seq_len(self):
+        batch, seq_len, n_heads, head_dim = 1, 8, 4, 32
+        cos, sin = _hf_freqs(seq_len, head_dim)
+        inputs = (torch.randn(batch, seq_len, n_heads, head_dim), cos, sin)
+        seq = Dim("seq", min=1, max=128)
+        dynamic_shapes = (
+            {0: None, 1: seq, 2: None, 3: None},
+            {0: None, 1: seq, 2: None},
+            {0: None, 1: seq, 2: None},
+        )
+        self._test_rope_bhsd(inputs, dynamic_shapes=dynamic_shapes)