Fold DyT alpha scalar into tanh LUT

backends/arm/_passes/__init__.py

@@ -1,3 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -114,6 +116,7 @@
 from .deduplicate_get_attr_pass import DeduplicateGetAttrPass  # noqa
 from .ensure_unique_output_nodes_pass import EnsureUniqueOutputNodesPass  # noqa
 from .exir_to_tosa_pass import ExirToTosaPass  # noqa
+from .fold_dyt_alpha_into_lut_pass import FoldDyTAlphaIntoLUTPass  # noqa
 from .fold_qdq_with_annotated_qparams_pass import (  # noqa
     FoldAndAnnotateQParamsPass,
     QuantizeClampArgumentsPass,

backends/arm/_passes/arm_pass_manager.py

@@ -208,6 +208,11 @@ class PassInsertions:
 _registered_pass_insertions: dict[type, PassInsertions] = {}
 
 
+_registered_pass_factories_before: dict[
+    type, list[Callable[[ExportedProgram], ExportPass]]
+] = {}
+
+
 def _graph_pass_name(graph_pass: Callable[[GraphModule], PassResult | None]) -> str:
     if isinstance(graph_pass, ExportPass):
         return ArmPass.get_name(graph_pass)
@@ -271,9 +276,21 @@ def register_pass_insertions_after(
     _registered_pass_insertions[target_pass_type].after_passes.extend(passes)
 
 
+def register_pass_factories_before(
+    target_pass_type: type,
+    factories: list[Callable[[ExportedProgram], ExportPass]],
+) -> None:
+    """Register factories instantiated with the backend ExportedProgram."""
+    registered = _registered_pass_factories_before.setdefault(target_pass_type, [])
+    for factory in factories:
+        if factory not in registered:
+            registered.append(factory)
+
+
 def clear_registered_pass_insertions() -> None:
-    """Clear all globally registered pass insertions."""
+    """Clear all globally registered pass insertions and pass factories."""
     _registered_pass_insertions.clear()
+    _registered_pass_factories_before.clear()
 
 
 class ArmPassManager(ExportedProgramPassManager):
@@ -430,6 +447,11 @@ def _configure_pass_insertions(self, exported_program: ExportedProgram) -> None:
                 self.insert_passes_before(pass_type, list(insertions.before_passes))
             if insertions.after_passes:
                 self.insert_passes_after(pass_type, list(insertions.after_passes))
+        for pass_type, factories in _registered_pass_factories_before.items():
+            self.insert_passes_before(
+                pass_type,
+                [factory(exported_program) for factory in factories],
+            )
 
     def add_passes(self, passes: Sequence[ExportPass | None]):
         for p in passes: