[GOID] Use unified printing for parameter types (#3191)

Small follow-up on #3188 to remove the manual Python -> MLIR type
conversion and manual string serialization of MLIR types for the
GraphOpID format. Instead the Python binding utilities are used and the
MLIR attribute printer handles the string serialization.

As a side-effect, more types are supported now that weren't present in
the table (e.g. uint).

Co-authored-by: paul0403 <paulhaochen.wang@gmail.com>

Author: dime10

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))
 
 

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<f64>": 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<f64>",
-    ir.F16Type: "f16",
-    ir.F32Type: "f32",
-    ir.F64Type: "f64",
-    (ir.ComplexType, ir.F32Type): "complex<f32>",
-    (ir.ComplexType, ir.F64Type): "complex<f64>",
-    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<f32>",
-    np.dtype("complex128"): "complex<f64>",
-}
-
 
 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):

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<f64>>"
-            ]
-        }
+        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]}, {})
 

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.