CUDA backend: lower complex rotary embeddings and free repeated outputs once

backends/aoti/aoti_backend.py

@@ -312,8 +312,8 @@ def preprocess(
 
         # ``run_decompositions`` retraces the complete ExportedProgram even
         # when none of the table's operators occur in the graph. The in-place
-        # passes above preserve graph inputs and outputs, so the existing graph
-        # signature remains valid when the decomposition table cannot apply.
+        # passes above keep the graph signature valid, so the program needs no
+        # retrace when the decomposition table cannot apply.
         if contains_any_call_fn_target_op(
             device_edge_program.graph_module, decomposition_table
         ):

backends/aoti/passes/replace_view_copy_with_view.py

@@ -33,8 +33,8 @@
     ops.edge.aten.split_copy.Tensor: ops.edge.aten.split.Tensor,
     # Edge turns ``detach`` into ``detach_copy``, which has no c-shim either.
     # Without autograd it is ``alias_copy``, which Inductor can decompose,
-    # unlike ``detach_copy``. A view would let two outputs share one handle,
-    # and ``clone`` fails Edge verification in dim order mode.
+    # unlike ``detach_copy``, and ``clone`` fails Edge verification in dim
+    # order mode.
     torch.ops.aten.detach_copy.default: torch.ops.aten.alias_copy.default,
     ops.edge.aten.detach_copy.default: ops.edge.aten.alias_copy.default,
 }

backends/cuda/BUCK

@@ -79,6 +79,7 @@ fbcode_target(
         "passes/__init__.py",
         "passes/move_cond_predicate_to_cpu.py",
         "passes/lower_offgraph_kv.py",
+        "passes/replace_complex_with_real.py",
         "passes/replace_int64_floordiv.py",
     ],
     visibility = [

backends/cuda/cuda_backend.py

@@ -39,6 +39,9 @@
 from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import (
     MoveCondPredicateToCpuPass,
 )
+from executorch.backends.cuda.passes.replace_complex_with_real import (
+    ReplaceComplexWithRealPass,
+)
 from executorch.backends.cuda.passes.replace_int64_floordiv import (
     ReplaceInt64FloorDivWithFloatPass,
 )
@@ -1013,7 +1016,11 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                         f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'."
                     )
                 triton_kernel_mode = mode
-        passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
+        passes = [
+            ReplaceComplexWithRealPass(),
+            MoveCondPredicateToCpuPass(),
+            ReplaceInt64FloorDivWithFloatPass(),
+        ]
         cell_layout = any(
             spec.key == OFFGRAPH_KV_COMPILE_SPEC
             and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"