🚧 Python Compiler

a.py

@@ -0,0 +1,105 @@
+import pennylane as qml
+from pennylane.tape import QuantumScript, QuantumScriptBatch
+from pennylane.transforms import cancel_inverses as pl_cancel_inverses
+from pennylane.typing import PostprocessingFn
+from catalyst.tracing.contexts import EvaluationContext
+from catalyst.passes.xdsl_plugin import getXDSLPluginAbsolutePath
+
+import catalyst
+import functools
+
+from dataclasses import dataclass
+from xdsl import context, passes
+from xdsl.utils import parse_pipeline
+from xdsl.dialects import builtin, func
+
+def xdsl_transform(_klass):
+
+    def identity_transform(tape):
+        return tape, lambda args: args[0]
+
+    identity_transform.__name__ = "xdsl_transform" + _klass.__name__
+    transform = qml.transform(identity_transform)
+    catalyst.from_plxpr.register_transform(transform, _klass.name, False)
+    from catalyst.python_compiler import register_pass
+    register_pass(_klass.name, lambda : _klass())
+
+    return transform
+
+
+from xdsl.rewriter import InsertPoint
+from xdsl import pattern_rewriter
+
+from catalyst.python_compiler import quantum
+
+
+self_inverses = ("PauliZ", "PauliX", "PauliY", "Hadamard", "Identity")
+
+
+def cancel_ops(rewriter, op, next_op):
+    rewriter.replace_all_uses_with(next_op.results[0], op.in_qubits[0])
+    rewriter.erase_op(next_op)
+    rewriter.erase_op(op)
+    owner = op.in_qubits[0].owner
+
+    if isinstance(owner, quantum.CustomOp) and owner.gate_name.data in self_inverses:
+        next_user = None
+
+        for use in owner.results[0].uses:
+            user = use.operation
+            if isinstance(user, quantum.CustomOp) and user.gate_name.data == owner.gate_name.data:
+                next_user = user
+                break
+
+        if next_user is not None:
+            cancel_ops(rewriter, owner, next_user)
+
+
+class DeepCancelInversesSingleQubitPattern(pattern_rewriter.RewritePattern):
+    @pattern_rewriter.op_type_rewrite_pattern
+    def match_and_rewrite(self, funcOp: func.FuncOp, rewriter: pattern_rewriter.PatternRewriter):
+        """Deep Cancel for Self Inverses"""
+        print(funcOp)
+        for op in funcOp.body.walk():
+            if not isinstance(op, quantum.CustomOp):
+                continue
+
+            if op.gate_name.data not in self_inverses:
+                continue
+
+            next_user = None
+            for use in op.results[0].uses:
+                user = use.operation
+                if isinstance(user, quantum.CustomOp) and user.gate_name.data == op.gate_name.data:
+                    next_user = user
+                    break
+
+            if next_user is not None:
+                cancel_ops(rewriter, op, next_user)
+
+
+@xdsl_transform
+class DeepCancelInversesSingleQubitPass(passes.ModulePass):
+    name = "deep-cancel-inverses-single-qubit"
+
+    def apply(self, ctx: context.MLContext, module: builtin.ModuleOp) -> None:
+        pattern_rewriter.PatternRewriteWalker(
+            pattern_rewriter.GreedyRewritePatternApplier([DeepCancelInversesSingleQubitPattern()])
+        ).rewrite_module(module)
+
+qml.capture.enable()
+
+@catalyst.qjit(pass_plugins=[getXDSLPluginAbsolutePath()])
+@DeepCancelInversesSingleQubitPass
+@qml.qnode(qml.device("lightning.qubit", wires=1))
+def captured_circuit(x: float):
+    qml.RX(x, wires=0)
+    qml.Hadamard(wires=0)
+    qml.Hadamard(wires=0)
+    return qml.state()
+
+
+qml.capture.disable()
+
+
+captured_circuit(0.0)

frontend/catalyst/compiler.py

@@ -467,6 +467,67 @@
 
         return output_object_name, out_IR
 
+    @debug_logger
+    def is_using_python_compiler(self):
+        xdsl_path = pathlib.Path("xdsl-does-not-use-a-real-path")
+        using_xdsl = False
+        if xdsl_path in self.options.pass_plugins:
+            plugins = self.options.pass_plugins
+            self.options.pass_plugins = tuple(elem for elem in plugins if elem != xdsl_path)
+            using_xdsl = True
+
+        if xdsl_path in self.options.dialect_plugins:
+            plugins = self.options.dialect_plugins
+            self.options.dialect_plugins = tuple(elem for elem in plugins if elem != xdsl_path)
+            using_xdsl = True
+
+        return using_xdsl
+
+    def run_python_compiler(self, mlir_module):
+
+        # TODO: check that xdsl is available to be imported
+        import xdsl
+        from xdsl.context import Context
+        from xdsl import passes
+        from xdsl.dialects import arith, builtin, func, scf, tensor, transform
+        from .python_compiler import quantum, ApplyTransformSequence
+        generic_assembly_format = mlir_module.operation.get_asm(binary=False, print_generic_op_form=True, assume_verified=True)
+        ctx = Context(allow_unregistered=True)
+        ctx.load_dialect(arith.Arith)
+        ctx.load_dialect(builtin.Builtin)
+        ctx.load_dialect(func.Func)
+        ctx.load_dialect(scf.Scf)
+        ctx.load_dialect(tensor.Tensor)
+        ctx.load_dialect(transform.Transform)
+        ctx.load_dialect(quantum.QuantumDialect)
+        # TODO: In order of importance
+        # TODO: Load gradient
+        # TODO: Load the stablehlo dialect? I kind of want to wait until all operations
+        # are properly represented. I am not sure how an incompletely known dialect would be handled
+        # but it is likely worse than an unregistered dialect.
+        # TODO: Load Catalyst
+        # TODO: Load ion/ppm/mbqc/zne...
+        module = xdsl.parser.Parser(ctx, generic_assembly_format).parse_module()
+        pipeline = passes.PipelinePass((ApplyTransformSequence(),))
+        pipeline.apply(ctx, module)
+
+        from jax._src.interpreters import mlir
+        from jaxlib.mlir.dialects import stablehlo
+        from jaxlib.mlir.ir import Context, Module
+        from xdsl.printer import Printer
+        import io
+
+        buffer = io.StringIO()
+        Printer(stream=buffer, print_generic_format=True).print(module)
+        with Context() as ctx:
+            ctx.allow_unregistered_dialects = True
+            ctx.append_dialect_registry(mlir.upstream_dialects)
+            stablehlo.register_dialect(ctx)
+            module = Module.parse(buffer.getvalue())
+        return module
+
+        # TODO: transform the program based on the transform dialect
+
     @debug_logger
     def run(self, mlir_module, *args, **kwargs):
         """Compile an MLIR module to a shared object.
@@ -483,6 +544,10 @@
             (str): filename of shared object
         """
 
+        python_compiler = self.is_using_python_compiler()
+        if python_compiler:
+            mlir_module = self.run_python_compiler(mlir_module)
+
         return self.run_from_ir(
             mlir_module.operation.get_asm(
                 binary=False, print_generic_op_form=False, assume_verified=True

frontend/catalyst/from_plxpr.py

@@ -201,11 +201,10 @@ def handle_qnode(
 }
 
 
-# This is our registration factory for PL transforms. The loop below iterates
-# across the map above and generates a custom handler for each transform.
-# In order to ensure early binding, we pass the PL plxpr transform and the
-# Catalyst pass as arguments whose default values are set by the loop.
-for pl_transform, (pass_name, decomposition) in transforms_to_passes.items():
+# pylint: disable-next=redefined-outer-name
+def register_transform(pl_transform, pass_name, decomposition):
+    """Register pennylane transforms and their conversion to Catalyst transforms"""
+
     # pylint: disable=unused-argument, too-many-arguments, cell-var-from-loop
     @WorkflowInterpreter.register_primitive(pl_transform._primitive)
     def handle_transform(
@@ -251,6 +250,14 @@ def wrapper(*args):
             return self.eval(inner_jaxpr, consts, *non_const_args)
 
 
+# This is our registration factory for PL transforms. The loop below iterates
+# across the map above and generates a custom handler for each transform.
+# In order to ensure early binding, we pass the PL plxpr transform and the
+# Catalyst pass as arguments whose default values are set by the loop.
+for pl_transform, (pass_name, decomposition) in transforms_to_passes.items():
+    register_transform(pl_transform, pass_name, decomposition)
+
+
 class QFuncPlxprInterpreter(PlxprInterpreter):
     """An interpreter that converts plxpr into catalyst-variant jaxpr.
 

frontend/catalyst/passes/xdsl_plugin/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2025 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.
+
+"""XDSL Plugin Interface"""
+
+# This file contains what looks like a plugin
+# but in reality it is just not a "real" MLIR plugin.
+# It just follows the same convention to be able to add passes
+# that are implemented in xDSL
+
+from pathlib import Path
+
+
+def getXDSLPluginAbsolutePath():
+    """Returns a fake path"""
+    return Path("xdsl-does-not-use-a-real-path")
+
+
+def name2pass(name):
+    return getXDSLPluginAbsolutePath(), name

frontend/catalyst/python_compiler/__init__.py

@@ -0,0 +1,120 @@
+from typing import Callable
+
+from dataclasses import dataclass
+from xdsl import context, passes
+from xdsl.utils import parse_pipeline
+from xdsl.dialects import builtin
+
+from dataclasses import dataclass
+
+from xdsl.dialects import builtin, transform
+from xdsl.interpreter import Interpreter
+from xdsl.interpreters.transform import TransformFunctions
+from xdsl.passes import Context, ModulePass
+from xdsl.utils.exceptions import PassFailedException
+
+from xdsl.interpreter import (
+    Interpreter,
+    InterpreterFunctions,
+    PythonValues,
+    ReturnedValues,
+    TerminatorValue,
+    impl,
+    impl_callable,
+    impl_terminator,
+    register_impls,
+)
+
+@register_impls
+class TransformFunctionsExt(TransformFunctions):
+    ctx: Context
+    passes: dict[str, Callable[[], type[ModulePass]]]
+
+    def __init__(
+        self, ctx: Context, available_passes: dict[str, Callable[[], type[ModulePass]]]
+    ):
+        self.ctx = ctx
+        self.passes = available_passes
+
+    @impl(transform.ApplyRegisteredPassOp)
+    def run_apply_registered_pass_op(
+        self,
+        interpreter: Interpreter,
+        op: transform.ApplyRegisteredPassOp,
+        args: PythonValues,
+    ) -> PythonValues:
+        pass_name = op.pass_name.data
+        requested_by_user = passes.PipelinePass.build_pipeline_tuples(
+            self.passes, parse_pipeline.parse_pipeline(pass_name)
+        )
+        # TODO: Switch between catalyst and xDSL
+
+        schedule = tuple(
+            pass_type.from_pass_spec(spec) for pass_type, spec in requested_by_user
+        )
+        pipeline = passes.PipelinePass(schedule)
+        pipeline.apply(self.ctx, args[0])
+        return (args[0],)
+
+from xdsl.transforms import get_all_passes
+
+updated_passes = get_all_passes()
+
+def register_pass(name, _callable):
+    global updated_passes
+    updated_passes[name] = _callable
+
+@dataclass(frozen=True)
+class TransformInterpreterPass(ModulePass):
+    """Transform dialect interpreter"""
+
+    name = "transform-interpreter"
+
+    entry_point: str = "__transform_main"
+
+    @staticmethod
+    def find_transform_entry_point(
+        root: builtin.ModuleOp, entry_point: str
+    ) -> transform.NamedSequenceOp:
+        for op in root.walk():
+            if (
+                isinstance(op, transform.NamedSequenceOp)
+                and op.sym_name.data == entry_point
+            ):
+                return op
+        raise PassFailedException(
+            f"{root} could not find a nested named sequence with name: {entry_point}"
+        )
+
+    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
+        schedule = TransformInterpreterPass.find_transform_entry_point(
+            op, self.entry_point
+        )
+        interpreter = Interpreter(op)
+        global updated_passes
+        interpreter.register_implementations(TransformFunctionsExt(ctx, updated_passes))
+        interpreter.call_op(schedule, (op,))
+
+@dataclass(frozen=True)
+class ApplyTransformSequence(passes.ModulePass):
+    name = "apply-transform-sequence"
+
+    def apply(self, ctx: context.MLContext, module: builtin.ModuleOp) -> None:
+        nested_modules = []
+        for region in module.regions:
+          for block in region.blocks:
+              for op in block.ops:
+                  if isinstance(op, builtin.ModuleOp):
+                      nested_modules.append(op)
+
+        pipeline = passes.PipelinePass((TransformInterpreterPass(),))
+        for op in nested_modules:
+            pipeline.apply(ctx, op)
+
+        for op in nested_modules:
+            for region in op.regions:
+                for block in region.blocks:
+                    for op in block.ops:
+                        if isinstance(op, builtin.ModuleOp) and op.get_attr_or_prop("transform.with_named_sequence"):
+                            block.erase_op(op)
+