Redefine ApplyRegisteredPass in the Transform dialect

pennylane/compiler/python_compiler/dialects/transform.py

@@ -0,0 +1,136 @@
+# 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.
+"""
+This file contains an updated version of the transform dialect.
+As of the time of writing, xDSL uses the MLIR released with LLVM's
+version 20.1.7. However, https://github.com/PennyLaneAI/catalyst/pull/1916
+will be updating MLIR where the transform dialect has the
+`apply_registered_pass` operation re-defined.
+
+See the following changelog on the above PR
+
+    Things related to transform.apply_registered_pass op:
+
+    It now takes in a dynamic_options
+
+    [MLIR][Transform] Allow ApplyRegisteredPassOp to take options as
+    a param llvm/llvm-project#142683. We don't need to use this as all our pass options are static.
+    https://github.com/llvm/llvm-project/pull/142683
+
+    The options it takes in are now dictionaries instead of strings
+    [MLIR][Transform] apply_registered_pass op's options as a dict llvm/llvm-project#143159
+    https://github.com/llvm/llvm-project/pull/143159
+
+This file will re-define the apply_registered_pass operation in xDSL
+and the transform dialect.
+
+Once xDSL moves to a newer version of MLIR, these changes should
+be contributed upstream.
+"""
+
+from xdsl.dialects.builtin import DictionaryAttr
+
+# pylint: disable=unused-wildcard-import,wildcard-import
+from xdsl.dialects.transform import *
+
+
+@irdl_op_definition
+# pylint: disable=function-redefined
+class ApplyRegisteredPassOp(IRDLOperation):
+    """
+    See external [documentation](https://mlir.llvm.org/docs/Dialects/Transform/#transformapply_registered_pass-transformapplyregisteredpassop).
+    """
+
+    name = "transform.apply_registered_pass"
+
+    options = prop_def(DictionaryAttr, default_value=DictionaryAttr({}))
+    pass_name = prop_def(StringAttr)
+    target = operand_def(TransformHandleType)
+    result = result_def(TransformHandleType)
+    # While this assembly format doesn't match
+    # the one in upstream MLIR,
+    # this is because xDSL currently lacks CustomDirectives
+    # https://mlir.llvm.org/docs/DefiningDialects/Operations/#custom-directives
+    # https://github.com/xdslproject/xdsl/pull/4829
+    # However, storing the property in the attribute should still work
+    # specially when parsing and printing in generic format.
+    # Which is how Catalyst and XDSL currently communicate at the moment.
+    # TODO: Add test.
+    assembly_format = "$pass_name `to` $target attr-dict `:` functional-type(operands, results)"
+    irdl_options = [ParsePropInAttrDict()]
+
+    def __init__(
+        self,
+        pass_name: str | StringAttr,
+        target: SSAValue,
+        options: dict | DictionaryAttr | None = None,
+    ):
+        if isinstance(pass_name, str):
+            pass_name = StringAttr(pass_name)
+
+        if isinstance(options, dict):
+            options = DictionaryAttr(options)
+
+        super().__init__(
+            properties={
+                "pass_name": pass_name,
+                "options": options,
+            },
+            operands=[target],
+            result_types=[target.type],
+        )
+
+
+# Copied over from xDSL's sources
+# the main difference will be the use
+# of a different ApplyRegisteredPassOp
+Transform = Dialect(
+    "transform",
+    [
+        ApplyRegisteredPassOp,
+        GetConsumersOfResultOp,
+        GetDefiningOp,
+        GetParentOp,
+        GetProducerOfOperandOp,
+        GetResultOp,
+        GetTypeOp,
+        IncludeOp,
+        MatchOperationEmptyOp,
+        MatchOperationNameOp,
+        MatchParamCmpIOp,
+        MergeHandlesOp,
+        ParamConstantOp,
+        SplitHandleOp,
+        SequenceOp,
+        YieldOp,
+        TileOp,
+        TileToForallOp,
+        SelectOp,
+        NamedSequenceOp,
+        CastOp,
+        MatchOp,
+    ],
+    [
+        # Types
+        AffineMapType,
+        AnyOpType,
+        AnyValueType,
+        AnyParamType,
+        OperationType,
+        ParamType,
+        TypeParamType,
+        # Attributes
+        FailurePropagationModeAttr,
+    ],
+)

pennylane/compiler/python_compiler/jax_utils.py

@@ -30,12 +30,12 @@
 from xdsl.dialects import scf as xscf
 from xdsl.dialects import stablehlo as xstablehlo
 from xdsl.dialects import tensor as xtensor
-from xdsl.dialects import transform as xtransform
 from xdsl.ir import Dialect as xDialect
 from xdsl.parser import Parser as xParser
 from xdsl.traits import SymbolTable as xSymbolTable
 
 from .dialects import MBQC, Quantum
+from .dialects import transform as xtransform
 
 JaxJittedFunction: TypeAlias = _jax.PjitFunction  # pylint: disable=c-extension-no-member
 

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

@@ -23,7 +23,7 @@
 
 from catalyst.compiler import _quantum_opt  # pylint: disable=protected-access
 from xdsl.context import Context
-from xdsl.dialects import builtin, transform
+from xdsl.dialects import builtin
 from xdsl.interpreter import Interpreter, PythonValues, impl, register_impls
 from xdsl.interpreters.transform import TransformFunctions
 from xdsl.parser import Parser
@@ -32,6 +32,8 @@
 from xdsl.rewriter import Rewriter
 from xdsl.utils.exceptions import PassFailedException
 
+from pennylane.compiler.python_compiler.dialects import transform
+
 
 # pylint: disable=too-few-public-methods
 @register_impls

tests/python_compiler/test_python_compiler.py

@@ -25,14 +25,16 @@
 jax = pytest.importorskip("jax")
 jaxlib = pytest.importorskip("jaxlib")
 xdsl = pytest.importorskip("xdsl")
+filecheck = pytest.importorskip("filecheck")
 
 from catalyst import CompileError
 from xdsl import passes
 from xdsl.context import Context
-from xdsl.dialects import builtin, transform
+from xdsl.dialects import builtin
 from xdsl.interpreters import Interpreter
 
 from pennylane.compiler.python_compiler import Compiler
+from pennylane.compiler.python_compiler.dialects import transform
 from pennylane.compiler.python_compiler.jax_utils import (
     jax_from_docstring,
     module,
@@ -205,5 +207,22 @@ def program():
     assert captured.out.strip() == "hello world"
 
 
+def test_transform_dialect_update(run_filecheck):
+
+    program = """
+        "builtin.module"() ({
+            "transform.named_sequence"() <{function_type = (!transform.any_op) -> (), sym_name = "__transform_main"}> ({
+            ^bb0(%arg0: !transform.any_op):
+                %0 = "transform.structured.match"(%arg0) <{ops = ["func.func"]}> : (!transform.any_op) -> !transform.any_op
+                // CHECK: "invalid-option"
+                %1 = "transform.apply_registered_pass"(%0) <{options = {"invalid-option" = 1 : i64}, pass_name = "canonicalize"}> : (!transform.any_op) -> !transform.any_op
+                "transform.yield"() : () -> ()
+            }) : () -> ()
+        }) {transform.with_named_sequence} : () -> ()
+    """
+
+    run_filecheck(program, ())
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])