diff --git a/frontend/catalyst/decomposition/graph_op_id.py b/frontend/catalyst/decomposition/graph_op_id.py index d9394c5d21..e92c3b6c44 100644 --- a/frontend/catalyst/decomposition/graph_op_id.py +++ b/frontend/catalyst/decomposition/graph_op_id.py @@ -14,7 +14,6 @@ """Python implementation of Graph Operator ID.""" -import contextlib from collections.abc import Mapping, Sequence from typing import Any @@ -29,25 +28,14 @@ replace_wires_with_placeholder_wires, ) from catalyst.from_plxpr.uid import generate_uid -from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval +from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval, mlir_build_context _SPECIAL_LOWERINGS = {} -@contextlib.contextmanager -def _attribute_context(): - """Provide an MLIR context and location for attribute construction.""" - if current := ir.Context.current: - with ir.Location.unknown(context=current): - yield current - else: - with ir.Context() as context, ir.Location.unknown(context=context): - yield context - - def format_static_data_dict_for_id(static_data): """Format the static-data group of a GraphOpID with MLIR's attribute printer.""" - with _attribute_context(): + with mlir_build_context(): return str(get_mlir_attribute_from_pyval(static_data)) diff --git a/frontend/catalyst/decomposition/type_utils.py b/frontend/catalyst/decomposition/type_utils.py index fe06cf4a44..3183183271 100644 --- a/frontend/catalyst/decomposition/type_utils.py +++ b/frontend/catalyst/decomposition/type_utils.py @@ -21,10 +21,13 @@ import jax.numpy as jnp import numpy as np import pennylane as qp +from jax._src.interpreters.mlir import dtype_to_ir_type from jax._src.lib.mlir import ir from jax.core import ShapedArray from pennylane.pytrees import flatten, unflatten +from catalyst.jax_extras.lowering import mlir_build_context + _MLIR_DTYPES_TO_PY_DTYPES = { "i1": jnp.bool_, "i8": jnp.int8, @@ -38,45 +41,23 @@ "complex": jnp.complex128, } -_PY_DTYPES_TO_MLIR_DTYPES = {v: k for k, v in _MLIR_DTYPES_TO_PY_DTYPES.items()} | { - float: "f64", - int: "i64", - complex: "complex", - ir.F16Type: "f16", - ir.F32Type: "f32", - ir.F64Type: "f64", - (ir.ComplexType, ir.F32Type): "complex", - (ir.ComplexType, ir.F64Type): "complex", - np.dtype("bool_"): "i1", - np.dtype("int8"): "i8", - np.dtype("int16"): "i16", - np.dtype("int32"): "i32", - np.dtype("int64"): "i64", - np.dtype("float16"): "f16", - np.dtype("float32"): "f32", - np.dtype("float64"): "f64", - np.dtype("complex64"): "complex", - np.dtype("complex128"): "complex", -} - def convert_item_to_mlir_type(item, is_special_lowering=False): - """Convert a string or PennyLane AbstractArray to an mlir type annotation.""" + """Convert a string or PennyLane AbstractArray to an mlir type annotation. + + The type is spelled by MLIR's own printer, applied to the type the value lowers to, which is + what ``printDynamicShape`` in mlir/lib/Quantum/IR/QuantumInterfaces.cpp does as well. One + printer spelling both sides is what keeps a rule compiled here findable by the + ``graph-decomposition`` pass. + """ if isinstance(item, str): return item - if item.shape == (): - if is_special_lowering: - return _PY_DTYPES_TO_MLIR_DTYPES[item.dtype] - return "tensor<" + _PY_DTYPES_TO_MLIR_DTYPES[item.dtype] + ">" - - return ( - "tensor<" - + "x".join(str(dim_size) for dim_size in item.shape) - + "x" - + _PY_DTYPES_TO_MLIR_DTYPES[item.dtype] - + ">" - ) + with mlir_build_context(): + element_type = dtype_to_ir_type(np.dtype(item.dtype)) + if is_special_lowering and item.shape == (): + return str(element_type) + return str(ir.RankedTensorType.get(item.shape, element_type)) def get_dummy_values_for_arg(arg): diff --git a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py index 4466410adc..455aacc78e 100644 --- a/frontend/catalyst/from_plxpr/qref_operator2_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_operator2_primitives.py @@ -17,6 +17,7 @@ """ # pylint: disable=unused-argument +import numpy as np import pennylane as qp from jax._src.lib.mlir import ir from jax.core import ShapedArray @@ -39,7 +40,10 @@ inject_new_rules_into_module, ) from catalyst.decomposition.graph_op_id import _SPECIAL_LOWERINGS, build_graph_op_id -from catalyst.decomposition.type_utils import get_dummy_values_for_arg +from catalyst.decomposition.type_utils import ( + convert_item_to_mlir_type, + get_dummy_values_for_arg, +) from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval from catalyst.jax_extras.patches import mock_attributes from catalyst.jax_primitives import ( @@ -318,11 +322,10 @@ def compile_decomp_rules( elif op_cls is qp.QubitUnitary: num_wires = wire_lens[0] matrix_size = 2**num_wires - dynamic_shape = { - qp.QubitUnitary.dynamic_argnames[0]: [ - f"tensor<{matrix_size}x{matrix_size}xcomplex>" - ] - } + matrix_type = convert_item_to_mlir_type( + ShapedArray((matrix_size, matrix_size), np.complex128) + ) + dynamic_shape = {qp.QubitUnitary.dynamic_argnames[0]: [matrix_type]} wire_argname = qp.QubitUnitary.wire_argnames[0] op_id = build_graph_op_id("QubitUnitary", dynamic_shape, {wire_argname: wire_lens[0]}, {}) diff --git a/frontend/catalyst/jax_extras/lowering.py b/frontend/catalyst/jax_extras/lowering.py index 2954d582d3..223a6b5f2c 100644 --- a/frontend/catalyst/jax_extras/lowering.py +++ b/frontend/catalyst/jax_extras/lowering.py @@ -15,6 +15,7 @@ from __future__ import annotations +import contextlib import dataclasses import logging import textwrap @@ -192,6 +193,17 @@ def custom_lower_jaxpr_to_module( return ctx.module, ctx.context +@contextlib.contextmanager +def mlir_build_context(): + """Provide an MLIR context and location for attribute and type construction.""" + if current := ir.Context.current: + with ir.Location.unknown(context=current): + yield current + else: + with ir.Context() as context, ir.Location.unknown(context=context): + yield context + + def get_mlir_attribute_from_pyval(value): """ Given a value of any type, construct an mlir attribute of corresponding type.