Add roundtrip testing capability to filecheck fixture (#8049)

### Before submitting

Please complete the following checklist when submitting a PR:

- [ ] All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to
the
      test directory!

- [ ] All new functions and code must be clearly commented and
documented.
If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [ ] Ensure that the test suite passes, by running `make test`.

- [ ] Add a new entry to the `doc/releases/changelog-dev.md` file,
summarizing the
      change, and including a link back to the PR.

- [ ] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


------------------------------------------------------------------------------------------------------------

**Context:**
xDSL uses round-trip lit tests for testing its dialect definitions. This
involves parsing an MLIR string to xDSL, printing in generic format, and
then parsing the generic format string again into xDSL before running
filecheck. We add the same ability here.

**Description of the Change:**
* Add a `roundtrip` boolean kwarg to `run_filecheck` to allow roundtrip
testing.
* Add a `verify` boolean kwarg to run `op.verify()` on the xDSL module
being tested.

**Benefits:**
More robust lit tests.

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-97491]

---------

Co-authored-by: Mehrdad Malek <[email protected]>

Author: mudit2812

doc/releases/changelog-dev.md

@@ -523,6 +523,10 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* Add capability for roundtrip testing and module verification to the Python compiler `run_filecheck` and
+`run_filecheck_qjit` fixtures.
+  [(#8049)](https://github.com/PennyLaneAI/pennylane/pull/8049)
+
 * Improve type hinting internally.
   [(#8086)](https://github.com/PennyLaneAI/pennylane/pull/8086)
 

pennylane/compiler/python_compiler/jax_utils.py

@@ -77,7 +77,8 @@ def __init__(
 
         extra_dialects = extra_dialects or ()
         for dialect in self.default_dialects + tuple(extra_dialects):
-            self.ctx.load_dialect(dialect)
+            if self.ctx.get_optional_dialect(dialect.name) is None:
+                self.ctx.load_dialect(dialect)
 
 
 def _module_inline(func: JaxJittedFunction, *args, **kwargs) -> jModule:

tests/python_compiler/conftest.py

@@ -28,6 +28,7 @@
     from xdsl.context import Context
     from xdsl.dialects import test
     from xdsl.passes import PassPipeline
+    from xdsl.printer import Printer
 
     from pennylane.compiler.python_compiler import Compiler
     from pennylane.compiler.python_compiler.jax_utils import (
@@ -38,7 +39,7 @@
     deps_available = False
 
 
-def _run_filecheck_impl(program_str, pipeline=()):
+def _run_filecheck_impl(program_str, pipeline=(), verify=False, roundtrip=False):
     """Run filecheck on an xDSL module, comparing it to a program string containing
     filecheck directives."""
     if not deps_available:
@@ -47,13 +48,27 @@ def _run_filecheck_impl(program_str, pipeline=()):
     ctx = Context(allow_unregistered=True)
     xdsl_module = QuantumParser(ctx, program_str, extra_dialects=(test.Test,)).parse_module()
 
+    if roundtrip:
+        # Print generic format
+        stream = StringIO()
+        Printer(stream=stream, print_generic_format=True).print_op(xdsl_module)
+        xdsl_module = QuantumParser(ctx, stream.getvalue()).parse_module()
+
+    if verify:
+        xdsl_module.verify()
+
     pipeline = PassPipeline(pipeline)
     pipeline.apply(ctx, xdsl_module)
 
+    if verify:
+        xdsl_module.verify()
+
+    stream = StringIO()
+    Printer(stream).print_op(xdsl_module)
     opts = parse_argv_options(["filecheck", __file__])
     matcher = Matcher(
         opts,
-        FInput("no-name", str(xdsl_module)),
+        FInput("no-name", stream.getvalue()),
         Parser(opts, StringIO(program_str), *pattern_for_opts(opts)),
     )
 
@@ -63,7 +78,23 @@ def _run_filecheck_impl(program_str, pipeline=()):
 
 @pytest.fixture(scope="function")
 def run_filecheck():
-    """Fixture to run filecheck on an xDSL module."""
+    """Fixture to run filecheck on an xDSL module.
+
+    This fixture uses FileCheck to verify the correctness of a parsed MLIR string. Testers
+    can provide a pass pipeline to transform the IR, and verify correctness by including
+    FileCheck directives as comments in the input program string.
+
+    Args:
+        program_str (str): The MLIR string containing the input program and FileCheck directives
+        pipeline (tuple[ModulePass]): A sequence containing all passes that should be applied
+            before running FileCheck
+        verify (bool): Whether or not to verify the IR after parsing and transforming.
+            ``False`` by default.
+        roundtrip (bool): Whether or not to use round-trip testing. This is useful for dialect
+            tests to verify that xDSL both parses and prints the IR correctly. If ``True``, we parse
+            the program string into an xDSL module, print it in generic format, and then parse the
+            generic program string back to an xDSL module. ``False`` by default.
+    """
     if not deps_available:
         pytest.skip("Cannot run lit tests without xDSL and filecheck.")
 
@@ -90,7 +121,7 @@ def _get_filecheck_directives(qjit_fn):
     return "\n".join(filecheck_directives)
 
 
-def _run_filecheck_qjit_impl(qjit_fn):
+def _run_filecheck_qjit_impl(qjit_fn, verify=False):
     """Run filecheck on a qjit-ed function, using FileCheck directives in its inline
     comments to assert correctness."""
     if not deps_available:
@@ -107,6 +138,9 @@ def _run_filecheck_qjit_impl(qjit_fn):
     )
     xdsl_module = parse_generic_to_xdsl_module(mod_str)
 
+    if verify:
+        xdsl_module.verify()
+
     opts = parse_argv_options(["filecheck", __file__])
     matcher = Matcher(
         opts,
@@ -127,6 +161,11 @@ def run_filecheck_qjit():
     output IR against FileCheck directives that may be present in the source
     function as inline comments.
 
+    Args:
+        qjit_fn (Callable): The QJIT object on which we want to run lit tests
+        verify (bool): Whether or not to verify the IR after parsing and transforming.
+            ``False`` by default.
+
     An example showing how to use the fixture is shown below. We apply the
     ``merge_rotations_pass`` and check that there is only one rotation in
     the final IR:

tests/python_compiler/dialects/test_catalyst_dialect.py

@@ -105,4 +105,4 @@ def test_assembly_format(run_filecheck):
     %callback_result = catalyst.callback_call @callback_func(%val) : (f64) -> f64
     """
 
-    run_filecheck(program)
+    run_filecheck(program, roundtrip=True)

tests/python_compiler/dialects/test_mbqc_dialect.py

@@ -98,7 +98,7 @@ def test_assembly_format(run_filecheck):
     %graph_reg = mbqc.graph_state_prep (%adj_matrix : tensor<6xi1>) [init "Hadamard", entangle "CZ"] : !quantum.reg
     """
 
-    run_filecheck(program)
+    run_filecheck(program, roundtrip=True)
 
 
 class TestMeasureInBasisOp:

tests/python_compiler/dialects/test_quantum_dialect.py

@@ -14,8 +14,6 @@
 
 """Unit test module for pennylane/compiler/python_compiler/dialects/quantum.py."""
 
-import io
-
 import pytest
 
 # pylint: disable=wrong-import-position
@@ -110,7 +108,7 @@ def test_all_attributes_names(attr):
     assert attr.name == expected_name
 
 
-def test_assembly_format():
+def test_assembly_format(run_filecheck):
     program = """
     // CHECK: quantum.alloc(1) : !quantum.reg
     %qreg_alloc_static = quantum.alloc(1) : !quantum.reg
@@ -172,30 +170,4 @@ def test_assembly_format():
     %mres2, %out_qubit2 = quantum.measure %qubit postselect 0 : i1, !quantum.bit
     """
 
-    ctx = xdsl.context.Context()
-    from xdsl.dialects import builtin, func, test
-
-    ctx.load_dialect(builtin.Builtin)
-    ctx.load_dialect(func.Func)
-    ctx.load_dialect(test.Test)
-    ctx.load_dialect(Quantum)
-
-    module = xdsl.parser.Parser(ctx, program).parse_module()
-
-    from filecheck.finput import FInput
-    from filecheck.matcher import Matcher
-    from filecheck.options import parse_argv_options
-    from filecheck.parser import Parser, pattern_for_opts
-
-    opts = parse_argv_options(["filecheck", __file__])
-    matcher = Matcher(
-        opts,
-        FInput("no-name", str(module)),
-        Parser(opts, io.StringIO(program), *pattern_for_opts(opts)),
-    )
-
-    assert matcher.run() == 0
-
-
-if __name__ == "__main__":
-    pytest.main(["-x", __file__])
+    run_filecheck(program, roundtrip=True)