add AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER to ai_edge_torch.config

PiperOrigin-RevId: 718950902

Author: chunnienc

ai_edge_torch/_config.py

@@ -65,5 +65,14 @@ def enable_group_norm_composite(self) -> bool:
   def enable_group_norm_composite(self, value: bool):
     os.environ["ENABLE_GROUP_NORM_COMPOSITE"] = "y" if value else "n"
 
+  @property
+  def layout_optimize_partitioner(self) -> str:
+    """The algorithm to use for layout optimization."""
+    return os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", "DEFAULT")
+
+  @layout_optimize_partitioner.setter
+  def layout_optimize_partitioner(self, value: str):
+    os.environ["AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER"] = str(value).upper()
+
 
 config = _Config()

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -18,6 +18,7 @@
 import os
 from typing import Union
 
+import ai_edge_torch
 from ai_edge_torch import fx_infra
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
@@ -261,10 +262,8 @@ def call(self, exported_program: torch.export.ExportedProgram):
     self.mark_const_nodes(exported_program)
 
     graph_module = exported_program.graph_module
-    partitioner = os.environ.get(
-        "AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None
-    )
-    if partitioner == "MINCUT":
+    partitioner = ai_edge_torch.config.layout_optimize_partitioner
+    if partitioner in ("MINCUT", "OPTIMAL"):
       graph_module = layout_partitioners.min_cut.partition(graph_module)
     elif partitioner == "GREEDY":
       graph_module = layout_partitioners.greedy.partition(graph_module)