add initial impl

Author: erick-xanadu

pennylane/compiler/python_compiler/impl.py

@@ -12,6 +12,8 @@
 
 from .quantum_dialect import QuantumDialect as Quantum
 
+from .transforms import ApplyTransformSequence
+
 class Compiler:
 
     def run(self, jmod: jaxModule) -> jaxModule:
@@ -33,7 +35,7 @@ def run(self, jmod: jaxModule) -> jaxModule:
         ctx.load_dialect(Quantum)
 
         xmod: builtin.ModuleOp = Parser(ctx, gentxtmod).parse_module()
-        pipeline = passes.PipelinePass((ApplyTransformSequence(),))
+        pipeline = PipelinePass((ApplyTransformSequence(),))
         # xmod is modified in place
         pipeline.apply(ctx, xmod)
 
@@ -43,6 +45,6 @@ def run(self, jmod: jaxModule) -> jaxModule:
             ctx.allow_unregistered_dialects = True
             ctx.append_dialect_registry(mlir.upstream_dialects)
             stablehlo.register_dialect(ctx)
-            newmod: jaxModule = Module.parse(buffer.getvalue())
+            newmod: jaxModule = jaxModule.parse(buffer.getvalue())
 
         return newmod

pennylane/compiler/python_compiler/transforms/__init__.py

@@ -1,5 +1,5 @@
-from apply_transform_sequence import ApplyTransformSequence
-from transform_interpreter import TransformInterpreterPass
+from .apply_transform_sequence import ApplyTransformSequence
+from .transform_interpreter import TransformInterpreterPass
 
 __all__ = [
         "ApplyTransformSequence",

pennylane/compiler/python_compiler/transforms/apply_transform_sequence.py

@@ -1,8 +1,10 @@
+from dataclasses import dataclass
+
 from xdsl.dialects import builtin
 from xdsl.context import MLContext
 from xdsl.passes import ModulePass, PipelinePass
 
-from transform_interpreter import TransformInterpreterPass
+from .transform_interpreter import TransformInterpreterPass
 
 @dataclass(frozen=True)
 class ApplyTransformSequence(ModulePass):
@@ -16,7 +18,7 @@ def apply(self, ctx: MLContext, module: builtin.ModuleOp) -> None:
                     if isinstance(op, builtin.ModuleOp):
                         nested_modules.append(op)
 
-        pipeline = PipelinePass((TransformInterpreterPass({}),))
+        pipeline = PipelinePass((TransformInterpreterPass(passes={}),))
         for op in nested_modules:
             pipeline.apply(ctx, op)
 

pennylane/compiler/python_compiler/transforms/transform_interpreter/__init__.py

@@ -0,0 +1,5 @@
+from .transform_interpreter_catalyst import TransformInterpreterPass
+
+__all__ = [
+    "TransformInterpreterPass"
+]

pennylane/compiler/python_compiler/transforms/transform_interpreter/interpreter/__init__.py

@@ -0,0 +1,5 @@
+from .impl import TransformFunctionsExt
+
+__all__ = [
+    "TransformFunctionsExt"
+]

pennylane/compiler/python_compiler/transforms/transform_interpreter/interpreter/impl.py

@@ -1,13 +1,15 @@
 import io
 from typing import Callable
 
-from xdsl.interpreter import Interpreter, PythonValues, register_impls
+from xdsl.interpreter import Interpreter, PythonValues, impl, register_impls
 from xdsl.interpreters.transform import TransformFunctions
+from xdsl.context import Context
 from xdsl.dialects import transform
 from xdsl.parser import Parser
 from xdsl.passes import ModulePass, PipelinePass
 from xdsl.printer import Printer
 from xdsl.rewriter import Rewriter
+from xdsl.utils import parse_pipeline
 
 from catalyst.compiler import _quantum_opt
 

pennylane/compiler/python_compiler/transforms/transform_interpreter/transform_interpreter_catalyst.py

@@ -1,19 +1,23 @@
 from dataclasses import dataclass
+from typing import Callable
 
+from xdsl.context import Context
 from xdsl.dialects import builtin, transform
 from xdsl.interpreters import Interpreter
 from xdsl.passes import ModulePass
 
-from interpreter import TransformFunctionsExt
+from .interpreter import TransformFunctionsExt
 
-@dataclass(frozen=True)
 class TransformInterpreterPass(ModulePass):
     """Transform dialect interpreter"""
 
+    passes: dict[str, Callable[[], type[ModulePass]]]
     name = "transform-interpreter"
 
     entry_point: str = "__transform_main"
-    passes: dict[str, Callable[[], type[ModulePass]]]
+
+    def __init__(self, passes):
+        self.passes = passes
 
     @staticmethod
     def find_transform_entry_point(
@@ -34,6 +38,6 @@ def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
             op, self.entry_point
         )
         interpreter = Interpreter(op)
-        interpreter.register_implementations(TransformFunctionsExt(ctx, passes))
+        interpreter.register_implementations(TransformFunctionsExt(ctx, self.passes))
         schedule.parent_op().detach()
         interpreter.call_op(schedule, (op,))