Decomp/precompiled rules

Makefile

@@ -122,7 +122,7 @@ frontend:
 	# versions of a package with the same version tag (e.g. 0.38-dev0).
 	$(PYTHON) -m pip uninstall -y pennylane
 	$(PYTHON) -m pip install -e . --extra-index-url https://test.pypi.org/simple $(PIP_VERBOSE_FLAG)
-	$(PYTHON) -m catalyst.utils.precompile_decomposition_rules
+	$(PYTHON) -m catalyst.decomposition.precompile_decomposition_rules
 	rm -r frontend/pennylane_catalyst.egg-info
 
 .PHONY: mlir llvm stablehlo enzyme dialects runtime oqc builtin-decomp-rules

doc/releases/changelog-dev.md

@@ -277,6 +277,7 @@
   [(#2983)](https://github.com/PennyLaneAI/catalyst/pull/2983)
   [(#3022)](https://github.com/PennyLaneAI/catalyst/pull/3022)
   [(#3039)](https://github.com/PennyLaneAI/catalyst/pull/3039)
+  [(#3046)](https://github.com/PennyLaneAI/catalyst/pull/3046)
 
 * The `graph-decomposition` pass eliminates three redundant IR manipulations:
   the cloning, removal, and re-insertion of user rules. This optimization is particularly

frontend/catalyst/decomposition/decomposition_rules.py

@@ -0,0 +1,206 @@
+# Copyright 2026 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module provides infrastructure for lowering decomposition rules via python.
+"""
+
+# pylint: disable=protected-access,bare-except
+
+import jax.numpy as jnp
+import pennylane as qp
+from jax._src.lib.mlir import ir
+from jaxlib.mlir.dialects.builtin import ModuleOp
+
+from catalyst.decomposition.type_utils import (
+    _MLIR_DTYPES_TO_PY_DTYPES,
+    _PY_DTYPES_TO_MLIR_DTYPES,
+    get_dummy_values_for_container,
+    mlir_stringify_type,
+)
+from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
+
+
+class GraphOpID:
+    """
+    Return the graph operator id for the operator2 instance `op`.
+
+    The FuncOp decomposition rules in the returned string satisfy the following requirements:
+        - Are named `{rule name}_{op graph ID}`.
+        - Are MLIR representations of the PennyLane decomposition rules associated with the
+          specified operator.
+        - Are instantiated with the static data provided, and all other data remains dynamic.
+        - Are self-contained, and do not contain any device initialization, setup/teardown etc.
+        - Are compatible with the `decompose-lowering` and `graph-decomposition` passes, meaning
+          the following:
+            - Their `target_gate` attribute is set to the provided graph operator ID
+            - They have a resources attribute containing an operations attribute which maps graph
+              operator IDs to counts of their occurrences in the rule.
+            - Their arguments are mappable to the operator they decompose via `decompose-lowering`.
+
+    Note that this function should not be updated without updating the corresponding method on the
+    DecomposableGate interface in mlir/lib/quantum/IR/QuantumInterfaces.cpp.
+    """
+
+    def __init__(self, op: qp.core.Operator2, uid=None):
+        """Create a new GraphOpId."""
+        assert isinstance(
+            op, qp.core.Operator2
+        ), "Graph-based decomposition expects an Operator2 instance"
+        self.op = op
+
+        self.operator_name = op.name
+        self.dynamic_shape = self.parse_dynamic_shape()
+        self.wire_lens = self.parse_wire_lens()
+        self.static_data = self.parse_static_data()
+        self.extra_data = uid
+
+    def parse_dynamic_shape(self) -> list:
+        """Return the dynamic shape as a list of dtypes."""
+        return list(self.op.dynamic_args.values())
+
+    def parse_wire_lens(self) -> list[int]:
+        """Return the length of each of the wire args."""
+        return list(map(len, self.op.wire_args.values()))
+
+    def parse_static_data(self) -> dict:
+        """Return a dictionary of names to static data values."""
+        return {
+            static_argname: getattr(self.op, static_argname)
+            for static_argname in self.op.compilable_argnames
+        }
+
+    def get_operator_name(self) -> str:
+        """Return the name of the operator."""
+        return self.operator_name
+
+    def get_dynamic_shape_id_format(self) -> str:
+        """Return the dynamic shape formatted for GraphOpId."""
+        return f"[{','.join(map(mlir_stringify_type, self.dynamic_shape))}]"
+
+    def get_wire_lens_id_format(self) -> str:
+        """Return the wire lengths formatted for GraphOpId."""
+        return f"[{','.join(map(str, self.wire_lens))}]"
+
+    def get_static_data_id_format(self) -> str:
+        """Return the static data formatted for GraphOpId."""
+        return f"{{{','.join(f'{k}:{v}' for k, v in self.static_data.items())}}}"
+
+    def getID(self) -> str:
+        """
+        Return the GraphOpId as a string.
+
+        NOTE: do not modify this method without also modifying the corresponding DecomposableGate
+        interface in MLIR.
+        """
+        ID_string = (
+            self.get_operator_name()
+            + self.get_dynamic_shape_id_format()
+            + self.get_wire_lens_id_format()
+            + self.get_static_data_id_format()
+        )
+        if self.extra_data:
+            ID_string += "[" + str(self.extra_data) + "]"
+        return ID_string
+
+
+def collect_resources_for_op(op_name, dummy_dynamic_args, dummy_wires, static_data):
+    """
+    Return resource data for all decomposition rules associated to op_name.
+
+    This includes a dictionary
+    """
+    decomp_rules = list(qp.decomposition.list_decomps(op_name))
+
+    # map rules to resource resources, in a more generic format
+    name_to_resource_ids = {}
+    name_to_resources = {}
+    for rule in decomp_rules:
+        # The `compute_resources` function's signature is the same as the Operator2 signature
+        # for the original op of the rule
+        resources = rule.compute_resources(*dummy_dynamic_args, *dummy_wires, **static_data)
+        name_to_resources[rule.name] = resources.gate_counts
+        name_to_resource_ids[rule.name] = {
+            GraphOpID(op).getID(): count for op, count in resources.gate_counts.items()
+        }
+
+    return name_to_resources, name_to_resource_ids, decomp_rules
+
+
+def compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data) -> ModuleOp:
+    """
+    Return a ModuleOp containing the decomposition rules for an operator instance.
+
+    The decomposition rules will be decorated with appropriate resource and target_gate attributes.
+    """
+    device = qp.device("null.qubit", wires=sum(wire_lens))
+    dummy_wires = tuple(jnp.array(range(length), dtype=int) for length in wire_lens)
+    dummy_dynamic_args = get_dummy_values_for_container(dynamic_shape)
+
+    _, name_to_resource_ids, decomp_rules = collect_resources_for_op(
+        op_name, dummy_dynamic_args, dummy_wires, static_data
+    )
+
+    def rule_to_subroutine(rule):
+        def decomp_rule(*args, **kwargs):
+            rule._impl(*args, **kwargs)
+
+        # keep the frontend name for readability, append target op_id for symbol uniqueness
+        decomp_rule.__name__ = rule._impl.__name__ + "_" + op_id
+
+        return qp.capture.subroutine(decomp_rule)
+
+    subroutines = [rule_to_subroutine(rule) for rule in decomp_rules]
+
+    @qp.qjit(
+        target="mlir",
+        capture=True,
+    )
+    @qp.qnode(device=device)
+    def circuit():
+        for subroutine in subroutines:
+            subroutine(*dummy_dynamic_args, wires=dummy_wires)
+
+    module = circuit.mlir_module
+
+    def update_funcop_attributes(op):
+        """Update the decomposition rule attributes if op is a decomposition rule.
+
+        For use with module.walk
+
+        This function updates the following attributes:
+            - Adds the `target_gate` attribute.
+            - Adds the `resources` attribute.
+        """
+        if op.name == "func.func":
+            rule_name = ir.StringAttr(op.attributes["sym_name"]).value.removesuffix("_" + op_id)
+            if rule_name in name_to_resource_ids:
+                op.attributes["resources"] = get_mlir_attribute_from_pyval(
+                    {"operations": name_to_resource_ids[rule_name]}
+                )
+                op.attributes["target_gate"] = ir.StringAttr.get(op_id)
+
+        return ir.WalkResult.ADVANCE
+
+    with module.context:
+        module.operation.walk(update_funcop_attributes)
+
+    return module
+
+
+def compile_decomposition_rules_wrapper(
+    op_name, op_id, dynamic_shape, wire_lens, static_data
+) -> str:
+    """Return a string MLIR module containing the decomposition rules for an operator instance."""
+    return str(compile_decomposition_rules(op_name, op_id, dynamic_shape, wire_lens, static_data))