diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 0282dc495a..95b064bcb0 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -21,12 +21,24 @@ * The graph-based decomposition system now supports **adjoint operators** for `Operator2`. [(#3120)](https://github.com/PennyLaneAI/catalyst/pull/3120) [(#3115)](https://github.com/PennyLaneAI/catalyst/pull/3115) + [(#3204)](https://github.com/PennyLaneAI/catalyst/pull/3204) For a target gate set, `Adjoint(Op)` is reached through any of three pathways: 1. Rules registered on the base `Op`, 2. Rules registered directly for `Adjoint(Op)`, and 3. Rules *synthesized by distribution* (`decompose(Adjoint(Op)) = adjoint(decompose(Op))`). + Pathway 2 now also covers the rules PennyLane registers with the *symbolic* operator's arguments, + i.e. `rule(base)` rather than the base op's `(*params, wires)`. This is how + `self_adjoint`, `adjoint_rotation` and other symbolic rules are written, so `Adjoint(H)`, `Adjoint(X)`, + `Adjoint(RZ)`, `Adjoint(Rot)`, ... now decompose straight back to their base operator instead of + falling through to the (much longer) distributed rules, or failing to solve at all when the base + op is the only member of the target `gate_set`. + + A rule registered for `Adjoint(Op)` must now take `base`; this is the convention PennyLane's own + graph calls such a rule with, and the one every `Operator2` rule in PennyLane already follows. A + rule written against the base op's parameters instead is skipped with a `RuleLoweringWarning`. + * The graph-based decomposition system now supports **controlled operators** for `Operator2`, including single control (`C(Op)`), multiple controls (`C(Op)`), and their composition with adjoint. diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py index f5bd9ddf8b..96291ff4aa 100644 --- a/frontend/catalyst/decomposition/decomposition_rules.py +++ b/frontend/catalyst/decomposition/decomposition_rules.py @@ -23,6 +23,7 @@ import jax.numpy as jnp import pennylane as qp from jax._src.lib.mlir import ir +from pennylane.core.operator import Operator2, abstractify from catalyst.compiler import _quantum_opt from catalyst.decomposition.graph_op_id import GraphOpID @@ -47,6 +48,26 @@ def _resources_have_measurement(gate_counts) -> bool: return any(isinstance(op, _NON_INVERTIBLE_RESOURCE_TYPES) for op in gate_counts) +def build_base_op(op_cls, kwargs, is_custom_op): + """Instantiate the base operator of a PennyLane's symbolic rule from prepared rule kwargs. + + Note that we have to pause capture here, because the base operator is only a carrier of + the traced parameters and wires, and the rule body is what decides whether (and how) it + gets applied. + + Args: + op_cls (type): the base operator's class + kwargs (dict): the prepared rule kwargs + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + + Returns: + Operator2: the base operator, built without binding its primitive + """ + args, kwargs = split_call_args(kwargs, is_custom_op) + with qp.capture.pause(): + return op_cls(*args, **kwargs) + + # Canonical nesting order for op-level modifiers, listed OUTERMOST first. The compiler's # ``wrapModifiers`` (mlir/lib/Quantum/IR/QuantumInterfaces.cpp) folds modifiers into a graphOpId # in this exact order: @@ -96,6 +117,16 @@ def wrap_modifier_id(op_id: str, modifier: str) -> str: optional ``[uid]``) follow it, matching the compiler's ``defaultGetGraphOpId``. This applies to any modifier (e.g. ``"Adjoint"``, ``"C"``), so nested ids compose as ``C(Adjoint(RX)){...}``. Extend callers here to support future op-level modifiers. + + Args: + op_id (str): the graphOpId to wrap + modifier (str): the modifier token, e.g. ``"Adjoint"``, ``"C"`` or ``"2C"`` + + Returns: + str: the modified graphOpId + + Raises: + ValueError: if the modifier would nest inside one that is canonically outer """ new_kind = _modifier_kind(modifier) inner_kind = _leading_modifier_kind(op_id) @@ -126,7 +157,18 @@ def name_wrap_adjoint(op_id: str) -> str: def name_unwrap_adjoint(op_name: str, op_id: str) -> str: - """Inverse of :func:`name_wrap_adjoint` given the base ``op_name``.""" + """Inverse of :func:`name_wrap_adjoint` given the base ``op_name``. + + Args: + op_name (str): the base operator's name + op_id (str): the adjoint graphOpId to unwrap + + Returns: + str: the base operator's graphOpId + + Raises: + ValueError: if ``op_id`` is not an adjoint id for ``op_name`` + """ prefix = f"Adjoint({op_name})" if not op_id.startswith(prefix): raise ValueError(f"{op_id!r} is not an adjoint id for base op {op_name!r}") @@ -137,7 +179,17 @@ def name_unwrap_control(op_name: str, op_id: str): """Inverse of control name-wrapping given the base ``op_name``. ``("RX", "2C(RX){...}")`` -> ``("RX{...}", 2)`` and ``("RX", "C(RX){...}")`` -> ``("RX{...}", 1)``. - Raises if ``op_id`` is not a control id for ``op_name``. + + Args: + op_name (str): the base operator's name + op_id (str): the controlled graphOpId to unwrap + + Returns: + str: the base operator's graphOpId + int: the number of controls + + Raises: + ValueError: if ``op_id`` is not a control id for ``op_name`` """ i = 0 while i < len(op_id) and op_id[i].isdigit(): @@ -151,6 +203,16 @@ def name_unwrap_control(op_name: str, op_id: str): def get_rule_strings_from_module(module: ir.Module) -> list[str]: + """Extract the decomposition rules held by a module as MLIR strings. + + Every FuncOp carrying a ``target_gate`` attribute is a decomposition rule. + + Args: + module (ir.Module): the module a rule-compiling qjit produced + + Returns: + list[str]: one string per rule, with the ``__builtin_`` prefix added to its name + """ raw_funcOps = [] def find_condition(op): @@ -193,8 +255,8 @@ def get_rules_from_module(module: ir.Module) -> str: Parse and modify decomposition rules from a ModuleOp. Args: - module: an MLIR module object; every FuncOp carrying a `target_gate` attribute is - extracted as a decomposition rule. + module (ir.Module): an MLIR module object; every FuncOp carrying a `target_gate` + attribute is extracted as a decomposition rule. Returns: str: The string representation of any decomposition rules from `module`, pre-pending the @@ -205,6 +267,14 @@ def get_rules_from_module(module: ir.Module) -> str: def inject_new_rules_into_module(module: ir.Module, decomp_rules: list[str]): + """Add decomposition rules to a module, skipping the ones it already holds. + + A rule counts as already held when both its ``target_gate`` and its ``resources`` match. + + Args: + module (ir.Module): the module to add the rules to + decomp_rules (list[str]): the rules to add, as MLIR strings + """ with ir.InsertionPoint(module.body): for decomp_rule in decomp_rules: decomp_rule_op = ir.Operation.parse(decomp_rule) @@ -240,6 +310,14 @@ def split_call_args(kwargs, is_custom_op): them by their real argnames, so they are passed positionally with only "wires" kept as keyword. Custom-op params are always scalar f64, so scalar ``0.0`` dummies are used (a shape-(1,) array would instead lower the gate through the general ``qref.operator`` path rather than ``qref.custom``). + + Args: + kwargs (dict): the prepared rule kwargs + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + + Returns: + tuple: the positional arguments to call the rule with + dict: the keyword arguments to call the rule with """ if is_custom_op: args = tuple(0.0 for key in kwargs if key != "wires") @@ -248,7 +326,19 @@ def split_call_args(kwargs, is_custom_op): def collect_resources_for_op(op_name, kwargs, is_custom_op=False, adjoint_resources=False): - """Return resource data for all decomposition rules associated to op_name.""" + """Return resource data for all decomposition rules associated to op_name. + + Args: + op_name (str): the operator's name, as PennyLane registers its rules + kwargs (dict): the arguments to compute the resources with + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + adjoint_resources (bool): whether to spell each produced id in its adjoint form + + Returns: + dict: rule name to the resources it produces + dict: rule name to the graphOpId of each resource + list: the rules considered + """ decomp_rules = list(qp.decomposition.list_decomps(op_name)) args, kwargs = split_call_args(kwargs, is_custom_op) @@ -278,6 +368,15 @@ def collect_resources_for_op(op_name, kwargs, is_custom_op=False, adjoint_resour def prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) -> dict: + """Build the dummy arguments an operator's decomposition rules are called with. + + Args: + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + + Returns: + dict: argument names to dummy values + """ kwargs = {} for wire_name, wire_len in wire_lens.items(): kwargs[wire_name] = jnp.array(range(wire_len), dtype=int) @@ -318,6 +417,21 @@ def compile_decomposition_rules( Note that ``wrap_adjoint`` and ``wrap_control`` may be combined to synthesize the nested modifier ``C(Adjoint(op_name))``: adjoint is applied innermost and control outermost (the canonical order matching the compiler's ``wrapModifiers``). + + Args: + op_name (str): the operator's name, as PennyLane registers its rules + op_id (str): the operator's graphOpId + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + wrap_adjoint (bool): whether to distribute the base rules over adjoint + wrap_control (bool): whether to distribute the base rules over control + n_ctrl (int): the number of controls to distribute over + + Returns: + ir.Operation: the ``builtin.module`` holding the rules """ kwargs = prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) extra_data = extra_data or {} @@ -325,7 +439,10 @@ def compile_decomposition_rules( device = qp.device("null.qubit", wires=n_base_wires + (n_ctrl if wrap_control else 0)) name_to_resources, name_to_resource_ids, decomp_rules = collect_resources_for_op( - op_name, kwargs | static_data | extra_data, is_custom_op, adjoint_resources=wrap_adjoint + op_name, + kwargs | static_data | extra_data, + is_custom_op, + adjoint_resources=wrap_adjoint, ) # TODO: The modified target id and the wrapped resource ids are derived here by string-wrapping @@ -392,11 +509,38 @@ def decomp_rule(*_args, _ctrl_wires=None, **_kwargs): call_args, call_kwargs = split_call_args(kwargs, is_custom_op) + return build_rule_module( + subroutines, device, call_args, call_kwargs, ctrl_wires, name_to_resource_ids, target_id + ) + + +# pylint: disable=too-many-arguments +def build_rule_module( + subroutines, device, call_args, call_kwargs, ctrl_wires, name_to_resource_ids, target_id +) -> ir.Operation: + """Trace ``subroutines`` into a module of standalone decomposition-rule functions. + + Args: + subroutines (list): the rule bodies, as captured subroutines + device (Device): the device to trace them on, sized for the operator's wires + call_args (tuple): positional arguments to call each subroutine with + call_kwargs (dict): keyword arguments to call each subroutine with + ctrl_wires: the control wires to pass, or None when the rules are not controlled + name_to_resource_ids (dict): rule name to the graphOpId of each resource + target_id (str): the graphOpId of the gate the rules decompose + + Returns: + ir.Operation: the module holding the rules + + Raises: + CompileError: if the rules could not be traced + """ + @qp.qjit(target="mlir", capture=True, collect_decomp_rules=False) @qp.qnode(device=device) def circuit(): for subroutine in subroutines: - if wrap_control: + if ctrl_wires is not None: subroutine(*call_args, _ctrl_wires=ctrl_wires, **call_kwargs) else: subroutine(*call_args, **call_kwargs) @@ -474,42 +618,218 @@ def re_privatize_rules(op): return inlined_module +def collect_symbolic_adjoint_resources(op_cls, op_name, kwargs, is_custom_op): + """Return resource data for the rules registered against ``Adjoint(op_name)``. + + Args: + op_cls (type): the base operator's class + op_name (str): the base operator's name + kwargs (dict): the arguments to build the base operator with + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + + Returns: + list: the rules considered + dict: the arguments the rules were probed with + dict: rule name to the resources it produces + dict: rule name to the graphOpId of each resource + """ + rules = list(qp.decomposition.list_decomps(f"Adjoint({op_name})")) + if not rules: + return [], {}, {}, {} + + # The base op is only a carrier of the dummy parameters and wires: the graph reasons about + # resources in terms of abstract operators, matching `_get_kwargs` in PennyLane's graph. + probe_args = {"base": abstractify(build_base_op(op_cls, kwargs, is_custom_op))} + + name_to_resources = {} + name_to_resource_ids = {} + for rule in rules: + try: + resources = rule.compute_resources(**probe_args) + name_to_resources[rule.name] = resources.gate_counts + # The rule body names the ops it produces itself, so unlike the distribution pathway + # these ids carry no added modifier. + # TODO: a resource op that is itself symbolic (e.g. a generic `Controlled(GlobalPhase)` + # rep) does not spell the compiler's canonical `C(GlobalPhase){...}` id here; such a + # rule registers an id the solver cannot match. + name_to_resource_ids[rule.name] = { + GraphOpID(op).getGraphOpId(): count for op, count in resources.gate_counts.items() + } + except Exception as e: # pylint: disable=broad-except + warnings.warn( + f"Failed to get resources for the {rule.name} decomposition rule: {e}", + category=RuleLoweringWarning, + ) + + return rules, probe_args, name_to_resources, name_to_resource_ids + + +# pylint: disable=too-many-arguments +def compile_registered_adjoint_rules( + op_name, + target_id, + dynamic_shape, + wire_lens, + static_data, + extra_data=None, + is_custom_op=False, + op_cls=None, +) -> ir.Operation | None: + """Return the module of rules registered against ``Adjoint(op_name)`` that follow PennyLane's + symbolic-argument convention, or None if there are none. + + These rules take a base operator instance rather than the base op's parameters, so the rule + body rebuilds the operator from its own traced arguments. + + Args: + op_name (str): the base operator's name + target_id (str): the adjoint graphOpId the rules decompose + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + op_cls (type[Operator2]): the base operator's class, required to rebuild it + + Returns: + ir.Operation or None: the module holding the rules, or None if there are none + + Raises: + ValueError: if ``op_cls`` is not given + """ + if op_cls is None: + raise ValueError( + f"The operator class of {op_name!r} is needed to lower the decomposition rules " + f"registered against {target_id}" + ) + assert issubclass(op_cls, Operator2), f"Expected an Operator2 subclass, got {op_cls}" + + extra_data = extra_data or {} + static_and_extra = static_data | extra_data + kwargs = prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) + device = qp.device("null.qubit", wires=sum(wire_lens.values())) + + rules, probe_args, name_to_resources, name_to_resource_ids = collect_symbolic_adjoint_resources( + op_cls, op_name, kwargs | static_and_extra, is_custom_op + ) + if not rules: + return None + + def rule_to_subroutine(rule): + def decomp_rule(*_args, **_kwargs): + with qp.capture.pause(): + base = op_cls(*_args, **_kwargs) + # TODO: Call the rule itself instead of its _impl after merging + # https://github.com/PennyLaneAI/pennylane/pull/10144 + rule._impl(base=base) + + decomp_rule_no_static_args = partial(decomp_rule, **static_and_extra) + decomp_rule_no_static_args.__name__ = rule.name + "_" + target_id + return qp.capture.subroutine(decomp_rule_no_static_args) + + subroutines = [] + for rule in rules: + if rule.name not in name_to_resource_ids: + continue + if _resources_have_measurement(name_to_resources[rule.name]): # pragma: no cover + warnings.warn( + f"Skipped the {rule.name} decomposition rule for {target_id}: it contains a " + "mid-circuit measurement, which is not supported with adjoint or control regions.", + category=RuleLoweringWarning, + ) + continue + if rule.is_applicable(**probe_args): + subroutines.append(rule_to_subroutine(rule)) + + if not subroutines: + return None + + call_args, call_kwargs = split_call_args(kwargs, is_custom_op) + + return build_rule_module( + subroutines, device, call_args, call_kwargs, None, name_to_resource_ids, target_id + ) + + +def registered_adjoint_rule_strings(op_name, target_id, **kwargs) -> list[str]: + """Return the rule strings from :func:`compile_registered_adjoint_rules`. + + A failure to lower them is reported as a warning and yields no rules. + + Args: + op_name (str): the base operator's name + target_id (str): the adjoint graphOpId the rules decompose + **kwargs: forwarded to :func:`compile_registered_adjoint_rules` + + Returns: + list[str]: the rules, as MLIR strings + """ + try: + module = compile_registered_adjoint_rules(op_name, target_id, **kwargs) + except Exception as e: # pylint: disable=broad-except # pragma: no cover + warnings.warn( + f"Failed to lower the registered adjoint decomposition rules for {target_id}: {e}", + category=RuleLoweringWarning, + ) + return [] + return get_rule_strings_from_module(module) if module is not None else [] + + def adjoint_variant_rule_strings( - op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=None, is_custom_op=False + op_name, + op_id, + dynamic_shape, + wire_lens, + static_data, + extra_data=None, + is_custom_op=False, + op_cls=None, ): """Return the rule strings whose ``target_gate`` is ``Adjoint(op_name)``. ``op_id`` is the *base* op's graphOpId (e.g. ``"S{...}"``). Two pathways contribute: - 1. rules registered directly against ``Adjoint(op_name)`` (``list_decomps("Adjoint(S)")``), and + 1. rules registered directly against ``Adjoint(op_name)`` (``list_decomps("Adjoint(S)")``), + which take the symbolic operator's arguments the way PennyLane writes them + (``self_adjoint``, ``adjoint_rotation``, ...), and 2. rules synthesized by distributing each base rule of ``op_name`` over adjoint (the ``wrap_adjoint`` pathway), dropping any whose body is non-invertible. Shared by the eager lowering-time closure (:func:`fetch_all_reachable_decomposition_rules_from_op`) and the compiler's on-demand loader (:func:`compile_decomposition_rules_wrapper`) so both build adjoint rules identically. + + Args: + op_name (str): the base operator's name + op_id (str): the base operator's graphOpId + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + op_cls (type): the base operator's class; pathway 1 is skipped without it + + Returns: + list[str]: the rules, as MLIR strings """ out = [] adj_name = f"Adjoint({op_name})" - # (1) Rules registered directly against Adjoint(op_name): - try: + adj_id = name_wrap_adjoint(op_id) + # (1) Rules registered directly against Adjoint(op_name). They take a base operator instance, + # so they need the operator's class. + # TODO: the on-demand decomp rules are skipped for now. + if op_cls is not None: out.extend( - get_rule_strings_from_module( - compile_decomposition_rules( - adj_name, - name_wrap_adjoint(op_id), - dynamic_shape, - wire_lens, - static_data, - extra_data=extra_data, - is_custom_op=is_custom_op, - ) + registered_adjoint_rule_strings( + op_name, + adj_id, + dynamic_shape=dynamic_shape, + wire_lens=wire_lens, + static_data=static_data, + extra_data=extra_data, + is_custom_op=is_custom_op, + op_cls=op_cls, ) ) - except Exception as e: # pylint: disable=broad-except - warnings.warn( - f"Failed to lower the decomposition rules for {adj_name}: {e}", - category=RuleLoweringWarning, - ) # (2) Rules for Adjoint(op_name) synthesized by adjointing each base rule of op_name: try: distributed = get_rule_strings_from_module( @@ -561,6 +881,19 @@ def control_variant_rule_strings( *adjointed* base rule (``wrap_adjoint`` + ``wrap_control``), so controlled-adjoint ops are reachable too. Distribution rules whose body is non-controllable are dropped. + + Args: + op_name (str): the base operator's name + op_id (str): the base operator's graphOpId + ctrl_counts (list[int]): the control counts to build rules for + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + + Returns: + list[str]: the rules, as MLIR strings """ out = [] for n in ctrl_counts: @@ -628,7 +961,20 @@ def compile_decomposition_rules_wrapper( extra_data=None, is_custom_op=False, ) -> str: - """Return a string MLIR module containing the decomposition rules for an operator instance.""" + """Return a string MLIR module containing the decomposition rules for an operator instance. + + Args: + op_name (str): the operator's name, as PennyLane registers its rules + op_id (str): the operator's graphOpId + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + + Returns: + str: the module holding the rules + """ return str( compile_decomposition_rules( op_name, @@ -668,6 +1014,18 @@ def compile_reachable_decomposition_rules_wrapper( to complete a path. :func:`fetch_all_reachable_decomposition_rules_from_op` builds that closure (base + adjoint-registered + distributed-adjoint rules, transitively) and each returned func keeps its own ``target_gate``, which is how the loader registers them. + + Args: + op_name (str): the operator's *base* name + op_id (str): the operator's full graphOpId, modifiers included + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + + Returns: + str: a module holding the whole reachable rule closure """ base_id = op_id n_ctrls = 0 @@ -702,13 +1060,37 @@ def fetch_all_reachable_decomposition_rules_from_op( extra_data=None, is_custom_op=False, n_ctrls=0, + op_cls=None, ): + """Return every decomposition rule reachable from an operator, as MLIR strings. + + Starting from the given operator, this walks the resources its rules produce and captures the + rules of each op it meets, together with their adjoint and controlled variants, until nothing + new turns up. + + Args: + op_name (str): the operator's name, as PennyLane registers its rules + op_id (str): the operator's graphOpId + dynamic_shape (dict): dynamic argument names to their MLIR types + wire_lens (dict): wire argument names to their lengths + static_data (dict): compiler-static argument names to their values + extra_data (dict): argument values the graphOpId identifies by UID instead of spelling + is_custom_op (bool): whether the operator lowers to ``qref.custom`` + n_ctrls (int): the number of controls on the operator instance being decomposed + op_cls (type): the operator's class; classes of the ops met along the way are taken from + the resources themselves + + Returns: + list[str]: the rules, as MLIR strings + """ extra_data = extra_data or {} queue = deque() start = (op_name, dynamic_shape, wire_lens, static_data, extra_data, is_custom_op) queue.append(start) visited = {op_id} # remember ops by their graph id + op_classes = {op_name: op_cls} if op_cls is not None else {} + # Control counts to synthesize `C(...)` rules for. A single control is always captured # proactively; a multi-controlled instance (`n_ctrls > 1`) additionally needs its own count. ctrl_counts = [1] if n_ctrls <= 1 else [1, n_ctrls] @@ -746,6 +1128,7 @@ def compile_variants( static_data, extra_data=extra_data, is_custom_op=is_custom_op, + op_cls=op_classes.get(name), ) ) out.extend( @@ -777,42 +1160,54 @@ def compile_variants( ) = queue.popleft() this_extra_data = this_extra_data or {} this_kwargs = prepare_dynamic_op_kwargs(this_dynamic_shape, this_wire_lens) - # Explore ops reachable through the rules of both this op and its adjoint: - for explore_name in _op_variant_names(this_name): - resources, _, _ = collect_resources_for_op( - explore_name, this_kwargs | this_static_data | this_extra_data, this_is_custom_op - ) - for _rule_name, resource in resources.items(): - try: - for op, _count in resource.items(): - graph_op_id = GraphOpID(op) - probe_id = graph_op_id.getGraphOpId() - if probe_id in visited: - continue - - visited.add(probe_id) - probe = ( - graph_op_id.get_operator_name(), - graph_op_id.dynamic_shape, - graph_op_id.wire_lens, - graph_op_id.static_data, - graph_op_id.extra_data, - graph_op_id.is_custom_op, - ) - queue.append(probe) - rules.extend(compile_variants(probe[0], probe_id, *probe[1:])) - except Exception as e: - warnings.warn( - f"Failed to lower the {_rule_name} decomposition rule for {this_name}: {e}", - category=RuleLoweringWarning, + all_kwargs = this_kwargs | this_static_data | this_extra_data + + # Explore the ops reachable through the rules of this op and of its adjoint. Keyed by + # (explored op, rule name): the same rule name may be registered against both. + resources = { + (this_name, name): res + for name, res in collect_resources_for_op(this_name, all_kwargs, this_is_custom_op)[ + 0 + ].items() + } + if ( + not this_name.startswith("Adjoint(") + and (this_op_cls := op_classes.get(this_name)) is not None + ): + resources |= { + (f"Adjoint({this_name})", name): res + for name, res in collect_symbolic_adjoint_resources( + this_op_cls, this_name, all_kwargs, this_is_custom_op + )[2].items() + } + + for (_, _rule_name), resource in resources.items(): + try: + for op, _count in resource.items(): + graph_op_id = GraphOpID(op) + probe_id = graph_op_id.getGraphOpId() + probe = ( + graph_op_id.get_operator_name(), + graph_op_id.dynamic_shape, + graph_op_id.wire_lens, + graph_op_id.static_data, + graph_op_id.extra_data, + graph_op_id.is_custom_op, ) - continue - return rules + # Remembered even for an op already visited: another op may reach it later and + # need its class to rebuild the base of a registered adjoint rule. + op_classes.setdefault(probe[0], type(op)) + if probe_id in visited: + continue -def _op_variant_names(op_name): - """Yield the operator names to explore for `op_name`: the op itself and, unless it is already - adjointed, its adjoint `Adjoint(op_name)`. Rules registered against both are collected.""" - yield op_name - if not op_name.startswith("Adjoint("): - yield f"Adjoint({op_name})" + visited.add(probe_id) + queue.append(probe) + rules.extend(compile_variants(probe[0], probe_id, *probe[1:])) + except Exception as e: + warnings.warn( + f"Failed to lower the {_rule_name} decomposition rule for {this_name}: {e}", + category=RuleLoweringWarning, + ) + continue + return rules diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 455aacc78e..451995ee08 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -255,6 +255,7 @@ def compile_decomp_rules( wire_lens={"wires": wire_lens[0]}, static_data={}, is_custom_op=True, + op_cls=op_cls, ) elif op_cls is qp.MultiRZ: @@ -268,6 +269,7 @@ def compile_decomp_rules( dynamic_shape=dynamic_shape, wire_lens={f"{wire_argname}": wire_lens[0]}, static_data={}, + op_cls=op_cls, ) elif op_cls is qp.PauliRot: @@ -287,6 +289,7 @@ def compile_decomp_rules( dynamic_shape=dynamic_shape, wire_lens={f"{wire_argname}": wire_lens[0]}, static_data=repack_static_data, + op_cls=op_cls, ) elif op_cls is qp.PCPhase: @@ -305,6 +308,7 @@ def compile_decomp_rules( dynamic_shape=dynamic_shape, wire_lens={f"{wire_argname}": wire_lens[0]}, static_data=repack_static_data, + op_cls=op_cls, ) elif op_cls is qp.GlobalPhase: @@ -317,6 +321,7 @@ def compile_decomp_rules( dynamic_shape=dynamic_shape, wire_lens={}, static_data={}, + op_cls=op_cls, ) elif op_cls is qp.QubitUnitary: @@ -335,6 +340,7 @@ def compile_decomp_rules( dynamic_shape=dynamic_shape, wire_lens={f"{wire_argname}": wire_lens[0]}, static_data={}, + op_cls=op_cls, ) else: @@ -421,6 +427,7 @@ def compile_decomp_rules( wire_lens=non_hybrid_wire_lens, static_data=repack_static_data, extra_data=extra_data, + op_cls=op_cls, ) inject_new_rules_into_module(module, decomp_rules) diff --git a/frontend/test/lit/test_operator2/test_lower_time_rules_adjoint.py b/frontend/test/lit/test_operator2/test_lower_time_rules_adjoint.py index eb16303205..2e6d567256 100644 --- a/frontend/test/lit/test_operator2/test_lower_time_rules_adjoint.py +++ b/frontend/test/lit/test_operator2/test_lower_time_rules_adjoint.py @@ -38,13 +38,15 @@ def base_rule(reg): def _adj_rule(): - def adj_resource_fn(reg): + """A rule for ``Adjoint(NoParams)``.""" + + def adj_resource_fn(base): return {SingleParam(x=Float, reg=Wire[2]): 2} @qp.register_resources(adj_resource_fn) - def adj_rule(reg): - SingleParam(x=0.2, reg=reg[0:2]) - SingleParam(x=0.3, reg=reg[0:2]) + def adj_rule(base): + SingleParam(x=0.2, reg=base.wires[0:2]) + SingleParam(x=0.3, reg=base.wires[0:2]) return adj_rule diff --git a/frontend/test/lit/test_operator2/test_on_demand_rules_adjoint.py b/frontend/test/lit/test_operator2/test_on_demand_rules_adjoint.py index 8dc4b87bb8..063a03b4ed 100644 --- a/frontend/test/lit/test_operator2/test_on_demand_rules_adjoint.py +++ b/frontend/test/lit/test_operator2/test_on_demand_rules_adjoint.py @@ -41,22 +41,24 @@ def base_rule(reg): def _adj_rule(): - def adj_resource_fn(reg): + """A rule for ``Adjoint(NoParams)``.""" + + def adj_resource_fn(base): return {SingleParam(x=Float, reg=Wire[2]): 2} @qp.register_resources(adj_resource_fn) - def adj_rule(reg): - SingleParam(x=0.2, reg=reg[0:2]) - SingleParam(x=0.3, reg=reg[0:2]) + def adj_rule(base): + SingleParam(x=0.2, reg=base.wires[0:2]) + SingleParam(x=0.3, reg=base.wires[0:2]) return adj_rule def test_on_demand_adjoint_id_routes_to_adjoint_rules(): """Requesting the rules for an ``Adjoint(NoParams)`` id yields the reachable closure with the - adjoint rules correctly routed: the base rule of NoParams, the rule registered on - Adjoint(NoParams), and the synthesized (adjointed base) rule. - Crucially, the adjoint target is never attached to the un-adjointed base body.""" + adjoint rules correctly routed: the base rule of ``NoParams`` and the synthesized + rule. Crucially, the adjoint target is never attached to the un-adjointed base body. + """ with qp.decomposition.local_decomps(): qp.add_decomps(NoParams, _base_rule()) qp.add_decomps("Adjoint(NoParams)", _adj_rule()) @@ -69,7 +71,7 @@ def test_on_demand_adjoint_id_routes_to_adjoint_rules(): # CHECK: module { # CHECK-DAG: func.func private @"__builtin_base_rule_NoParams{}{reg:2}{}"{{.*}}"SingleParam{{.*}}target_gate = "NoParams{}{reg:2}{}" - # CHECK-DAG: func.func private @"__builtin_adj_rule_Adjoint(NoParams){}{reg:2}{}"{{.*}}"SingleParam{{.*}} = 2 : i64{{.*}}target_gate = "Adjoint(NoParams){}{reg:2}{}" + # CHECK-NOT: __builtin_adj_rule_Adjoint(NoParams) # CHECK-DAG: func.func private @"__builtin_base_rule_Adjoint(NoParams){}{reg:2}{}"{{.*}}"Adjoint(SingleParam{{.*}}target_gate = "Adjoint(NoParams){}{reg:2}{}" # CHECK: qref.adjoint diff --git a/frontend/test/lit/test_operator2/test_symbolic_registered_rules.py b/frontend/test/lit/test_operator2/test_symbolic_registered_rules.py new file mode 100644 index 0000000000..8ae0ab1a1a --- /dev/null +++ b/frontend/test/lit/test_operator2/test_symbolic_registered_rules.py @@ -0,0 +1,79 @@ +# 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. + +""" +Test the lowering of decomposition rules registered against ``Adjoint(Op)`` and written against the +*symbolic* operator's arguments following PennyLane's conventions. +""" + +# RUN: %PYTHON %s | FileCheck %s + +# pylint: disable = line-too-long,unused-argument,missing-function-docstring + +import pennylane as qp +from operator2_dummy_gates import NoParams, SingleParam +from pennylane.typing import Float, Wire + +from catalyst.decomposition.decomposition_rules import ( + fetch_all_reachable_decomposition_rules_from_op, +) + + +def _base_rule(): + def base_resource_fn(reg): + return {SingleParam(x=Float, reg=Wire[2]): 1} + + @qp.register_resources(base_resource_fn) + def base_rule(reg): + SingleParam(x=0.1, reg=reg[0:2]) + + return base_rule + + +def _self_adjoint_rule(): + """Test the lowering into a rule for ``Adjoint(NoParams)`` that applies the + bare base op without regions.""" + + @qp.register_resources(lambda base: {qp.core.abstractify(base): 1}) + def self_adjoint(base): + qp.apply(base) + + return self_adjoint + + +def test_registered_self_adjoint_rule_targets_the_adjoint_op(): + """Test the lowering into a rule for ``Adjoint(NoParams)`` that applies the + bare base op without regions.""" + with qp.decomposition.local_decomps(): + qp.add_decomps(NoParams, _base_rule()) + qp.add_decomps("Adjoint(NoParams)", _self_adjoint_rule()) + + print( + "\n".join( + fetch_all_reachable_decomposition_rules_from_op( + op_name="NoParams", + op_id="NoParams{}{reg:2}{}", + dynamic_shape={}, + wire_lens={"reg": 2}, + static_data={}, + op_cls=NoParams, + ) + ) + ) + + # CHECK-DAG: func.func private @"__builtin_self_adjoint_Adjoint(NoParams){}{reg:2}{}"(%arg0: !qref.reg<2>, %arg1: tensor<2xi64>){{.*}}"NoParams{}{reg:2}{}" = 1 : i64{{.*}}target_gate = "Adjoint(NoParams){}{reg:2}{}" + # CHECK-DAG: func.func private @"__builtin_base_rule_NoParams{}{reg:2}{}"{{.*}}target_gate = "NoParams{}{reg:2}{}" + + +test_registered_self_adjoint_rule_targets_the_adjoint_op() diff --git a/frontend/test/pytest/test_decomposition.py b/frontend/test/pytest/test_decomposition.py index f66c59192e..4f5e5c45fe 100644 --- a/frontend/test/pytest/test_decomposition.py +++ b/frontend/test/pytest/test_decomposition.py @@ -45,6 +45,8 @@ _modifier_kind, compile_decomposition_rules_wrapper, compile_reachable_decomposition_rules_wrapper, + compile_registered_adjoint_rules, + get_rule_strings_from_module, name_unwrap_adjoint, name_unwrap_control, name_wrap_adjoint, @@ -592,5 +594,68 @@ def test_wrap_modifier_id_rejects_non_canonical(self, op_id): wrap_modifier_id(op_id, "Adjoint") +class TestSymbolicRules: + """Tests for the rules registered against a symbolic operator that take the symbolic + op's args; following the convention in PennyLane.""" + + def test_self_adjoint_rule_is_lowered(self): + """Test ``self_adjoint`` rule on ``Adjoint(Hadamard)``.""" + + module = compile_registered_adjoint_rules( + "Hadamard", "Adjoint(Hadamard){}{wires:1}{}", {}, {"wires": 1}, {}, op_cls=qp.Hadamard + ) + (rule,) = get_rule_strings_from_module(module) + + assert 'target_gate = "Adjoint(Hadamard){}{wires:1}{}"' in rule + assert 'resources = {operations = {"Hadamard{}{wires:1}{}" = 1 : i64}}' in rule + assert "qref.adjoint" not in rule + assert "(%arg0: !qref.reg<1>, %arg1: tensor<1xi64>)" in rule + assert rule.count('gate_name = "Hadamard"') == 1 + + def test_adjoint_rotation_rule_is_lowered(self): + """Test ``adjoint_rotation`` reads the angle off the base operator.""" + + module = compile_registered_adjoint_rules( + "RZ", + "Adjoint(RZ){0:[f64]}{wires:1}{}", + {"0": ["f64"]}, + {"wires": 1}, + {}, + is_custom_op=True, + op_cls=qp.RZ, + ) + (rule,) = get_rule_strings_from_module(module) + + assert 'target_gate = "Adjoint(RZ){0:[f64]}{wires:1}{}"' in rule + assert 'resources = {operations = {"RZ{0:[f64]}{wires:1}{}" = 1 : i64}}' in rule + assert "qref.adjoint" not in rule + assert "stablehlo.negate" in rule + + def test_no_registered_symbolic_rules(self): + """Test an op with no symbolic rules registered against its adjoint yields no module.""" + + with local_decomps(): + assert ( + compile_registered_adjoint_rules( + "NoParams", + "Adjoint(NoParams){}{reg:2}{}", + {}, + {"reg": 2}, + {}, + op_cls=NoParams, + ) + is None + ) + + def test_missing_op_class_raises(self): + """Test lowering cannot proceed without the base operator's class: these rules take a base + operator instance, which the operator's name alone cannot produce.""" + + with pytest.raises(ValueError, match="operator class of 'Hadamard' is needed"): + compile_registered_adjoint_rules( + "Hadamard", "Adjoint(Hadamard){}{wires:1}{}", {}, {"wires": 1}, {} + ) + + if __name__ == "__main__": pytest.main(["-x", __file__])