Memoize compile_decomposition_rules per operator variant

doc/releases/changelog-dev.md

@@ -55,6 +55,14 @@
 
 <h3>Improvements 🛠</h3>
 
+* Compiled decomposition rules are now memoized per operator variant and registered rule set, so
+  each unique operator and modifier variant is compiled only once. Under program capture, the
+  on-demand rule loader previously rebuilt (and recompiled) the same variants many times over,
+  making rule compilation the dominant cost of decomposition-heavy programs. The cache is keyed on
+  the registered decomposition rules as well as the operator variant, so changing the registered
+  rules (e.g. across `qml.decomposition.local_decomps` blocks) never returns a stale module.
+  [(#3174)](https://github.com/PennyLaneAI/catalyst/pull/3174)
+
 * Add the `XMEM_REPLY_BRAM` memory type and use it to allocate reply buffers in dedicated BRAM.
   [(#3148)](https://github.com/PennyLaneAI/catalyst/pull/3148)
 

frontend/catalyst/decomposition/decomposition_rules.py

@@ -270,6 +270,55 @@ def collect_resources_for_op(op_name, kwargs, is_custom_op=False, adjoint_resour
     return name_to_resources, name_to_resource_ids, decomp_rules
 
 
+# Memoization for `compile_decomposition_rules`. Keyed by the fully-resolved operator instance,
+# modifier variant, and a fingerprint of the registered rule set. Values are
+# ``(result_module, pinned_rules)`` tuples: the (context-owning) result module is only read (walked
+# and cloned) downstream, never mutated, so sharing it across callers is safe, and the pinned rule
+# objects keep the fingerprinted ids from being recycled while the entry lives.
+_COMPILE_DECOMP_CACHE = {}
+
+
+def _hashable(value):
+    """Return a stable, hashable representation of a (possibly nested dict/list) value."""
+    if isinstance(value, dict):
+        return tuple(sorted((k, _hashable(v)) for k, v in value.items()))
+    if isinstance(value, (list, tuple)):
+        return tuple(_hashable(v) for v in value)
+    try:
+        hash(value)
+        return value
+    except TypeError:
+        return repr(value)
+
+
+def _compile_decomp_cache_key(
+    op_name,
+    op_id,
+    dynamic_shape,
+    wire_lens,
+    static_data,
+    extra_data,
+    is_custom_op,
+    wrap_adjoint,
+    wrap_control,
+    n_ctrl,
+    rule_ids,
+):
+    return (
+        op_name,
+        op_id,
+        _hashable(dynamic_shape),
+        _hashable(wire_lens),
+        _hashable(static_data),
+        _hashable(extra_data),
+        is_custom_op,
+        wrap_adjoint,
+        wrap_control,
+        n_ctrl,
+        rule_ids,
+    )
+
+
 def prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) -> dict:
     kwargs = {}
     for wire_name, wire_len in wire_lens.items():
@@ -312,6 +361,37 @@ def compile_decomposition_rules(
     ``<n>C(Adjoint(op_name))``: adjoint is applied innermost and control outermost (the canonical
     order matching the compiler's ``wrapModifiers``).
     """
+    # The on-demand rule loader rebuilds the full reachable-rule closure for every operator the
+    # compiler asks about, which recompiles the same operator variants many times over. Compiling
+    # a single variant is cheap (~tens of ms), but the closure explosion makes this the dominant
+    # cost of decomposition-heavy programs. Memoize the (context-owning) result module keyed by the
+    # fully-resolved operator instance and modifier variant so each unique variant is compiled once.
+    #
+    # The compiled module is fully determined by the operator variant *and* the decomposition rules
+    # currently registered for ``op_name`` (``list_decomps``). That rule set can differ between
+    # programs (e.g. across ``qp.decomposition.local_decomps`` blocks), so it is part of the key: a
+    # variant compiled under one rule set must not be reused under a different one. We fingerprint
+    # the rule set by the identities of its rule objects and pin those objects in the cache value so
+    # their ids cannot be recycled (which could otherwise alias distinct rule sets) while the entry
+    # is alive.
+    registered_rules = tuple(qp.decomposition.list_decomps(op_name))
+    _cache_key = _compile_decomp_cache_key(
+        op_name,
+        op_id,
+        dynamic_shape,
+        wire_lens,
+        static_data,
+        extra_data,
+        is_custom_op,
+        wrap_adjoint,
+        wrap_control,
+        n_ctrl,
+        tuple(id(rule) for rule in registered_rules),
+    )
+    _cached = _COMPILE_DECOMP_CACHE.get(_cache_key)
+    if _cached is not None:
+        return _cached[0]
+
     kwargs = prepare_dynamic_op_kwargs(dynamic_shape, wire_lens)
     extra_data = extra_data or {}
     n_base_wires = sum(wire_lens.values())
@@ -451,6 +531,8 @@ def re_privatize_rules(op):
     with inlined_module.context, ir.Location.unknown():
         inlined_module.operation.walk(re_privatize_rules)
 
+    # Store the module together with the pinned rule objects (see the fingerprint note above).
+    _COMPILE_DECOMP_CACHE[_cache_key] = (inlined_module, registered_rules)
     return inlined_module
 
 

frontend/test/pytest/test_decomposition.py

@@ -37,10 +37,12 @@
 from catalyst import qjit
 from catalyst.decomposition import GraphOpID, RuleLoweringWarning
 from catalyst.decomposition.decomposition_rules import (
+    _COMPILE_DECOMP_CACHE,
     _MODIFIER_CANONICAL_ORDER,
     _control_modifier,
     _leading_modifier_kind,
     _modifier_kind,
+    compile_decomposition_rules,
     compile_decomposition_rules_wrapper,
     compile_reachable_decomposition_rules_wrapper,
     name_unwrap_adjoint,
@@ -501,5 +503,68 @@ def test_wrap_modifier_id_rejects_non_canonical(self, op_id):
             wrap_modifier_id(op_id, "Adjoint")
 
 
+class TestCompileDecompositionRuleCache:
+    """``compile_decomposition_rules`` memoizes its compiled module per operator variant."""
+
+    def test_same_variant_compiled_once(self):
+        """Compiling the same operator variant twice returns the identical cached module."""
+        _COMPILE_DECOMP_CACHE.clear()
+        args = ("RZ", "RZ{0:[f64]}{wires:2}{}", {"0": ["f64"]}, {"wires": 1}, {})
+
+        first = compile_decomposition_rules(*args, is_custom_op=True)
+        assert len(_COMPILE_DECOMP_CACHE) == 1
+        second = compile_decomposition_rules(*args, is_custom_op=True)
+
+        assert first is second
+        assert len(_COMPILE_DECOMP_CACHE) == 1
+
+    def test_distinct_variants_not_shared(self):
+        """Distinct operator variants get independent cache entries and modules."""
+        _COMPILE_DECOMP_CACHE.clear()
+
+        rz = compile_decomposition_rules(
+            "RZ", "RZ{0:[f64]}{wires:2}{}", {"0": ["f64"]}, {"wires": 1}, {}, is_custom_op=True
+        )
+        rx = compile_decomposition_rules(
+            "RX", "RX{0:[f64]}{wires:2}{}", {"0": ["f64"]}, {"wires": 1}, {}, is_custom_op=True
+        )
+
+        assert rz is not rx
+        assert len(_COMPILE_DECOMP_CACHE) == 2
+
+    def test_same_variant_distinct_rule_sets_not_shared(self):
+        """The same operator variant compiled under two different registered rule sets gets
+        independent cache entries: the compiled module depends on the registered rules, not just the
+        operator variant, so the rule set is part of the cache key."""
+        _COMPILE_DECOMP_CACHE.clear()
+        args = ("NoParams", "NoParams{}{reg:2}{}", {}, {"reg": 2}, {})
+
+        def rule_a_resources(reg):
+            return {SingleParam(x=Float, reg=Wire[2]): 1}
+
+        @register_resources(rule_a_resources)
+        def rule_a(reg):
+            SingleParam(x=0.1, reg=reg[0:2])
+
+        def rule_b_resources(reg):
+            return {SingleParam(x=Float, reg=Wire[2]): 2}
+
+        @register_resources(rule_b_resources)
+        def rule_b(reg):
+            SingleParam(x=0.1, reg=reg[0:2])
+            SingleParam(x=0.2, reg=reg[0:2])
+
+        with local_decomps():
+            add_decomps(NoParams, rule_a)
+            first = compile_decomposition_rules(*args)
+
+        with local_decomps():
+            add_decomps(NoParams, rule_b)
+            second = compile_decomposition_rules(*args)
+
+        assert first is not second
+        assert len(_COMPILE_DECOMP_CACHE) == 2
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])