Decomp/lower time rules

frontend/catalyst/decomposition/decomposition_rules.py

@@ -18,6 +18,7 @@
 
 # pylint: disable=protected-access,bare-except
 
+from collections import deque
 from functools import partial
 
 import jax.numpy as jnp
@@ -27,12 +28,12 @@
 from pennylane.pytrees import flatten
 
 from catalyst.decomposition.type_utils import (
+    convert_types_to_mlir_strings,
+    format_dynamic_params_for_id,
     get_dummy_values_for_container,
-    mlir_stringify_type,
     post_process_concretize_leaves,
     replace_abstract_wires_with_concrete_wires,
 )
-from catalyst.from_plxpr.qref_operator2_primitives import _is_custom_op
 from catalyst.from_plxpr.uid import generate_uid
 from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
 
@@ -77,10 +78,7 @@ def __init__(self, op: qp.core.Operator2):
 
     def parse_dynamic_shape(self) -> dict:
         """Return the dynamic shape as a dictionary of dtypes from the dynamic arg names."""
-        return {
-            argname: mlir_stringify_type(argtype)
-            for argname, argtype in sorted(self.op.dynamic_args.items())
-        }
+        return {argname: argtype for argname, argtype in sorted(self.op.dynamic_args.items())}
 
     def parse_wire_lens(self) -> dict:
         """Return the length of each of the wire args as a dictionary from the wire arg names."""
@@ -131,9 +129,12 @@ def get_operator_name(self) -> str:
         """Return the name of the operator."""
         return self.operator_name
 
+    def get_dynamic_shape(self) -> dict:
+        return self.dynamic_shape
+
     def get_dynamic_shape_id_format(self) -> str:
         """Return the dynamic shape formatted for GraphOpId."""
-        return "{" + ",".join(f"{name}:{shape}" for name, shape in self.dynamic_shape.items()) + "}"
+        return format_dynamic_params_for_id(convert_types_to_mlir_strings(self.dynamic_shape))
 
     def get_wire_lens_id_format(self) -> str:
         """Return the wire lengths formatted for GraphOpId."""
@@ -161,6 +162,44 @@ def getID(self) -> str:
         return ID_string
 
 
+def get_rule_funcs_from_module(module: ir.Module) -> list[ir.Operation]:
+    funcOps = []
+
+    def find_condition(op):
+        if op.name == "func.func":
+            if "target_gate" in op.attributes:
+                old_attr = op.attributes["sym_name"]
+                op.attributes["sym_name"] = ir.StringAttr.get(
+                    "__builtin_" + old_attr.value.strip('"'), context=old_attr.context
+                )
+                funcOps.append(op)
+                return ir.WalkResult.SKIP
+        return ir.WalkResult.ADVANCE
+
+    module.operation.walk(find_condition)
+    return funcOps
+
+
+def get_rules_from_module_as_list(module: ir.Module) -> list[str]:
+    funcOps = get_rule_funcs_from_module(module)
+    return [str(funcOp) for funcOp in funcOps]
+
+
+def get_rules_from_module(module: ir.Module) -> str:
+    """
+    Parse and modify decomposition rules from a ModuleOp.
+
+    Args:
+        module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted
+
+    Returns:
+        str: The string representation of any decomposition rules from `module`, pre-pending the
+             `__builtin_` prefix to their names.
+    """
+    funcOps = get_rule_funcs_from_module(module)
+    return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else ""
+
+
 def collect_resources_for_op(op_name, kwargs, is_custom_op=False):
     """Return resource data for all decomposition rules associated to op_name."""
     decomp_rules = list(qp.decomposition.list_decomps(op_name))
@@ -184,6 +223,13 @@ def collect_resources_for_op(op_name, kwargs, is_custom_op=False):
     return name_to_resources, name_to_resource_ids, decomp_rules
 
 
+def prepare_dynamic_op_kwargs(dynamic_shape, wire_lens) -> dict:
+    kwargs = {}
+    for wire_name, wire_len in wire_lens.items():
+        kwargs[wire_name] = jnp.array(range(wire_len), dtype=int)
+    return kwargs | get_dummy_values_for_container(dynamic_shape)
+
+
 def compile_decomposition_rules(
     op_name,
     op_id,
@@ -198,15 +244,9 @@ def compile_decomposition_rules(
 
     The decomposition rules will be decorated with appropriate resource and target_gate attributes.
     """
-    kwargs = {}
+    kwargs = prepare_dynamic_op_kwargs(dynamic_shape, wire_lens)
     extra_data = extra_data or {}
-
     device = qp.device("null.qubit", wires=sum(wire_lens.values()))
-    for wire_name, wire_len in wire_lens.items():
-        kwargs[wire_name] = jnp.array(range(wire_len), dtype=int)
-    for arg_name, arg_shape in dynamic_shape.items():
-        kwargs[arg_name] = get_dummy_values_for_container(arg_shape)
-
     _, name_to_resource_ids, decomp_rules = collect_resources_for_op(
         op_name, kwargs | static_data | extra_data, is_custom_op
     )
@@ -228,6 +268,8 @@ def decomp_rule(*_args, **_kwargs):
 
     subroutines = [rule_to_subroutine(rule) for rule in decomp_rules]
 
+    # TODO: reconcretify abstracted hybrid ops
+
     @qp.qjit(
         target="mlir",
         capture=True,
@@ -285,3 +327,50 @@ def compile_decomposition_rules_wrapper(
             is_custom_op=is_custom_op,
         )
     )
+
+
+def fetch_all_reachable_decomposition_rules_from_op(
+    op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=None
+):
+    queue = deque()
+    start = (op_name, dynamic_shape, wire_lens, static_data, extra_data)
+    queue.append(start)
+    visited = [start]
+    rules = get_rules_from_module_as_list(
+        compile_decomposition_rules(
+            op_name, op_id, dynamic_shape, wire_lens, static_data, extra_data=extra_data
+        )
+    )
+
+    while len(queue) != 0:
+        this_name, this_dynamic_shape, this_wire_lens, this_static_data, this_extra_data = (
+            queue.popleft()
+        )
+        this_extra_data = this_extra_data or {}
+        this_kwargs = prepare_dynamic_op_kwargs(this_dynamic_shape, this_wire_lens)
+        resources, _, _ = collect_resources_for_op(
+            this_name, this_kwargs | this_static_data | this_extra_data
+        )
+        for _rule_name, resource in resources.items():
+            for op, _count in resource.items():
+                graph_op_id = GraphOpID(op)
+                probe = (
+                    graph_op_id.get_operator_name(),
+                    convert_types_to_mlir_strings(graph_op_id.get_dynamic_shape()),
+                    graph_op_id.wire_lens,
+                    graph_op_id.static_data,
+                    graph_op_id.extra_data,
+                )
+                if not probe in visited:
+                    visited.append(probe)
+                    queue.append(probe)
+                    module = compile_decomposition_rules(
+                        probe[0],
+                        graph_op_id.getID(),
+                        probe[1],
+                        probe[2],
+                        probe[3],
+                        probe[4],
+                    )
+                    rules.extend(get_rules_from_module_as_list(module))
+    return rules

frontend/catalyst/decomposition/precompile_decomposition_rules.py

@@ -21,7 +21,11 @@
 from pennylane.operation import Operator, Operator2
 
 from catalyst.compiler import _quantum_opt
-from catalyst.decomposition.decomposition_rules import GraphOpID, compile_decomposition_rules
+from catalyst.decomposition.decomposition_rules import (
+    GraphOpID,
+    compile_decomposition_rules,
+    get_rules_from_module,
+)
 from catalyst.utils.runtime_environment import BYTECODE_FILE_PATH
 
 # TODO: Uncomment dynamic size wires ops once they are supported
@@ -71,40 +75,6 @@
 }
 
 
-def get_rule_funcs_from_module(module: ir.Module) -> list[ir.Operation]:
-    funcOps = []
-
-    def find_condition(op):
-        if op.name == "func.func":
-            if "target_gate" in op.attributes:
-                old_attr = op.attributes["sym_name"]
-                op.attributes["sym_name"] = ir.StringAttr.get(
-                    "__builtin_" + old_attr.value.strip('"'), context=old_attr.context
-                )
-                funcOps.append(op)
-                return ir.WalkResult.SKIP
-        return ir.WalkResult.ADVANCE
-
-    module.operation.walk(find_condition)
-    return funcOps
-
-
-def get_rules_from_module(module: ir.Module) -> str:
-    """
-    Parse and modify decomposition rules from a ModuleOp.
-
-    Args:
-        module: an MLIR module object containing a FuncOp named `rule_wrapper` to be extracted
-
-    Returns:
-        str: The string representation of any decomposition rules from `module`, pre-pending the
-             `__builtin_` prefix to their names.
-    """
-    funcOps = get_rule_funcs_from_module(module)
-
-    return "\n".join(str(funcOp) for funcOp in funcOps) if funcOps else ""
-
-
 def get_abstract_args(op_class: type[Operator]) -> list[type]:
     """
     Create jax-compatible abstract args for catalyst DecompositionRules that apply to op_class.

frontend/catalyst/decomposition/type_utils.py

@@ -18,6 +18,7 @@
 
 import jax.numpy as jnp
 import pennylane as qp
+from jax._src.lib.mlir import ir
 from jax.core import ShapedArray
 
 _MLIR_DTYPES_TO_PY_DTYPES = {
@@ -33,42 +34,98 @@
     "complex<f128>": jnp.complex128,
 }
 
-_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()}
+_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} | {
+    (ir.IntegerType, 1): "i1",
+    (ir.IntegerType, 8): "i8",
+    (ir.IntegerType, 16): "i16",
+    (ir.IntegerType, 32): "i32",
+    (ir.IntegerType, 64): "i64",
+    ir.F16Type: "f16",
+    ir.F32Type: "f32",
+    ir.F64Type: "f64",
+    (ir.ComplexType, ir.F64Type): "complex<f64>",
+}
 
 
-def _stringify_shaped_type(shape: tuple, dim: int, element_type) -> str:
-    """Return a string representation of the given shaped data type."""
-    if dim + 1 == len(shape):
-        inner_content = _PY_DTYPES_TO_MLIR_DTYPES[element_type]
-    else:
-        inner_content = _stringify_shaped_type(shape, dim + 1, element_type)
-    length = shape[dim]
-    return f"[{','.join([inner_content] * length)}]"
-
-
-def mlir_stringify_type(dtype: qp.typing.AbstractArray):
-    """Return a string representation of the given data type."""
-    assert isinstance(
-        dtype, qp.typing.AbstractArray
-    ), f"Expected an AbstractArray to stringify, got {dtype}"
-    element_type = dtype.dtype.type
-    if dtype.shape == ():
-        return f"[{_PY_DTYPES_TO_MLIR_DTYPES[element_type]}]"
-    else:
-        return _stringify_shaped_type(dtype.shape, 0, element_type)
+def get_mlir_tensor_type_map_key(mlir_type):
+    if isinstance(mlir_type, ir.ComplexType):
+        return (type(mlir_type), type(mlir_type.element_type))
+    if isinstance(mlir_type, ir.IntegerType):
+        return (type(mlir_type), mlir_type.width)
+    return type(mlir_type)
+
+
+def convert_shaped_type_to_mlir_string(shaped_type, current_dim=0):
+    """Convert a shape of arbitrary dimension to a string with MLIR type strings for values."""
+    if isinstance(shaped_type, (ShapedArray, qp.typing.AbstractArray)):
+        if current_dim == shaped_type.ndim:
+            return _PY_DTYPES_TO_MLIR_DTYPES[shaped_type.dtype.type]
+
+        return [
+            convert_shaped_type_to_mlir_string(shaped_type, current_dim + 1)
+        ] * shaped_type.shape[current_dim]
+    elif isinstance(shaped_type, ir.RankedTensorType):
+        if current_dim == shaped_type.rank:
+            return _PY_DTYPES_TO_MLIR_DTYPES[get_mlir_tensor_type_map_key(shaped_type.element_type)]
+
+        return [
+            convert_shaped_type_to_mlir_string(shaped_type, current_dim + 1)
+        ] * shaped_type.shape[current_dim]
+
+
+def convert_types_to_mlir_strings(d: dict) -> dict:
+    """Convert the values of a dictionary to MLIR type strings."""
+
+    def handle_item(item):
+        if isinstance(item, type):
+            if item in _PY_DTYPES_TO_MLIR_DTYPES:
+                return _PY_DTYPES_TO_MLIR_DTYPES[item]
+            raise TypeError(
+                f"encountered unknown type {type(item)} of item {item} when converting to mlir strings."
+            )
+        elif type(item) in _PY_DTYPES_TO_MLIR_DTYPES:
+            return _PY_DTYPES_TO_MLIR_DTYPES[type(item)]
+        elif isinstance(item, str):
+            return item
+        elif isinstance(item, (list, tuple)):
+            return [handle_item(i) for i in item]
+        elif isinstance(item, (ShapedArray, qp.typing.AbstractArray)):
+            if item.shape == ():
+                return [_PY_DTYPES_TO_MLIR_DTYPES[item.dtype.type]]
+            return convert_shaped_type_to_mlir_string(item)
+        elif isinstance(item, ir.RankedTensorType):
+            if len(item.shape) == 0:
+                return [_PY_DTYPES_TO_MLIR_DTYPES[get_mlir_tensor_type_map_key(item.element_type)]]
+            return convert_shaped_type_to_mlir_string(item)
+        else:
+            raise TypeError(
+                f"encountered unknown type {type(item)} of item {item} when converting to mlir strings."
+            )
+
+    return {k: handle_item(v) for k, v in d.items()}
+
+
+def format_dynamic_params_for_id(d):
+    """Format a structure for ID, after calling convert_types_to_mlir_string on it."""
+
+    def handle_item(item):
+        if isinstance(item, str):
+            return item
+        elif isinstance(item, list):
+            return "[" + ",".join(handle_item(i) for i in item) + "]"
+
+    return "{" + ",".join(k + ":" + handle_item(v) for k, v in d.items()) + "}"
 
 
-def get_dummy_values_for_container(container):
+def get_dummy_values_for_container(dictionary):
     """
-    Given a container of python or MLIR types, replace the types with corresponding dummy values.
+    Given a dictionary of python or MLIR types, replace the types with corresponding dummy values.
 
-    Each item in the container must be representible as an MLIR tensor with at most one layer of
+    Each item in the dictionary must be representable as an MLIR tensor with at most one layer of
     nesting, i.e. cannot be nested and all elements must be of the same type.
     Ex.
     [[float, float], [int, int, int], [int32, int32, int32, int32]]
     """
-    if isinstance(container, str):
-        return jnp.zeros((), dtype=_MLIR_DTYPES_TO_PY_DTYPES[container])
 
     def handle_item(item):
         if isinstance(item, (list, tuple)):
@@ -85,7 +142,7 @@ def handle_item(item):
                     f"Unexpected type in container when creating dummy values: {type(item)}"
                 )
 
-    return tuple(handle_item(item) for item in container)
+    return {k: handle_item(v) for k, v in dictionary.items()}
 
 
 def replace_abstract_wires_with_concrete_wires(node):